Research map / Large language models / Instruction and fine-tuning
Multi-Token Prediction Needs Registers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes MuToR, which interleaves learnable register tokens that each predict a future token at a sampled offset, enabling multi-token prediction during fine-tuning, pretraining, and PEFT without architectural changes or inference cost.
Problem
Multi-token prediction improves pretraining but its gains have not consistently transferred to fine-tuning; existing methods add extra output heads/layers that must be trained from scratch and are hard to combine with pretrained LLMs.
Approach
Insert interleaved learnable register tokens (sharing one embedding) into the training sequence, with attention masks that hide registers from regular tokens and position ids set to the target offset, so each register predicts the token d steps ahead; registers are discarded at inference. A 2D offset variant extends this to image token grids.
Key results
- On GSM8K with Gemma 2B, MuToR reaches 42.10% vs 38.87% for next-token fine-tuning (3-seed average).
- On ImageNet 256x256 with LlamaGen-B at 360K steps, MuToR-2D reaches FID 5.09 vs 6.18 for next-token.
- With LoRA on Gemma 2B, LoRA-MuToR reaches 68.11% on 1M-GSM vs 66.11% for LoRA-next-token, matching full fine-tuning next-token (66.09%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace next-token fine-tuning with MuToR | Gemma 2B · Mathematical reasoning: GSM8K, 1M-GSM, 2M-GSM, MATH500, and AQUA-RAT | accuracy (%): 38.87 -> 42.10; 66.09 -> 68.33; 69.02 -> 70.56; 26.73 -> 28.13; 40.16 -> 41.73 | Helped |
| Replace next-token fine-tuning with MuToR | Llama 3 8B · Mathematical reasoning: GSM8K, 1M-GSM, 2M-GSM, and MATH500 | accuracy (%): 66.41 -> 67.85; 85.74 -> 87.03; 87.33 -> 87.64; 41.4 -> 43.2 | Helped |
| Replace next-token fine-tuning with MuToR | Gemma 2B · SAMSum and DialogSum | ROUGE-L: 43.23 -> 44.09; 38.77 -> 39.48 | Helped |
| Use 2D rather than 1D prediction offsets (ablation) | MuToR image generation · ImageNet 256×256, 100K, 200K, and 360K iterations | FID: 7.01 -> 6.57; 6.43 -> 5.65; 5.79 -> 5.09 | Helped |
| Increase maximum language lookahead offset to d_max=4 (ablation) | MuToR · GSM8K and 1M-GSM, Gemma 2B | accuracy (%): GSM8K: 41.93 (d_max=2), 41.60 (d_max=3), 42.10 (d_max=4), 39.90 (d_max=6); 1M-GSM: 67.15, 68.01, 68.33, 68.16 | Mixed |
| Use distinct register embeddings per offset instead of a shared embedding (ablation) | MuToR, Gemma 2B · GSM8K and 1M-GSM | accuracy (%): 42.10 -> 41.85; 68.33 -> 68.18 | Hurt |
| Increase the number of image registers from 80 to 256 (ablation) | MuToR-2D · ImageNet 256×256, 100K, 200K, and 360K iterations | FID: 6.87 -> 6.57; 5.89 -> 5.65; 5.03 -> 5.09 | Mixed |
| Set d_max_2D to 4 rather than 3 or 5 (ablation) | MuToR-2D with 80 registers · ImageNet 256×256, 100K, 200K, and 360K iterations | FID: 100K: 7.20, 6.87, 7.08; 200K: 6.19, 5.89, 6.21; 360K: 5.56, 5.03, 5.53 | Helped |
| Train Next-Token for 10 rather than 5 epochs (ablation) | Gemma 2B · GSM8K, 1M-GSM, and MATH500 | accuracy (%): 38.87 -> 37.91; 66.09 -> 65.07; 26.73 -> 27.07 | Mixed |
| Train Multi-Token for 10 rather than 5 epochs (ablation) | Gemma 2B · GSM8K, 1M-GSM, and MATH500 | accuracy (%): 40.66 -> 39.98; 66.69 -> 64.92; 26.87 -> 26.73 | Hurt |
| Train Next-Token for 5 rather than 3 epochs (ablation) | Gemma 2B · SAMSum and DialogSum | ROUGE-L: SAMSum: 43.23 -> 43.32; DialogSum: 38.77 -> 38.25 | Mixed |
| Train Multi-Token for 5 rather than 3 epochs (ablation) | Gemma 2B · SAMSum and DialogSum | ROUGE-L: 43.50 -> 42.48; 39.25 -> 38.97 | Hurt |
Surprising findings
- Multi-Token baseline performance degrades as d_max increases, barely beating next-token on 1M-GSM at larger horizons.
- Shared register embedding across offsets performs slightly better than per-offset embeddings.
Limitations
- Register placement is uniform or random and not optimized; authors note this leaves room for improvement.
- Optimal d_max and loss weight a are task-dependent and tuned per setting.
- Llama 3 8B results are from a single seed due to compute cost.
Experimental setup
Datasets: OpenMathInstruct-2 GSM8K AQUA-RAT SAMSum DialogSum ImageNet star-graph path-finding
Benchmarks: GSM8K MATH500 AQUA-RAT SAMSum DialogSum ImageNet 256x256 generation
Metrics: exact-match accuracy ROUGE-1 ROUGE-2 ROUGE-L FID IS Precision Recall solve rate
Baselines: Next-Token prediction Multi-Token prediction LoRA-Next-Token LlamaGen Next-Token
Builds on: multi-token prediction (Gloeckle et al.) register tokens (Darcet et al.) RoPE LoRA LlamaGen
Closest papers by meaning
- Better & Faster Large Language Models via Multi-token Prediction (ICML 2024 · similarity 0.89)
- Beyond Next Token Prediction: Patch-Level Training for Large Language Models (ICLR 2025 · similarity 0.86)
- Leviathan: Decoupling Input and Output Representations in Language Models (NeurIPS 2026 · similarity 0.86)
- On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.86)
- Not All Tokens Are What You Need for Pretraining (NeurIPS 2024 · similarity 0.86)
- Efficient Training-Free Multi-Token Prediction via Embedding-Space Probing (ICML 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.