[Paper] 모멘텀 제어: 저랭크 근사화를 통한 옵티마이저 상태 재고

Source: arXiv - 2602.24283v1

Overview

오늘날 대규모 언어 모델을 학습하는 데는 Adam 및 그 변형과 같은 적응형 옵티마이저가 크게 의존합니다. 이러한 방법은 수렴 속도를 높이는 반면, 매 파라미터마다 1차 및 2차 “모멘텀” 벡터를 저장해야 하므로 메모리 사용량이 증가하고 주어진 GPU에서 학습할 수 있는 모델 크기가 제한됩니다. 논문 *“Taming Momentum: Rethinking Optimizer States Through Low‑Rank Approximation”*은 이러한 모멘텀 버퍼를 온라인 선형 회귀 모델의 파라미터로 보는 새로운 관점을 제시하고, 이를 활용해 저‑랭크 행렬 분해로 압축합니다. 그 결과 옵티마이저 LoRA‑Pre는 메모리 오버헤드를 크게 줄이면서도 사전 학습 및 파인‑튜닝 작업 모두에서 최첨단 성능을 제공합니다.

주요 기여

• EMA 모멘텀 재해석: Adam‑style 옵티마이저에서 사용되는 지수 이동 평균이 온라인 그래디언트 흐름을 통한 선형 회귀기 훈련과 수학적으로 동등함을 보여준다.
• 저랭크 옵티마이저 설계 (LoRA‑Pre): 전체 모멘텀 행렬을 압축된 저랭크 부분공간으로 분해하여 옵티마이저 상태 크기를 크게 줄인다.
• Llama 계열에 대한 실증적 검증: 60 M에서 1 B 파라미터까지 다양한 모델에서 일관된 성능 향상을 보여주며, 모든 베이스라인 중 최고의 결과를 달성한다.
• 랭크 효율성: 랭크를 1/8만 사용하여(즉, 저장 파라미터가 훨씬 적음) 베이스라인 성능과 동등하거나 능가한다.
• 파인튜닝 우수성: 동일한 랭크 예산 하에서 Llama‑3.1‑8B에서는 3.14 포인트, Llama‑2‑7B에서는 6.17 포인트 만큼 인기 있는 효율적 파인튜닝 방법(예: 표준 LoRA)을 능가한다.
• 오픈소스 공개: 전체 구현 코드와 학습 스크립트가 GitHub에 공개되어 있다.

방법론

1. EMA를 온라인 선형 회귀로 보기

   • 이 논문은 Adam의 1차 모멘텀의 핵심인 그래디언트의 지수 이동 평균(EMA)을 온라인 선형 회귀 문제의 해로 표현하는 것으로 시작합니다. 이 관점에서 각 옵티마이저 단계는 과거 그래디언트로부터 현재 그래디언트를 예측하려는 작은 선형 모델을 업데이트하는 과정과 같습니다.

2. 모멘텀 행렬의 저‑랭크 근사

   • 전체 크기의 모멘텀 행렬(크기 = 파라미터 수 × 히든 차원)을 저장하는 대신, LoRA‑Pre는 이를 두 개의 작은 행렬 U와 V(랭크‑r 근사)로 분해합니다. U·Vᵀ의 곱은 실시간으로 모멘텀 추정치를 재구성합니다.
   • U와 V를 업데이트하는 방식은 EMA 해석에서 도출된 동일한 온라인 그래디언트 흐름 역학을 따르며, 저‑랭크 표현이 원래 옵티마이저 역학에 충실하도록 보장합니다.

3. 기존 학습 파이프라인과의 통합

   • LoRA‑Pre는 Adam/Muon의 드롭‑인 대체물입니다: 옵티마이저 API는 변하지 않으며, 내부 상태 처리만 다릅니다.
   • 사전 학습 단계에서 랭크 r은 하이퍼파라미터로 취급됩니다; 저자들은 매우 낮은 랭크(예: 1 B 모델에 대해 r = 8)만으로도 충분함을 발견했습니다.

4. 평가 프로토콜

   • Llama 계열 모델(60 M, 160 M, 410 M, 1 B)을 표준 언어 모델링 코퍼스에서 사전 학습합니다.
   • 파인‑튜닝 실험에서는 Alpaca, MMLU와 같은 다운스트림 벤치마크를 사용해 LoRA, QLoRA 및 기타 파라미터 효율적 방법들과 비교합니다.

Results & Findings

• Pre‑training: LoRA‑Pre는 훨씬 적은 옵티마이저 상태를 사용함에도 불구하고 전체 메모리 Adam 기준을 지속적으로 능가합니다.
• Fine‑tuning: 동일한 낮은 차원(rank)에서도 LoRA‑Pre는 표준 LoRA보다 3.14점(Llama‑3.1‑8B) 및 6.17점(Llama‑2‑7B) 더 높은 성능을 보입니다.
• Rank Efficiency: 실험 결과 전체 차원의 약 1/8 정도를 초과하는 rank에서는 수익이 감소함을 보여주며, 대부분의 유용한 모멘텀 정보가 저차원 부분공간에 존재함을 확인했습니다.

실용적 시사점

• 범용 하드웨어에서 확장: 옵티마이저 메모리를 원래 크기의 약 12 %로 줄임으로써, 개발자는 동일한 GPU 메모리 예산으로 더 큰 모델을 학습하거나 단일 노드에 여러 실험을 배치할 수 있습니다.
• 더 빠른 반복 주기: 작은 옵티마이저 상태는 GPU와 호스트 메모리 간 데이터 이동을 감소시켜, 특히 옵티마이저 동기화가 병목인 다중 GPU 파이프라인에서 약간의 속도 향상으로 이어질 수 있습니다.
• 파라미터 효율적인 파인튜닝: LoRA‑Pre는 기존 PEFT(파라미터 효율 파인튜닝) 프레임워크와 결합될 수 있어, “이중 효율” 향상을 제공합니다—저랭크 어댑터 가중치 와 저랭크 옵티마이저 상태.
• 단순화된 인프라: 상태 크기가 이미 작기 때문에 맞춤형 체크포인트 트릭(예: 옵티마이저 상태 샤딩)이 필요 없습니다. 이는 GPU 메모리 GB당 요금을 부과하는 클라우드 플랫폼에 배포를 용이하게 합니다.
• 다른 분야에 대한 잠재력: EMA‑as‑online‑regressor 통찰은 언어 모델에만 국한되지 않으며, Adam 계열 옵티마이저를 사용하는 대규모 학습(비전, 음성, 강화 학습)에서도 메모리 절감을 위해 LoRA‑Pre를 적용할 수 있습니다.

제한 사항 및 향후 작업

• Rank selection is still heuristic: 논문은 경험적 가이드를 제공하지만, 모델 또는 레이어별 최적 랭크를 자동으로 선택하는 방법이 있으면 접근 방식을 더 플러그‑앤‑플레이하게 만들 수 있습니다.
• Compatibility with second‑order moments: LoRA‑Pre는 1차 모멘트(EMA)에 초점을 맞춥니다. Adam의 분산 항(2차 모멘트)에 저랭크 아이디어를 확장하면 추가적인 절감 효과를 얻을 수 있지만, 이는 향후 연구 과제로 남겨져 있습니다.
• Benchmarks limited to Llama family: 저자들은 1 B 파라미터까지의 Llama 모델을 평가했습니다; 수십 억 파라미터 규모의 실제 대형 모델 및 다른 아키텍처(e.g., GPT‑NeoX, T5)에서 테스트하면 보편적 적용 가능성 주장을 강화할 수 있습니다.
• Potential numerical stability concerns: 저랭크 분해는 매우 깊거나 희소하게 업데이트되는 레이어에서 조건수 문제를 일으킬 수 있습니다; 논문에서는 랭크를 너무 낮게 설정하면 가끔 수렴 실패가 발생한다는 점을 언급하며, 강인한 보호 장치가 필요함을 시사합니다.

LoRA‑Pre를 직접 사용해 보고 싶다면, 저자들이 코드를 https://github.com/mrflogs/LoRA-Pre 에서 오픈소스로 공개했습니다. 기존 PyTorch 학습 스크립트에 적용하고 원하는 랭크를 설정하면, 모델 품질을 손상시키지 않으면서 즉시 메모리 절감 효과를 확인할 수 있습니다.

저자

• Zhengbo Wang
• Jian Liang
• Ran He
• Zilei Wang
• Tieniu Tan

논문 정보

• arXiv ID: 2602.24283v1
• Categories: cs.LG, cs.AI, cs.CL
• Published: 2026년 2월 27일
• PDF: Download PDF