Soft-MTP, 멀티토큰 예측의 노이즈 그래디언트를 제거하다

배경

멀티토큰 예측(MTP)은 언어 모델 사전학습에서 널리 사용되는 구성 요소가 됐다. 모델이 다음 토큰뿐 아니라 더 먼 위치의 토큰까지 예측하도록 하면 다운스트림 성능을 높일 수 있고, 추측 디코딩을 위한 더 강한 초안 예측도 기대할 수 있다. 그러나 예측 헤드를 추가하는 것만으로 효과가 보장되지는 않는다. 공유 모델 본체로 어떤 그래디언트를 전달하느냐가 학습 품질에 큰 영향을 준다.

표준 MTP의 문제

이번 연구는 일반적인 MTP 교차엔트로피 그래디언트에 숨은 문제를 분석했다. 이 그래디언트는 두 부분으로 정확히 분해할 수 있다. 하나는 KL 발산 형태의 자기 증류 신호다. 이는 각 MTP 헤드가 모델 자체의 다음 토큰 예측 분포를 따르도록 만든다. 본 모델과 초안 헤드의 분포를 맞추기 때문에 추측 디코딩에 유용하다. 다른 하나는 신호대잡음비가 낮은 잔차다.

이 잔차는 단순한 여분의 항이 아니다. 공유 트렁크로 잡음이 전달되면서 다음 토큰 학습과 멀티토큰 학습을 동시에 방해할 수 있다. 연구는 모델 규모가 커질수록 이런 간섭이 더 두드러진다고 설명한다. 따라서 표준 MTP 손실을 그대로 추가하는 방식은 실제로 유용한 감독 신호와 잡음을 함께 역전파할 수 있다.

Soft-MTP의 방식

연구진은 ‘KL 그래디언트 복원’ 기법을 제안했다. NTP와 MTP가 선형 출력 헤드를 공유한다는 점을 이용하면, 기존 교차엔트로피 그래디언트를 조합해 KL 항에 해당하는 그래디언트를 정확히 복원할 수 있다. 직통 그래디언트 라우팅을 사용하므로 어휘 크기의 로짓 행렬을 별도로 만들 필요가 없다.

Soft-MTP는 표준 MTP에 한 줄 수준의 변경으로 적용할 수 있다. 기존 NTP와 MTP 교차엔트로피 계산을 재사용하고, 융합 교차엔트로피 커널과도 호환된다. 추가 파라미터가 없으며 측정 가능한 계산 오버헤드도 발생하지 않는다. DeepSeek 방식과 Meta 방식 MTP를 모두 지원하고, 167M부터 1.2B 파라미터 규모의 평가에서 표준 MTP보다 나은 결과를 보였다고 연구는 보고한다.

의미와 영향

이 결과가 주는 broader한 교훈은 보조 학습 목표에서 용량을 늘리는 것만큼 그래디언트를 선별하는 일이 중요할 수 있다는 점이다. MTP의 핵심 역할이 미래 토큰 헤드와 본 모델의 분포를 정렬하는 것이라면, 손실의 모든 성분을 동일하게 역전파하기보다 KL 자기 증류 신호에 집중하는 편이 합리적일 수 있다.

실무 측면에서도 접근성이 높다. 새로운 아키텍처나 학습 파라미터, 최적화된 교차엔트로피 연산자의 교체가 필요하지 않기 때문이다. 이미 MTP 사전학습 파이프라인을 보유한 팀이라면 전체 구조를 바꾸지 않고 학습 목표를 가볍게 수정하는 방식으로 검토할 수 있다. 다만 제공된 소재에는 전체 실험 표가 포함되지 않았으므로, 더 큰 모델과 더 많은 예측 헤드, 추가 과제에서의 효과는 논문 전문을 통해 확인해야 한다.

출처：Hugging Face Daily Papers