MTP 학습에서 잡음 성분을 걷어내 추론과 정확도를 함께 개선한다

무엇인가

이 논문은 다중 토큰 예측(MTP) 사전학습의 표준 교차엔트로피 학습에 결함이 있다고 주장한다. MTP는 다음 토큰 예측(NTP)보다 긴 지평의 감독을 주고 추론 시 스페큘러티브 디코딩 초안으로 쓰이지만, 표준 MTP 그래디언트는 정확히 두 부분으로 분해된다. 위치 정렬된 NTP 분포 q_t(k)=p_{t+k}(0)로 MTP 분포 p_t(k)를 끌어당기는 KL 자기증류 성분 g_KL과, 데이터의 원-핫 타깃에 대한 NTP 교차엔트로피 그래디언트인 잔차 g_res다. 저자들은 g_res가 신호대잡음비가 낮고 공유 트렁크의 MTP 그래디언트 중 약 80%를 차지하며, 학습 내내 NTP 업데이트 방향과 반대 정렬되어 NTP와 MTP 학습을 모두 방해한다고 보고한다. 이 간섭은 모델 규모가 커질수록 심해진다.

어떻게 동작하나

Soft-MTP의 핵심은 g_res를 없애고 g_KL만 정확히 복원하는 KL 그래디언트 복구 트릭이다. 두 분포가 같은 선형 출력 헤드 Wθ를 공유하면, 교차엔트로피 그래디언트의 조합으로 순전파 KL의 그래디언트를 정확히 얻을 수 있다. 구현은 NTP 히든 상태를 h~_t=h_t(0)+Σ[sg(h_{t-k}(k))-h_{t-k}(k)]로 증강한다. 이 항은 수치적으로 0이라 NTP 예측은 바뀌지 않지만, 역전파 때 MTP 히든 상태 h_t(k)에 -g_res를 주입한다. 따라서 증강 NTP 손실과 기존 MTP 교차엔트로피 손실을 합치면 표준 MTP 그래디언트 g_hard에서 g_res를 뺀 g_KL이 된다. 별도 어휘 크기 로짓 행렬이나 KL 계산 없이 기존 융합 교차엔트로피 커널을 재사용하는 한 줄 코드 변경이며, DeepSeek 스타일과 Meta 스타일 MTP를 모두 지원한다.

무엇과 다른가

Soft-MTP는 MTP 헤드와 공유 트렁크에는 g_KL만 보내지만, 공유 출력 헤드 Wθ에는 MTP 교차엔트로피 그래디언트를 계속 보낸다. 이 차이가 중요하다. KL-MTP처럼 출력 헤드까지 자기증류 그래디언트로 업데이트하면, KL을 줄이는 지름길로 MTP와 NTP 히든 상태가 다른 방향에 덜 민감하게 읽히도록 헤드를 만들 수 있다. 실제로 로짓 맵의 유효 랭크는 NTP 496, 표준 MTP 131, KL-MTP 321, Soft-MTP 437로 보고된다. Soft-MTP는 데이터에 맞는 교차엔트로피 감독을 헤드에 유지하면서 잡음 잔차만 트렁크에서 제거한다.

어떻게 쓰나

실험은 LLaMA 스타일 모델 167M, 332M, 689M, 1.2B 파라미터 규모에서 DeepSeek 스타일 순차 MTP 헤드 3개(K=3)로 진행됐다. 각 MTP 헤드는 트렁크 비임베딩 파라미터의 3.2%로 고정하고 λ_MTP=1.0을 썼다. FineWeb-Edu에서 배치 0.5M 토큰, 시퀀스 길이 2048, 60,000스텝, 약 31.5B 토큰을 학습했으며 가장 큰 모델 기준 약 1.3배 Chinchilla-optimal이다. Muon 옵티마이저와 update-RMS 리스케일링을 사용했고, Meta 스타일 MTP 헤드 실험은 부록 C에 있다. 결과는 모든 규모에서 표준 MTP보다 낮은 NTP 퍼플렉시티, 더 빠른 자기 스페큘러티브 디코딩, 더 높은 제로샷 정확도를 보였다. MTP 퍼플렉시티는 689M부터 더 낮아졌다. 다섯 도메인 평균으로 수용 길이는 1.8~2.2%, 속도향상은 1.4~1.5% 개선됐고 MTP 분포는 NTP 검증기와의 KL 및 TV 거리가 모든 규모에서 더 작았다. 제로샷 평균 정확도는 표준 MTP 대비 167M과 332M에서 약 0.1점, 689M과 1.2B에서 약 0.5점 높았으며 689M과 1.2B에서는 NTP까지 추월했다. 1.2B에서는 WikiText와 LAMBADA 퍼플렉시티 및 평균 정확도가 세 모델 중 가장 좋았다.

전제와 한계

1.2B 모델에서 MTP 헤드를 버리고 NTP만으로 파인튜닝한 전이 실험도 있다. 수학은 OpenMathInstruct-2 4.2B 토큰으로 학습해 GSM8K와 MATH-500을, 코드는 Stack-Edu 파이썬 부분집합 21.0B 토큰으로 학습해 HumanEval과 MBPP를 평가했다. Soft-MTP는 여덟 칸 중 일곱 칸에서 표준 MTP를 앞섰고 NTP 아래로 떨어지지 않은 유일한 목적함수였다. 코드에서는 pass@64 기준 HumanEval +3.1%p, MBPP +1.2%p, 수학에서는 GSM8K +1.1점, MATH-500 +1.8점이었다. 332M 모델을 7.3B 토큰으로 학습한 스윕에서도 λ_MTP 0.25~1.0과 K 3~8 전 구간에서 NTP 및 MTP 손실 파레토 우위를 보였다.

비용 비교에서 Soft-MTP는 1.2B 기준 표준 MTP와 같은 matmul FLOPs, MFU, 융합 CE 피크 메모리를 유지했고 처리량 차이는 0.07% 이내였다. 반면 명시적 KL-MTP는 FLOPs 1.5% 증가, 처리량 5.6% 감소, 피크 메모리 21GB 증가가 있었다. Soft-MTP는 대부분 지표에서 KL-MTP와 동등하거나 우세했지만, 수학 파인튜닝에서는 KL-MTP가 앞섰고 FineWeb-Edu 스페큘러티브 디코딩에서는 KL-MTP가 최대 0.01 앞섰다. Soft-MTP는 1.2B에서 모든 퍼플렉시티 지표와 평균 제로샷 정확도에서 KL-MTP와 동등하거나 우세했고, 코드 파인튜닝 여섯 지표에서는 모두 앞섰다.

실무에서는 이미 MTP 헤드를 쓰는 사전학습 코드베이스에서 기존 MTP 손실에 Soft-MTP 라우팅을 추가하는 방식으로 가장 먼저 검증할 만하다. 추가 파라미터가 없고 측정 가능한 오버헤드가 없으며 융합 교차엔트로피 커널을 유지하기 때문이다. 다만 사후 학습용 드래프터만 바꾸는 기법이 아니라 사전학습 그래디언트 라우팅을 바꾸는 방법이므로, 자체 규모와 병렬화 환경에서 NTP 및 MTP 손실, 스페큘러티브 수용률, 다운스트림 정확도를 함께 확인해야 한다.

저자들이 밝힌 전제는 NTP 히든 상태와 MTP 히든 상태가 같은 장치에 있고 NTP 손실에서 MTP 히든 상태로의 역전파 경로가 필요하다는 점이다. 데이터 병렬, ZeRO, 전문가 병렬은 추가 통신이 없지만 시퀀스 샤딩과 컨텍스트 병렬에서는 샤드 경계마다 최대 K개의 히든 상태 그래디언트 교환이 필요하고, 파이프라인 병렬에서는 MTP 모듈이 출력 헤드와 같은 스테이지에 있어야 한다. 또한 공유 출력 헤드는 정확한 KL이 아니라 교차엔트로피로 업데이트되며, 실험은 167M~1.2B 규모에 한정되고 최전선 규모는 호환성 논의에 그친다.