Source-linked AI summary
Direct Multi-Token Decoding
Xuan Luo, Weizhi Wang, Xifeng Yan
TL;DR
Decoder-only models may underuse computation because their layers can specialize, raising the question of whether hidden states can support multiple future tokens. DMTD reuses late layers in fixed multi-token cycles, and at cycle length four reports up to a 2× speedup while retaining 96.3% of vanilla performance.
Problem
The authors ask whether underused computation and layer specialization in decoder-only models can support generating multiple future tokens without repeatedly traversing early and middle layers.
Source: §1 · pp. 1–2Method
DMTD uses fixed multi-token cycles, performing one full forward pass and then reusing late layers to decode consecutive tokens without added parameters or post-generation verification.
Source: §1 · p. 2Results
At cycle length 4, DMTD achieves up to a 2× inference speedup while retaining 96.3% of vanilla performance.
Source: §1 · p. 2Takeaways & Limitations
DMTD offers a simple direction for accelerating LLM inference, and its performance may improve with larger training datasets.
Source: §6 · p. 10Only 1.5B tokens were used versus Qwen3’s approximately 36T-token pretraining dataset, preventing full continual training and a direct experimental comparison with speculative decoding.
Source: §5 · p. 10
Abstract
from arXiv · showhide
Decoder-only transformers have become the standard architecture for large language models (LLMs) due to their strong performance. Recent studies suggest that, in pre-trained LLMs, early, middle, and late layers may serve distinct roles: Early layers focus on understanding the input context, middle layers handle task-specific processing, and late layers convert abstract representations into output tokens. We hypothesize that once representations have been processed by the early and middle layers, the resulting hidden states may encapsulate sufficient information to support the generation of multiple tokens using only the late layers, eliminating the need to repeatedly traverse the early and middle layers. We refer to this inference paradigm as Direct Multi-Token Decoding (DMTD). Unlike speculative decoding, our method introduces no additional parameters, auxiliary routines, or post-generation verification. Despite being trained on a limited dataset, a fine-tuned DMTD Qwen3-4B model has already demonstrated promising results, achieving up to a 2x speedup with only minor performance loss. Moreover, as shown in our scaling analysis, its performance is expected to further improve with larger training datasets.
1 Introduction
Decoder-only layers may specialize into context encoding, task-specific processing, and token decoding, motivating DMTD’s reuse of late layers for multi-token generation. The paper reports up to 2× inference speedup at cycle length four after end-to-end fine-tuning, with minor performance degradation.
1 Introduction
- Decoder-only layers may specialize in encoding context, handling reasoning and task-specific processing, and generating token-level predictions.The paper labels these stages encoding, thinking, and decoding layers, while noting that their boundaries are unclear.Source: §1 · p. 1
- Dynamic layer skipping on simple tasks such as string copy suggests spare computation that DMTD seeks to repurpose for future-token information.The motivating question is whether later tokens can use that information by attending only through decoding layers.Source: §1 · pp. 1–2
- DMTD performs one full forward pass per cycle, then reuses late layers to decode multiple tokens consecutively.Unlike speculative decoding, it adds no parameters, LM heads, or post-processing routines.Source: §1 · p. 2
- DMTD is trained end-to-end from a pretrained LLM and fine-tuned on approximately 1.5B tokens; performance improves continuously as training data increases.The authors identify large-scale continued pretraining followed by post-training as a way to further exploit the architecture.Source: §1 · p. 2
- 100% of vanilla performance at two-token cycles, 98.4% at three-token cycles, and 96.3% at four-token cycles; cycle length four yields up to 2× speedup.The evaluation reuses the last 8 layers of the 36-layer Qwen3-4B model.Source: §1 · p. 2
2 Method
DMTD combines cyclical masking for multi-future-token training with cycle-based inference that reuses late layers after an initial full pass and refills missing KV-cache entries across cycles.
2.1 Parallel Training with Cyclical Masking
- An alternative training setup uses encoding representations instead of input embeddings to form masked states, reusing encoding and decoding layers during multi-token decoding.Source: §2.1 · p. 3
- Cyclical masking trains multiple future-token predictions in one sequence by reusing shared prefix representations instead of creating divergent sequences.For cycle length 3, the mask follows a repeating 1, 0, 0 pattern; the training pipeline uses one forward pass and adds no parameters.Source: §2.1 · pp. 3–4, Figure 2 · p. 3
2.2 Multi-Token Decoding with cyclical refilling
- Inference cycles begin with a full-layer pass, then use only late layers for subsequent tokens and refill missing KV-cache entries at the next cycle.The refilling mechanism restores entries for earlier tokens so the full context remains available during generation.Source: §2.2 · p. 4
- DMTD processes fewer layers per token despite vanilla-transformer-equivalent computational load, and larger cycles or a smaller decoding-layer proportion reduce PLT.PLT averages layers processed per generated token; inference is memory-bound, so processing fewer layers per token improves efficiency.Source: §2.2 · pp. 4–5
3 Experiments
Experiments evaluate DMTD’s benchmark retention, layer allocation, inference speed, and scaling across training data, model sizes, and inference cycle lengths. Results show speedups alongside quality trade-offs that vary with decoding configuration and model size.
3.1 Implementation Details
- DMTD uses supervised fine-tuning initialized from a pretrained LLM, with evaluation spanning five benchmarks and varied prompting settings.The implementation uses the last eight layers of Qwen3-4B as decoding layers; evaluations include ARC-Easy, ARC-Challenge, WinoGrande, GSM8K, and CoQA.Source: §3.1 · p. 5
3.2 Decoding Cycle Length
- At cycle length 4, DMTD retains 96.3% of vanilla performance, while cycle length 6 retains 82.1%.The evaluation uses the same cycle length for training and inference; the authors hypothesize limited hidden-state dimensionality constrains longer prediction horizons.Source: §3.2 · p. 5, Table 1 · p. 6
3.3 Impact of Encoding and Decoding Layer Allocation
- E16D0 reaches 75.2% overall performance, while E8D8 reaches 99.8%, comparable to E0D16 at 100.1%.The results indicate that reusing decoding layers is important; scaling encoding layers alone performs worse than configurations emphasizing decoding layers.Source: §3.3 · p. 6, Table 2 · p. 6
3.4 Inference Speedup
- Compared with the vanilla model, MTD4 achieves 2.15× speedup at batch size 1, declining to 1.77× at batch size 8; gains generally increase with cycle length, particularly at lower batch sizes.Source: §3.4 · pp. 6–7
3.5 Scaling with Training Data
- Cross-entropy loss decreases as training data increases across all three model sizes, with log-linear fits having slopes of -0.179, -0.191, and -0.178 for 0.6B, 1.7B, and 4B.The corresponding R² values are 0.966, 0.972, and 0.994; the experiments use supervised fine-tuning on limited data.Source: §3.5 · p. 7
3.6 Impact of Model Scale on Direct Multi-Token Decoding
- Under E0D8MTD3, Qwen3-4B retains 98.4% of vanilla performance, compared with 72.6% for Qwen3-0.6B.The authors hypothesize that larger models encode richer anticipatory information; reusing eight layers is 22.2% of Qwen3-4B’s layers versus approximately 28.6% for the 28-layer models.Source: §3.6 · pp. 7–8, Table 4 · p. 8, Figure 5 · p. 8
3.7 Inference Cycle Length
- A model trained at cycle length 3 retains 98.4% performance at length 3 and 92.2% at length 4, with sharper declines at lengths 5 and 6.The reported drops are particularly evident on reasoning-intensive tasks such as GSM8K.Source: §3.7 · pp. 8–9, Table 5 · p. 9
4 Related Works
The related work traces decoder-only models’ layered specialization and prior multi-token decoding approaches, motivating DMTD’s reuse of late layers without auxiliary models or heads.
4.1 Large Language Models
- Encoder-decoder models established separate context encoding and token generation, while T5 [26] and FLAN [39] extended them to text-to-text and instruction-tuned tasks.Decoder-only models later became mainstream for scalable, prompt-driven generation.Source: §4.1 · p. 9
- Prior studies describe early layers as encoding context, middle layers as handling reasoning and task processing, and late layers as specializing in token prediction.MoD and FlexiDepth [22] explore layer skipping; their irregular patterns can hinder acceleration in memory-bound settings, motivating DMTD’s cyclic reuse of late layers.Source: §4.1 · p. 9
4.2 Multi-token Prediction
- Speculative decoding [13], Medusa [5], and EAGLE [17] [16] [18] generate multiple candidate tokens using drafts, decoding heads, or feature-level sampling, with verification or related mechanisms.These inference methods address autoregressive decoding’s memory bottlenecks.Source: §4.2 · p. 10
- Multi-token prediction also appears as an auxiliary or core training objective, while DMTD reuses existing late layers without auxiliary models or heads.The cited training approaches predict multiple future tokens during pre-training or use independent heads; DMTD instead uses cyclical decoding.Source: §4.2 · p. 10
5 Limitations
The evaluation is constrained by unavailable Qwen3 pre-training data: the authors used 1.5B tokens rather than the approximately 36T-token training dataset and did not directly compare DMTD with speculative decoding.
5 Limitations
- Only 1.5B tokens were used, versus Qwen3’s approximately 36T-token pre-training dataset, preventing full continual training of DMTD.The dataset is not publicly available, so a fully developed DMTD could not be assessed.Source: §5 · p. 10
- The authors do not provide a direct experimental comparison between DMTD and speculative decoding, leaving that evaluation to future work.Source: §5 · p. 10
- Performance is expected to improve with additional training data, and speculative decoding can still be applied to tokens generated by DMTD.The authors identify MoE settings and large batch sizes as contexts for further investigation.Source: §5 · p. 10
6 Conclusion
The conclusion presents DMTD as a fixed-cycle approach that repurposes underused computation for sustained multi-token generation, with feasibility shown experimentally and improvement expected from larger training datasets.
6 Conclusion
- DMTD repurposes underutilized computation in pre-trained LLMs into fixed cycles of sustained multi-token generation without additional parameters or post-generation verification.Source: §6 · p. 10
- The experiments demonstrate the feasibility of DMTD, while the conclusion says performance may improve with larger training datasets.Source: §6 · p. 10
- The authors frame DMTD as a direction for accelerating LLM inference.Source: §6 · p. 10