Training for the Model You Return: Improving Optimization for Iterate-Averaged Language Models

Iterate average(EMA) 기반 LM 반환을 전제로, 최적 제어 이론에서 도출한 PACE 옵티마이저가 학습 궤적을 직접 개선하는 연구.

현대 LM 파이프라인은 학습 iterate의 EMA를 반환하는 것이 표준인데, 이 논문은 'EMA로 반환할 것이라면 학습 자체를 어떻게 바꿔야 하는가'라는 근본 질문을 던진다. 저자들은 이를 확률적 2차 모델에서의 최적 제어 문제로 정식화하고, AdamW 위에 얹는 가벼운 래퍼인 PACE를 도출했다. PACE는 현재 가중치를 EMA 쪽으로 clipped per-coordinate pullback으로 당기면서, iterate average의 오차를 최소화한다. convex 수렴 보장(standard SCO rate)과 2차 설정에서 임의로 큰 개선을 증명했으며, SmolLM2-1.7B, Qwen3-1.7B, Gemma3-1B fine-tuning과 GPT-2 pretraining에서 광범위한 하이퍼파라미터에 걸쳐 EMA 단독 및 AdamW를 일관되게 능가했다. Schedule-Free와도 경쟁력을 보였다.

핵심 요약

• iterate-average 최적화를 continuous-time stochastic quadratic 최적 제어 문제로 정식화
• PACE: AdamW 위에 얹는 lightweight wrapper, 추가 메모리 = model weights 1본
• convex 설정에서 standard SCO 수렴률 보장, 2차 설정에서는 임의로 큰 한정 오차 개선 가능
• SmolLM2/Qwen3/Gemma3 fine-tuning 3종 일관된 개선, GPT-2 FineWeb pretraining에서도 유효
• Schedule-Free와 경쟁적이면서 LR decay 불필요 — constant LR으로 WSD의 모든 token budget에서 우위
• pullback strength c, EMA power κ, update frequency에 대한 광범위한 ablation으로 견고성 입증

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.