Looped transformers adapt iteration depth per token for efficient reasoning
T-LoopFormer: Token-Level Elastic-Depth Looped Transformers for Latent Reasoning With Dynamic Routing
Artificial Intelligence
Summary
Transformers are computer programs that help understand language and solve reasoning problems. Usually, these programs do the same amount of work for every word or part of a sentence, which can waste computing time. The authors introduce a new method where each word decides how much processing it needs, so simpler words get less attention and harder ones get more. They also made a smart way to remember previous work to speed up the process. Their approach improves performance, speeds up processing, and uses less computing power.
What this means in practice
- •For natural language processing engineers: Design transformer models that allocate compute per token dynamically to improve efficiency in language understanding and reasoning tasks.
- •For machine learning infrastructure teams: Implement recursion-wise caching mechanisms to reduce redundant computations and lower latency during autoregressive model inference.
- •For chatbot developers: Create chatbots that provide faster, higher-quality responses by using token-level adaptive depth in transformer models.$Commercial implications: This paper enables chatbot products with improved response speed and accuracy by adaptive compute allocation per token.
Authors
Mingqian Yu, Wenpeng Zhang, Peilin Zhao
Abstract
Looped Transformers have recently demonstrated strong performance in both reasoning and language tasks by reusing a shared set of parameters across multiple iterations, achieving parameter efficiency without sacrificing representational power. Besides, looped Transformers perform inference directly in the latent space (latent reasoning) to reduce the number of tokens consumed during inference, thereby achieving improved sample efficiency. However, these models typically apply a fixed recursion depth uniformly to every token, leading to suboptimal compute allocation and leaving significant efficiency gains on the table. In this work, we propose \textbf{dynamic token-choice routing} for looped transformers, enabling each token to adaptively determine its own number of loop iterations based on its hidden state. We use a dynamic router to decide whether a token should continue recursing or exit early, allowing simple tokens to bypass unnecessary computation while hard tokens receive deeper processing. To ensure that this adaptive mechanism does not compromise decoding efficiency, we further introduce recursion-wise KV caching, which maintains an independent key-value cache for each recursion loop. This design ensures that tokens at different depths only attend to their corresponding cached states, effectively eliminating redundant computations for exited tokens and enabling fast autoregressive decoding. Extensive experiments show that T-LoopFormer reaches the sota performance under the same parameters on PPL and 10 zero-shot reasoning tasks, even surpassing the base model at 24x FLOPs and our model could reach the lowest inference latency, which validate the effectiveness of token-choice router and recursion-wise KV cache. Code: https://github.com/YuMingQian1234/T-LoopFormer.