Multi-Token Prediction Needs Registers

TL;DR

MuToR introduces register tokens that are interleaved with regular tokens during training to predict future targets, enabling scalable multi-token prediction without architectural changes at inference. By using a shared register embedding and a carefully designed attention mask, registers contribute a richer supervisory signal while leaving inference identical to standard autoregressive decoding. The approach achieves consistent gains in language modeling tasks (including mathematical reasoning and summarization) and in autoregressive image generation, and it integrates well with parameter-efficient fine-tuning approaches like LoRA. This method broadens the lookahead horizons available to pretraining and fine-tuning, offering a practically efficient path to improved generative modeling across modalities.

Abstract

Multi-token prediction has emerged as a promising objective for improving language model pretraining, but its benefits have not consistently generalized to other settings such as fine-tuning. In this paper, we propose MuToR, a simple and effective approach to multi-token prediction that interleaves learnable register tokens into the input sequence, each tasked with predicting future targets. Compared to existing methods, MuToR offers several key advantages: it introduces only a negligible number of additional parameters, requires no architectural changes--ensuring compatibility with off-the-shelf pretrained language models--and remains aligned with the next-token pretraining objective, making it especially well-suited for supervised fine-tuning. Moreover, it naturally supports scalable prediction horizons. We demonstrate the effectiveness and versatility of MuToR across a range of use cases, including supervised fine-tuning, parameter-efficient fine-tuning (PEFT), and pretraining, on challenging generative tasks in both language and vision domains. Our code will be available at: https://github.com/nasosger/MuToR.

Figures (4)

• Figure 1: Next-token prediction vs. Multi-token prediction with registers (MuToR). The transformer block represents any decoder-only autoregressive model, with colored lines indicating permitted attention connections between tokens. Left: Standard next-token prediction, where each $x_t$ predicts $x_{t+1}$ conditioned on preceding tokens. Right: MuToR interleaves register tokens $r_d$ to predict tokens $d$ steps ahead ($x_{t+d}$), conditioned only on previous regular tokens. Register tokens are assigned position ids (e.g., $t+d-1$ for $r_d$ targeting $x_{t+d}$) that mimic next-token prediction. Regular tokens follow the standard next-token prediction formulation, unaffected by the registers.
• Figure 2: MuToR's attention mask. Each cell indicates whether the row can attend to the column.
• Figure 3: The 2D neighborhood of possible prediction targets (depicted in red) for a register token. The register $r_d$ is inserted after $x_7$, and $d_{\mathrm{max\_2D}}$ is set to 3.
• Figure 4: Solve rate (%) of finetuned GPT2-L model on different star graph configurations.