Multi-Token Prediction(MTP) 학습 기법: 한 번에 여러 토큰을 예측하고 드래프터를 공짜로 얻는 방법
요약
| 항목 | 내용 |
|---|---|
| 핵심 아이디어 | 학습 시 N개의 예측 헤드를 추가해 t+1, t+2, …, t+N 토큰을 동시에 예측 |
| 추론 시 효과 | 드래프트 토큰을 별도 모델 없이 MTP 헤드에서 생성 → Speculative Decoding 무료 제공 |
| 채택 사례 | DeepSeek-V3/R1, Kimi K3, Nemotron-Puzzle, Gemma 4 |
| 학습 비용 | 헤드당 약 0.1 헤드 크기의 추가 파라미터, 손실 가중치 조절로 주 손실 품질 유지 |
| 현재 상태 | vLLM 0.8+, SGLang 0.5+에서 MTP 기반 Speculative Decoding 지원 |
왜 Multi-Token Prediction인가
표준 언어 모델 학습은 다음 토큰 하나만 예측하는 Next-Token Prediction(NTP)을 사용한다. 각 위치에서 하나의 소프트맥스 분포를 계산하고, 그것만으로 손실을 역전파한다. 이 설계는 단순하고 강력하지만 두 가지 낭비가 있다.
첫째, 긴 시퀀스의 모든 정보가 마지막 숨겨진 상태 하나를 통해서만 흘러야 한다. 위치 t에서 모델이 문장 전체의 장기 구조를 파악하더라도, 그 지식은 t+1 토큰을 맞추는 데만 쓰인다. 더 먼 미래를 직접 예측하는 학습 신호가 없다.
둘째, 추론 속도는 Auto-regressive의 한계에 갇힌다. 한 번에 토큰 하나씩 생성하므로 GPU 연산이 메모리 대역폭 병목에 묶인다. Speculative Decoding이 해결책으로 부상했지만, 별도의 드래프트 모델을 훈련·관리해야 한다.
MTP(Multi-Token Prediction)는 이 두 문제를 동시에 공략한다.
학습 구조: 공유 트렁크 + N개 예측 헤드
총 손실은 헤드별 손실의 가중 합이다. Meta AI 논문(arXiv:2404.19737)의 원래 설계에서는 모든 λ를 1.0으로 뒀지만, DeepSeek-V3는 헤드 수가 늘수록 λ를 줄이는 방식을 채택했다.
중요한 구현 디테일:
헤드마다 트렁크 출력 hₜ에서 독립적으로 투영을 배운다. 각 헤드는 작은 2-레이어 MLP + 언임베딩 행렬로 구성되며, 트렁크 파라미터는 공유된다. 헤드당 추가 파라미터는 전체 모델 크기의 약 0.1~0.3% 수준이다.
왜 더 먼 토큰을 예측하면 표현이 좋아지는가
직관은 간단하다. 모델이 "다음 토큰 하나"만 맞추면 지역적(local) 패턴만 학습해도 충분하다. 하지만 t+5나 t+10을 동시에 맞춰야 한다면, 숨겨진 상태 hₜ는 더 먼 문맥 구조를 인코딩해야 한다.
Meta AI 실험 결과:
- 코딩 벤치마크(HumanEval, MBPP)에서 MTP 훈련 모델이 NTP 동규모 모델 대비 일관되게 우위
- 특히 긴 함수 생성, 반복 구조가 있는 코드에서 차이가 두드러짐
- N=4까지는 품질 이득이 꾸준히 증가했으며, 그 이상은 한계 이득이 줄어듦
추론 시 무료 Speculative Decoding
MTP의 실용적 핵심이다.
표준 Speculative Decoding은 두 모델을 필요로 한다.
- 드래프트 모델: 작은 모델이 k개 토큰을 빠르게 생성
- 검증 모델: 큰 모델이 드래프트를 배치로 검증, 불일치 위치 이후를 거부
MTP 모델은 추론 시 헤드 2, 3, …, N을 드래프트 소스로 활용할 수 있다.
실제 속도 이득: DeepSeek-V3 기준 MTP 2헤드 설정에서 배치 추론 처리량이 약 1.8×~2.1× 향상됐다고 보고됐다. 수락률은 입력 도메인에 따라 다르며, 코드 생성에서 더 높다.
제약: 수락률은 헤드가 예측한 분포가 실제 분포와 얼마나 일치하는지에 달려 있다. N이 클수록 먼 헤드의 수락률이 떨어져 평균 수락 길이가 수렴한다.
실제 채택 사례
DeepSeek-V3 / DeepSeek-R1
DeepSeek-V3 기술 보고서(2024년 12월)에 MTP 구조가 공개됐다. 671B 파라미터 MoE 모델에서 N=2 (헤드 2개)를 사용했다. 학습 시에는 두 헤드를 모두 활성화하고, 추론 시에는 두 번째 헤드를 드래프터로 전환했다. 단일 A100 노드에서 토큰당 FLOP 증가 없이 디코딩 속도가 유의미하게 상승했다고 명시했다.
Kimi K3
月之暗面(Moonshot AI)의 Kimi K3 보고서(arXiv:2607.24653, 2026년 7월)는 MTP를 N=4까지 확장했다. 코딩 추론(AIME, Codeforces)에서의 이득을 강조했고, 헤드 손실 가중치 λ를 거리에 따라 지수적으로 감소시켰다.
Nemotron-Puzzle
NVIDIA의 Nemotron-Puzzle(2026년)은 MTP를 퍼즐 형태 데이터 믹스와 결합했다. 중간 정답이 여러 단계에 걸쳐 예측되어야 하는 수학 문제에서 MTP가 긴 의존성 학습에 특히 효과적이라고 보고했다.
Gemma 4
Google DeepMind의 Gemma 4(2026년)는 N=4 MTP 헤드를 사용하며, vLLM / Ollama / MLX에서 MTP 기반 스펙 디코딩이 기본 활성화 상태다.
구현 시 고려 사항
헤드 수 선택
N=2가 비용 대비 효과의 최적점으로 자주 보고된다. N=4는 코딩 집중 모델에서 추가 이득이 있으나 검증 배치 구성이 복잡해진다. N≥8은 수락률 감소로 실질 이득이 줄어든다.
손실 가중치 튜닝
λ를 너무 높게 설정하면 주 NTP 손실이 묻혀 기본 언어 품질이 낮아질 수 있다. DeepSeek-V3는 λ₂=0.1로 보수적인 값을 사용했다.
추론 프레임워크 지원
vLLM 0.8+ (--enable-chunked-prefill, --speculative-model mtp 설정 없이 모델 구조에서 자동 인식), SGLang 0.5+에서 MTP 헤드를 드래프터로 자동 활용한다. TensorRT-LLM도 DeepSeek 시리즈를 위해 MTP 경로를 구현했다.
단점
- 학습 중 메모리 증가: 헤드 수 × 배치 × 시퀀스 길이만큼 추가 활성화 메모리
- Gradient 흐름이 트렁크 쪽으로 혼재되어 학습률 스케줄 조정이 필요할 수 있음
- 추론 프레임워크가 MTP를 이해하지 못하면 헤드를 그냥 버리게 됨 (속도 이득 없음)
요점 정리
- 학습 시: N개 예측 헤드를 공유 트렁크에 붙여 t+1, …, t+N을 동시에 예측한다. 추가 파라미터는 극소, 표현력은 향상.
- 추론 시: 헤드 2~N을 드래프터로 재활용해 별도 모델 없이 Speculative Decoding을 구현한다.
- N=2가 대부분의 프로덕션 배포에서 비용-효과 최적점이다.
- 코딩·추론 집중 모델에서 이득이 더 두드러지며, 수용 정도는 도메인에 따라 달라진다.
References
- Gloeckle et al., "Better & Faster Large Language Models via Multi-token Prediction," arXiv:2404.19737 (Meta AI, 2024). https://arxiv.org/abs/2404.19737
- DeepSeek-AI, "DeepSeek-V3 Technical Report," arXiv:2412.19437 (2024). https://arxiv.org/abs/2412.19437
- Moonshot AI, "Kimi K3 Technical Report," arXiv:2607.24653 (2026). https://arxiv.org/abs/2607.24653
- vLLM Documentation — Speculative Decoding with MTP. https://docs.vllm.ai/en/latest/features/spec_decode.html
- SGLang Documentation — Multi-Token Prediction. https://sglang.readthedocs.io/en/latest/speculative_decoding.html
- NVIDIA, "Nemotron-Puzzle Technical Report," 2026. https://research.nvidia.com/nemotron-puzzle