기술 요약: 이차 모델(Quadratic Model)에 대한 옹호

문제 정의

대규모 언어 모델(LLM)의 최적화 이론은 비볼록(non-convex) 손실 지형의 복잡성을 탐색하기 위해 종종 이상적인 모델에 의존합니다. 그러나 이론적 다루기 쉬움(tractability)과 경험적 정확성 사이에는 상당한 격차가 존재합니다. 수학적으로 단순한 모델은 실제 역학을 설명하는 데 실패하는 경우가 많고, 특정 문제 구조를 포착하는 모델은 이론적으로 다루기 어려운 경우가 빈번합니다. 본 연구가 다루는 핵심 질문은 다음과 같습니다: 가장 단순한 최적화 모델인 이차 모델(quadratic model)이 실제로 LLM의 훈련 역학을 어느 정도까지 포착하는가?

방법론

저자들은 30억 개의 토큰으로 학습된 150M 파라미터 트랜스포머를 사용하여 이차 모델에 대한 엄격한 "스트레스 테스트"를 수행합니다. 방법론은 세 가지 주요 구성 요소로 이루어집니다.

1. 국소 이차 근사 (Section 3)

저자들은 중간 체크포인트(훈련의 매 10%마다)에서 국소 테일러 전개(local Taylor expansions)의 예측력을 평가합니다. 이들은 두 가지 대리 손실(proxy losses)을 정의합니다:

• Prox (선형화): 원래의 크로스 엔트로피 목적 함수를 유지하면서 참조점 $\theta_{ref}$ 주변에서 모델 $f_\theta$를 선형화합니다. 이는 뉴럴 탠전트 커널(NTK) 체제 하에서의 다중 클래스 로지스틱 회귀와 동일합니다.
• Quad (이차): Prox 손실을 2차까지 테일러 전개하여 볼록한 이차 목적 함수를 생성합니다. 결정적으로, 저자들은 이차 목적 함수가 볼록성(양의 준정부호)을 유지하고 음의 곡률 방향에서 유계가 벗어나는 것을 방지하기 위해 전체 헤시안($H$) 대신 가우스-뉴턴 행렬($G$)을 사용합니다.

이들은 각 체크포인트에서 시작하여 전체 토큰 예산의 10% 동안 이 대리 모델들을 훈련시키고, 그 결과로 나타나는 최적화 궤적을 실제 LLM 훈련 역학에 대해 비교합니다.

2. 헤시안 스펙트럼 추정 (Section 4)

이러한 국소 이차 문제의 구조를 분석하기 위해, 저자들은 훈련 궤적을 따라 헤시안 및 가우스-뉴턴 행렬(원래 형태 및 Adam 프리컨디셔닝 적용 형태 모두)의 스펙트럼을 추정합니다.

• 기법: 매우 깊은 프로브($m=1200$)와 엄격한 오차 범위(Gauss-Radau bounds)를 사용하는 **확률적 랜초스 쿼드러처(Stochastic Lanczos Quadrature, SLQ)**를 채택합니다.
• 규모: 배치 크기 $B \in \{1, 64, 1024\}$에 대해 훈련의 10%, 50%, 100% 지점의 체크포인트에서 프로브를 실행합니다.
• 분석: 스펙트럼 감쇠(거듭제곱 법칙), 네트워크 레이어 전반의 고유벡터 분포, 그리고 그래디언트와 고유벡터 간의 정렬(소스 조건)을 조사합니다.

3. 안정성의 경계(Edge of Stability, EOS) 분석 (Section 5)

저자들은 LLM 훈련이 "안정성의 경계"에서 작동하는지 경험적으로 테스트합니다.

• 설정: Adam 프리컨디셔너를 고정하고 국소 이차 대리 모델 상에서 프리컨디셔닝된 SGD를 실행합니다.
• 섭동(Perturbation): 학습률($\eta$)과 배치 크기($B$) 승수를 격자 형태로 변화시켜 안정성 경계를 결정합니다.
• 지표: 손실이 임계값(100)을 초과하면 발산하는 것으로 간주합니다. 이들은 평균 제곱 안정성(heavy tails에 민감함)과 리아푸노프 안정성(일반적인 동작)을 구분합니다.

주요 결과

1. 이차 모델의 예측력

• 일치 구간: 국소 이차 및 선형화 모델은 전체 훈련 예산의 최대 10% 동안 실제 최적화 역학을 정확하게 예측합니다.
• 시간적 정확도: 근사 품질은 훈련 초기 단계보다 중후반 단계에서 현저히 더 높습니다.
• Prox vs. Quad: 선형화 모델(Prox)은 초기 훈련에서 이차 모델(Quad)보다 실제 손실과 더 잘 일치하는 경향이 있는 반면, Quad는 훈련의 60% 이후에 높은 정확도로 수렴합니다.

2. 헤시안 스펙트럼의 구조

• Head-Tail 분리: 스펙트럼은 인덱스 $i \approx 8192$(어휘 사전 크기) 부근에서 급격한 전이를 보입니다. "Head"(상위 고유값들)는 "뉴럴 콜랩스(neural collapse)" 현상과 일치하게 **언임베딩 레이어(unembedding layer)**에 의해 지배됩니다.
• 보편적 Tail: 스펙트럼의 "Tail"(인덱스 $> 10^5$)은 보편적인 거듭제곱 법칙을 따릅니다. 이는 서로 다른 배치 크기에 걸쳐 동일하며, Adam 프리컨디셔너의 영향을 받지 않습니다.
• 지수(Exponents):
  • 용량 지수 ($\alpha$): 가우스-뉴턴 테일(Gauss-Newton tail)에 대해 $\alpha \approx 0.96$, 헤시안 테일(Hessian tail)에 대해 $\alpha \approx 0.65$로 추정됩니다.
  • 소스 지수 ($\beta$): 소스 조건(그래디언트 정렬)은 가우스-뉴턴의 경우 $\beta \in [0.1, 0.2]$, 헤시안의 경우 $\beta \in [0.3, 0.6]$를 나타냅니다. 이는 $\beta > 1$이라는 기존 문헌의 일반적인 가정과 상충하며, 문제가 임계 시간 이후 지수적 편향 감소를 갖는 유한 차원처럼 작동함을 시사합니다.
• 음의 고유값: 프리컨디셔닝된 헤시안은 특히 큰 배치 크기에서 큰 음의 고유값을 보이며, 이는 MLP 업-프로젝션(up-projection) 레이어에 국한되어 나타납니다. 저자들은 이것이 "죽은 뉴런(dead neurons)" 및 희귀 토큰 효과와 연관되어 있다고 가설을 세웁니다.

3. 안정성의 경계 역학

• 확률적 vs. 결정론적:
  • 작은 및 중간 배치 크기($B=1, 64$)에서 훈련은 **확률적 안정성의 경계(stochastic edge of stability)**에서 발생합니다. 안정성 경계는 $(\eta, B)$ 공간에서 대각선 형태를 따르며, 이는 불안정성이 노이즈(stochasticity)에 의해 유발됨을 나타냅니다.
  • 큰 배치 크기($B=1024$)에서 훈련은 풀 배치 그래디언트 디센트 동작과 일치하는 평탄한 컷오프를 특징으로 하는 **결정론적 안정성의 경계(deterministic edge of stability)**에서 작동합니다.
• 경계 근접성: 모든 설정에서 훈련은 안정성 경계로부터 2배 이내의 범위 내에서 작동합니다.

의의 및 주장

본 논문은 이차 모델이 단순한 이론적 장난감이 아니라 **검증된 실험적 대리 모델(proxy)**임을 주장합니다.

• 이론적 다루기 쉬움: 이러한 발견은 복잡한 LLM 최적화가 국소적으로 구조화된 이차 문제의 연속체로 이해될 수 있음을 시사합니다. 이를 통해 연구자들은 이차 최적화 이론(예: 안정성 분석, 학습률 스케줄링)의 풍부한 도구 상자를 활용하여 LLM 훈련을 추론할 수 있습니다.
• 가정의 재검토: 결과는 소스 조건의 합산 가능성($\beta > 1$) 및 프리컨디셔닝 하에서의 헤시안 테일의 보편성에 대한 기존 문헌의 가정을 반박합니다.
• 실질적 함의: 작은 배치 크기에서의 "확률적 안정성의 경계" 식별은 왜 특정 하이퍼파라미터 구성에서 훈련이 발산하는지에 대한 정교한 이해를 제공하며, 이산화 오류(discretization errors)와 확률적 노이즈(stochastic noise)를 구분해 줍니다.

저자들은 전역적 손실 지형은 비볼록하지만, 국소 이차 근사가 대규모 모델의 훈련을 분석하고 설계하기 위한 강력하고 다루기 쉬운 프레임워크를 제공한다고 결론짓습니다.