연구진이 MTP 학습의 결함을 찾아내고 한 줄 수정으로 해결했다

Multi-Token Prediction Done Right

Multi-Token Prediction(MTP)은 언어모델 사전학습의 표준 요소가 됐지만, 표준 MTP 그래디언트에는 모델 규모가 커질수록 next-token 학습과 multi-token 학습을 모두 방해하는 잔차 성분이 섞여 있다. 저자들은 MTP cross-entropy 그래디언트를 KL self-distillation 항과 신호대잡음비가 낮은 잔차로 정확히 분해하고, 두 분포가 선형 출력 헤드를 공유할 때 straight-through gradient routing으로 KL 그래디언트를 정확히 복원하는 기법을 제안한다. 이를 적용한 Soft-MTP는 코드 한 줄 수정으로 잔차를 제거하며, 167M에서 1.2B 파라미터까지 모든 규모에서 더 낮은 NTP perplexity와 더 빠른 speculative decoding, 더 높은 zero-shot 정확도를 달성한다.