LOOPMTP: A Looped Transformer Guided by Latent Multi-Token Prediction

Summary (Overview)

• Core contribution: LOOPMTP is a novel looped transformer architecture that addresses two key weaknesses of existing looped models—latent overthinking (where correct early predictions get corrupted by later iterations) and undifferentiated computation (where successive iterations produce increasingly similar representations)
• Key innovation: The method applies soft MTP alignment in latent space, steering each iteration's hidden state toward the embedding of the token it should anticipate (t steps ahead), using cosine similarity rather than expensive vocabulary-sized projections
• Architectural elements: Includes an MTP-guided looped block with iteration-index embeddings, per-iteration normalization, input fusion via concatenation-projection, and a novel Loop-LNS (layer-norm scaling) with fixed 1/T scaling factor
• Results: Achieves up to 8.1% relative improvement over a parameter-matched non-looped baseline, outperforms the state-of-the-art LoopFormer in 27 of 28 matched comparisons, and demonstrates stable training up to 15 loops
• Practical application: Shows effectiveness in training small domain-specific math experts, achieving 19.03% GSM8K accuracy vs. 7.05% for the non-looped baseline at equal parameter count (~260M)

Introduction and Theoretical Foundation

Background: The Need for Parameter-Efficient Reasoning

Large language models (LLMs) demonstrate strong reasoning capabilities, but these are typically tied to scale—specifically depth. Standard transformers stack more layers to improve reasoning, but this increases parameter count. Two orthogonal approaches exist:

1. Multi-token prediction (MTP): Modifies the pretraining objective to predict multiple future tokens simultaneously, encouraging the model to "think ahead"
2. Looped transformers: Reuse the same stack of transformer blocks multiple times, achieving effective depth without adding parameters

The Problem with Existing Looped Models

Current looped models suffer from two critical issues:

Latent overthinking: Each iteration overwrites the previous iteration's hidden representation, discarding potentially useful information. Predictions that are already correct after early iterations can be corrupted by later iterations.

Undifferentiated computation: As the number of iterations grows, hidden representations become increasingly similar across iterations. Successive iterations perform redundant work, wasting computation.

LOOPMTP's Solution

LOOPMTP bridges MTP and looping through a structural correspondence: a model that loops T times can anticipate T future tokens. This provides dense, forward-looking supervision at every iteration, addressing both weaknesses simultaneously.

Methodology

3.1 Notation and Problem Setup

Given an input token sequence $\mathbf{u} = (u_1, \ldots, u_S) \in \mathcal{V}^S$:

Token embedding:

$$\mathbf{x}^{(0)} = \operatorname{Embed}(\mathbf{u}) \in \mathbb{R}^{S \times d} \tag{1}$$

Recursive application of the shared backbone $f_\theta$ for $t = 1, \ldots, T$:

$$\mathbf{x}^{(t)} = f_\theta\left(\mathbf{x}^{(t-1)}, \mathbf{x}^{(0)}, t\right) \tag{2}$$

3.2 MTP-Guided Looped Block

The block processes inputs through three operations:

Iteration-index embedding and per-iteration normalization:

$$\mathbf{h}^{(t-1)} = \mathrm{LN}_t^{\text{prev}}\left(\left[\frac{t-1}{T} \| \mathbf{x}^{(t-1)}\right]\right) \in \mathbb{R}^{S \times (d+1)} \tag{3}$$ $$\mathbf{e} = \mathrm{LN}^{\mathrm{tok}}\left(\mathbf{x}^{(0)}\right) \in \mathbb{R}^{S \times d} \tag{4}$$

Input fusion via concatenation and linear projection:

$$\mathbf{y}^{(t)} = \mathrm{P}\left(\left[\mathbf{e} \| \mathbf{h}^{(t-1)}\right]\right) \in \mathbb{R}^{S \times d} \tag{5}$$

Backbone application:

$$\mathbf{x}^{(t)} = f_\theta\left(\mathbf{y}^{(t)}\right) = \operatorname{Layer}_L \circ \cdots \circ \operatorname{Layer}_1\left(\mathbf{y}^{(t)}\right) \tag{6}$$

Loop-LNS (Looped Layer-Norm Scaling): A key innovation—uses a fixed scaling factor of $1/T$ rather than running-depth alternatives:

$$\mathbf{x} = \mathbf{x} + \operatorname{Attn}\left(\frac{1}{T}\operatorname{RMSNorm}(\mathbf{x})\right) \tag{7}$$ $$\mathbf{x} = \mathbf{x} + \operatorname{FFN}\left(\frac{1}{T}\operatorname{RMSNorm}(\mathbf{x})\right) \tag{8}$$

The fixed factor is chosen because progressive counters (scaling by $\frac{1}{(t-1)\cdot L + \ell}$) produce tiny factors for later iterations, neutralizing their gradient contribution.

Per-Iteration MTP Alignment: For iteration $t \geq 2$, position $i$ is aligned to the output embedding of token $u_{i+t}$ (the token $t$ steps ahead). The first iteration ($t=1$) is deliberately unconstrained to serve as a rich representation for subsequent iterations.

3.3 Aggregation via Gating

Instead of discarding intermediate iterations, LOOPMTP aggregates them with a content-conditional gate:

Gate computation (shared linear layer $W_g \in \mathbb{R}^{d \times d}$):

$$\mathbf{g}_i^{(t)} = \text{softplus}\left(W_g\mathbf{x}_i^{(t)} + \beta_t \cdot \mathbf{1}_d\right) \in \mathbb{R}^d \tag{9}$$

Normalization across iterations:

$$\tilde{\mathbf{g}}_i^{(t)} = \frac{\mathbf{g}_i^{(t)}}{\sum_{s=1}^{T} \mathbf{g}_i^{(s)} + \varepsilon} \tag{10}$$

Weighted aggregation:

$$\mathbf{z}_i = \sum_{t=1}^{T} \tilde{\mathbf{g}}_i^{(t)} \odot \mathbf{x}_i^{(t)} \tag{11}$$

The gating adds < 0.5% parameters. Initialization uses a strong prior toward the first iteration (vanilla transformer behavior), with later iterations free to take over during training.

3.4 Training Objective

The total loss combines three terms:

Main NTP loss (applied to aggregated representation):

$$\mathcal{L}_{\mathrm{NTP}} = -\frac{1}{S-1} \sum_{i=1}^{S-1} \log \mathbf{p}_i^{\mathrm{agg}}[u_{i+1}] \tag{12}$$

Hidden-state soft MTP alignment (using stop-gradient on targets):

$$\mathcal{L}_{\mathrm{align}}^{(t)} = \frac{1}{S-t} \sum_{i=1}^{S-t} \left(1 - \cos\left(\mathbf{x}_i^{(t)}, \operatorname{sg}[E_{u_{i+t}}]\right)\right) \tag{13}$$ $$\mathcal{L}_{\text{align}} = \frac{1}{T-1} \sum_{t=2}^{T} \mathcal{L}_{\text{align}}^{(t)} \tag{14}$$

Ponder regularizer (KL divergence toward uniform prior over iterations):

$$\mathcal{L}_{\text{ponder}} = \frac{1}{S} \sum_{i=1}^{S} \mathrm{KL}(G_i \| Q) \tag{15}$$

Final loss:

$$\mathcal{L} = \mathcal{L}_{\mathrm{NTP}} + \lambda_{\text{align}} \mathcal{L}_{\text{align}} + \lambda_{\text{ponder}} \mathcal{L}_{\text{ponder}} \tag{16}$$

Empirical Validation / Results

Experimental Setup

• Model: GPT-2-style decoder-only transformer with $L=12$ layers, $d=1024$ embedding dimension, 32 attention heads, FFN hidden dimension 4096
• Baselines: Non-looped (with more parameters, ~266M vs. 260M) and LoopFormer (state-of-the-art looped model, ~280M)
• Training: 6.8B tokens from Nemotron-CC-v2 and Nemotron-CC-Math-v1, Muon/AdamW optimizers, peak LR $1.9 \times 10^{-3}$
• Evaluation: Perplexity on FineWeb-Edu and OpenWebText; accuracy on 7 commonsense benchmarks; BPB on QA, math, and code suites

Key Results (Table 1)

Key findings:

• LOOPMTP outperforms non-looped baseline by up to 8.08% on general tasks and 7.5% on BPB metrics
• Outperforms LoopFormer in 27 of 28 matched comparisons (4 loop counts × 7 benchmarks)
• Perplexity improves monotonically with loop count; QA/math/code BPB improves through $T=7$ with slight regression at $T=9$
• Training remains stable unlike LoopFormer (which required per-loop grid search and still diverged in some runs)

MTP Auxiliary Objective Impact

Performance effects: Runs with $\lambda_{\text{align}} > 0$ show clear improvements, particularly on math tasks. Without MTP, increasing loop count causes math performance to degrade monotonically.

Representation effects (Figure 2, Bottom): With MTP, the second iteration already maps to a distinctly different subspace; without MTP, representations change gradually with high cross-iteration similarities, confirming undifferentiated computation.

Ground-truth retrieval (Figure 3): The w/ MTP model achieves ground-truth rank up to 35.6× better than the w/o MTP model, confirming alignment is actually achieved and transfers to the LM head's coordinate system.

Gating Mechanism Impact

• All (learnable) consistently outperforms Only last, Only last (learnable), and All (uniform) variants
• Gate values increase with iteration depth (from ≈0.17 for early iterations to ≈0.26 for late ones)
• Later iterations show wider distributions, indicating more input-dependent (selective) blending

Domain-Specific Expert Results

Training on math data with up to $T=15$ loops:

• GSM8K accuracy: 19.03% vs. 7.05% for non-looped baseline (≈170% relative improvement)
• The slope of −169 in the BPB-accuracy relationship indicates even small BPB reductions yield meaningful accuracy gains

Theoretical and Practical Implications

Theoretical Contributions

1. Addresses latent overthinking: The MTP alignment signal provides forward-looking supervision that prevents later iterations from corrupting earlier correct predictions
2. Addresses undifferentiated computation: MTP encourages each iteration to produce a distinct representation, increasing expressiveness
3. Soft alignment is effective: Cosine similarity in latent space avoids vocabulary-sized projections while still providing meaningful guidance

Practical Implications

1. Parameter efficiency: LOOPMTP achieves better performance with fewer parameters (260M vs. 266M non-looped, 280M LoopFormer)
2. Training stability: Unlike LoopFormer, LOOPMTP trains stably without per-loop hyperparameter tuning
3. Domain expertise: Can train small, specialized models suitable for memory-constrained deployment (e.g., healthcare on local hardware)
4. Scaling potential: The relationship between loop count and performance suggests further gains may be possible with more loops

Conclusion and Future Directions

LOOPMTP demonstrates that combining looped architectures with soft MTP supervision creates a powerful, parameter-efficient approach to reasoning. The method achieves up to 8.1% relative gains over parameter-matched baselines, wins 27 of 28 matched comparisons against the state-of-the-art, and enables small domain-specific experts with substantial accuracy improvements.

Future work:

• Investigating scaling behavior and establishing scaling laws for LOOPMTP
• Verifying transfer of domain-expert benefits to other specialized domains beyond math
• Characterizing the effective upper bound on useful depth for looped architectures

Limitations:

• Domain-expert results cover only math (single seed due to computational constraints)
• Gains diminish and occasionally become non-monotonic as loop count grows
• Full characterization of the effective upper bound on useful looping depth remains open