2025-11-16T03:28:12.300331

The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton

Abreu, Vyas, Kakade et al.
Recent efforts to accelerate LLM pretraining have focused on computationally-efficient approximations that exploit second-order structure. This raises a key question for large-scale training: how much performance is forfeited by these approximations? To probe this question, we establish a practical upper bound on iteration complexity by applying full Gauss-Newton (GN) preconditioning to transformer models of up to 150M parameters. Our experiments show that full GN updates yield substantial gains over existing optimizers, achieving a 5.4x reduction in training iterations compared to strong baselines like SOAP and Muon. Furthermore, we find that a precise layerwise GN preconditioner, which ignores cross-layer information, nearly matches the performance of the full GN method. Collectively, our results suggest: (1) the GN approximation is highly effective for preconditioning, implying higher-order loss terms may not be critical for convergence speed; (2) the layerwise Hessian structure contains sufficient information to achieve most of these potential gains; and (3) a significant performance gap exists between current approximate methods and an idealized layerwise oracle.
academic

LLM을 위한 이계 최적화의 잠재력: 완전 가우스-뉴턴 연구

기본 정보

  • 논문 ID: 2510.09378
  • 제목: The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
  • 저자: Natalie Abreu (Harvard), Nikhil Vyas (Harvard/OpenAI), Sham Kakade (Harvard), Depen Morwani (Harvard)
  • 분류: cs.LG cs.AI
  • 발표 시간: 2025년 10월 10일 (arXiv 사전인쇄본)
  • 논문 링크: https://arxiv.org/abs/2510.09378

초록

본 논문은 대규모 언어 모델(LLM) 사전학습에서 기존 이계 최적화 방법의 계산 효율적 근사가 얼마나 많은 성능을 손실하는지 연구했습니다. 저자들은 150M 매개변수의 Transformer 모델에 완전 가우스-뉴턴(GN) 전조건화를 적용하여 반복 복잡도의 실용적 상한을 확립했습니다. 실험 결과, 완전 GN 업데이트는 SOAP 및 Muon 등의 강력한 기준선 대비 5.4배의 훈련 반복 감소를 달성했습니다. 또한 계층 간 정보를 무시하는 정확한 계층별 GN 전조건화기는 완전 GN 방법의 성능에 거의 근접했습니다.

연구 배경 및 동기

문제 정의

LLM의 계산 요구량이 지속적으로 증가함에 따라, 최적화 방법의 개선이 훈련 효율성 향상의 핵심 전략이 되었습니다. 전통적인 일계 방법(SGD 및 Adam 등)은 광범위하게 사용되지만, 이계 방법은 이론적으로 더 빠른 수렴 속도와 더 나은 대규모 배치 확장성을 제공합니다.

연구 동기

  1. 기존 이계 방법의 한계: 현재의 이계 최적화기(Shampoo, SOAP, Muon)는 계산 가능성을 유지하기 위해 Hessian의 근사를 모두 사용하지만, 이러한 근사가 얼마나 많은 성능을 손실하는지는 명확하지 않습니다.
  2. 이론과 실제의 격차: 이계 방법이 이론적으로 우수하지만, 완전 Hessian의 저장 및 계산 비용이 너무 높아 실제 응용에서는 근사 방법을 사용해야 합니다.
  3. 핵심 연구 질문: "LLM에서 이계 최적화의 기본 성능 극한은 무엇인가? Hessian의 어떤 구조적 특성이 이러한 극한을 달성하는 데 필요한가?"

핵심 기여

  1. 성능 상한 확립: 완전 가우스-뉴턴 방법을 통해 이계 최적화의 실용적 성능 상한을 확립하여 반복 복잡도에서 SOAP 대비 5.4배 향상을 달성했습니다.
  2. 핵심 구조 규명: 계층별 Hessian 구조가 대부분의 성능 향상을 달성하기 위한 충분한 정보를 포함하며, 계층 간 곡률 정보의 중요성이 제한적임을 발견했습니다.
  3. 이론적 통찰: GN 근사가 전조건화에 매우 효과적임을 증명하여, 고계 손실 항이 수렴 속도에 중요하지 않을 수 있음을 시사합니다.
  4. 배치 크기 확장: 임계 배치 크기를 크게 확장하여 거의 최적에 가까운 확장 성능을 보여줍니다.

방법 상세 설명

작업 정의

모델 매개변수 θ, 입력 x 및 레이블 y가 주어질 때, 손실 함수 L(f(θ,x), y)를 정의합니다. 목표는 기대 손실을 최소화하며, 반복 복잡도(목표 손실에 도달하는 데 필요한 단계 수)에 중점을 둡니다.

가우스-뉴턴 방법 원리

수학적 기초

완전 Hessian 행렬은 다음과 같이 분해됩니다:

∇²θL(θ) = ∇θf(θ)ᵀ∇²zL(θ)∇θf(θ) + Σₐ(δL/δzₐ)∇²θ[f(θ)]ₐ

여기서 첫 번째 항은 가우스-뉴턴 행렬 G이고, 두 번째 항은 모델의 곡률입니다.

알고리즘 구현

알고리즘 1: 가우스-뉴턴 방법

  1. 모델의 일계 테일러 전개: f⁽¹⁾θₜ(θ,x) := f(θₜ,x) + ∇f(θₜ,x)ᵀ(θ-θₜ)
  2. 손실 볼록화: L̃θₜ(θ) := (1/b)Σ₍ₓ,ᵧ₎∈B ℓ(f⁽¹⁾θₜ(θ,x), y)
  3. 이계 테일러 근사 구성: L̃⁽²⁾θₜ(θ)
  4. 최소제곱 문제 해결: θ̂ = argminθ L̃⁽²⁾θₜ(θ)
  5. 선 탐색: θₜ₊₁ ← θₜ + α*(θ̂ - θₜ)

메모리 가능한 구현

Hessian 행렬의 명시적 저장을 피하기 위해 Jacobian-벡터 곱셈(JVP)을 사용하여 기능적으로 동등한 방법을 구현합니다. 핵심 아이디어는 손실 함수 L의 이계 테일러 근사와 모델 f의 일계 테일러 근사를 최적화하는 것입니다.

변형 방법

GN-근접-선형 방법

선형화된 모델의 손실을 직접 최소화: θ* = argminθ L̃θₜ(θ), 고계 손실 항의 영향을 연구하는 데 사용됩니다.

계층별 가우스-뉴턴

각 계층 l에 대해 독립적으로:

  1. 해당 계층의 일계 테일러 전개 계산: f⁽¹⁾θₗ,ₜ(θₗ)
  2. 해결: θₗ,ₜ₊₁ = argminθₗ L̃⁽²⁾θₗ,ₜ(θₗ)
  3. 모든 계층의 업데이트를 병합하고 선 탐색 적용

실험 설정

데이터셋 및 모델

  • 모델: 45M 및 150M 매개변수의 LLaMA 아키텍처
  • 데이터셋: C4 데이터셋
  • 시퀀스 길이: 1024

기준선 방법

  • AdamW: 가장 널리 사용되는 LLM 최적화기
  • Muon: Newton-Schulz 직교화를 사용하는 방법
  • SOAP: Shampoo의 최신 변형

실험 구성

  • 내부 최적화기: 최소제곱 문제 해결을 위해 Muon 사용
  • 배치 크기: 그래디언트 누적으로 제어, bᵢₙₙₑᵣ = 32(45M) / 128(150M)
  • 학습률 스케줄: 전역 코사인, 전역+내부 코사인, 상수+내부 코사인 세 가지 전략
  • 정규화: 가중치 감소, 선 탐색 등 다양한 전략

실험 결과

주요 결과

반복 복잡도

손실 3.25에 도달하는 실험에서:

  • 가우스-뉴턴: 54단계
  • SOAP: 292단계 (5.4배 차이)
  • Muon: 약 16배 차이
  • 계층별 GN: 78단계 (1.4배 차이만)

배치 크기 확장

고정 3B 토큰 훈련에서:

  • 가우스-뉴턴은 120M 배치 크기에서도 양호한 성능 유지(손실 3.45)
  • AdamW는 동일 배치 크기에서 심각한 성능 저하(손실 >4.4)
  • 임계 배치 크기가 크게 확장되어 거의 최적 확장 추세에 근접

절제 실험

GN vs GN-근접-선형

두 방법의 성능이 거의 동일하여 고계 손실 항이 성능 향상에 제한적 기여함을 나타냅니다.

완전 GN vs 계층별 GN

계층별 방법이 대부분의 설정에서 완전 GN 성능에 근접하여 계층 간 곡률 정보의 중요성이 제한적임을 보여줍니다.

주요 발견

  1. 학습률 스케줄의 중요성: 전역 코사인 스케줄이 중소 배치에서 최고 성능 발휘
  2. 선 탐색의 필요성: GN 방법의 안정적 수렴에 필수적
  3. 내부 최적화기 선택: Muon이 내부 최적화기로서 AdamW보다 우수

관련 연구

이계 최적화 방법

  • Hessian-자유 최적화: Martens(2010)가 제안한 켤레 기울기 방법
  • 대각선 Hessian 근사: AdaHessian, Sophia 등 경량 방법
  • 계층별 근사: Shampoo 계열 방법의 핵심 아이디어

LLM 최적화기 발전

  • 전통 방법: SGD, Adam 계열
  • 현대 이계 방법: AlgoPerf 경쟁에서 28% 우수한 Shampoo
  • 특화 방법: LLM을 위해 설계된 Muon, SOAP 등

결론 및 논의

주요 결론

  1. 성능 상한 확립: 완전 GN 방법이 이계 최적화의 명확한 성능 목표 제공
  2. 구조 중요성: 계층별 Hessian 구조가 대부분의 향상을 달성하기 위한 충분한 정보 포함
  3. 근사 효과: 현재 근사 방법과 이상적 계층별 예언기 간 상당한 성능 격차 존재

한계

  1. 계산 오버헤드: 현재 구현이 표준 훈련보다 4-5배 느림
  2. 규모 제한: 실험이 150M 매개변수 모델로만 제한됨
  3. 실용성: 주로 분석 도구로서 직접적인 실용 최적화기가 아님

향후 방향

  1. 효율적 구현: 계산 효율적인 정확 이계 방법 개발
  2. 더 나은 근사: 계층별 Hessian 근사 방법 개선
  3. 규모 확장: 더 큰 모델에서 발견 검증

심층 평가

장점

  1. 이론적 깊이: 이계 최적화 성능 극한에 대한 중요한 이론적 통찰 제공
  2. 실험 엄밀성: 광범위한 하이퍼매개변수 탐색 및 다양한 정규화 전략
  3. 실용적 가치: 기존 이계 방법 개선을 위한 명확한 목표 제시
  4. 방법론 혁신: JVP를 영리하게 사용하여 명시적 Hessian 저장 회피

부족한 점

  1. 계산 비용: 높은 계산 오버헤드가 실제 응용 제한
  2. 규모 한계: 진정한 대규모 LLM에서 검증되지 않음
  3. 이론 분석: 계층별 근사가 왜 그렇게 효과적인지에 대한 심층 이론 설명 부족

영향력

  1. 학술 기여: 이계 최적화 연구에 중요한 기준 제공
  2. 실제 지침: 기존 방법 개선 방향 제시
  3. 방법론적 가치: 이계 방법 평가를 위한 새로운 프레임워크 확립

적용 시나리오

  • 이계 최적화 방법의 이론 분석
  • 새로운 최적화 알고리즘의 성능 기준
  • 대규모 배치 훈련 시나리오의 최적화 선택

참고문헌

본 논문은 최적화 분야의 중요한 연구를 인용하며, 다음을 포함합니다:

  • Martens (2010): Hessian-자유 최적화의 개척적 연구
  • Gupta et al. (2018): Shampoo 최적화기
  • Jordan et al. (2024): Muon 최적화기
  • Vyas et al. (2025): SOAP 최적화기

종합 평가: 이는 엄밀한 실험을 통해 LLM 훈련에서 이계 최적화의 성능 상한을 확립하고 해당 분야에 중요한 이론적 통찰과 실제 지침을 제공하는 고품질 연구 논문입니다. 계산 비용과 규모 제한이 있음에도 불구하고, 학술적 가치와 향후 연구에 대한 지도 의의는 상당합니다.