Multi-Token Prediction Done Right

Paper: Jayden Teoh (2026). "Multi-Token Prediction Done Right". arXiv preprint, arXiv:2610.38331v1, 2026-10-28. 원문 · PDF.

판본: arXiv:2610.38331v1, 2026-10-28, 23쪽 읽기 버전: 2026-11-01

1. 여러 토큰을 한꺼번에 예측하는 학습

LLM은 보통 다음 토큰 하나를 맞히도록 학습합니다. 이를 **NTP(next-token prediction, 다음 토큰 예측)**라고 합니다. 최근 DeepSeek-V3, Nemotron 3, GLM-5, Kimi K3 같은 모델은 여기에 **MTP(multi-token prediction, 다중 토큰 예측)**를 더합니다. 본체 위에 작은 head를 몇 개 더 달아 두 번째, 세 번째 뒤의 토큰까지 미리 맞히게 하는 방식입니다(Gloeckle et al., 2024).

MTP의 쓸모는 두 가지입니다. 먼 앞을 내다보게 해 학습 신호를 보완하고, 추론 때는 MTP head가 다음 토큰들을 미리 써 두면 본체가 한 번에 검사해 받아들이는 **speculative decoding(추측 디코딩)**으로 생성을 빠르게 합니다.

2. 표준 MTP의 학습 신호에는 잡음이 섞여 있습니다

추측 디코딩에서 MTP가 쓴 초안은 본체의 NTP 예측과 맞아야 받아들여집니다. 그런데 표준 MTP head는 본체의 예측이 아니라 데이터의 실제 정답 토큰을 맞히도록 학습합니다.

저자는 표준 MTP의 학습 방향(gradient)을 정확히 두 조각으로 나눕니다.

• KL 항(자기 증류 신호): MTP head의 예측을 같은 토큰에 대한 본체의 NTP 예측 쪽으로 끌어당깁니다. 추측 디코딩에 필요한 바로 그 방향입니다.
• 잔차(residual): 나머지 부분입니다. 저자의 측정에서 잔차는 신호보다 잡음이 훨씬 많았습니다.

원 논문 Figure 3(논문 PDF 7쪽): (a–b) 파란 선(KL 항)이 빨간 선(잔차)보다 신호 대 잡음비가 높습니다. (c) 그런데 본체에 도달하는 MTP 학습 신호의 약 80%는 잔차가 차지합니다. 원문 PDF에서 보기.

잔차의 **신호 대 잡음비(SNR)**는 KL 항보다 약 4배 낮았고, 이 격차는 167M 모델의 4.1배에서 1.2B 모델의 4.6배로 커졌습니다. 그런데도 잔차는 본체로 가는 MTP 신호의 약 80%를 차지했습니다.

3. 잔차는 본래의 다음 토큰 학습까지 방해합니다

원 논문 Figure 4(논문 PDF 8쪽): (a) 각 성분 방향으로 한 걸음 움직였을 때 NTP loss의 변화로, 파란 KL 항은 loss를 낮추고 빨간 잔차는 높입니다. (b–c) 학습 중 NTP·MTP loss에서 Soft-MTP(파란 선)가 표준 MTP(검은 선)보다 대체로 낮고, 오른쪽 숫자는 마지막 차이입니다. 원문 PDF에서 보기.

KL 항 방향으로 움직이면 NTP loss가 일관되게 줄었고, 잔차 방향으로 움직이면 일관되게 늘었습니다. 이 패턴은 학습 내내, 모든 모델 크기에서 유지됐습니다.

4. Soft-MTP: 코드 한 줄로 잔차만 빼냅니다

KL 항만 쓰려면 보통 MTP와 NTP의 어휘 전체 확률 분포를 둘 다 메모리에 만들어 비교해야 합니다. 이러면 메모리가 많이 들고, 확률 분포를 만들지 않고 손실을 계산하는 빠른 fused cross-entropy 커널을 쓸 수 없습니다.

저자의 KL gradient-recovery trick은 이 문제를 피합니다. NTP와 MTP가 같은 출력층을 공유한다는 점을 이용해, 이미 계산하는 두 cross-entropy의 gradient를 조합해서 KL 항만 정확히 되살립니다. 구현은 NTP hidden state의 값은 그대로 두고 MTP 쪽 gradient가 NTP 경로로 흐르게 하는 한 줄입니다.

원 논문 Figure 1(논문 PDF 1쪽): 왼쪽은 straight-through gradient routing의 흐름, 가운데 파란 상자가 바뀌는 코드 한 줄 h0[:, k:] += hk.detach() - hk, 오른쪽은 표준 MTP 대비 NTP·MTP loss 곡선입니다. 원문 PDF에서 보기.

이렇게 만든 Soft-MTP는 파라미터를 하나도 늘리지 않으며, DeepSeek식 순차 MTP 모듈과 Meta식 병렬 head를 모두 지원합니다.

5. 실험 설정

167M부터 1.2B까지 네 크기의 LLaMA 계열 모델을 NTP, 표준 MTP, Soft-MTP로 각각 사전학습했습니다. FineWeb-Edu 데이터로 60,000 step, 약 31.5B 토큰을 학습했으며 이는 가장 큰 모델 기준 Chinchilla 최적량의 약 1.3배입니다. MTP head는 3개이고, 각 head는 본체(임베딩 제외) 파라미터의 3.2% 크기입니다. 비교를 위해 KL 발산을 직접 계산해 최소화하는 KL-MTP도 학습했습니다.

6. 결과: 모든 크기에서 표준 MTP보다 좋았습니다

• 생성 속도: 다섯 도메인 평균으로 Soft-MTP는 표준 MTP보다 받아들여지는 초안 길이가 1.8–2.2%, 속도 향상이 1.4–1.5% 더 컸습니다(모델 크기별 범위).
• 언어 모델 품질: 평균 zero-shot 정확도가 모든 크기에서 표준 MTP보다 높았습니다. 차이는 167M·332M에서 약 0.1점, 689M·1.2B에서 약 0.5점으로 작습니다. MTP loss는 689M부터 표준 MTP보다 낮아졌고, 그보다 작은 모델에서는 오히려 조금 높았습니다.
• 미세조정 후 추론: 1.2B 모델을 수학·코드로 미세조정하면 Soft-MTP가 8개 지표 중 7개에서 표준 MTP보다 높았습니다. HumanEval pass@64에서 +3.1점, MATH-500에서 +1.8점이었고, NTP 기준선보다 낮아진 지표가 하나도 없는 유일한 방식이었습니다.

1.2B에서 세 MTP 방식을 함께 비교하면 다음과 같습니다(원문 Table 6 일부).

지표	NTP	표준 MTP	KL-MTP	Soft-MTP
FineWeb-Edu 퍼플렉서티 ↓	10.57	10.89	10.74	10.74
평균 zero-shot 정확도 ↑	56.75	56.80	57.22	57.31
초안 수용 길이 ↑	–	1.67	1.71	1.70
GSM8K ↑	47.69	48.60	50.34	49.66
MATH-500 ↑	24.80	23.20	26.00	25.00
HumanEval pass@64 ↑	39.02	40.85	41.46	43.90

KL-MTP와 Soft-MTP는 성능이 비슷하며, 수학에서는 KL-MTP가 앞섰습니다. 차이는 비용에서 났습니다. Soft-MTP는 표준 MTP와 계산량·메모리가 같고 처리량 차이가 0.07% 이내였지만, KL-MTP는 계산량이 1.5% 늘고 처리량이 5.6% 줄며 최대 메모리가 21GB 늘었습니다.

7. 큰 규모 학습에서도 쓸 수 있습니다

이 trick은 NTP hidden state와 MTP hidden state가 같은 장치에 있기만 하면 됩니다. 데이터 병렬, ZeRO, expert 병렬은 그대로 쓸 수 있습니다. pipeline 병렬에서는 MTP 모듈이 출력층과 같은 단계에 있어야 하고, 시퀀스를 나눠 처리하는 방식에서는 경계마다 최대 3개(head 수만큼)의 hidden-state gradient를 주고받아야 합니다. 이 비용은 시퀀스 길이와 무관합니다.

8. 저자가 밝힌 한계

• MTP가 앞으로도 계속 쓰일지는 불확실합니다. DeepSeek-V4.1-Flash는 MTP 대신 사전학습 뒤 따로 학습하는 추측 디코더 DSpark를 택했습니다. 다만 MTP로 미리 내다보는 표현을 익힌 적 없는 모델에서 이런 사후 학습 drafter가 얼마나 잘 작동하는지는 열린 문제라고 저자는 씁니다.
• 공유 출력층은 정확한 KL gradient가 아니라 cross-entropy로 갱신됩니다. 정확한 KL gradient는 MTP head와 본체만 받습니다. 저자는 출력층에서 cross-entropy를 유지하는 것이 KL-MTP의 지름길(출력층이 두 hidden state의 차이에 둔감해지는 현상)을 막는다고 설명하며, 출력 logit map의 effective rank는 NTP 496, Soft-MTP 437, KL-MTP 321, 표준 MTP 131이었습니다.
• trick의 적용 범위는 teacher와 student가 선형 출력층을 공유하는 설정입니다.

References

Gloeckle, F., Idrissi, B. Y., Rozière, B., Lopez-Paz, D., & Synnaeve, G. (2024). Better & faster large language models via multi-token prediction. Proceedings of the 41st International Conference on Machine Learning.

Teoh, J. (2026). Multi-token prediction done right (Version 1, arXiv:2610.38331) [Preprint]. arXiv. https://arxiv.org/abs/2610.38331v1