메모리 절감 LLM 사전학습을 위한 최소주의 옵티마이저 SCALE

초록

본 논문은 대규모 언어 모델(LLM) 사전학습에서 메모리 사용량을 크게 줄이면서 Adam 수준의 성능을 달성할 수 있는 최소주의 옵티마이저인 SCALE을 제안한다. 핵심 아이디어는 (1) 출력 차원(컬럼) 기준으로 그래디언트를 정규화해 스케일링 효과를 제공하고, (2) 변동성이 가장 큰 마지막 레이어에만 1차 모멘텀을 적용하는 것이다. 실험 결과, 60M‑7B 규모의 LLaMA 모델에 대해 기존 Adam 대비 35‑45% 메모리만 사용하면서 동일하거나 더 낮은 퍼플렉시티를 기록했으며, GaLore, Fira, APOLLO 등 메모리 효율적 최신 옵티마이저들을 모두 능가하였다.

상세 분석

SCALE은 두 가지 설계 선택을 통해 Adam의 핵심 이점을 최소한의 상태(state)만으로 재현한다. 첫 번째는 컬럼‑와이즈 그래디언트 정규화이다. LLM의 각 레이어는 입력 차원(d_in)과 출력 차원(d_out)으로 구성된 행렬 형태의 가중치를 갖는데, 기존 Adam은 원소별 2‑노름(v_t)으로 정규화한다. 저자는 출력 차원(컬럼) 별로 L2‑노름을 나누어 정규화함으로써, 특히 LM‑head와 같은 출력이 큰 레이어에서 그래디언트 스케일을 균일화하고, 행(row) 정규화에서 발생하는 극단적인 값 폭을 방지한다. 실험적으로 컬럼 정규화는 행 정규화보다 2‑3배 빠른 연산 속도와 더 낮은 퍼플렉시티를 보였으며, SVD 기반의 특잇값 정규화보다 메모리와 시간 효율이 뛰어나다. 두 번째는 마지막 레이어 전용 1차 모멘텀 적용이다. 변동성이 가장 큰 LM‑head 레이어에만 EMA(β₁) 모멘텀을 유지하고, 나머지 레이어는 순수 SGD 업데이트를 수행한다. 이는 모멘텀 상태를 전체 파라미터에 걸쳐 저장할 필요가 없으므로 메모리 오버헤드가 거의 0에 가깝다. 또한 마지막 레이어에만 모멘텀을 두면 학습 초기에 급격한 파라미터 변동을 완화하면서도 전체 모델의 수렴 속도는 유지된다.

메모리 측면에서 SCALE은 Adam 대비 **35‑45%**의 메모리만 사용한다. 구체적으로 1B 모델에서는 전체 메모리 중 10%만 추가로 필요하고, 7B 모델에서는 2% 수준이다. 이는 Adam이 파라미터당 2개의 추가 상태(m, v)를 저장하는 반면, SCALE은 마지막 레이어에만 1개의 모멘텀 상태를 저장하기 때문이다.

실험에서는 LLaMA 60M, 130M, 350M, 1B, 7B 모델을 C4 데이터셋으로 사전학습시켰으며, 퍼플렉시티와 메모리 사용량을 Adam, GaLore, Fira, APOLLO, Muon 등과 비교하였다. 모든 규모에서 SCALE은 Adam과 거의 동등하거나 더 낮은 퍼플렉시티를 기록했으며, 특히 7B 모델에서는 APOLLO와 Muon을 능가하는 결과를 보였다. 추가적인 Ablation 실험에서는 (i) 컬럼 정규화 없이 마지막 레이어 모멘텀만 적용했을 때 성능이 급격히 떨어짐을, (ii) 마지막 레이어 모멘텀 없이 컬럼 정규화만 적용했을 때도 Adam 수준에 도달하지 못함을 확인하였다. 이는 두 요소가 상호 보완적으로 작용한다는 것을 의미한다.

한계점으로는 현재 구현이 GPU 메모리 최적화에 초점을 맞추었으며, CPU‑기반 대규모 분산 학습 환경에서의 효율성은 추가 검증이 필요하다. 또한 마지막 레이어가 여러 개 존재하는 모델(예: 멀티‑헤드 어텐션)에서는 어느 레이어에 모멘텀을 적용할지에 대한 정책이 필요할 수 있다.

전반적으로 SCALE은 “적은 상태, 충분한 정규화”라는 원칙을 통해 메모리 제약이 심한 대규모 LLM 사전학습에 실용적인 솔루션을 제공한다.