COAP: Memory-Efficient Training with Correlation-Aware Gradient Projection
This addresses memory constraints for researchers and practitioners training large models in vision, language, and multimodal domains, offering an incremental improvement over existing low-rank gradient projection methods.
The paper tackles the problem of high memory usage in training large neural networks by proposing COAP, a memory-efficient method that reduces optimizer memory by 61% for LLaMA-1B with minimal time cost and achieves 81% reduction with 4x speedup for LLaVA-v1.5-7B fine-tuning while maintaining or improving performance.
Training large-scale neural networks in vision, and multimodal domains demands substantial memory resources, primarily due to the storage of optimizer states. While LoRA, a popular parameter-efficient method, reduces memory usage, it often suffers from suboptimal performance due to the constraints of low-rank updates. Low-rank gradient projection methods (e.g., GaLore, Flora) reduce optimizer memory by projecting gradients and moment estimates into low-rank spaces via singular value decomposition or random projection. However, they fail to account for inter-projection correlation, causing performance degradation, and their projection strategies often incur high computational costs. In this paper, we present COAP (Correlation-Aware Gradient Projection), a memory-efficient method that minimizes computational overhead while maintaining training performance. Evaluated across various vision, language, and multimodal tasks, COAP outperforms existing methods in both training speed and model performance. For LLaMA-1B, it reduces optimizer memory by 61% with only 2% additional time cost, achieving the same PPL as AdamW. With 8-bit quantization, COAP cuts optimizer memory by 81% and achieves 4x speedup over GaLore for LLaVA-v1.5-7B fine-tuning, while delivering higher accuracy.