Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training
Sophia: 언어 모델 사전학습을 위한 확장 가능한 확률적 2차 최적화기
Hong Liu, Zhiyuan Li, David Hall, et al. (2023)
헤시안의 대각 추정을 경량화하여 2차 곡률 정보를 활용하는 Sophia 옵티마이저를 제안하여, Adam 대비 약 2배 빠른 수렴 속도로 LLM 사전학습을 가능하게 하면서 계산 오버헤드를 최소화했다.
배경
대규모 언어 모델 사전학습은 Adam 옵티마이저에 의존하고 있으나, Adam은 1차 모멘트(그래디언트)와 2차 모멘트(그래디언트 제곱)만 사용하여 곡률(curvature) 정보를 간접적으로만 활용한다. 뉴턴 방법 등 2차 최적화기는 곡률을 직접 활용하여 이론적으로 더 빠른 수렴이 가능하지만, 헤시안 행렬의 계산과 역행렬이 O(d²) 이상의 비용을 요구하여 수십억 파라미터 모델에는 적용 불가능했다. 확률적 2차 방법(K-FAC, Shampoo 등)도 LLM 규모에서는 메모리와 계산 오버헤드가 과도했다.
핵심 아이디어
Sophia의 핵심은 헤시안의 대각 원소를 확률적으로 추정하여, 파라미터별 곡률에 반비례하는 학습률을 적용하는 것이다. 곡률이 큰 방향(날카로운 방향)에서는 작은 스텝을, 곡률이 작은 방향(평탄한 방향)에서는 큰 스텝을 밟는다. 헤시안 대각 추정을 위해 두 가지 경량 방법을 제안한다: (1) Hutchinson의 확률적 추정자(Sophia-H) - 랜덤 벡터와 헤시안-벡터 곱으로 대각 원소를 추정, (2) 가우스-뉴턴-바르톨디 추정자(Sophia-G) - 미니배치별 그래디언트의 제곱으로 일반화된 가우스-뉴턴 행렬의 대각을 추정. 또한 업데이트 크기를 클리핑하여 비볼록 최적화에서의 안정성을 보장한다.
방법론
각 스텝에서 1차 모멘트(EMA of gradients)를 계산하고, k 스텝마다(기본 k=10) 헤시안 대각 추정치 h_t를 갱신한다. 업데이트 규칙은 θ_{t+1} = θ_t - η × clip(m_t / max(h_t, γ), λ)로, m_t는 그래디언트의 지수이동평균, h_t는 헤시안 대각 추정의 지수이동평균이며, γ는 0 나누기를 방지하는 상수, λ는 클리핑 경계이다. Sophia-G의 경우 h_t = E[g_B ⊙ g_B]로 미니배치 그래디언트의 원소별 제곱의 기대값을 사용하여, 추가 역전파 없이 추정이 가능하다.
주요 결과
GPT-2 규모(125M~770M)의 모델 사전학습에서 Sophia는 Adam 대비 동일 검증 손실에 도달하는 데 약 50% 적은 스텝(약 2배 빠른 수렴)을 보였다. 벽시계 시간(wall-clock time) 기준으로도 약 2배 빠른 학습을 달성했는데, Sophia-G의 경우 계산 오버헤드가 스텝당 약 2%에 불과하기 때문이다. 다운스트림 태스크(SuperGLUE, HellaSwag 등)에서도 Sophia로 학습한 모델이 Adam 대비 동등하거나 우수한 성능을 보였다.
임팩트
Sophia는 2차 최적화의 이점을 LLM 규모에서 실용적으로 활용할 수 있음을 보여준 선구적 연구이다. Adam이 10년 가까이 지배적이던 딥러닝 최적화 분야에서, 곡률 정보의 활용이라는 새로운 방향을 제시했다. 이후 Muon, Shampoo의 LLM 적용 등 2차 최적화 연구가 활성화되는 계기가 되었으며, LLM 사전학습 비용 절감이라는 실용적 측면에서도 큰 관심을 받고 있다.