기술 요약: 토큰 기하학(Token Geometry)과 Ember 옵티마이저

문제 정의

현재의 대규모 언어 모델(LLM)은 임베딩 테이블과 LM-head 행렬(통칭 "토큰 인터페이스")을 최적화하기 위해 Adam 옵티마이저를 사용합니다. Adam은 우수한 성능을 보이지만, 분산 학습 시 상당한 엔지니어링 부담을 초박합니다. 구체적으로, Adam은 모든 파라미터에 대해 1차 및 2차 모멘트 버퍼를 유지하므로, 어휘 크기 $V$와 임베딩 차원 $D$에 대해 $O(2VD)$의 옵티마이저 상태 복잡도를 가집니다. 대규모 모델의 경우, 이 상태는 기가바이트 단위의 VRAM을 소비하며, 이는 모델의 연산 병목이 아닌 파라미터에 대해서도 복잡한 샤딩 전략(예: ZeRO/FSDP)을 필요로 하게 만듭니다.

또한, 본 논문은 토큰 인터페이스에 Adam을 표준적으로 적용하는 것이 이러한 파라미터들의 특정한 기하학적 구조를 간과하고 있다고 주장합니다. 밀집 선형 레이어(dense linear layers)와 달리, 토큰 인터페이스는 희소한 그래디언트(매 단계마다 활성화된 토크의 행만 업데이트됨)를 통해 업데이트되며, 크로스 엔트로피(cross-entropy) 손실 함수의 기저에 깔린 손실 경관(loss landscape)은 표준 옵티마이저가 효율적으로 활용하지 못하는 독특한 특성을 가질 수 있습니다.

방법론: Ember 옵티마이저

저자들은 임베딩 및 LM-head 행렬을 위해 설계된 경량 옵티마이저인 Ember를 소개합니다. 핵심 방법론은 정보 기하학(Information Geometry), 특히 **피셔 정보 메트릭(Fisher Information metric)**에 근거합니다.

이론적 기초

본 논문은 크로스 엔트로피 손실에 대해 피셔 메트릭(그래디언트의 기대 외적)이 쿨백-라이블러 발산(Kullback–Leibler divergence)의 정확한 2차 항이라고 상정합니다. 최적점 근처에서 피셔 메트릭은 헤시안(Hessian)을 근사합니다.

• 제곱근 피셔(Square-Root Fisher): 전체 뉴턴 단계(Hessian 역행렬 계산)를 적용하는 것은 계산 비용이 많이 들고 불안정하므로, 저자들은 그래디언트를 피셔 메트릭의 제곱근으로 컨디셔닝하는 것만으로도 충분하다고 주장합니다. 이는 효과적으로 그래디언트를 표준 편차로 정규화하는 것(즉, "z-score" 변환)과 같습니다.
• 그래디언트 희소성(Gradient Sparsity): 토큰 인터페이스에서 토큰 $i$에 대한 그래디언트는 해당 토큰이 선택될 때만 0이 아닙니다. 본 논문은 행 단위 그래디언트 에너지(row-wise gradient energy)가 토큰의 참여 확률($p_i$)에 대한 불편 추정치(unbiased estimator)임을 관찰했습니다.

알고-리드믹 설계

Ember는 $O(2VD)$의 메모리 풋프린트를$O(V + D)$ 방식으로 대체하기 위해 2차 모멘트 추정치를 인수분해합니다:

1. 1차 모멘트 제거: 저자들은 Adam의 1차 모멘트(모멘텀) EMA를 완전히 제거하는 것이 성능 면에서 유리하며 메모리를 절약한다는 것을 발견했습니다.
2. 인수분해된 2차 모멘트: 전체 $V \times D$ 2차 모멘트 버퍼를 유지하는 대신, Ember는 두 개의 1D 벡터를 유지합니다:
   • $r_t$: 행 단위 2차 모멘트 추정치 (토큰 빈도/에너지를 추적).
   • $c_t$: 열 단위 2차 모멘트 추정치 (특징별 에너지를 추적).
3. 프리컨디셔닝(Preconditioning): 업데이트 규칙은 경량 외적 프리컨디셔너를 형성합니다. 파라미터 $(i, j)$에 대한 유효 분모는 행과 열 추정치의 기하 평균($\sqrt{\hat{r}_i \hat{c}_j}$)으로부터 도출됩니다.
4. 편향 수정(Bias Correction): 프리컨더를 형성하기 전에 행 및 열 추정치에 표준 편향 수정이 적용됩니다.

결과적으로 이 업데이트 규칙은 그래디언트를 피셔 메트릭의 제곱근 추정치로 나누는 효과를 가지며, 이는 토큰 빈도와 무관하게 정규화되면서 빈번한 토큰에 대한 분산을 줄여줍니다.

주요 기여

1. 메모리 효율성: Ember는 토큰 인터페이스에 대한 옵티마이저 상태를 $O(2VD)$에서$O(V + D)$로 줄입니다. 이를 통해 임베딩 테이블에 대한 옵티마이저 상태를 샤딩할 필요를 없애주며, 거대 모델의 경우에도 전체 상태가 단일 GPU에 들어갈 수 있게 합니다.
2. 기하학적 통찰: 본 논문은 Transformer를 위한 경쟁력 있는 옵티마이저들(Adam, Lion, Muon 포함)이 그래디언트를 제곱근 피셔 메트릭(또는 그 행렬 유사체)으로 컨디셔닝하는 프레임워크 아래에서 통합될 수 있다는 이론적 관점을 제공합니다.
3. 최적화 궤적 분석: 저자들은 토큰 궤적의 특이값 분해(SVD)를 통해, 임베딩 토큰들이 놀랍게도 단순한 **1D 레이(1D ray, 단항식 기저)**를 따라 최적화된다는 것을 입증했습니다. 이는 신경망 최적화가 매우 비볼록(non-convex)하고 혼돈스러운 경관을 탐색한다는 기존의 견해에 도전하며, 토큰 인터페이스의 경우 매끄럽고 선형적인 수렴 경로를 가짐을 시사합니다.
4. 오픈 소스 구현: 기존 ZeRO/FSDP 설정과 호환되는 분산 구현체가 공개되어 즉각적인 도입이 가능합니다.

결과

저자들은 다양한 환경에서 경험적 증거를 제시합니다:

• 지도 미세 조정(SFT): GPT-2 및 Pythia 모델에서 Ember는 AdamW와 대등하거나 약간 더 나은 검증 손실(validation loss)을 보이면서, 옵티마이저 상태를 기가바이트에서 킬로바이트 단위로 줄였습니다 (예: Pythia-2.8B의 경우 2.06 GB $\to$ 413 KB).
• 강화 학습(RL): GRPO를 사용한 Post-training Qwen2.5-3B에서, Ember는 다양한 롤아웃 배치 크기 전반에 걸쳐 AdamW와 대등한 성능을 달면서도 4,000배 적은 옵티마이저 상태를 사용했습니다.
• 사전 학습(Pretraining): FineWeb 데이터셋에서 Ember는 AdamW와 대등한 성능을 보였으며, 특히 낮은 배치 크기 환경에서 Adafactor보다 뛰어난 성능을 보였습니다. 이는 Rank-1 그래디언트 분해가 희소하거나 노이즈가 섞인 그래디언트에 대한 안정성을 제공하기 때문입니다.
• 스케일링 법칙(Scaling Laws): Ember는 모델 크기와 배치 크기에 따라 효과적으로 스케일링됩니다. 모델 크기가 커질수록 Ember와 AdamW 사이의 성능 격차가 사라지는데, 이는 Adam의 전체 2차 모멘트 버퍼가 대규모 환경에서 Ember의 Rank-1 인수분해에 의해 잘 근사될 수 있음을 시사합니다.
• 안정성: Ember는 AdamW에서 관찰되는 작은 배치 크기에서의 초기 학습 손실 스파이크 현상을 보이지 않습니다.

의의 및 주장

본 논문은 Ember가 "파레토 최적(Pareto frontier)" 개선을 제공한다고 주장하며, 극적인 메모리 오버헤드 감소와 동시에 최첨단 성능을 달 수 있다고 말합니다. 그 의의는 두 가지 측면에서 나타납니다:

1. 실용적 엔지니어링: 임베딩 옵티마이저 상태를 샤딩할 필요를 제거함으로써, Ember는 분산 학습 설정을 단순화하고, Adam의 상태가 제약 조건이었던 최대 7B 파라미터 이상의 모델에 대해서도 단일 장치에서의 빠른 반복 학습을 가능하게 합니다.
2. 이론적 이해: 토큰 최적화 궤적이 깨끗한 1D 단항식 기저로 분해된다는 발견은, "비볼록성"의 특성이 토큰 인터페이스 자체의 고유한 속성이 아니라 옵티마이저나 파라미터화의 산물일 수 있음을 시사합니다. 이는 Transformer를 위한 최적의 옵티마이저를 찾는 탐색 공간이 생각보다 좁으며, 제곱근 피셔 메트릭을 중심으로 형성되어 있음을 의미합니다.

저자들은 Ember가 메모리가 제한된 학습을 위한 실용적인 도구이기도 하지만, LLM이 어떻게 학습하는지에 대한 새로운 관점을 제공하는 이론적 틀을 갖추고 있으며, 특히 토큰 표현을 위한 최적화 경로의 단순성을 강조하고 있다고 결론짓습니다.