멀티토큰 예측, KL 신호 복원으로 한 줄 개선

핵심 요약

• Multi-Token Prediction(MTP)은 언어모델 사전학습에서 다음토큰 예측과 추론 속도를 모두 개선하지만, 표준 MTP 훈련에 결함이 있어 모델 규모가 커질수록 그래디언트 간섭 문제가 심해진다.
• MTP 교차엔트로피 그래디언트는 수학적으로 분해되는데, 유용한 KL 자기증류 항과 신호대잡음비가 낮은 잔여항으로 나뉜다.
• 저자들은 KL 그래디언트 복원 트릭을 제시했다: 두 분포가 선형 출력층을 공유할 때, 스트레이트스루 그래디언트 라우팅으로 교차엔트로피 그래디언트를 조합하여 KL 발산의 그래디언트를 정확히 복원할 수 있다.
• 이 방식은 어휘 크기의 로짓 행렬이 필요 없고 표준 융합 교차엔트로피 커널과 호환되므로 오버헤드가 없다.
• Soft-MTP는 기존 MTP와 NTP 교차엔트로피 계산을 재사용하면서 한 줄 코드 변경으로 잡음 있는 잔여항을 제거하고 KL 신호만 역전파한다.
• Soft-MTP는 DeepSeek 스타일과 Meta 스타일 MTP를 모두 지원하며 추가 파라미터와 측정 가능한 오버헤드가 없다.
• 167M에서 1.2B 파라미터 범위에서 Soft-MTP는 표준 MTP 대비 더 낮은 NTP 퍼플렉시티, 더 빠른 추론, 더 높은 제로샷 정확도를 보이며, 689M 이상에서는 더 낮은 MTP 퍼플렉시티도 달성한다.
• 1.2B 규모에서 수학과 코드 추론 세밀조정 후 더 강한 성능을 보인다.
• Soft-MTP는 직접 KL 발산 최소화와 대비해도 대부분의 지표에서 동등하거나 능가한다.