LoopMTP: A looped transformer guided by latent multi-token prediction

2026-08-04

AI Summary

This work addresses the tendency of existing recurrent Transformers to suffer from overthinking and computational redundancy due to the absence of cross-iteration guidance signals. To mitigate this, the authors propose integrating recurrence with multi-token prediction (MTP), aligning the number of recurrence steps with future token prediction horizons in latent space—such that the t-th recurrence step directly predicts the token t steps ahead—thereby providing dense lookahead supervision. A lightweight gating mechanism is further introduced to preserve useful information across iterations. This approach achieves, for the first time, an explicit alignment between recurrence depth and prediction horizon in latent space, significantly alleviating overthinking and enhancing inference efficiency. Experiments demonstrate up to an 8.1% relative improvement in average accuracy over non-recurrent baselines, with stable training observed even with up to 15 recurrence steps.

Abstract

Looped transformers have emerged as a parameter-efficient alternative to scaling depth for strong reasoning. By reusing one stack of layers across $T$ iterations, they attain the effective depth and reasoning capabilities of larger models at a fixed parameter count. Yet existing approaches suffer from latent overthinking and undifferentiated computation, largely because intermediate representations receive no guidance across loops. Multi-token prediction (MTP) supplies exactly the dense, forward-looking supervision the loop is missing. We propose \textsc{LoopMTP}, which links the two through a structural correspondence in latent space: a model that loops $T$ times can anticipate $T$ future tokens. \textsc{LoopMTP} realizes this by softly aligning the hidden state of loop $t$ with the embedding of the token $t$ steps ahead, while a lightweight gate preserves useful information across iterations. \textsc{LoopMTP} improves average accuracy by up to 8.1\% (relative) over the non-looped baseline, with training remaining stable for up to 15 loops.

Problem

• looped transformers
• latent overthinking
• undifferentiated computation
• multi-token prediction
• intermediate representations

Innovation

• Looped Transformer
• Multi-token Prediction
• Latent Alignment
• Parameter Efficiency
• Iterative Reasoning