다중 토큰 예측(Multi-Token Prediction; MTP)

GPT-2부터 Llama 3까지 모든 자기회귀(autoregressive) LLM은 위치마다 하나의 손실 함수(loss)로 학습합니다. 바로 다음 토큰(token)을 예측하는 것입니다. DeepSeek-V3는 위치마다 두 번째 손실 함수를 추가했습니다. 그다음 토큰까지 예측합니다. 671B 모델에 추가된 14B 매개변수(parameter)는 기울기 흐름(gradient flow)을 통해 본체 모델(main model)로 증류(distill)되었고, 학습된 MTP 헤드(head)는 추론 시 추측 디코딩(speculative decoding)의 초안 생성기(drafter)로 재사용되어 80% 이상의 수용률(acceptance rate)을 보였습니다. 1.8배의 생성 처리량(generation throughput)이 사실상 공짜로 따라왔습니다. 이 강의에서는 DeepSeek 기술 보고서(technical report)의 순차적(sequential) MTP 모듈을 만들고, 손실 함수와 공유 헤드(shared-head) 매개변수 배치를 계산하며, Gloeckle et al.의 원래 병렬(parallel) MTP가 인과 사슬(causal chain)을 깨뜨린 반면 DeepSeek MTP가 이를 어떻게 유지하는지 설명합니다.

유형: Build 언어: Python (표준 라이브러리) 선수 지식: Phase 10 · 04 (소형 GPT 사전학습), Phase 10 · 15 (추측 디코딩) 예상 시간: 약 60분

학습 목표

• MTP 학습 목표(training objective)를 설명하고 예측 깊이(prediction depth) 전체의 결합 손실(joint loss)을 유도합니다.
• Gloeckle et al.의 병렬 MTP 헤드(2024)와 DeepSeek-V3의 순차적 MTP 모듈의 차이를 설명하고, 순차적 설계가 인과 사슬을 보존하는 이유를 설명합니다.
• 사전학습(pre-training) 진행에 MTP 모듈을 추가할 때의 매개변수 부담(parameter overhead)과 메모리 부담(memory overhead)을 계산합니다.
• MTP 모듈 하나를 처음부터 구현합니다. 공유 임베딩(shared embedding), 깊이별 트랜스포머 블록(per-depth transformer block), 사영(projection), 공유 출력 헤드(shared output head)를 포함합니다.

문제

다음 토큰 예측(next-token prediction)은 표준 LLM 학습 목표입니다. 모든 은닉 상태(hidden state)는 정확히 하나만 예측하도록 지도(supervision)됩니다. 바로 다음 토큰입니다. 이는 놀라울 정도로 약한 신호입니다. 시퀀스(sequence)에 담긴 정보 대부분은 토큰 하나를 넘어서 이어집니다. 구조, 일관성(coherence), 사실성(factuality), 산술 흐름(arithmetic flow)이 모두 그렇습니다. 모델은 수조 개 토큰(trillion token)에 걸친 수많은 한 토큰 단위 신호(one-token signal)를 누적하면서 이를 배워야 합니다.

MTP는 이렇게 묻습니다. 모든 은닉 상태가 여러 미래 토큰을 동시에 예측하도록 지도된다면 어떨까요? Gloeckle et al.(Meta, 2024)은 이것이 도움이 된다는 것을 보였습니다. 그 구현은 백본(backbone) 위에 여러 개의 독립 출력 헤드(independent output head)를 두고, 각 헤드가 서로 다른 오프셋(offset)을 예측했습니다. 병렬이고 단순했지만, 헤드들은 계층적 정제(hierarchical refinement) 없이 같은 은닉 상태를 보았습니다. 또한 예측이 인과적으로 사슬(causally chain)을 이루지 않았기 때문에 추측 디코딩에는 사용할 수 없었습니다.

DeepSeek-V3(2024년 12월)는 각 예측 깊이에서 인과 사슬을 유지하도록 MTP를 순차적 모듈로 다시 설계했습니다. 모델은 h_i^(0)에서 t+1을 예측하고, 그다음에는 h_i^(0)와 E(t+1) 임베딩을 결합한 새로운 은닉 상태 h_i^(1)에서 t+2를 예측합니다. 이런 식으로 이어집니다. 각 깊이는 자기만의 작은 트랜스포머 블록(transformer block)을 가집니다. 공유 임베딩과 공유 출력 헤드를 사용해 매개변수 부담을 적정한 수준으로 유지합니다. DeepSeek-V3 규모에서는 671B 본체 모델 가중치(main-model weight) 위에 MTP 모듈 전체가 14B 추가 매개변수를 차지했습니다. 이 2%의 부담으로 더 촘촘한 학습 신호(training signal)와 추론에서 곧바로 쓸 수 있는 추측 디코딩 초안(speculative-decoding draft)을 함께 얻은 셈입니다.

이 강의에서는 단일 MTP 모듈과 깊이 D짜리 손실(D-depth loss)을 처음부터 만듭니다. 수식은 깔끔하고 구현은 150줄입니다.

개념

순차적 MTP 조리법

DeepSeek-V3는 본체 모델 위에 D개의 MTP 모듈을 추가합니다. 각 모듈 k(k = 1..D)는 깊이 k의 토큰을 예측합니다. 다시 말해 위치 i까지의 접두부(prefix)가 주어졌을 때 t_{i+k}를 예측합니다.

모듈 k는 다음으로 구성됩니다.

• 자체 어텐션(attention)과 MLP를 가진 트랜스포머 블록 T_k.
• 이전 깊이의 은닉 상태와 다음 깊이의 정답 토큰(ground-truth token) 임베딩을 결합하는 사영 행렬(projection matrix) M_k.
• 공유 임베딩 E. 본체 모델과 동일한 임베딩입니다.
• 공유 출력 헤드 Out. 본체 모델과 동일한 헤드입니다.

학습 시 위치 i까지의 접두부에 대해 깊이별 은닉 상태는 다음과 같습니다.

h_i^(0) = main model backbone at position i
h_i^(k) = T_k( M_k * concat(RMSNorm(h_i^(k-1)), RMSNorm(E(t_{i+k}))) )   for k >= 1

깊이별 예측은 다음과 같습니다.

logits_{i+k} = Out(h_i^(k-1))   for k = 1..D

깊이별 손실은 정답 t_{i+k}에 대한 교차 엔트로피(cross-entropy)입니다.

L_k = CE(logits_{i+k}, t_{i+k})

깊이 전체의 결합 손실은 다음과 같습니다.

L_MTP = (lambda / D) * sum_{k=1..D} L_k

lambda는 작은 가중치 계수(weighting factor)입니다. DeepSeek-V3는 학습 초반 10%에는 0.3을, 이후에는 0.1을 사용합니다. 전체 학습 손실은 L_main + L_MTP입니다.

왜 병렬이 아니라 순차여야 하는가

Gloeckle의 원래 병렬 MTP는 D개의 출력 헤드를 두고, 각각을 h_i^(0)에 직접 적용했습니다. 각 헤드는 같은 백본 은닉 상태에서 t_{i+k}를 예측합니다. 이 방식도 학습은 되지만 예측끼리 서로 조건화(conditioning)되지 않습니다. head_1의 출력을 사용해 head_2를 도울 수 없습니다. 헤드가 병렬로 점화하기 때문입니다.

DeepSeek-V3의 순차적 설계는 h_i^(k-1)과 실제 다음 토큰 임베딩 E(t_{i+k})로부터 h_i^(k)를 만듭니다. 이것이 인과 사슬을 보존합니다. t_{i+k+1}을 예측하려면 깊이 k+1의 모듈이 t_{i+k}에 무엇이 있었는지를 봅니다. 이는 자기회귀 디코더(autoregressive decoder)가 자기 출력을 다시 입력으로 소비하는 방식과 구조적으로 동일합니다. 따라서 MTP 모듈은 추측 디코딩의 초안 생성기로 그대로 사용할 수 있습니다.

추론 시에는 h_i^(k-1)과 초안으로 뽑힌(drafted) t_{i+k}를 모듈 k+1에 넣고 t_{i+k+1} 예측을 얻은 뒤 반복합니다. 이는 학습된 MTP 모듈을 초안 네트워크(draft network)로 사용하는 EAGLE 방식 초안과 정확히 동일합니다. DeepSeek-V3는 첫 번째 MTP 모듈에서 80% 이상의 수용률과 약 1.8배의 속도 향상(speedup)을 보고했습니다.

매개변수 회계

은닉 차원이 h, 어휘(vocabulary)가 V인 모델을 생각합니다.

• 본체 모델(Main model): 수십억에서 수천억 매개변수와 V * h 크기의 출력 헤드 하나.
• 공유 출력 헤드: 본체 모델의 헤드를 재사용합니다. 추가 매개변수가 없습니다.
• 공유 임베딩: 본체 모델의 임베딩을 재사용합니다. 추가 매개변수가 없습니다.
• MTP 모듈 하나당:
  • 사영 M_k: (2h) * h = 2h^2.
  • 트랜스포머 블록 T_k: 어텐션(MHA 기준 4h^2)과 MLP(SwiGLU 비율 8/3 기준 보통 8h^2). 블록 하나당 약 12h^2.

모듈 하나당 추가량은 약 ~14h^2입니다. DeepSeek-V3의 h = 7168, D = 1 모듈이라면 논문 수준의 계산으로는 ~14 * 7168^2 = ~720M 매개변수입니다. DeepSeek-V3가 보고한 값은 14B입니다. 차이는 대부분 MTP 모듈의 전문가 층(expert layer)도 전문가 혼합(Mixture of Experts; MoE) 구조를 상속하기 때문입니다.

추측 디코딩의 보상

사전학습 중 MTP 모듈은 학습 속도를 약 10% 늦춥니다. 전방 계산(forward compute)이 늘고 추가 손실이 생기기 때문입니다. 보상은 두 가지입니다.

1. 더 촘촘한 학습 신호. 각 은닉 상태는 D+1개의 지도 목표(supervision target)를 봅니다. DeepSeek-V3의 절제 실험(ablation)에서 MMLU, GSM8K, MATH, HumanEval에 대해 일관되게 몇 퍼센트포인트(percentage point) 향상이 측정되었습니다.

2. 추론에서 공짜로 얻는 추측 디코딩 초안. MTP 모듈은 이미 다음 몇 개 토큰을 예측하도록 학습되어 있습니다. 초안 네트워크로 재사용하면 80% 이상의 수용률을 보입니다. 이 수준이면 N=3 또는 N=5 추측 디코딩이 1.8배의 처리량을 제공합니다. 10%의 학습 시간 비용은 추론을 처음 운영하는 순간 회수됩니다.

EAGLE과의 관계

EAGLE은 사전학습 이후에 작은 초안 모델(draft model)을 별도로 학습합니다. MTP는 초안을 사전학습 안에 녹여 둡니다. 두 접근은 비슷한 수용률로 수렴하지만 파이프라인(pipeline)이 다릅니다.

직접 만들기

code/main.py는 단일 MTP 모듈을 처음부터 끝까지(end-to-end) 만듭니다. 공유 임베딩, 사영, 트랜스포머 블록, 공유 출력 헤드를 포함합니다. 이어서 짧은 합성 시퀀스(synthetic sequence)에서 깊이별 교차 엔트로피 손실을 계산하고 구성 요소별 매개변수 개수를 출력합니다. 장난감용 어휘(toy vocabulary)는 32개 토큰으로 두어 숫자가 읽기 쉽게 만듭니다.

Step 1: 공유 임베딩 표

하나의 vocab_size x hidden 표(table)가 본체 모델과 모든 깊이의 모든 MTP 모듈에서 사용됩니다. 두 번째 사본이 아니라 말 그대로 동일한 텐서(tensor)입니다.

Step 2: 깊이별 결합

def combine(prev_hidden, next_token_embed, M_k):
    # 특징 차원 방향으로 concat한 뒤 hidden 차원으로 사영한다.
    concat = rms_norm(prev_hidden) + rms_norm(next_token_embed)  # 벡터 덧셈으로 대체
    projected = matvec(M_k, concat)
    return projected

실제 DeepSeek-V3는 RMSNorm을 적용한 두 벡터(vector)를 [2h]로 이어 붙인(concatenate) 뒤 h x 2h 행렬로 사영합니다. 장난감 구현에서는 표준 라이브러리로 간결하게 보여 주기 위해 벡터 덧셈으로 대체합니다.

Step 3: 깊이 k의 트랜스포머 블록

자기 어텐션(self-attention)과 MLP로 구성합니다. 장난감 구현에서는 구조가 잘 드러나도록 1개 층의 선형 어텐션 블록(linear attention block)과 SwiGLU MLP를 사용합니다. NumPy는 사용하지 않습니다.

Step 4: 공유 출력 헤드

본체 모델의 출력 사영을 재사용합니다. 어휘 위의 로짓(logits)을 만듭니다.

Step 5: 깊이별 손실

softmax(logits)와 오프셋 k의 정답 토큰 사이의 교차 엔트로피입니다. lambda / D 스케일링 계수로 깊이 전체를 합산(aggregate)합니다.

Step 6: 매개변수 회계

전체 매개변수 개수, 공유(임베딩, 헤드) 개수, 모듈 하나당 추가 개수를 출력합니다. 본체 모델 크기 대비 MTP 추가 비율을 보여 줍니다.

사용해보기

MTP는 DeepSeek-V3(2024년 12월)와 DeepSeek-R1 계열에 통합되어 있습니다. 추론 시에는 다음과 같이 사용됩니다.

• DeepSeek 자체 서빙 스택(serving stack)은 MTP 모듈을 추측 디코더(speculative decoder)로 그대로 소비합니다.
• 2026년 4월 기준 vLLM과 SGLang은 DeepSeek-V3 MTP 통합 경로(integration path)를 갖고 있습니다.
• AMD의 ROCm SGLang 튜토리얼은 V3 체크포인트(checkpoint)에서 1.8배의 속도 향상을 측정한 구체적인 MTP 추측 디코딩 설정을 보여 줍니다.

새 사전학습 진행에서 MTP를 사용하기 좋은 경우는 다음과 같습니다.

• 전체 사전학습 파이프라인을 통제하고 더 촘촘한 학습 신호를 확보하고 싶을 때.
• 모델을 대규모로 서빙할 것을 미리 알고 있고, 추측 디코딩을 사실상 공짜로 얻고 싶을 때.
• 은닉 차원이 최소 4096 이상일 때. 1B 규모에서는 부담이 이득보다 더 큽니다.

사용하지 않는 편이 나은 경우도 있습니다.

• 기존에 사전학습된 밀집 모델(dense model)을 미세 조정(fine-tuning)할 때. MTP 모듈이 학습되어 있지 않습니다.
• 깨끗한 기준선(baseline)과 비교해야 하는 연구용 모델. MTP는 구조(architecture)를 바꿉니다.

산출물 만들기

이 강의는 outputs/skill-mtp-planner.md를 만듭니다. 사전학습 진행 명세(모델 크기, 데이터, 연산 자원)가 주어지면 MTP 통합 계획(integration plan)을 반환합니다. 예측 깊이 D, lambda 스케줄, 메모리 부담, 추론 시 추측 디코딩 배선(wiring)을 포함합니다.

연습문제

1. (쉬움) code/main.py를 실행합니다. 합성 신호가 강해질수록 깊이별 손실이 단조롭게 감소하는지 확인합니다. 합성 데이터를 고정 패턴(fixed pattern)으로 바꾸고, 깊이-1과 깊이-2 손실이 모두 수렴하는지 확인합니다.

2. (중간) D=1 MTP 모듈을 가진 밀집 70B 모델(은닉 8192, 80개 층)의 매개변수 부담을 계산합니다. DeepSeek-V3가 보고한 14B 부담과 비교합니다. DeepSeek의 수치가 더 큰 이유를 설명합니다. MTP 트랜스포머 블록이 같은 MoE 구조를 상속해 모듈당 매개변수 개수를 키우기 때문입니다.

3. (중간) 장난감 구현에서 D=2를 구현합니다. h^(1)을 받아 t_{i+2}를 예측하는 두 번째 MTP 모듈을 추가합니다. 결합 손실과 매개변수 회계가 DeepSeek 논문의 식 19-21과 일치하는지 확인합니다.

4. (어려움) 장난감 구현을 병렬 MTP(Gloeckle 방식)로 바꿉니다. 본체 은닉 상태 위에 D개의 출력 헤드를 두고 각각 다른 오프셋을 예측하게 합니다. 같은 합성 신호에서 깊이별 손실을 순차적 구현과 비교합니다. 순차적 구현은 중간 예측에 조건화하므로 k > 1에서 더 낮은 깊이-k 손실을 보여야 합니다.

5. (어려움) 학습된 MTP 모듈을 EAGLE 방식 초안으로 사용합니다. 추론에서 모듈 k를 호출해 t_{i+k}를 제안합니다. 보류된 시퀀스(held-out sequence)에서 본체 모델 예측과 비교해 초안 토큰의 수용률을 측정합니다. 장난감 구현에서 50% 이상을 달성하면 MTP를 초안 생성기로 활용하는(MTP-as-draft) 성질을 경험적으로 재현한 것입니다.

핵심 용어

더 읽을거리

• DeepSeek-AI — DeepSeek-V3 Technical Report (arXiv:2412.19437) — 순차적 MTP의 전체 설명입니다. Section 2.2에 결합 손실 식과 추론에서의 1.8배 속도 향상이 포함됩니다.
• Gloeckle et al. — Better & Faster Large Language Models via Multi-token Prediction (arXiv:2404.19737) — DeepSeek 설계가 개선한 병렬 MTP 기준선(baseline)입니다.
• DeepSeek-V3 model card on Hugging Face — 총 685B(본체 671B + MTP 14B)와 배포 노트입니다.
• Leviathan et al. — Fast Inference from Transformers via Speculative Decoding (arXiv:2211.17192) — MTP가 들어가는 추측 디코딩 프레임워크입니다.
• Li et al. — EAGLE-3 (arXiv:2503.01840) — MTP와 경쟁하는 EAGLE의 2025년 초안 구조입니다.