모멘텀 제어: 저차원 근사화를 통한 옵티마이저 상태 재고

Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

Adam 및 Muon과 같은 현대적인 최적화 알고리즘은 대규모 언어 모델 훈련에 필수적이지만, 이들은 1차 및 2차 모멘텀에 의존하여 상당한 메모리 오버헤드를 발생시키며, 이는 확장성과 계산 효율성을 저해합니다. 본 연구에서는 이러한 모멘텀에 사용되는 지수 이동 평균(EMA)을 온라인 그래디언트 흐름을 통한 선형 회귀 학습으로 재해석합니다. 이러한 동등성을 바탕으로, 효율적인 사전 훈련을 위해 설계된 새로운 저차원 최적화 알고리즘인 LoRA-Pre를 소개합니다. 구체적으로, LoRA-Pre는 온라인 선형 학습 모델 내에서 전체 모멘텀 행렬을 작은 저차원 부분 공간으로 분해하여 최적화기의 메모리 공간을 줄임으로써, 최적화 성능을 유지하면서 메모리 효율성을 향상시킵니다. Llama 아키텍처 기반의 모델을 60M에서 1B 파라미터까지 확장하여 LoRA-Pre의 효과를 실험적으로 검증했습니다. LoRA-Pre는 모든 모델 크기에서 최고의 성능을 달성했습니다. 주목할 만한 점은 LoRA-Pre가 뛰어난 저차원 효율성을 보여주며, 기준 방법보다 1/8 수준의 낮은 차원만으로도 유사하거나 더 나은 결과를 얻는다는 것입니다. 사전 훈련 외에도, LoRA-Pre의 효과를 미세 조정 시나리오에서도 평가했습니다. 동일한 차원을 사용했을 때, LoRA-Pre는 모든 효율적인 미세 조정 기준 성능보다 꾸준히 우수한 성능을 보였습니다. 특히, 표준 LoRA와 비교하여 LoRA-Pre는 Llama-3.1-8B 모델에서 3.14 포인트, Llama-2-7B 모델에서 6.17 포인트의 상당한 성능 향상을 달성하여, 사전 훈련 및 미세 조정 모두에서 본 연구의 접근 방식이 효과적임을 입증합니다. 본 연구의 코드는 다음 주소에서 공개적으로 이용할 수 있습니다: https://github.com/mrflogs/LoRA-Pre.

Modern optimizers like Adam and Muon are central to training large language models, but their reliance on first- and second-order momenta introduces significant memory overhead, which constrains scalability and computational efficiency. In this work, we reframe the exponential moving average (EMA) used in these momenta as the training of a linear regressor via online gradient flow. Building on this equivalence, we introduce LoRA-Pre, a novel low-rank optimizer designed for efficient pre-training. Specifically, LoRA-Pre reduces the optimizer's memory footprint by decomposing the full momentum matrix into a compact low-rank subspace within the online linear learner, thereby maintaining optimization performance while improving memory efficiency. We empirically validate LoRA-Pre's efficacy by pre-training models from the Llama architecture family, scaling from 60M to 1B parameters. LoRA-Pre achieves the highest performance across all model sizes. Notably, LoRA-Pre demonstrates remarkable rank efficiency, achieving comparable or superior results using only 1/8 the rank of baseline methods. Beyond pre-training, we evaluate LoRA-Pre's effectiveness in fine-tuning scenarios. With the same rank, LoRA-Pre consistently outperforms all efficient fine-tuning baselines. Specifically, compared to standard LoRA, LoRA-Pre achieves substantial improvements of 3.14 points on Llama-3.1-8B and 6.17 points on Llama-2-7B, validating our approach's effectiveness across both pre-training and fine-tuning paradigms. Our code is publicly available at https://github.com/mrflogs/LoRA-Pre.