AdaMTP: An Adaptive Training Paradigm for Multi-Token Prediction
AdaMTP uses entropy-based boundary detection and dynamic masking to improve multi-token prediction accuracy and speed.
Key Findings
Methodology
AdaMTP leverages token-level predictive entropy estimated by a pretrained base model to identify semantic boundaries marked by sudden entropy surges. The sequence is segmented into variable-length chunks, with each token assigned an adaptive prediction depth based on its position within these chunks. During training, a dynamic masking mechanism suppresses loss contributions for predictions crossing boundaries, reducing noisy gradients. The training involves a two-stage process: initial head warm-up with frozen backbone, followed by joint fine-tuning with LoRA, where the loss is masked according to the adaptive depths. During inference, the model supports fixed and adaptive decoding modes, enabling faster and more reliable multi-token generation.
Key Results
- On Llama-3.1-8B, Qwen-2.5-7B, and Gemma-3-12B, AdaMTP outperforms NTP and standard MTP across tasks like math reasoning and code generation, with average improvements of around 2x and speedups up to 2.75x.
- In benchmarks such as GSM8K, HumanEval, and MMLU, AdaMTP achieves higher accuracy and significantly faster inference, especially in large-batch scenarios.
- By filtering out high-entropy boundary predictions, AdaMTP reduces gradient noise, leading to more stable representations and improved core model capabilities.
Significance
This work addresses the core limitation of fixed prediction horizons in multi-token models, introducing an adaptive strategy that aligns prediction scope with sequence predictability. It effectively mitigates gradient interference at semantic boundaries, enhancing both performance and efficiency. The approach is architecture-agnostic, applicable across diverse models and tasks, and paves the way for more robust, scalable language models capable of long-range reasoning and generation, with broad implications for NLP, code synthesis, and AI-assisted reasoning.
Technical Contribution
The paper introduces an entropy-based boundary detection algorithm that segments sequences into variable-length chunks, assigning each token an adaptive prediction depth. It incorporates a dynamic masking mechanism during training to prevent cross-boundary prediction loss, combined with a two-stage training pipeline involving head warm-up and LoRA fine-tuning. The method is validated across three large models, demonstrating consistent improvements in accuracy and inference speed, establishing a new paradigm for adaptive multi-token training.
Novelty
This is the first work to utilize token-level entropy surges as semantic boundary indicators for adaptive prediction horizon adjustment. Unlike fixed-range multi-token models, AdaMTP dynamically modulates prediction depth based on sequence predictability, effectively reducing noisy gradients and improving robustness. Its innovative boundary detection and masking strategies mark a significant departure from prior fixed-horizon approaches, offering a new perspective on sequence modeling.
Limitations
- The boundary detection relies on the accuracy of entropy estimation, which may be affected by model quality and data complexity. In noisy or ambiguous sequences, boundary identification might be less reliable.
- The added complexity of dynamic masking and boundary segmentation increases training and implementation difficulty, requiring careful hyperparameter tuning.
- The method's effectiveness in multilingual or multimodal settings remains to be explored, as current experiments focus on English text and code.
Future Work
Future research could explore integrating reinforcement learning to optimize dynamic prediction depths further, extend boundary detection to multimodal data, and develop more robust entropy estimation techniques. Additionally, applying this framework to real-time systems and multi-task learning scenarios could unlock broader applications and performance gains.
AI Executive Summary
Large language models (LLMs) have revolutionized NLP, achieving remarkable success in tasks like reasoning, code generation, and conversational AI. However, their autoregressive nature limits inference speed, especially for long sequences. Multi-token prediction (MTP) emerged as a promising solution, enabling models to generate multiple tokens simultaneously, thus accelerating inference. Yet, fixed prediction horizons across sequences fail to account for the inherent variability in natural language and code, where some parts are highly predictable while others involve semantic shifts. Forcing auxiliary heads to predict across these unpredictable boundaries introduces noisy, conflicting gradients, impairing the model's core capabilities.
To address this, AdaMTP introduces an adaptive training paradigm that dynamically aligns the prediction horizon with the sequence's intrinsic predictability. It employs an entropy-based segmentation algorithm, using a pretrained base model to estimate token-level predictive entropy. Sudden surges in entropy mark semantic boundaries, prompting the model to segment sequences into variable-length chunks. Each token is assigned an adaptive prediction depth based on its position within these chunks. During training, a dynamic masking mechanism suppresses loss contributions from predictions crossing boundaries, effectively reducing gradient noise and representation interference.
Extensive experiments across three large models—Llama-3.1-8B, Qwen-2.5-7B, and Gemma-3-12B—demonstrate that AdaMTP consistently outperforms standard NTP and fixed-horizon MTP in both accuracy and inference speed. On benchmarks like GSM8K, HumanEval, and MMLU, the method achieves an average performance boost of around 2x, with speedups up to 2.75x. The approach effectively mitigates the gradient interference caused by high-entropy boundaries, leading to more robust and efficient models.
This research marks a significant step toward adaptive, context-aware sequence modeling, enabling large models to generate longer, more accurate outputs faster. Its architecture-agnostic design and empirical validation suggest broad applicability in NLP, code synthesis, and beyond. Future work may focus on refining boundary detection, extending to multimodal data, and integrating reinforcement learning for further optimization, promising a new paradigm in scalable AI systems.
Deep Dive
Abstract
Multi-Token Prediction (MTP) has emerged as an effective paradigm that augments a shared Large Language Model backbone with auxiliary heads, training the model to predict several future tokens in parallel to enrich its supervision signal and accelerate inference. However, existing training frameworks adopt a rigid, fixed-length prediction horizon, disregarding the highly non-uniform information density of natural language and code. Forcing the auxiliary heads to predict across high-entropy semantic boundaries injects noisy, conflicting training signals; because these heads share the backbone's latent representations, the resulting gradients backpropagate and interfere with the model's core capabilities. We propose AdaMTP, an adaptive training paradigm that dynamically aligns the prediction horizon with the intrinsic predictability of the sequence. At its core, an entropy-based segmentation algorithm leverages the base model to detect sudden surges in uncertainty as semantic boundaries, partitioning sequences into variable-length groups. Each token is assigned an adaptive prediction depth, and a dynamically masked MTP objective suppresses the loss for predictions that cross these boundaries, attenuating the noisy gradients that degrade the backbone. Across mathematical reasoning, code generation, and general benchmarks on three backbones (Llama-3.1-8B, Qwen-2.5-7B, Gemma-3-12B), AdaMTP consistently outperforms standard MTP in both task performance and inference speedup.