AdaMTP: An Adaptive Training Paradigm for Multi-Token Prediction

2026-08-01

AI Summary

This work addresses a key limitation of existing Multi-Token Prediction (MTP) methods, which employ a fixed prediction length and thereby overlook the non-uniform information density in natural language and code, often introducing noisy gradients across semantic boundaries and impairing model training. To mitigate this, the authors propose an Adaptive MTP paradigm that leverages an entropy-based segmentation algorithm to identify semantic boundaries, dynamically assigns prediction depth per token, and masks cross-boundary losses to suppress detrimental gradient updates. Integrated with a large language model backbone (e.g., Llama-3.1-8B), a dynamic masking objective, and a shared auxiliary head architecture, the proposed approach consistently outperforms standard MTP across mathematical reasoning, code generation, and general-purpose benchmarks, achieving simultaneous gains in both performance and inference speed.

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.

Problem

• Multi-Token Prediction
• prediction horizon
• semantic boundaries
• information density
• gradient interference

Innovation

• Adaptive Training
• Multi-Token Prediction
• Entropy-based Segmentation
• Dynamic Masking
• Semantic Boundaries