기초 모델 규모에서 헤시안 스펙트럼을 정확히 측정하는 실용적 방법

초록

본 논문은 Fully Sharded Data Parallel(FSDP) 환경에서 파라미터를 모으지 않고도 중앙 차분 방식의 헤시안-벡터 곱(HVP)을 구현한다. 이를 이용해 100 B 파라미터 규모의 오픈소스 언어 모델에 대해 확률적 Lanczos Quadrature(SLQ)를 수행해 실제 헤시안 스펙트럼을 추정한다. FP32와 BF16에서의 유한 차분 편향·수치 잡음, Krylov 안정성 분석을 통해 실용적인 ε 스텝 크기를 제시하고, 전체 연산 비용이 1차 학습 대비 상수 배 수준임을 실험적으로 검증한다. 또한 널리 쓰이는 블록 대각 근사법이 중규모 모델에서도 순위‑1 상대 오차와 방향 정렬 실패를 보이며, 기존 근사의 질을 크게 과대평가하고 있음을 밝힌다.

상세 분석

이 연구는 두 가지 핵심 기술적 공헌을 제시한다. 첫째, FSDP‑sharded 파라미터에 직접 적용 가능한 중앙 차분 기반 HVP를 설계하였다. 알고리즘 1에 따르면 각 GPU는 로컬 파라미터 샤드를 ±εv 만큼 변형하고, 동일한 데이터 로더 슬라이스에 대해 두 번의 순방향·역전파를 수행한다. 이렇게 얻은 로컬 그라디언트 차이를 2ε 로 나누면 샤드‑레벨의 Hv = H v 가 된다. 파라미터 전체를 모으는 all‑gather 없이 구현되므로 통신 비용은 기존 2배 gradient pass 수준에 머문다.

정리 3.1은 부동소수점 연산 오차 ε_mach 를 고려한 근사 오차 상한을 제시한다.
‖eHv − Hv‖ ≤ ε_mach² · C₁ + O(ε_mach)·‖∇L‖,
여기서 C₁ = (1/6)‖∇(D³_v L)‖. 최적 ε* ≈ (ε_mach ‖∇L‖ ‖∇(D³_v L)‖)^{1/3} 로 선택하면 오차는 O(ε_mach^{2/3}) 수준으로 감소한다. FP32(ε_mach≈1.2·10⁻⁷)에서는 ε*≈10⁻³~10⁻², BF16(ε_mach≈3.9·10⁻³)에서는 ε*≈10⁻¹ 정도가 실험적으로 확인되었다.

두 번째 공헌은 이러한 근사 HVP를 Lanczos 알고리즘에 연결한 이론적 분석이다. 정리 4.1은 유한 차분 오차가 독립이고 무편향일 때, m 단계 Lanczos가 생성하는 삼대각 행렬 ˜T_m 은 실제 T_m 에 작은 교란 ΔT_m 을 더한 형태이며,
E‖ΔT_m‖₂ ≲ ‖H‖₂ · σ_f^{1/2} · ‖D⁴f‖^{1/2} · r̄ · η · P,
와 같은 상한을 갖는다. 여기서 η는 Krylov 벡터의 평균 ℓ∞‑노름 역수, r̄는 재직교화 단계 수, P는 파라미터 수이다. 결과적으로 Ritz 값 λ̃_i 와 실제 고유값 λ_i 사이의 차이는 O(‖H‖₂ · …) 로 제한되며, 재직교화가 부족할 경우 “ghost eigenvalue” 가 발생해 RMS 스플리팅 Δλ_ghost ≈ 2‖H‖₂ · … 로 나타난다. 최적 ε* 를 사용하면 기존 Paige의 결과와 일치한다는 점도 강조된다.

시스템 측면에서는 HVP 한 번이 2 · T_grad + T_vec 로 표현되며, 여기서 T_vec 은 로컬 AXPY 연산과 스칼라 all‑reduce 비용이다. Lanczos 한 단계는 T_HvP + (2+r)·T_scalar + O((1+r)·P·γ) 로 계산된다. 전체 SLQ는 s · m 단계 Lanczos 호출에 비례하므로, 실험에서는 s=20, m=30 정도에서도 전체 실행 시간이 1차 학습 대비 1.3~1.7배에 머물렀다.

마지막으로, 실제 7 B, 30 B, 100 B 모델에 대해 스펙트럼 밀도와 고유값 분포를 추정한 결과, 블록‑대각 근사(예: K‑FAC, EK‑FAC)가 전체 헤시안과의 코사인 유사도가 0.2 이하이며, 상대 오차가 O(1) 수준임을 확인했다. 이는 기존 연구가 제시한 “근사는 충분히 정확하다”는 가정을 크게 뒤흔든다.

요약하면, 이 논문은 대규모 분산 학습 환경에서 정확한 헤시안 정보를 얻을 수 있는 실용적 파이프라인을 제시하고, 수치적·이론적 분석을 통해 최적 파라미터와 비용 모델을 제시함으로써 향후 1조 파라미터 시대에도 2차 정보 활용이 가능함을 증명한다.