기술 요약: 반환되는 모델을 위한 학습 (Training for the Model You Return)

문제 정의

현대적인 언어 모델(LM) 파이프라인은 대개 최종 반복값(final iterate) 자체가 아니라, 일반적으로 훈련 반복의 지수 이동 평균(Exponential Moving Average, EMA)인 평균화된 모델을 반환합니다. 반복 평균화(iterate averaging)가 최적화를 안정화하고 일반화 성능을 향상시킨다는 점은 잘 알려져 있으나, 근본적인 질문이 남아 있습니다: 최종 출력물이 반복 평균 추정치라는 점을 고려할 때, 이 특정 추정치의 성능을 최적화하기 위해 훈련 과정 자체를 어떻게 수정해야 하는가?

기존 연구들은 평균화가 더 공격적인 학습률을 허용한다는 관찰을 통해 이를 암묵적으로 다루어 왔습니다. 그러나 본 논문은 직접적인 접근 방식을 취합니다. 즉, 반환되는 반복 평균 추정치의 오차를 최소화하도록 훈련 궤적을 명시적으로 제어하는 최적화 알고리즘을 설계하고자 합니다.

방법론

이론적 정식화

저자들은 이 문제를 연속 시간 확률 이차 모델(continuous-time stochastic quadratic model) 내의 **최적 제어 문제(optimal-control problem)**로 프레임화합니다.

• 역학(Dynamics): 파라미터 궤적 $\theta_t$는 (노이즈가 섞인 경사 하강법을 나타내는) Ornstein–Uhlenbeck 과정에 의해 생성되는 확률 미분 방정식(SDE)을 따르며, 여기에 제어 입력 $u_t$가 추가됩니다.
• 목표: 제어 개입의 크기에 대한 이차 펜널티(규제 매개변수 $\lambda$)를 부과하면서, 반환된 시간 평균 추정치 $\hat{\theta}_{EMA, T}$와 실제 최적점 $\mu^*$ 사이의 평균 제곱 오차를 최소화하는 제어 전략 $u_t$를 찾는 것입니다.
• 유도: 저자들은 선형-이차-가우시안(LQG) 제어 이론과 확실성 등가 원리(princance of certainty equivalence)를 사용하여 최적 제어기를 유도합니다. 정확한 해는 궤적에 기반한 최적점의 사후 평균에 의존하지만, 이는 대규모 LM에서는 다루기 어렵습니다(intractable).

실용적 근사: PACE

해결책을 구현 가능하게 만들기 위해, 저자들은 일련의 휴리스틱 근사법을 도입합니다.

1. 추정치 대체: 알 수 없는 최적점의 사후 평균을 현재의 반복 평균인 $\hat{\theta}_{EMA, t}$로 대체합니다.
2. 시간 근사: 현재 시간 $t$가 전체 훈련 시간 $T$보다 훨씬 작다($t \ll T$)고 가정하여, 시간 의존적 계수들을 단순화합니다.
3. 이산화 및 프리컨디셔닝: 연속 제어를 이산화하고, 알려지지 않은 헤시안(Hessian) 행렬을 AdamW에서 사용되는 대각 프리컨디셔너(이차 모멘트 추정치 $v_t$)로 대체합니다.

이를 통해 **PACE (Pullback Averaging Control for Efficient Optimization)**가 도출됩니다. 이는 AdamW를 감싸는 가벼운 래퍼(wrapper)이며, 라이브 가중치를 EMA 가중치 쪽으로 이동시키는 "풀백(pullback)" 항을 추가합니다:
$$\theta_{k+1} = \theta_k - \eta g_k + C(\hat{\theta}_{EMA, k} - \theta_k)$$
여기 여기서 $C$는 Adam 프리컨디셔너로부터 유도된 대각 행렬이며, 풀백 강도는 업데이트가 표준 스텝과 풀백의 볼록 결합(convex combination)으로 유지되도록 클리핑됩니다.

주요 기여

1. 최적 제어 정식화: 본 논문은 최종 반복값이 아닌, 반복 평균 추정치의 오차를 최소화하기 위해 특별히 설계된 최적화 알고리즘에 대한 최초의 원리적 유도를 제공합니다.
2. 이론적 보장:
   • 볼록(Convex) 설정: 저자들은 PACE의 정형화된 버전이 편향되지 않은 확률적 그래디언트를 사용하는 모든 볼록 손실 함수에 대해, EMA 강도에 따른 상수 계수 차이를 제외하면 표준 $O(1/\sqrt{T})$ 속도로 수렴함을 증명합니다. 이는 이 방법이 표준 SGD보다 결코 크게 나쁘지 않음을 보장합니다.
   • 이차(Quadratic) 설정: 유도에 사용된 이차 설정에서, 제어된 역학은 제어되지 않은 역학(표준 SGD와 EMA)에 비해 반복 평균 추정치의 한계 제곱 오차를 엄격하게 개선합니다. 이 개선 효과는 특정 이차 인스턴스에서 임의로 커질 수 있습니다.
3. 실증적 검증: 방법론은 다음 모델 및 데이터셋에 대해 평가되었습니다:
   • 미세 조정(Fine-tuning): SmolLM2-1.7B, Qwen3-1.7B, Gemma3-1B 모델을 smol-smoltalk 데이터셋에 적용.
   • 사전 학습(Pretraining): GPT-2 (124M)를 Chinchilla-optimal 토큰 예산에 따라 FineWeb에서 학습.
   • 베이스라인: PACE를 vanilla AdamW, EMA 평가를 포함한 AdamW, 그리고 Schedule-Free 알고리즘과 비교했습니다.

결과

• 성능: PACE는 테스트된 모든 모델(1B–2B 파라미터)과 데이터셋에서 vanilla AdamW 및 EMA 평가를 포함한 AdamW보다 엄격하게 우수한 성능을 보였습니다.
• 강건성: 개선 효과는 다양한 학습률, 감쇠 스케줄(constant, cosine, Warmup-Stable-Decay) 및 하이퍼파라미터(풀백 강도 $c$, EMA 지수 $\kappa$) 범위에서 견고하게 나타났습니다.
• Schedule-Free와의 비교: PACE는 학습률 감소의 필요성을 제거하기 위해 반복 평균을 사용하는 Schedule-Free 알고리즘(Defazio et al., 2024)과 경쟁 가능한 수준입니다. 그러나 PACE는 다른 업데이트 규칙을 통해 이를 달성합니다.
• 학습률 감소(Learning Rate Decay): PACE는 일정한 학습률에서도 잘 작동하며, 이는 학습률 감소가 일반적으로 필요한 영역에서도 훈련을 안정화할 수 있음을 시사합니다. 토큰 예산 비교 시, PACE의 constant-LR 궤적은 WSD 분기만큼 효과적으로 손실 지형의 "리버 플로어(river floor)"를 추적하며, 종종 이를 능가하기도 합니다.

의의 및 한계

의의:
본 논문은 최종 반복값이 아닌 *반환되는 모델(평균)*을 위해 최적화하는 것이 실질적인 성능 향상으로 이어진다는 것을 보여줍니다. 훈련 궤적을 제어 문제로 다룸으로써, PACE는 복잡한 스케줄 튜닝 없이도 훈련을 안정화하고 일반화 성능을 높일 수 있는 이론적 근거를 갖춘 방법을 제공합니다. 저자들은 PACE가 일반적인 볼록 설정에서 EMA 요인으로 인해 이론적 수렴 보장이 약간 느려질 수 있음에도 불구하고, 이차 설정에서는 표준 방식보다 개선된 이론적 수렴성을 제공한다고 언급했습니다.

한계:

• 이론적 가정: 유도는 가산 노이즈가 있는 단순화된 이차 모델에 의존합니다. 일반적인 볼록 함수에 대한 수렴은 증명되었으나, 모멘텀과 적응형 프리컨디셔닝(AdamW의 핵심)을 이론적 분석에 엄격하게 통합하는 것은 여전히 열린 과제로 남아 있습니다.
• 규모(Scale): 실험은 미세 조정의 경우 최대 2B 파라미터, 사전 학습의 경우 124M 파라미터 규모로 제한되었습니다. 70B 이상의 대규모 규모에서의 효용성은 아직 검증되지 않았습니다.
• 메모리 오버헤드: PACE는 훈련 중에 모델 가중치의 추가 복사본(EMA 상태)을 유지해야 하므로 무시할 수 없는 메모리 오버헤드를 발생시키지만, 실제로는 큰 배치 사이즈를 통해 이를 완화합니다.
• 시드 민감도: 컴퓨팅 자원의 제약으로 인해 주요 실행에는 단일 시드가 사용되었으나, 멀티 시드 강건성 연구를 통해 방법론의 순위 안정성을 확인했습니다.

저자들은 PACE가 현대 LM의 실제 배포 전략(평균된 모델이 최종 제품이 되는 방식)과 최적화 알고리즘을 정렬하는 데 있어 유망한 단계임을 결론짓습니다.