MLorc: Momentum Low-rank Compression for Memory Efficient Large Language Model Adaptation
Wei Shen, Zhang Yaxiang, Minhui Huang, Mengfan Xu, Jiawei Zhang, Cong Shen
TL;DR
MLorc tackles the memory bottleneck of fine-tuning large language models by compressing and reconstructing the momentum of matrix parameters using Randomized SVD, enabling full-parameter updates with reduced memory. The method preserves training dynamics more faithfully than gradient-based approaches and provides convergence guarantees for the Lion optimizer. Empirically, MLorc outperforms LoRA, GaLore, and LDAdamW across NLG and NLU tasks, closely matching or exceeding full fine-tuning at small ranks (e.g., r = 4) while maintaining competitive runtime and memory. This yields a practical pathway to memory-efficient, high-quality fine-tuning of large models and suggests potential extensions to pre-training and larger architectures.
Abstract
With increasing size of large language models (LLMs), full-parameter fine-tuning imposes substantial memory demands. To alleviate this, we propose a novel memory-efficient training paradigm called Momentum Low-rank compression (MLorc). The key idea of MLorc is to compress and reconstruct the momentum of matrix parameters during training to reduce memory consumption. Compared to LoRA, MLorc avoids enforcing a fixed-rank constraint on weight update matrices and thus enables full-parameter learning. Compared to GaLore, MLorc directly compress the momentum rather than gradients, thereby better preserving the training dynamics of full-parameter fine-tuning. We provide a theoretical guarantee for its convergence under mild assumptions. Empirically, MLorc consistently outperforms other memory-efficient training methods, matches or even exceeds the performance of full fine-tuning at small ranks (e.g., $r=4$), and generalizes well across different optimizers -- all while not compromising time or memory efficiency.