Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

이 논문은 Adam 및 Muon 같은 최적화기의 메모리 오버헤드를 줄이기 위해 모멘텀 행렬을 저랭크 부분 공간으로 분해하는 새로운 최적화기 LoRA-Pre 를 제안하고, 사전 학습과 미세 조정 모두에서 기존 방법보다 뛰어난 성능과 메모리 효율성을 입증했습니다.

핵심

1. 문제: 거대한 '학습 노트'가 메모리를 차지해요

대형 인공지능 모델을 훈련시킨다는 건, 수천만 권의 책을 읽으며 배운 내용을 기억하는 것과 비슷합니다.
기존의 최적화 방법 (Adam 같은 것들) 은 이 학습 과정에서 **매우 방대한 양의 '학습 노트 (기억 상태)'**를 저장합니다.

• 비유: 학생이 시험을 보기 위해 모든 교과서를 다 외우려다 보니, 책가방 (컴퓨터 메모리) 이 너무 커져서 더 이상 다른 책을 넣을 공간이 없어진 상황입니다.
• 결과: 더 큰 모델을 만들려면 더 큰 가방이 필요하고, 이는 비용과 시간을 엄청나게 잡아먹습니다.

2. 핵심 아이디어: "기억은 단순한 '요약'일 뿐이다"

저자들은 이 '학습 노트'를 다시 바라봤습니다.

• 기존 생각: "기억 (Momentum) 은 과거의 모든 정보를 정확히 저장해야 하는 무거운 데이터야."
• 새로운 통찰 (이 논문의 핵심): "아니야, 기억은 사실 **과거의 흐름을 예측하는 '간단한 공학자 (선형 회귀)'**가 하는 일이야."

비유:
과거의 시험 문제를 다 외울 필요 없이, "최근 출제 경향을 보면 이런 문제가 나올 거야"라고 핵심만 요약해서 기억해도 충분하다는 거죠. 이 논리는 "기억을 저장하는 것"을 "과거 데이터를 바탕으로 미래를 예측하는 간단한 수학적 모델"로 바꾼 것입니다.

3. 해결책: "LoRA-Pre" - 메모리 효율적인 '요약 노트'

이 새로운 통찰을 바탕으로, 저자들은 메모리를 아끼는 방법을 개발했습니다. 바로 LoRA-Pre입니다.

• 기존 방식: 모든 과거 정보를 고해상도 사진 (고차원 행렬) 으로 저장. → 메모리 폭탄.
• LoRA-Pre 방식: 고해상도 사진을 **저해상도 스케치 (저랭크 행렬)**로 변환해서 저장.
  • 비유: 4K 화질의 거대한 지도를 저장하는 대신, 핵심 길만 표시한 작은 주머니 지도를 들고 다니는 것과 같습니다.
  • 특이점: 이 '주머니 지도'는 정적인 것이 아니라, 매 순간 새로운 정보를 받으면 실시간으로 스스로 업데이트됩니다. (기존 방법들은 지도를 주기적으로 다시 그려야 해서 오차가 생겼지만, 이 방법은 매순간 바로바로 고쳐서 정확도를 유지합니다.)

4. 왜 이것이 대단한가요? (실험 결과)

이 논문은 이 방법이 두 가지 상황에서 탁월한 성과를 보였다고 말합니다.

1. 처음부터 배우기 (Pre-training):

   • 60 만 개부터 10 억 개까지 다양한 크기의 모델을 훈련시켰습니다.
   • 결과: 메모리는 훨씬 적게 쓰면서, 기존 최고 성능을 내는 방법들보다 더 좋은 점수를 받았습니다.
   • 비유: "작은 가방만 들고도, 큰 가방을 든 사람보다 더 멀리, 더 빠르게 달릴 수 있다"는 뜻입니다.

2. 이미 배운 것을 다듬기 (Fine-tuning):

   • 이미 훈련된 모델을 특정 작업 (수학 문제 풀기 등) 에 맞게 수정할 때도 효과가 있었습니다.
   • 결과: 기존에 유명한 'LoRA' 방법보다 훨씬 더 높은 정확도를 기록했습니다. (예: Llama-3.1 모델에서 3 점 이상, Llama-2 에서 6 점 이상 향상)

5. 요약: 이 기술이 가져오는 변화

• 기존: "기억하려면 많은 공간이 필요해. 그래서 큰 서버가 필요해."
• LoRA-Pre: "기억은 '요약'으로 충분해. 작은 메모리만으로도 똑똑하게 학습할 수 있어."

한 줄 평:
이 논문은 인공지능을 훈련시킬 때 **"무거운 기억을 버리고, 스마트한 요약으로 대체하자"**고 제안하며, 이를 통해 작은 컴퓨터로도 거대한 인공지능을 키울 수 있는 새로운 길을 열었습니다. 마치 거대한 도서관의 모든 책을 다 읽지 않고도, 핵심 요약본만 정독해서 똑똑한 사람이 될 수 있게 해주는 기술이라고 생각하시면 됩니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 학습에는 Adam 이나 Muon 과 같은 현대적인 옵티마이저가 필수적입니다. 그러나 이러한 옵티마이저는 1 차 및 2 차 모멘트 (기울기의 이동 평균) 를 유지해야 하므로, 모델 파라미터 수에 비례하여 엄청난 메모리 오버헤드가 발생합니다. 이는 LLM 의 확장성 (Scalability) 과 계산 효율성을 제한하는 주요 병목 현상입니다.

기존의 저랭크 최적화 방법들 (예: GaLore 등) 은 주기적인 부분공간 (subspace) 업데이트를 통해 기울기를 투영하는 방식을 사용하지만, 이는 부분공간 적응의 지연으로 인해 오차 누적 (error accumulation) 이 발생하고 최적화 성능이 저하되는 문제가 있습니다. 따라서 메모리 효율성을 높이면서도 최적화 동역학을 유지할 수 있는 새로운 접근법이 필요합니다.

2. 제안 방법론: LoRA-Pre (Methodology)

저자들은 모멘텀 업데이트와 온라인 선형 회귀 (Online Linear Regression) 사이의 수학적 동치성을 발견하고, 이를 바탕으로 LoRA-Pre라는 새로운 저랭크 옵티마이저를 제안했습니다.

핵심 통찰 (Key Insight)

기존 옵티마이저의 지수 이동 평균 (EMA) 기반 모멘텀 업데이트 식은, 온라인 기울기 흐름 (Online Gradient Flow) 을 통해 선형 회귀 모델을 학습하는 과정과 수학적으로 동일하다는 것을 증명했습니다.

• 기존 식: $m_{t+1} = \beta \cdot m_t + (1-\beta) \cdot g_t$
• 재해석: 이는 모멘텀 $m$을 파라미터로 하는 선형 회귀 모델이 현재 기울기 $g$에 대해 손실 함수 $L(m; g) = \frac{1}{2}\|m - g\|_F^2$를 최소화하는 과정과 동일합니다.

LoRA-Pre 의 작동 원리

이 통찰을 바탕으로, 전체 모멘텀 행렬을 저랭크 행렬 두 개의 곱으로 분해하여 메모리를 압축합니다.

1. 1 차 모멘텀 압축: 전체 모멘텀 행렬 $m \in \mathbb{R}^{p \times q}$를 $m = m_B \cdot m_A$ ($m_B \in \mathbb{R}^{p \times r}, m_A \in \mathbb{R}^{r \times q}, r \ll \min(p,q)$) 로 분해합니다.
2. 최적화: 분해된 행렬 $m_B, m_A$에 대해 선형 회귀 목적 함수를 최소화하는 뉴턴 방법 (Newton's Method) 을 적용하여 **폐쇄형 업데이트 규칙 (Closed-form update rules)**을 유도했습니다. 이는 역전파 (Backpropagation) 없이도 효율적으로 업데이트가 가능함을 의미합니다.
3. 2 차 모멘텀 압축: Adam 의 2 차 모멘트 ($v$) 는 양수 제약이 있으므로, $v = (v_B \cdot v_A) \circ 2$ (Hadamard product) 형태로 재파라미터화하여 양수성을 보장하면서 저랭크 구조를 유지합니다.
4. 적용: 이 프레임워크는 Adam 과 Muon 옵티마이저 모두에 적용 가능합니다.

3. 주요 기여 (Key Contributions)

1. 이론적 연결 고리 확립: EMA 기반 모멘텀 업데이트가 온라인 기울기 흐름을 통한 선형 회귀 학습과 수학적으로 동치임을 증명했습니다.
2. LoRA-Pre 제안: 위 통찰을 바탕으로 모멘텀 행렬을 저랭크 성분으로 분해하여 메모리 효율성을 극대화하는 새로운 옵티마이저를 개발했습니다. Adam 및 Muon 기반의 변형 모델을 제공하며, 이론적으로 유도된 업데이트 규칙을 통해 메모리 사용량을 크게 줄이면서도 최적화 동역학을 유지합니다.
3. 광범위한 실험적 검증: 60M 에서 1B 파라미터까지의 다양한 크기의 Llama 모델에 대한 사전 학습 (Pre-training) 과 Llama-2/3.1 모델에 대한 미세 조정 (Fine-tuning) 실험을 통해 방법론의 유효성을 입증했습니다.

4. 실험 결과 (Results)

사전 학습 (Pre-training) 성능

• 모델: Llama 아키텍처 (60M, 130M, 350M, 1B) 를 C4 데이터셋에서 학습.
• 결과: LoRA-Pre (Adam 및 Muon 버전) 은 모든 모델 크기에서 기존 저랭크 옵티마이저 (GaLore, Fira, LoRA 등) 보다 최고 또는 차상위 성능을 기록했습니다.
  • 특히 130M, 350M, 1B 모델에서 기존 최강 효율적 베이스라인보다 0.81~2.45 퍼플렉시티 (Perplexity) 포인트 향상.
  • Fira 와 같은 투영 기반 방법들은 큰 모델에서 오차 누적으로 인해 성능이 저하되는 반면, LoRA-Pre 는 일관된 성능을 보였습니다.

미세 조정 (Fine-tuning) 성능

• 모델: Llama-3.1-8B 및 Llama-2-7B 를 MetaMathQA 데이터셋으로 미세 조정 후 GSM8K, MATH-500 평가.
• 결과: 동일한 랭크 (Rank) 설정에서 기존 효율적 미세 조정 방법 (LoRA, DoRA, GaLore 등) 을 모두 압도했습니다.
  • Llama-3.1-8B: 표준 LoRA 대비 평균 3.14 점 향상.
  • Llama-2-7B: 표준 LoRA 대비 평균 6.17 점 향상.

랭크 효율성 (Rank Efficiency)

• LoRA-Pre 는 매우 낮은 랭크에서도 우수한 성능을 발휘합니다.
• 60M 모델: 랭크 16 에서 GaLore 의 랭크 128 과 동급 성능 (8 배 효율성 향상).
• 130M 모델: 랭크 16 에서 GaLore 의 랭크 256 과 동급 성능 (16 배 효율성 향상).
• 이는 LoRA-Pre 가 지속적인 부분공간 적응 (Continuous subspace adaptation) 메커니즘을 통해 오차 누적을 방지하고, 작은 부분공간으로도 기울기 구조를 효과적으로 포착하기 때문입니다.

5. 의의 및 결론 (Significance)

이 논문은 옵티마이저 상태 (Optimizer States) 를 단순한 메모리 부담이 아닌, 온라인 회귀 문제로 재해석함으로써 메모리 효율적인 최적화의 새로운 지평을 열었습니다.

• 메모리 효율성: 저랭크 분해를 통해 옵티마이저 상태의 메모리 사용량을 획기적으로 줄여, 고사양 GPU 없이도 대규모 모델의 학습을 가능하게 합니다.
• 범용성: Adam 뿐만 아니라 Muon 과 같은 다양한 모멘텀 기반 옵티마이저에 적용 가능하며, 사전 학습과 미세 조정 모두에서 우수한 성능을 입증했습니다.
• 실용성: 기존 저랭크 방법들의 한계였던 주기적 업데이트로 인한 불안정성을 해결하여, 더 안정적이고 효율적인 대규모 언어 모델 학습 파이프라인을 제공합니다.

결론적으로, LoRA-Pre 는 메모리 제약이 있는 환경에서도 대규모 모델의 학습과 적응을 가능하게 하는 강력한 솔루션으로 평가됩니다.