Looped State-Space Language Models with Adaptive Exit-State Selection
Summary
This paper explores looped (recurrent) state-space language models using Mamba and hybrid Mamba-Transformer backbones, showing they outperform non-looped baselines on reasoning tasks and remain competitive under iso-parameter and iso-FLOPs pretraining, with adaptive exit-state selection improving intermediate-depth performance.
View Cached Full Text
Cached at: 07/14/26, 04:18 AM
# Looped State-Space Language Models with Adaptive Exit-State Selection
Source: [https://arxiv.org/html/2607.10110](https://arxiv.org/html/2607.10110)
Zhenxuan YuRikkyo UniversityYutaka MatsuoThe University of TokyoTakeshi Kojima22footnotemark:2The University of TokyoYusuke Iwasawa22footnotemark:2The University of Tokyo
###### Abstract
Recent work on looped language models suggests that many reasoning problems benefit from greater computational depth rather than from additional independent parameters\. Existing studies, however, focus almost exclusively on Transformer backbones, leaving open whether this principle also applies to state\-space language models\. We investigate Looped Mamba and Looped Hybrid Mamba–Transformer architectures, which repeatedly apply a shared Mamba \(or hybrid\) block to introduce explicit finite\-depth recurrent computation\. On two controlled reasoning tasks—Mano \(modular\-arithmetic manipulation\) and p\-hop induction—Looped Mamba consistently outperforms parameter\-matched non\-looped baselines and, in several settings, matches or exceeds non\-looped models of equal effective depth\. We then extend the study to language model pre\-training under matched iso\-parameter and iso\-FLOPs protocols, which jointly disentangle the effects of parameter sharing and effective depth: looped models remain competitive on downstream benchmarks with substantially fewer distinct parameters, although deeper non\-looped models retain an advantage in validation perplexity under strict iso\-FLOPs comparisons\. Finally, we adapt Ouro’s two\-stage exit gate to Looped Mamba for threshold\-controlled selection among recurrent\-step outputs\. Since all recurrent steps are still executed, the selected exit step represents prediction depth rather than reduced wall\-clock computation\. At the scales studied, adaptive exit\-state selection improves downstream performance at intermediate depths, while actual inference\-time savings require additional state\-handling mechanisms\.111Code will be released publicly\.
Looped State\-Space Language Models with Adaptive Exit\-State Selection
Zhenxuan Yu††thanks:E\-mail:24wr006m@rikkyo\.ac\.jpRikkyo UniversityYutaka Matsuo††thanks:Email:\{t\.kojima, matsuo, iwasawa\}@weblab\.t\.u\-tokyo\.ac\.jpThe University of TokyoTakeshi Kojima22footnotemark:2The University of TokyoYusuke Iwasawa22footnotemark:2The University of Tokyo
## 1Introduction
Large Language Models \(LLMs\) have achieved remarkable progress in language understanding, text generation, and complex reasoning\(DeepSeek\-AIet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib1); Yanget al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib6)\)\. However, these advances have largely been driven by scaling model size, training data, and compute\-intensive training and test\-time inference\(Hoffmannet al\.,[2022](https://arxiv.org/html/2607.10110#bib.bib7); Snellet al\.,[2024](https://arxiv.org/html/2607.10110#bib.bib2)\)\. This scaling paradigm introduces substantial challenges for practical deployment: model weights and intermediate activations increase memory consumption, while Transformer\-based self\-attention requires maintaining large key\-value caches during long\-context inference, leading to rapidly growing inference costs as context length increases\(Beltagyet al\.,[2020](https://arxiv.org/html/2607.10110#bib.bib5)\)\. Improving model capability under fixed parameter, FLOP, or memory budgets has therefore become a central problem in the design of efficient LLM architectures\.
State Space Models \(SSMs\), particularly Mamba\-2\(Gu and Dao,[2024](https://arxiv.org/html/2607.10110#bib.bib3); Dao and Gu,[2024](https://arxiv.org/html/2607.10110#bib.bib11)\), provide an efficient alternative to attention with favorable long\-context scaling, and Hybrid Mamba\-Transformer designs\(Lieberet al\.,[2024](https://arxiv.org/html/2607.10110#bib.bib4); Waleffeet al\.,[2024](https://arxiv.org/html/2607.10110#bib.bib8); Blakemanet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib12)\)have further emerged as practical backbones for modern LLMs\. Orthogonally, Looped Transformers reuse a shared block across depth to gain effective computational depth without extra parameters, and recent work shows that this benefits reasoning\(Dehghaniet al\.,[2019](https://arxiv.org/html/2607.10110#bib.bib9); Saunshiet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib13)\)\.
Our core idea is to share the parameters of Mamba or hybrid blocks across layers and apply them repeatedly through explicit finite\-depth loops, constructing a deeper effective computational path without increasing the number of distinct trainable parameters\. This design allows us to disentangle the effects of parameter count, compute budget, and effective depth\.
We validate Looped Mamba on controlled synthetic reasoning tasks \(Mano and p\-hop induction\) that require compositional computation and recursive retrieval, and then extend to language\-model pre\-training under controlled iso\-parameter and iso\-FLOPs protocols, which together let us separate the contributions of parameter sharing, additional compute, and effective depth\. We further evaluate downstream reasoning benchmarks and compare Looped Mamba with Looped Transformer and Looped Hybrid Mamba\-Transformer\.
We also adapt the post\-training Exit Gate mechanism from Ouro\(Zhuet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib14)\)to Looped Mamba, enabling threshold\-controlled adaptive exit\-state selection\. Empirically, the exit\-gated variant achieves higher downstream performance than fixed\-depth inference at comparable selected prediction depths, while introducing only minimal additional parameters\.
The main contributions of this work are as follows\. First, we present a controlled empirical study of explicit finite\-depth looping for Mamba\-based and Hybrid Mamba\-Transformer language models, under matched iso\-parameter and iso\-FLOPs protocols against both Mamba and Transformer baselines; the two protocols jointly disentangle the effects of parameter sharing, compute budget, and effective depth\. Second, we present empirical results on controlled reasoning tasks \(Mano, p\-hop induction\) and on language\-model pre\-training with downstream benchmarks, showing that recurrent depth improves compositional and recursive computation under fixed parameter budgets\. Third, we introduce an adaptive exit\-state selection mechanism for looped state\-space language models, adapting the post\-training Exit Gate of Ouro to Mamba and hybrid backbones, and show that threshold\-controlled adaptive exit\-state selection can improve downstream performance at comparable selected prediction depths\.
## 2Related Work
#### State Space Models and Mamba\.
State Space Models \(SSMs\), particularly Mamba and Mamba\-2, offer a promising alternative to the efficiency bottlenecks of pure Transformer architectures\(Gu and Dao,[2024](https://arxiv.org/html/2607.10110#bib.bib3)\)\. Prior work has shown that Mamba\-style models can match or outperform Transformers at small to medium scales in language modeling, while Mamba\-2 further establishes a connection between SSMs and attention through structured state space duality and introduces a more efficient core layer design\(Dao and Gu,[2024](https://arxiv.org/html/2607.10110#bib.bib11)\)\. Owing to their favorable sequence\-modeling and memory\-efficiency characteristics compared with standard attention, Mamba\-based architectures provide an attractive path toward efficient long\-sequence language modeling and inference\. Meanwhile, Hybrid Mamba\-Transformer\(Lieberet al\.,[2024](https://arxiv.org/html/2607.10110#bib.bib4); Waleffeet al\.,[2024](https://arxiv.org/html/2607.10110#bib.bib8)\)architectures have begun to appear in recent LLM systems; for example, Nemotron 3 combines Mamba, Transformer, and Mixture\-of\-Experts components to improve throughput and support long\-context modeling\(Blakemanet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib12)\)\. These developments suggest that Mamba and hybrid Mamba\-Transformer designs are becoming important components of modern efficient LLM architectures\.
#### Recursive Mamba models\.
A closely related recent study also explores Mamba\-2 in a recursive reasoning setting\.Wang and Reid \([2026](https://arxiv.org/html/2607.10110#bib.bib17)\)replace the attention blocks in the Tiny Recursive Model \(TRM\)\(Jolicoeur\-Martineau,[2025](https://arxiv.org/html/2607.10110#bib.bib16)\)with Mamba\-2 attention\-hybrid blocks and evaluate the resulting model on ARC\-AGI\. Their results show that Mamba\-2 hybrid operators can preserve, and in some settings improve, recursive reasoning capability within a tiny task\-specialized model\. However, their focus is on a small TRM\-style architecture with approximately 7M parameters trained specifically for ARC\-AGI\-style abstract reasoning\. In contrast, we study general\-purpose pre\-trained Mamba\-based language models at 50M–370M scale and evaluate on standard downstream benchmarks\.
#### Implicit and self\-iterated language models\.
Closely related work bySchöneet al\.\([2025](https://arxiv.org/html/2607.10110#bib.bib15)\)studies implicit SSMs that self\-iterate toward a fixed point; their iterations act as a fixed\-point solver with phantom\-gradient training, whereas we study explicit, gradient\-tracked finite\-depth loops over an entire Mamba \(or hybrid\) backbone, enabling controlled iso\-parameter and iso\-FLOPs comparisons\.
#### Other Loop models\.
Looped or recurrent\-depth architectures have a long history in sequence modeling\. Universal Transformers apply the same transformation recurrently over sequence representations, combining the parallelism and global receptive field of Transformers with a recurrent inductive bias; they also introduce adaptive computation mechanisms for dynamic halting\(Dehghaniet al\.,[2019](https://arxiv.org/html/2607.10110#bib.bib9)\)\. More recently,Saunshiet al\.\([2025](https://arxiv.org/html/2607.10110#bib.bib13)\)study Looped Transformers for reasoning and show that akk\-layer Transformer loopedLLtimes can nearly match, and sometimes outperform, a non\-loopedkLkL\-layer model on synthetic and downstream reasoning tasks, while using far fewer distinct parameters\. These results suggest that many reasoning tasks may require effective depth more than independent parameters\. However, these studies focus primarily on Transformer backbones\. In contrast, we investigate whether the same looped\-depth principle applies to Mamba and Hybrid Mamba\-Transformer architectures\. Building on this line of work, Ouro\(Zhuet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib14)\)further introduces an entropy\-regularized exit gate that learns a token\-level distribution over recurrent steps for adaptive exit\-state selection on Looped Transformers; we adapt this mechanism to Looped Mamba and Hybrid Mamba\-Transformer backbones and study its selected\-depth–performance trade\-off\.
#### Concurrent work on looped subquadratic models\.
Concurrently with this work, LT2\(Denget al\.,[2026](https://arxiv.org/html/2607.10110#bib.bib19)\)studies looped architectures with linear and sparse attention, and additionally pre\-trains looped Mamba\-2 models at the 0\.6B–1\.3B scale, observing that purely Mamba\-2\-based looped models trail attention\-based looped models on downstream tasks and can exhibit training instabilities\. Notably, LT2 deliberately avoids adaptive computation at scale\. Our study is complementary along three axes: \(i\) we perform tightly controlled iso\-parameter*and*iso\-FLOPs comparisons that disentangle parameter sharing from effective depth, including on synthetic compositional\-reasoning tasks; \(ii\) we study Hybrid Mamba–Transformer backbones in which attention layers are looped jointly with Mamba layers; and \(iii\) we introduce an adaptive exit\-state selection mechanism for looped state\-space models, which LT2 leaves unaddressed\.
## 3Proposed Method
### 3\.1Overall Architecture
#### Looped Mamba models\.
We study the standard looped\-model construction with Mamba backbones, following the notation of looped models\(Saunshiet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib13)\)\. For any sequence\-to\-sequence functionff, let
f\(R\)=f∘f∘⋯∘f⏟Rtimesf^\{\(R\)\}=\\underbrace\{f\\circ f\\circ\\cdots\\circ f\}\_\{R\\ \\mathrm\{times\}\}denote the function obtained by applyingffrecurrently forRRsteps\. In this work,ffcorresponds to anNN\-layer Mamba backboneMθM\_\{\\theta\}, and we denote the resulting Looped Mamba model byN⊗RN\\otimes R, whereNNis the number of independently parameterized Mamba layers andRRis the number of recurrent steps\. The effective depth is thereforeNRNR\. Since the sameNN\-layer backbone is reused across recurrent steps, the parameter budget scales withNN, while the FLOP budget scales withNRNR\. Thus,N⊗1N\\otimes 1corresponds to a standard non\-looped model,N⊗RN\\otimes Ris iso\-parameter withN⊗1N\\otimes 1, and it is iso\-FLOPs with the non\-looped modelNR⊗1NR\\otimes 1\.
Given input tokensxx, we compute
h\(0\)\\displaystyle h^\{\(0\)\}=emb\(x\),\\displaystyle=\\mathrm\{emb\}\(x\),h\(r\)\\displaystyle h^\{\(r\)\}=Mθ\(h\(r−1\)\),r=1,…,R\.\\displaystyle=M\_\{\\theta\}\(h^\{\(r\-1\)\}\),\\quad r=1,\\ldots,R\.and obtain the final language\-model distribution from
pθ\(R\)\(⋅∣x≤t\)=softmax\(lmhead\(ht\(R\)\)\)\.p\_\{\\theta\}^\{\(R\)\}\(\\cdot\\mid x\_\{\\leq t\}\)=\\mathrm\{softmax\}\\\!\\left\(\\mathrm\{lmhead\}\(h\_\{t\}^\{\(R\)\}\)\\right\)\.During pre\-training, we optimize only the final recurrent step for a sequence of lengthTT:
ℒfinal=𝔼x\[∑t=1T−1−logpθ\(R\)\(xt\+1∣x≤t\)\]\.\\mathcal\{L\}\_\{\\mathrm\{final\}\}=\\mathbb\{E\}\_\{x\}\\left\[\\sum\_\{t=1\}^\{T\-1\}\-\\log p\_\{\\theta\}^\{\(R\)\}\(x\_\{t\+1\}\\mid x\_\{\\leq t\}\)\\right\]\.
#### Looped Hybrid Mamba\-Transformer models\.
We also extend the same looped construction to Hybrid Mamba\-Transformer backbones\. In this setting, the recurrent backboneMθM\_\{\\theta\}is still anNN\-layer parameterized block, but a small subset of its layers is implemented as causal self\-attention while the remaining layers are Mamba\-2 layers\. The full attention configurations are provided in Appendix[E\.4](https://arxiv.org/html/2607.10110#A5.SS4)\. For the 140M hybrid model, we replace the layers with indices\{5,10,15,20\}\\\{5,10,15,20\\\}with causal attention layers; for the 370M hybrid model, we replace the layers with indices\{6,18,30,42\}\\\{6,18,30,42\\\}\. Importantly, looping is applied to the entire hybrid backbone rather than only to the Mamba layers:
h\(r\)=Mθhybrid\(h\(r−1\)\),r=1,…,R\.h^\{\(r\)\}=M\_\{\\theta\}^\{\\mathrm\{hybrid\}\}\(h^\{\(r\-1\)\}\),\\qquad r=1,\\ldots,R\.This choice follows the standard looped\-model formulation, where the loop is applied to a complete sequence\-to\-sequence block rather than to a single submodule\. It also allows both Mamba state\-space updates and attention\-based token interactions to be revisited at every recurrent step\. Thus, each loop performs another round of interaction between the Mamba and attention components, enabling iterative refinement of the hidden representation\. This design allows us to evaluate whether recurrent parameter sharing remains effective when a small number of attention layers are inserted into an otherwise Mamba\-based backbone\.
#### Three\-stage training pipeline\.
We train Looped Mamba models in three stages\. First, we pre\-train the base looped model with the final\-loop objective above, without an exit gate\. Second, we continue training with an entropy\-regularized adaptive\-computation objective adapted from Ouro\(Zhuet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib14)\), jointly using the outputs of all recurrent steps to learn a token\-level distribution over exit depths\. Third, we freeze the Mamba backbone and the language\-modeling head, and train only the exit\-gate parameters with a focused adaptive gate objective that encourages continuation only when the next recurrent step improves the token\-level prediction loss\. The detailed objectives for the two exit\-gate stages are provided in Appendix[B](https://arxiv.org/html/2607.10110#A2)\.
#### Exit\-gated inference\.
When the exit gate is enabled, each token’s exit step is selected as the smallestrrwhose cumulative exit probability exceeds a thresholdqq\(Appendix[B](https://arxiv.org/html/2607.10110#A2)\)\. Unlike attention\-based looped models, looped SSMs impose a*state\-continuity*constraint: the recurrent scan at stepr\+1r\{\+\}1consumes the step\-\(r\+1\)\(r\{\+\}1\)SSM state of every preceding position, so a token cannot simply be dropped from later recurrent steps without affecting all subsequent tokens\. We therefore adopt an*output\-selection*semantics in this work: the state chain of allRRrecurrent steps is preserved, and the gate selects, per token, which recurrent step’s hidden state is passed to the language\-modeling head\. Under this semantics, the average exit step reported in our experiments \(Iter\) measures the average depth selected for prediction rather than a wall\-clock compute reduction; realizing actual decode\-time savings requires additional state\-handling policies for exited tokens \(e\.g\., state pass\-through\), which we leave to future work\.
## 4Looped Models on Simple Reasoning Tasks
#### Looped model notation\.
We use theN⊗RN\\otimes Rnotation defined in Section[3](https://arxiv.org/html/2607.10110#S3), whereNNdenotes the number of independently parameterized Mamba layers andRRdenotes the number of recurrent steps\. This notation is used consistently for all simple reasoning experiments below\.
### 4\.1Mano Task
#### Setup\.
We first evaluate the knowledge manipulation capability of Looped Mamba on the Mano task\(Allen\-Zhu,[2025](https://arxiv.org/html/2607.10110#bib.bib18)\)\. Mano is a synthetic modular\-arithmetic reasoning task where the model must directly output the final answer without producing any explicit intermediate steps\. Solving this task requires the model to apply arithmetic rules modulo 23 and to parse the tree structure encoded by the prefix expression in order to compose multiple intermediate computations\.
To thoroughly evaluate manipulation capability, we consider test accuracy across different difficulty levels, determined by the maximum expression lengthLL, i\.e\., the number of operations in each sample\. During training, expression lengthsℓ\\ellare sampled from\[1,L\]\[1,L\]; during evaluation, we test only on the hardest expressions withℓ=L\\ell=L\. We use three difficulty levels,L∈\{10,16,24\}L\\in\\\{10,16,24\\\}, and compare non\-looped Mamba baselines with Looped Mamba models\. Specifically, we train4⊗14\\otimes 1,6⊗16\\otimes 1, and12⊗112\\otimes 1as iso\-parameter baselines, and compare them with looped models4⊗64\\otimes 6,6⊗46\\otimes 4, and12⊗212\\otimes 2, respectively\. We additionally include24⊗124\\otimes 1as a non\-looped iso\-FLOPs reference for all three looped models\. All three looped models have effective depth 24 and are therefore iso\-FLOPs with the24⊗124\\otimes 1non\-looped baseline, while using only1/61/6,1/41/4, and1/21/2of its parameterized layers, respectively\. More details are provided in Appendix[C](https://arxiv.org/html/2607.10110#A3)\.
#### Results\.
Table[1](https://arxiv.org/html/2607.10110#S4.T1)reports the exact\-match accuracy on the Mano task\. Under the iso\-parameter comparison, Looped Mamba consistently outperforms its non\-looped counterpart across all model sizes and difficulty levels\. For example, on the hardest settingL=24L=24, the4⊗14\\otimes 1,6⊗16\\otimes 1, and12⊗112\\otimes 1baselines achieve only 7\.50%, 8\.00%, and 22\.70% accuracy, respectively, whereas their looped counterparts improve to 99\.35%, 59\.85%, and 38\.40%\. These results indicate that recurrently reusing Mamba layers substantially improves the model’s ability to parse arithmetic tree structures and compose modular computations, without increasing the number of trainable parameters\.
We further compare models under the same effective depth of 24\. In this iso\-FLOPs setting, the best looped model,4⊗64\\otimes 6, outperforms the24⊗124\\otimes 1non\-looped baseline on all difficulty levels, while using only one sixth of the parameterized layers\. In particular, onL=24L=24,4⊗64\\otimes 6improves accuracy from 37\.45% to 99\.35%\. This comparison shows that the gain is not simply due to using more computation: with matched effective depth, parameter sharing through recurrent loops can provide a stronger inductive bias for knowledge manipulation than increasing the number of unshared layers\. Moreover, among the effective\-depth\-24 looped models, performance improves as the number of recurrent steps increases, from12⊗212\\otimes 2to6⊗46\\otimes 4and further to4⊗64\\otimes 6\. This suggests that, for Mano, repeated application of a smaller Mamba block is particularly effective for learning reusable algorithmic computation\.
Table 1:Accuracy \(%\) of looped/non\-looped Mamba models on the Mano task\. Bold indicates the better result within each iso\-parameter group\.ModelParams /FLOPsL=10L=10L=16L=16L=24L=24Baseline modelBase\(24⊗1\)\(24\\otimes 1\)24×/24×24\\times\\,/\\,24\\times96\.1089\.4537\.454\-layer modelBase\(4⊗1\)\(4\\otimes 1\)4×/4×4\\times\\,/\\,4\\times55\.6012\.857\.50Loop\(4⊗6\)\(4\\otimes 6\)4×/24×4\\times\\,/\\,24\\times99\.2199\.6199\.356\-layer modelBase\(6⊗1\)\(6\\otimes 1\)6×/6×6\\times\\,/\\,6\\times86\.3318\.358\.00Loop\(6⊗4\)\(6\\otimes 4\)6×/24×6\\times\\,/\\,24\\times97\.9082\.4259\.8512\-layer modelBase\(12⊗1\)\(12\\otimes 1\)12×/12×12\\times\\,/\\,12\\times92\.9759\.7022\.70Loop\(12⊗2\)\(12\\otimes 2\)12×/24×12\\times\\,/\\,24\\times95\.7082\.0338\.40
### 4\.2p\-hop Task
#### Setup\.
We further evaluate the recursive retrieval capability of Looped Mamba on the p\-hop induction task\. This task is a multi\-step generalization of the induction\-head problem\. Given a sequencev=\(v1,…,vn\)v=\(v\_\{1\},\\ldots,v\_\{n\}\), 1\-hop induction starts from the last tokenvnv\_\{n\}, finds its previous occurrence in the sequence, and outputs the token immediately following that occurrence\. The p\-hop task recursively applies this operationpptimes, requiring the model to repeatedly backtrack through the context and retrieve the final answer\. We replace the Transformer backbone in the original p\-hop experiments with a Mamba\-2 backbone while keeping the same looped\-model protocol\. We use an alphabet of size 4, sequence lengthn=256n=256, and evaluate two difficulty levels,p=16p=16andp=32p=32\. Following the Looped Transformer setup, we compare non\-looped baselines, iso\-parameter looped models, and iso\-FLOPs baselines\. Since the answer is one of four alphabet tokens, random guessing obtains roughly 25% accuracy\.
#### Results\.
Table[2](https://arxiv.org/html/2607.10110#S5.T2)reports the results of Looped Mamba on p\-hop induction\. Shallow non\-looped Mamba models struggle on this task: for instance, the1⊗11\\otimes 1model obtains only 38\.32% and 41\.08% accuracy forp=16p=16andp=32p=32, respectively, while the2⊗12\\otimes 1model reaches 63\.51% and 67\.22%\. In contrast, adding recurrent loops dramatically improves performance\. The1⊗61\\otimes 6model already reaches 99\.51% and 98\.78%, and the1⊗121\\otimes 12model further improves to 99\.70% and 99\.89%\. Similarly,2⊗62\\otimes 6,3⊗43\\otimes 4, and4⊗34\\otimes 3all achieve near\-perfect accuracy under both difficulty settings\.
These results suggest that p\-hop induction primarily requires recursive computational depth rather than a large number of distinct parameters\. Even one\- or two\-layer Mamba\-2 backbones, when given sufficient effective depth via looping, can learn recursive backtracking and answer retrieval\. Consistent with the findings for Looped Transformers, Looped Mamba substantially outperforms the iso\-parameter baseline and matches or exceeds the performance of deeper non\-looped baselines\. This indicates that Mamba’s selective state\-update mechanism can also benefit from recurrent parameter sharing and can perform multi\-step latent retrieval with far fewer parameters\.
## 5Language Modeling with Looped Models
### 5\.1Setup
#### Models
We pre\-trained looped language models based on LLaMA\-2 TransformerTouvronet al\.\([2023](https://arxiv.org/html/2607.10110#bib.bib31)\)and Mamba\-2 architectures, with model sizes ranging from 50M to 370M parameters\. All models follow the GPT\-like autoregressive decoder pre\-training methodologyBrownet al\.\([2020](https://arxiv.org/html/2607.10110#bib.bib20)\)\. In Table[5](https://arxiv.org/html/2607.10110#A1.T5), we compare LLaMA\-2 Transformer and Mamba\-2 models under iso\-parameter and iso\-FLOPs settings\. In Table[4](https://arxiv.org/html/2607.10110#S5.T4), we further evaluate LLaMA\-2, Mamba\-2, Mamba\-2 Hybrid, and exit\-gate variants\.
#### Training settings
Models in Table[5](https://arxiv.org/html/2607.10110#A1.T5)are pre\-trained for 100B tokens, corresponding to 50,000 iterations, from the Nemotron\-ClimbMix datasetDiaoet al\.\([2025](https://arxiv.org/html/2607.10110#bib.bib21)\)\. For the models in Table[4](https://arxiv.org/html/2607.10110#S5.T4), we use a three\-stage training procedure\. First, all models are pre\-trained for 150B tokens, corresponding to 75,000 iterations, from the same dataset\. Second, the exit\-gate variants are continued with Exit\-Gate Stage 1 training for 5\.2B tokens, corresponding to 2,500 iterations\. Third, Exit\-Gate Stage 2 further trains the gate for 1\.2B tokens, corresponding to 600 iterations\. We use the Mistral\-7B\-v0\.1 tokenizerJianget al\.\([2023](https://arxiv.org/html/2607.10110#bib.bib22)\)with a vocabulary size of 32,000 and set the context length to 2048 tokens\. All experiments are conducted on 8 H100 GPUs with an effective global batch size of 1024 sequences\. Architecture\-specific configurations, optimizer settings, and detailed results are provided in Appendix[E](https://arxiv.org/html/2607.10110#A5)\.
#### Evaluation
For downstream evaluation, we report performance on a set of commonsense and reasoning benchmarks, including PIQA\(Bisket al\.,[2020](https://arxiv.org/html/2607.10110#bib.bib24)\), BoolQ\(Clarket al\.,[2019](https://arxiv.org/html/2607.10110#bib.bib23)\), WinoGrande\(Sakaguchiet al\.,[2021](https://arxiv.org/html/2607.10110#bib.bib26)\), OpenBookQA\(Mihaylovet al\.,[2018](https://arxiv.org/html/2607.10110#bib.bib28)\), HellaSwag\(Zellerset al\.,[2019](https://arxiv.org/html/2607.10110#bib.bib25)\), ARC\-Challenge, and ARC\-Easy\(Clarket al\.,[2018](https://arxiv.org/html/2607.10110#bib.bib27)\)\. We follow a fixed few\-shot evaluation protocol: PIQA, BoolQ, WinoGrande, and OpenBookQA are evaluated with 5\-shot prompting; HellaSwag with 10\-shot prompting; ARC\-Challenge with 25\-shot prompting; and ARC\-Easy with 8\-shot prompting\.
Table 2:Accuracy \(%\) of looped/non\-looped Mamba models on the p\-hop induction task\.ModelParams /FLOPs𝒑=𝟏𝟔,\\boldsymbol\{p\{=\}16,\}𝒏=𝟐𝟓𝟔\\boldsymbol\{n\{=\}256\}𝒑=𝟑𝟐,\\boldsymbol\{p\{=\}32,\}𝒏=𝟐𝟓𝟔\\boldsymbol\{n\{=\}256\}Base\(12⊗1\)\(12\\otimes 1\)12×/12×12\\times/12\\times99\.9599\.871\-layer modelBase\(1⊗1\)\(1\\otimes 1\)1×/1×1\\times/1\\times38\.3241\.08Loop\(1⊗6\)\(1\\otimes 6\)1×/6×1\\times/6\\times99\.5198\.78Loop\(1⊗12\)\(1\\otimes 12\)1×/12×1\\times/12\\times99\.7099\.892\-layer modelBase\(2⊗1\)\(2\\otimes 1\)2×/2×2\\times/2\\times63\.5167\.22Loop\(2⊗3\)\(2\\otimes 3\)2×/6×2\\times/6\\times99\.5598\.39Loop\(2⊗6\)\(2\\otimes 6\)2×/12×2\\times/12\\times99\.9399\.843\-layer modelBase\(3⊗1\)\(3\\otimes 1\)3×/3×3\\times/3\\times87\.1877\.38Loop\(3⊗2\)\(3\\otimes 2\)3×/6×3\\times/6\\times99\.7498\.55Loop\(3⊗4\)\(3\\otimes 4\)3×/12×3\\times/12\\times99\.9299\.804\-layer modelBase\(4⊗1\)\(4\\otimes 1\)4×/4×4\\times/4\\times97\.9989\.39Loop\(4⊗3\)\(4\\otimes 3\)4×/12×4\\times/12\\times99\.8999\.846\-layer modelBase\(6⊗1\)\(6\\otimes 1\)6×/6×6\\times/6\\times99\.8199\.38Loop\(6⊗2\)\(6\\otimes 2\)6×/12×6\\times/12\\times99\.9799\.91
### 5\.2iso\-FLOPs Evaluation
Table[3](https://arxiv.org/html/2607.10110#S5.T3)reports validation perplexity and downstream performance\. The table supports two complementary comparisons\. First, under the iso\-parameter setting, looped models consistently improve validation perplexity over their non\-looped counterparts with the same number of parameterized layers, for both Mamba\-2 and LLaMA\-2 backbones\. This indicates that recurrently reusing the same backbone improves language modeling performance when the trainable parameter budget is fixed\. For the LLaMA\-2 Transformer backbone, each parameterized layer contains both causal self\-attention and a gated MLP, whereas each Mamba\-2 layer uses a Mamba\-2 mixer without an external MLP\. In Table[3](https://arxiv.org/html/2607.10110#S5.T3), both backbones are compared using the same layer countsN∈\{6,12,24\}N\\in\\\{6,12,24\\\}and the same loop notationN⊗RN\\otimes R\. Thus, the comparison aligns the number of parameterized layers and effective depth, while architectural details such as attention, MLPs, and Mamba mixing remain backbone\-specific\.
Second, under the strict iso\-FLOPs setting, the deeper non\-looped24⊗124\\otimes 1models still achieve the best validation perplexity\. Although looped models match the effective depth of24⊗124\\otimes 1, they use substantially fewer distinct parameters and do not fully close the perplexity gap\. This behavior is consistent with prior observations for Looped Transformers, where looped language models improve over iso\-parameter baselines but remain worse than iso\-FLOPs non\-looped models in perplexity\(Saunshiet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib13)\)\.
Notably, downstream performance does not follow validation perplexity exactly\. Despite their worse perplexity relative to the24⊗124\\otimes 1baselines, looped models remain competitive on the downstream average and in some cases outperform the deeper non\-looped baseline\. This suggests that recurrent depth may provide a parameter\-efficient inductive bias that is not fully captured by perplexity alone, although the strength of this effect depends on the backbone and task\.
Table 3:iso\-FLOPs evaluation of Mamba\-2 and LLaMA\-2 Transformer models trained on 100B tokens\. PPL denotes validation perplexity, and Avg\. denotes the average score over seven downstream tasks\. Full per\-task results are reported in Appendix[A](https://arxiv.org/html/2607.10110#A1)\.ModelParams / FLOPsPPLAvg\.Mamba\-2Base\(6⊗1\)\(6\\otimes 1\)6×/6×6\\times/6\\times21\.9635\.31Loop\(6⊗4\)\(6\\otimes 4\)6×/24×6\\times/24\\times19\.9837\.39Base\(12⊗1\)\(12\\otimes 1\)12×/12×12\\times/12\\times18\.6638\.12Loop\(12⊗2\)\(12\\otimes 2\)12×/24×12\\times/24\\times18\.2138\.02Base\(24⊗1\)\(24\\otimes 1\)24×/24×24\\times/24\\times16\.1537\.57LLaMA\-2 TransformerBase\(6⊗1\)\(6\\otimes 1\)6×/6×6\\times/6\\times21\.4838\.19Loop\(6⊗4\)\(6\\otimes 4\)6×/24×6\\times/24\\times19\.7538\.45Base\(12⊗1\)\(12\\otimes 1\)12×/12×12\\times/12\\times18\.5338\.71Loop\(12⊗2\)\(12\\otimes 2\)12×/24×12\\times/24\\times18\.0639\.04Base\(24⊗1\)\(24\\otimes 1\)24×/24×24\\times/24\\times16\.4037\.21
### 5\.3iso\-Parameter Evaluation and Adaptive Exit
Table 4:Downstream evaluation of Mamba\-2, Mamba\-2 Hybrid, and LLaMA\-2 Transformer models\. PPL denotes validation perplexity, and Avg\. denotes the average score over the seven downstream tasks shown in the table\. Iter denotes the average exit step selected by the gate \(Section[3](https://arxiv.org/html/2607.10110#S3)\), averaged across downstream evaluation tasks\. Bold numbers indicate the best result within each model group; ties are broken by a smaller Iter\.ModelExitqqIterPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.140M LLaMA\-2llama2\-140m\-baselineno–1\.0015\.9650\.9856\.5150\.5927\.6025\.4829\.4423\.9137\.79llama2\-140m\-loop4no–4\.0014\.7551\.0960\.8351\.7025\.6025\.8827\.5626\.1838\.41140M Mamba\-2Mamba\-2\-140m\-baselineno–1\.0015\.8049\.9555\.3851\.2228\.0026\.1129\.4425\.5537\.95Mamba\-2\-140m\-loop2no–2\.0015\.4151\.5239\.6649\.0128\.0026\.0627\.8226\.9435\.57Mamba\-2\-140m\-loop4no–4\.0014\.6852\.3458\.8748\.2227\.0025\.2128\.0724\.8337\.79Mamba\-2\-140m\-loop4\-exitgate\-stage1yes0\.42\.8215\.5752\.5062\.2350\.6724\.6026\.1228\.8425\.0438\.57Mamba\-2\-140m\-loop4\-exitgate\-stage2yes0\.83\.0015\.5452\.9962\.2350\.2026\.0026\.0128\.9225\.7238\.87140M Mamba\-2 HybridMamba\-2\-hybrid\-140m\-baselineno–1\.0015\.2751\.0937\.8350\.2827\.0025\.7128\.5026\.3535\.25Mamba\-2\-hybrid\-140m\-loop4no–4\.0014\.3751\.3145\.8747\.4326\.8025\.0728\.9226\.0935\.93Mamba\-2\-hybrid\-140m\-loop4\-exitgate\-stage1yes0\.42\.8115\.0751\.7454\.7749\.9626\.8025\.2928\.1625\.3437\.44Mamba\-2\-hybrid\-140m\-loop4\-exitgate\-stage2yes0\.83\.2915\.0750\.3854\.8049\.2526\.4025\.3927\.7326\.6837\.23370M Mamba\-2Mamba\-2\-370m\-baselineno–1\.0012\.7248\.2061\.3049\.6027\.2024\.6028\.9025\.7037\.93Mamba\-2\-370m\-loop4no–4\.0012\.0651\.9061\.9648\.1527\.8024\.7229\.1025\.9338\.51Mamba\-2\-370m\-loop4\-exitgate\-stage1yes0\.603\.4012\.5452\.0759\.2449\.8824\.4025\.0926\.9626\.4737\.73Mamba\-2\-370m\-loop4\-exitgate\-stage2yes0\.803\.1812\.5552\.3457\.6150\.5924\.8024\.7726\.8826\.0937\.58370M Mamba\-2 HybridMamba\-2\-hybrid\-370m\-baselineno–1\.0012\.2252\.3960\.0650\.5926\.2024\.9529\.3526\.0138\.51Mamba\-2\-hybrid\-370m\-loop4no–4\.0011\.5951\.7461\.2249\.3328\.0025\.9729\.8626\.2638\.91Mamba\-2\-hybrid\-370m\-loop4\-exitgate\-stage1yes0\.603\.3412\.1251\.0339\.5450\.6727\.6026\.3630\.6325\.8435\.95Mamba\-2\-hybrid\-370m\-loop4\-exitgate\-stage2yes0\.803\.0112\.1251\.0939\.4550\.0427\.4025\.8430\.2926\.0935\.74
\(a\)140M
\(b\)140M \(hybrid\)
\(c\)370M
\(d\)370M \(hybrid\)
Figure 1:Accuracy–selected\-depth trade\-off of exit\-gated Looped Mamba and Looped Mamba Hybrid models under different thresholdsqq\. The x\-axis denotes the average exit step selected by the gate, and the y\-axis denotes the downstream average accuracy\. Exit\-gated models trace a threshold\-controlled trade\-off curve, while fixed\-depth inference corresponds to discrete forced\-loop points\. “Fixed Exit Depth” denotes forced fixed\-depth inference of the looped model, and “R1 Baseline” denotes the non\-looped \(R=1R\{=\}1\) baseline\.In addition to the iso\-FLOPs comparison above, we evaluate looped models under an iso\-parameter setting\. Here, the number of independently parameterized layers is fixed, while the number of recurrent loops is varied\. Thus, anN⊗RN\\otimes Rmodel has the same number of trainable parameters as itsN⊗1N\\otimes 1baseline, but uses a larger effective computation depth at inference time\. This setting allows us to examine whether recurrent computation improves language modeling and downstream performance under a fixed parameter budget\. Note that the 140M and 370M models use different backbone depths: the 140M models contain 24 parameterized layers, whereas the 370M models contain 48 parameterized layers\. Consequently, loop4 corresponds to an effective depth of 96 layers for 140M models and 192 layers for 370M models\. Therefore, the 140M and 370M results should be interpreted as separate iso\-parameter comparisons within each model scale, rather than as directly depth\-matched comparisons across scales\.
Table[4](https://arxiv.org/html/2607.10110#S5.T4)reports results for LLaMA\-2 Transformer, Mamba\-2, and Mamba\-2 Hybrid backbones\. For exit\-gated models, we select a representative threshold fromq∈\{0\.2,0\.4,0\.6,0\.8,1\.0\}q\\in\\\{0\.2,0\.4,0\.6,0\.8,1\.0\\\}\. We choose the threshold that yields an average selected exit step below four while achieving the highest downstream average among such lower\-selected\-depth configurations\. The full threshold sweep is reported in Appendix[G](https://arxiv.org/html/2607.10110#A7)\.
The iso\-parameter results show that recurrent depth is generally useful, especially for validation perplexity\. For example, the 140M LLaMA\-2 loop4 model improves Avg\. from 37\.79 to 38\.41 over the non\-looped baseline, and fixed\-loop Mamba\-2 variants also reduce PPL\. However, downstream performance is not monotonic with the number of loops\. The 140M Mamba\-2 loop2 model improves PPL but suffers a large drop on BoolQ, and the loop4 model achieves the best PPL among fixed\-loop variants while its downstream average remains close to the non\-looped baseline\. This suggests that validation perplexity alone is insufficient to determine the optimal recurrent depth for downstream use\.
The fixed\-loop inference ablation in Table[12](https://arxiv.org/html/2607.10110#A6.T12)further shows that, under the final\-loop pre\-training objective, intermediate recurrent steps are poorly calibrated with respect to the language\-modeling head: forcing fewer thanRRloops at inference time yields severely degraded perplexity, while different downstream tasks prefer different effective depths withinr≤Rr\\leq R\. Naively truncating a fixed\-loop model at inference time is therefore not a viable way to save computation, and no single forced depth is optimal for all tasks\. This motivates adaptive exit\-state selection instead of always using the maximum recurrent depth\.
We therefore attach a lightweight exit gate to pretrained fixed\-loop checkpoints and continue training with the adaptive objectives described in Appendix[B](https://arxiv.org/html/2607.10110#A2)\. At inference time, the gate predicts a distribution over recurrent exit steps for each token, and the thresholdqqcontrols the selected\-depth–performance trade\-off\. Smallerqqvalues select earlier recurrent\-step outputs, while larger values select later recurrent\-step outputs\. Thus, exit\-gated inference can exploit intermediate recurrent states and avoid forcing all tokens to use the same fixed depth\.
For the 140M Mamba\-2 model, exit gating improves both downstream performance and exit\-state selection effectiveness\. The fixed loop4 model obtains an Avg\. of 37\.79 with a fixed prediction depth of 4\. In contrast, the Stage 1 exit gate reaches 38\.57 Avg\. with an average selected exit step of 2\.82 atq=0\.4q=0\.4, and the Stage 2 gate further improves to 38\.87 Avg\. with an average selected exit step of 3\.00 atq=0\.8q=0\.8\. A similar pattern appears for the 140M Mamba\-2 Hybrid model: although fixed loop4 achieves the best PPL, its Avg\. is only 35\.93, while Stage 1 exit gating improves Avg\. to 37\.44 with an average selected exit step of 2\.81\. These results indicate that adaptive exit can outperform fixed full\-loop inference by selecting more suitable recurrent depths\.
For the 370M models, the effect is different\. The best downstream average remains concentrated near fixed loop4: 370M Mamba\-2 reaches 38\.51 Avg\., and 370M Mamba\-2 Hybrid reaches 38\.91 Avg\. The corresponding exit\-gated variants do not surpass these peak fixed\-loop4 scores\. This suggests that larger models may rely more strongly on the final recurrent step, or that the current post\-training gate is not yet sufficiently calibrated for larger\-scale models\.
Nevertheless, Figure[1](https://arxiv.org/html/2607.10110#S5.F1)shows that exit gating provides a controllable selected\-depth–performance frontier\. For the 140M models, the exit\-gated variants can exceed fixed\-depth baselines at smaller average selected exit steps\. For the 370M models, although the exit\-gated variants do not surpass the peak fixed\-loop4 scores, they remain advantageous in the intermediate\-selected\-depth regime\. At average selected exit steps near three, the exit\-gated 370M Mamba\-2 and Hybrid models outperform forced fixed\-depth inference withr=3r=3by 2\.89 and 1\.20 percentage points, respectively\. We note that forced fixed\-depth inference atr<Rr<Ris inherently miscalibrated under the final\-loop objective \(Appendix[F](https://arxiv.org/html/2607.10110#A6)\); this comparison therefore shows that gate training recovers usable intermediate\-depth predictions, rather than demonstrating a reduction in executed computation\. Overall, these results suggest that recurrent depth can support adaptive selection among recurrent\-step outputs rather than a fixed maximum depth that is always fully used\.
## 6Conclusion
We presented Looped Mamba, a recurrent\-depth extension of Mamba\-based language models that repeatedly applies a shared Mamba or Hybrid Mamba–Transformer backbone to increase effective depth without adding distinct trainable parameters\. On Mano and p\-hop induction, Looped Mamba consistently improves over iso\-parameter non\-looped baselines and, in several reasoning settings, matches or exceeds deeper models with comparable effective depth, suggesting that recurrent parameter sharing benefits compositional computation and recursive retrieval\. In language\-model pre\-training, recurrent depth improves validation perplexity under fixed parameter budgets, although deeper non\-looped models still perform better under strict iso\-FLOPs comparisons\. Nevertheless, looped models remain competitive on downstream benchmarks\. Finally, exit\-gated Looped Mamba enables threshold\-controlled adaptive exit\-state selection and improves performance through adaptive selection among recurrent\-step outputs, especially at the 140M scale\. Overall, at the scales we study, recurrent depth offers a promising, parameter\-efficient scaling axis for Mamba and hybrid language models\.
## Limitations
This work has several limitations\. First, our experiments are limited to models up to 370M parameters, and we have not yet validated the scaling behavior of Looped Mamba at billion\-parameter or larger scales\. Second, our language\-model pre\-training budget is 100B–150B tokens, which is modest compared with modern large\-scale LLM training regimes\. Thus, our results should be interpreted as controlled empirical evidence rather than a complete characterization of large\-scale performance\. Third, Mano and p\-hop induction are synthetic reasoning tasks\. Although they isolate compositional computation and recursive retrieval, they may not fully represent open\-ended real\-world reasoning\. Finally, our Exit Gate experiments use fixed global thresholdsqqto control the selected recurrent\-step output, but do not reduce the number of recurrent steps actually executed\. Realizing and optimizing compute savings per token, example, or task requires additional state\-handling mechanisms\. Future work should evaluate larger models, longer training runs, more realistic reasoning benchmarks, and more fine\-grained adaptive\-compute strategies\.
## References
- Z\. Allen\-Zhu \(2025\)Physics of language models: part 4\.1, architecture design and the magic of canon layers\.Note:SSRN preprintAvailable at SSRNCited by:[§4\.1](https://arxiv.org/html/2607.10110#S4.SS1.SSS0.Px1.p1.1)\.
- I\. Beltagy, M\. E\. Peters, and A\. Cohan \(2020\)Longformer: the long\-document transformer\.External Links:2004\.05150,[Link](https://arxiv.org/abs/2004.05150)Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p1.1)\.
- Y\. Bisk, R\. Zellers, J\. Gao, Y\. Choi,et al\.\(2020\)Piqa: reasoning about physical commonsense in natural language\.InProceedings of the AAAI conference on artificial intelligence,Vol\.34,pp\. 7432–7439\.Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px3.p1.1)\.
- A\. Blakeman, A\. Grattafiori, A\. Basant, A\. Gupta, A\. Khattar, A\. Renduchintala, A\. Vavre, A\. Shukla, A\. Bercovich, A\. Ficek,et al\.\(2025\)NVIDIA nemotron 3: efficient and open intelligence\.arXiv preprint arXiv:2512\.20856\.Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p2.1),[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px1.p1.1)\.
- T\. Brown, B\. Mann, N\. Ryder, M\. Subbiah, J\. D\. Kaplan, P\. Dhariwal, A\. Neelakantan, P\. Shyam, G\. Sastry, A\. Askell, S\. Agarwal, A\. Herbert\-Voss, G\. Krueger, T\. Henighan, R\. Child, A\. Ramesh, D\. Ziegler, J\. Wu, C\. Winter, C\. Hesse, M\. Chen, E\. Sigler, M\. Litwin, S\. Gray, B\. Chess, J\. Clark, C\. Berner, S\. McCandlish, A\. Radford, I\. Sutskever, and D\. Amodei \(2020\)Language models are few\-shot learners\.InAdvances in Neural Information Processing Systems,H\. Larochelle, M\. Ranzato, R\. Hadsell, M\.F\. Balcan, and H\. Lin \(Eds\.\),Vol\.33,pp\. 1877–1901\.External Links:[Link](https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf)Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px1.p1.1)\.
- C\. Clark, K\. Lee, M\. Chang, T\. Kwiatkowski, M\. Collins, and K\. Toutanova \(2019\)BoolQ: exploring the surprising difficulty of natural yes/no questions\.arXiv preprint arXiv:1905\.10044\.Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px3.p1.1)\.
- P\. Clark, I\. Cowhey, O\. Etzioni, T\. Khot, A\. Sabharwal, C\. Schoenick, and O\. Tafjord \(2018\)Think you have solved question answering? try arc, the ai2 reasoning challenge\.arXiv preprint arXiv:1803\.05457\.Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px3.p1.1)\.
- T\. Dao and A\. Gu \(2024\)Transformers are ssms: generalized models and efficient algorithms through structured state space duality\.arXiv preprint arXiv:2405\.21060\.Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p2.1),[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px1.p1.1)\.
- DeepSeek\-AI, A\. Liu, B\. Feng, B\. Xue, B\. Wang, B\. Wu, C\. Lu, C\. Zhao, C\. Deng, C\. Zhang, C\. Ruan, D\. Dai, D\. Guo, D\. Yang, D\. Chen, D\. Ji, E\. Li, F\. Lin, F\. Dai, F\. Luo, G\. Hao, G\. Chen, G\. Li, H\. Zhang, H\. Bao, H\. Xu, H\. Wang, H\. Zhang, H\. Ding, H\. Xin, H\. Gao, H\. Li, H\. Qu, J\. L\. Cai, J\. Liang, J\. Guo, J\. Ni, J\. Li, J\. Wang, J\. Chen, J\. Chen, J\. Yuan, J\. Qiu, J\. Li, J\. Song, K\. Dong, K\. Hu, K\. Gao, K\. Guan, K\. Huang, K\. Yu, L\. Wang, L\. Zhang, L\. Xu, L\. Xia, L\. Zhao, L\. Wang, L\. Zhang, M\. Li, M\. Wang, M\. Zhang, M\. Zhang, M\. Tang, M\. Li, N\. Tian, P\. Huang, P\. Wang, P\. Zhang, Q\. Wang, Q\. Zhu, Q\. Chen, Q\. Du, R\. J\. Chen, R\. L\. Jin, R\. Ge, R\. Zhang, R\. Pan, R\. Wang, R\. Xu, R\. Zhang, R\. Chen, S\. S\. Li, S\. Lu, S\. Zhou, S\. Chen, S\. Wu, S\. Ye, S\. Ye, S\. Ma, S\. Wang, S\. Zhou, S\. Yu, S\. Zhou, S\. Pan, T\. Wang, T\. Yun, T\. Pei, T\. Sun, W\. L\. Xiao, W\. Zeng, W\. Zhao, W\. An, W\. Liu, W\. Liang, W\. Gao, W\. Yu, W\. Zhang, X\. Q\. Li, X\. Jin, X\. Wang, X\. Bi, X\. Liu, X\. Wang, X\. Shen, X\. Chen, X\. Zhang, X\. Chen, X\. Nie, X\. Sun, X\. Wang, X\. Cheng, X\. Liu, X\. Xie, X\. Liu, X\. Yu, X\. Song, X\. Shan, X\. Zhou, X\. Yang, X\. Li, X\. Su, X\. Lin, Y\. K\. Li, Y\. Q\. Wang, Y\. X\. Wei, Y\. X\. Zhu, Y\. Zhang, Y\. Xu, Y\. Xu, Y\. Huang, Y\. Li, Y\. Zhao, Y\. Sun, Y\. Li, Y\. Wang, Y\. Yu, Y\. Zheng, Y\. Zhang, Y\. Shi, Y\. Xiong, Y\. He, Y\. Tang, Y\. Piao, Y\. Wang, Y\. Tan, Y\. Ma, Y\. Liu, Y\. Guo, Y\. Wu, Y\. Ou, Y\. Zhu, Y\. Wang, Y\. Gong, Y\. Zou, Y\. He, Y\. Zha, Y\. Xiong, Y\. Ma, Y\. Yan, Y\. Luo, Y\. You, Y\. Liu, Y\. Zhou, Z\. F\. Wu, Z\. Z\. Ren, Z\. Ren, Z\. Sha, Z\. Fu, Z\. Xu, Z\. Huang, Z\. Zhang, Z\. Xie, Z\. Zhang, Z\. Hao, Z\. Gou, Z\. Ma, Z\. Yan, Z\. Shao, Z\. Xu, Z\. Wu, Z\. Zhang, Z\. Li, Z\. Gu, Z\. Zhu, Z\. Liu, Z\. Li, Z\. Xie, Z\. Song, Z\. Gao, and Z\. Pan \(2025\)DeepSeek\-v3 technical report\.External Links:2412\.19437,[Link](https://arxiv.org/abs/2412.19437)Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p1.1)\.
- M\. Dehghani, S\. Gouws, O\. Vinyals, J\. Uszkoreit, and Ł\. Kaiser \(2019\)Universal transformers\.External Links:1807\.03819,[Link](https://arxiv.org/abs/1807.03819)Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p2.1),[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px4.p1.3)\.
- C\. Deng, Y\. Zhang, R\. Zhu, Y\. Xu, J\. Liu, T\. S\. E\. Ng, and H\. Chen \(2026\)LT2: linear\-time looped transformers\.External Links:2605\.20670,[Link](https://arxiv.org/abs/2605.20670)Cited by:[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px5.p1.1)\.
- S\. Diao, Y\. Yang, Y\. Fu, X\. Dong, D\. Su, M\. Kliegl, Z\. Chen, P\. Belcak, Y\. Suhara, H\. Yin, M\. Patwary, C\. Lin, J\. Kautz, and P\. Molchanov \(2025\)CLIMB: clustering\-based iterative data mixture bootstrapping for language model pre\-training\.arXiv preprint\.External Links:[Link](https://arxiv.org/abs/2504.13161)Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px2.p1.1)\.
- A\. Gu and T\. Dao \(2024\)Mamba: linear\-time sequence modeling with selective state spaces\.External Links:2312\.00752,[Link](https://arxiv.org/abs/2312.00752)Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p2.1),[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px1.p1.1)\.
- J\. Hoffmann, S\. Borgeaud, A\. Mensch, E\. Buchatskaya, T\. Cai, E\. Rutherford, D\. de Las Casas, L\. A\. Hendricks, J\. Welbl, A\. Clark, T\. Hennigan, E\. Noland, K\. Millican, G\. van den Driessche, B\. Damoc, A\. Guy, S\. Osindero, K\. Simonyan, E\. Elsen, J\. W\. Rae, O\. Vinyals, and L\. Sifre \(2022\)Training compute\-optimal large language models\.External Links:2203\.15556,[Link](https://arxiv.org/abs/2203.15556)Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p1.1)\.
- A\. Q\. Jiang, A\. Sablayrolles, A\. Mensch, C\. Bamford, D\. S\. Chaplot, D\. de las Casas, F\. Bressand, G\. Lengyel, G\. Lample, L\. Saulnier, L\. R\. Lavaud, M\. Lachaux, P\. Stock, T\. L\. Scao, T\. Lavril, T\. Wang, T\. Lacroix, and W\. E\. Sayed \(2023\)Mistral 7b\.External Links:2310\.06825,[Link](https://arxiv.org/abs/2310.06825)Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px2.p1.1)\.
- A\. Jolicoeur\-Martineau \(2025\)Less is more: recursive reasoning with tiny networks\.External Links:2510\.04871,[Link](https://arxiv.org/abs/2510.04871)Cited by:[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px2.p1.1)\.
- O\. Lieber, B\. Lenz, H\. Bata, G\. Cohen, J\. Osin, I\. Dalmedigos, E\. Safahi, S\. Meirom, Y\. Belinkov, S\. Shalev\-Shwartz, O\. Abend, R\. Alon, T\. Asida, A\. Bergman, R\. Glozman, M\. Gokhman, A\. Manevich, N\. Ratner, N\. Rozen, E\. Shwartz, M\. Zusman, and Y\. Shoham \(2024\)Jamba: a hybrid transformer\-mamba language model\.External Links:2403\.19887,[Link](https://arxiv.org/abs/2403.19887)Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p2.1),[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px1.p1.1)\.
- T\. Mihaylov, P\. Clark, T\. Khot, and A\. Sabharwal \(2018\)Can a suit of armor conduct electricity? a new dataset for open book question answering\.InEMNLP,Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px3.p1.1)\.
- K\. Sakaguchi, R\. L\. Bras, C\. Bhagavatula, and Y\. Choi \(2021\)Winogrande: an adversarial winograd schema challenge at scale\.Communications of the ACM64\(9\),pp\. 99–106\.Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px3.p1.1)\.
- N\. Saunshi, N\. Dikkala, Z\. Li, S\. Kumar, and S\. J Reddi \(2025\)Reasoning with latent thoughts: on the power of looped transformers\.InInternational Conference on Learning Representations,Vol\.2025,pp\. 14855–14881\.Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p2.1),[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px4.p1.3),[§3\.1](https://arxiv.org/html/2607.10110#S3.SS1.SSS0.Px1.p1.1),[§5\.2](https://arxiv.org/html/2607.10110#S5.SS2.p2.2)\.
- M\. Schöne, B\. Rahmani, H\. Kremer, F\. Falck, H\. Ballani, and J\. Gladrow \(2025\)Implicit language models are rnns: balancing parallelization and expressivity\.External Links:2502\.07827,[Link](https://arxiv.org/abs/2502.07827)Cited by:[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px3.p1.1)\.
- C\. Snell, J\. Lee, K\. Xu, and A\. Kumar \(2024\)Scaling llm test\-time compute optimally can be more effective than scaling model parameters\.External Links:2408\.03314,[Link](https://arxiv.org/abs/2408.03314)Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p1.1)\.
- J\. Su, Y\. Lu, S\. Pan, A\. Murtadha, B\. Wen, and Y\. Liu \(2023\)RoFormer: enhanced transformer with rotary position embedding\.External Links:2104\.09864,[Link](https://arxiv.org/abs/2104.09864)Cited by:[§E\.4](https://arxiv.org/html/2607.10110#A5.SS4.SSS0.Px1.p1.1)\.
- H\. Touvron, L\. Martin, K\. Stone, P\. Albert, A\. Almahairi, Y\. Babaei, N\. Bashlykov, S\. Batra, P\. Bhargava, S\. Bhosale, D\. Bikel, L\. Blecher, C\. C\. Ferrer, M\. Chen, G\. Cucurull, D\. Esiobu, J\. Fernandes, J\. Fu, W\. Fu, B\. Fuller, C\. Gao, V\. Goswami, N\. Goyal, A\. Hartshorn, S\. Hosseini, R\. Hou, H\. Inan, M\. Kardas, V\. Kerkez, M\. Khabsa, I\. Kloumann, A\. Korenev, P\. S\. Koura, M\. Lachaux, T\. Lavril, J\. Lee, D\. Liskovich, Y\. Lu, Y\. Mao, X\. Martinet, T\. Mihaylov, P\. Mishra, I\. Molybog, Y\. Nie, A\. Poulton, J\. Reizenstein, R\. Rungta, K\. Saladi, A\. Schelten, R\. Silva, E\. M\. Smith, R\. Subramanian, X\. E\. Tan, B\. Tang, R\. Taylor, A\. Williams, J\. X\. Kuan, P\. Xu, Z\. Yan, I\. Zarov, Y\. Zhang, A\. Fan, M\. Kambadur, S\. Narang, A\. Rodriguez, R\. Stojnic, S\. Edunov, and T\. Scialom \(2023\)Llama 2: open foundation and fine\-tuned chat models\.External Links:2307\.09288,[Link](https://arxiv.org/abs/2307.09288)Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px1.p1.1)\.
- R\. Waleffe, W\. Byeon, D\. Riach, B\. Norick, V\. Korthikanti, T\. Dao, A\. Gu, A\. Hatamizadeh, S\. Singh, D\. Narayanan,et al\.\(2024\)An empirical study of mamba\-based language models\.arXiv preprint arXiv:2406\.07887\.Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p2.1),[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px1.p1.1)\.
- W\. Wang and F\. Reid \(2026\)Tiny recursive reasoning with mamba\-2 attention hybrid\.External Links:2602\.12078,[Link](https://arxiv.org/abs/2602.12078)Cited by:[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px2.p1.1)\.
- A\. Yang, A\. Li, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Gao, C\. Huang, C\. Lv, C\. Zheng, D\. Liu, F\. Zhou, F\. Huang, F\. Hu, H\. Ge, H\. Wei, H\. Lin, J\. Tang, J\. Yang, J\. Tu, J\. Zhang, J\. Yang, J\. Yang, J\. Zhou, J\. Zhou, J\. Lin, K\. Dang, K\. Bao, K\. Yang, L\. Yu, L\. Deng, M\. Li, M\. Xue, M\. Li, P\. Zhang, P\. Wang, Q\. Zhu, R\. Men, R\. Gao, S\. Liu, S\. Luo, T\. Li, T\. Tang, W\. Yin, X\. Ren, X\. Wang, X\. Zhang, X\. Ren, Y\. Fan, Y\. Su, Y\. Zhang, Y\. Zhang, Y\. Wan, Y\. Liu, Z\. Wang, Z\. Cui, Z\. Zhang, Z\. Zhou, and Z\. Qiu \(2025\)Qwen3 technical report\.External Links:2505\.09388,[Link](https://arxiv.org/abs/2505.09388)Cited by:[§1](https://arxiv.org/html/2607.10110#S1.p1.1)\.
- R\. Zellers, A\. Holtzman, Y\. Bisk, A\. Farhadi, and Y\. Choi \(2019\)Hellaswag: can a machine really finish your sentence?\.arXiv preprint arXiv:1905\.07830\.Cited by:[§5\.1](https://arxiv.org/html/2607.10110#S5.SS1.SSS0.Px3.p1.1)\.
- R\. Zhu, Z\. Wang, K\. Hua, T\. Zhang, Z\. Li, H\. Que, B\. Wei, Z\. Wen, F\. Yin, H\. Xing,et al\.\(2025\)Scaling latent reasoning via looped language models\.arXiv preprint arXiv:2510\.25741\.Cited by:[Appendix B](https://arxiv.org/html/2607.10110#A2.SS0.SSS0.Px1.p1.3),[Appendix B](https://arxiv.org/html/2607.10110#A2.SS0.SSS0.Px2.p1.1),[§1](https://arxiv.org/html/2607.10110#S1.p5.1),[§2](https://arxiv.org/html/2607.10110#S2.SS0.SSS0.Px4.p1.3),[§3\.1](https://arxiv.org/html/2607.10110#S3.SS1.SSS0.Px3.p1.1)\.
## Appendix AFull Results for iso\-FLOPs Evaluation
Table[5](https://arxiv.org/html/2607.10110#A1.T5)presents the full iso\-FLOPs evaluation results for Mamba\-2 and LLaMA\-2\-style Transformer models trained on 100B tokens\. PPL denotes validation perplexity, while Avg\. denotes the average score across the seven downstream tasks reported in the table\. These per\-task results complement the main\-text analysis by showing how each architecture and looping configuration performs under matched parameter and compute budgets\.
Table 5:iso\-FLOPs evaluation of Mamba\-2 and LLaMA2\-style Transformer models trained on 100B tokens\. PPL denotes validation perplexity, and Avg\. denotes the average score over the seven downstream tasks shown in the table\.Arch\.ModelParams / FLOPsPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.Mamba\-2Mamba\-2Base\(6⊗1\)\(6\\otimes 1\)6×/6×6\\times/6\\times21\.9651\.5239\.5750\.2826\.8026\.3627\.3925\.2535\.31Mamba\-2Loop\(6⊗4\)\(6\\otimes 4\)6×/24×6\\times/24\\times19\.9852\.7750\.6151\.1428\.0025\.9127\.9925\.2937\.39Mamba\-2Base\(12⊗1\)\(12\\otimes 1\)12×/12×12\\times/12\\times18\.6652\.8861\.9348\.8624\.2026\.0227\.3925\.5938\.12Mamba\-2Loop\(12⊗2\)\(12\\otimes 2\)12×/24×12\\times/24\\times18\.2149\.8962\.1149\.4923\.4025\.8028\.5826\.8538\.02Mamba\-2Base\(24⊗1\)\(24\\otimes 1\)24×/24×24\\times/24\\times16\.1551\.9057\.3748\.0726\.4025\.4628\.4125\.3837\.57LLaMA\-2TransformerBase\(6⊗1\)\(6\\otimes 1\)6×/6×6\\times/6\\times21\.4852\.9459\.3047\.9925\.8025\.7528\.9226\.6438\.19TransformerLoop\(6⊗4\)\(6\\otimes 4\)6×/24×6\\times/24\\times19\.7552\.6757\.9849\.6427\.8026\.1728\.2426\.6438\.45TransformerBase\(12⊗1\)\(12\\otimes 1\)12×/12×12\\times/12\\times18\.5351\.4762\.0550\.5926\.0026\.0229\.0125\.8438\.71TransformerLoop\(12⊗2\)\(12\\otimes 2\)12×/24×12\\times/24\\times18\.0652\.7262\.1750\.5928\.6025\.6227\.0526\.5639\.04TransformerBase\(24⊗1\)\(24\\otimes 1\)24×/24×24\\times/24\\times16\.4050\.9249\.5451\.1428\.8026\.0228\.7525\.2937\.21
## Appendix BDetails of Two\-Stage Exit\-Gate Training
#### Exit\-Gate Stage 1: adaptive exit\-state selection\.
Following the entropy\-regularized adaptive computation objective of Ouro\(Zhuet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib14)\), we attach an exit gate to the Looped Mamba model\. Unlike the final\-loop objective used in base pre\-training, this stage uses the outputs from all recurrent steps\. The exit gate operates at the token level: for token positioniiand recurrent steprrwithr=1,…,R−1r=1,\\ldots,R\-1, it predicts an instantaneous exit probability
λi\(r\)=σ\(Linearϕ\(hi\(r\)\)\)∈\(0,1\),\\lambda\_\{i\}^\{\(r\)\}=\\sigma\\\!\\left\(\\mathrm\{Linear\}\_\{\\phi\}\(h\_\{i\}^\{\(r\)\}\)\\right\)\\in\(0,1\),wherehi\(r\)h\_\{i\}^\{\(r\)\}is the hidden state of tokeniiafter therr\-th recurrent step,ϕ\\phidenotes the gate parameters, andσ\(⋅\)\\sigma\(\\cdot\)is the sigmoid function\. The final recurrent step does not require a separate gate, since it receives all remaining probability mass\.
For each token position, we define the survival probability as the probability of not exiting before a given recurrent step:
Si\(0\)=1,Si\(r\)=∏j=1r\(1−λi\(j\)\)\.S\_\{i\}^\{\(0\)\}=1,\\qquad S\_\{i\}^\{\(r\)\}=\\prod\_\{j=1\}^\{r\}\\left\(1\-\\lambda\_\{i\}^\{\(j\)\}\\right\)\.The probability that tokeniiexits exactly at recurrent steprris
πi\(r\)=pϕ\(zi=r∣x\)=λi\(r\)Si\(r−1\),r=1,…,R−1\.\\begin\{split\}\\pi\_\{i\}^\{\(r\)\}&=p\_\{\\phi\}\(z\_\{i\}=r\\mid x\)\\\\ &=\\lambda\_\{i\}^\{\(r\)\}S\_\{i\}^\{\(r\-1\)\},\\qquad r=1,\\ldots,R\-1\.\\end\{split\}\(1\)whereziz\_\{i\}denotes the latent exit step for tokenii\. For the final recurrent step, we assign the remaining probability mass:
πi\(R\)=pϕ\(zi=R∣x\)=Si\(R−1\)\.\\pi\_\{i\}^\{\(R\)\}=p\_\{\\phi\}\(z\_\{i\}=R\\mid x\)=S\_\{i\}^\{\(R\-1\)\}\.This defines a valid token\-level distribution over recurrent steps:
∑r=1Rπi\(r\)=1\.\\sum\_\{r=1\}^\{R\}\\pi\_\{i\}^\{\(r\)\}=1\.
Let
ℓi\(r\)=−logpθ\(r\)\(xi\+1∣x≤i\)\\ell\_\{i\}^\{\(r\)\}=\-\\log p\_\{\\theta\}^\{\(r\)\}\(x\_\{i\+1\}\\mid x\_\{\\leq i\}\)be the token\-level next\-token loss at recurrent steprr\. To avoid collapse to a single recurrent depth, we add an entropy regularizer over the token\-level exit distribution:
Hi=−∑r=1Rπi\(r\)logπi\(r\)\.H\_\{i\}=\-\\sum\_\{r=1\}^\{R\}\\pi\_\{i\}^\{\(r\)\}\\log\\pi\_\{i\}^\{\(r\)\}\.The Stage 1 objective is therefore
ℒstage1=𝔼x\[1M∑i=1M\(∑r=1Rπi\(r\)ℓi\(r\)−βHi\)\],\\mathcal\{L\}\_\{\\mathrm\{stage1\}\}=\\mathbb\{E\}\_\{x\}\\left\[\\frac\{1\}\{M\}\\sum\_\{i=1\}^\{M\}\\left\(\\sum\_\{r=1\}^\{R\}\\pi\_\{i\}^\{\(r\)\}\\ell\_\{i\}^\{\(r\)\}\-\\beta H\_\{i\}\\right\)\\right\],whereβ\\betacontrols the strength of entropy regularization, andMMdenotes the number of valid next\-token prediction positions\. This objective encourages the gate to explore different recurrent depths while weighting each recurrent\-step loss by the probability of exiting at that step\.
At inference time, we use the cumulative exit probability
CDFi\(r∣x\)=∑j=1rπi\(j\)\\mathrm\{CDF\}\_\{i\}\(r\\mid x\)=\\sum\_\{j=1\}^\{r\}\\pi\_\{i\}^\{\(j\)\}to determine the selected output step\. Given a thresholdq∈\[0,1\]q\\in\[0,1\], tokeniiselects recurrent step
zi\(q\)=min\{r∈\{1,…,R\}:CDFi\(r∣x\)≥q\}\.z\_\{i\}\(q\)=\\min\\left\\\{r\\in\\\{1,\\ldots,R\\\}:\\mathrm\{CDF\}\_\{i\}\(r\\mid x\)\\geq q\\right\\\}\.The thresholdqqcontrols the selected\-depth–accuracy trade\-off: smaller values favor earlier recurrent\-step outputs, while larger values select later recurrent\-step outputs\.
#### Exit\-Gate Stage 2: focused exit\-gate training\.
Following the focused adaptive gate training stage of Ouro\(Zhuet al\.,[2025](https://arxiv.org/html/2607.10110#bib.bib14)\), we further calibrate the exit gate after entropy\-regularized training\. In this stage, the Mamba backbone and the language\-modeling head are frozen, and only the exit\-gate parametersϕ\\phiare updated\. The goal is to train the gate to decide whether continuing to the next recurrent step is useful, based on the realized improvement in next\-token loss\.
We first detach the per\-step token losses from the computation graph:
ℓ¯i\(r\)=sg\(ℓi\(r\)\),\\bar\{\\ell\}\_\{i\}^\{\(r\)\}=\\mathrm\{sg\}\\\!\\left\(\\ell\_\{i\}^\{\(r\)\}\\right\),wheresg\(⋅\)\\mathrm\{sg\}\(\\cdot\)denotes the stop\-gradient operator\. For each non\-final recurrent stepr∈\{1,…,R−1\}r\\in\\\{1,\\ldots,R\-1\\\}, we define the improvement obtained by continuing from steprrto stepr\+1r\+1as
Ii\(r\)=max\(0,ℓ¯i\(r\)−ℓ¯i\(r\+1\)\)\.I\_\{i\}^\{\(r\)\}=\\max\\left\(0,\\,\\bar\{\\ell\}\_\{i\}^\{\(r\)\}\-\\bar\{\\ell\}\_\{i\}^\{\(r\+1\)\}\\right\)\.A largerIi\(r\)I\_\{i\}^\{\(r\)\}means that the additional recurrent step reduces the loss, suggesting that the model should continue rather than exit\.
We convert this improvement into a soft continuation target:
wi\(r\)=σ\(κ\(Ii\(r\)−γ\)\),w\_\{i\}^\{\(r\)\}=\\sigma\\\!\\left\(\\kappa\\left\(I\_\{i\}^\{\(r\)\}\-\\gamma\\right\)\\right\),whereκ\\kappacontrols the sharpness of the target andγ\\gammais the minimum improvement threshold\. When the improvement is larger thanγ\\gamma,wi\(r\)w\_\{i\}^\{\(r\)\}is close to 1 and the gate is encouraged to continue; otherwise, it is encouraged to exit\.
Sinceλi\(r\)\\lambda\_\{i\}^\{\(r\)\}is the predicted exit probability, the predicted continuation probability is1−λi\(r\)1\-\\lambda\_\{i\}^\{\(r\)\}\. We train the gate by matching this predicted continuation probability towi\(r\)w\_\{i\}^\{\(r\)\}using binary cross\-entropy:
ℒadapt\(r\)=−1M∑i=1M\[wi\(r\)log\(1−λi\(r\)\)\+\(1−wi\(r\)\)logλi\(r\)\]\.\\begin\{split\}\\mathcal\{L\}\_\{\\mathrm\{adapt\}\}^\{\(r\)\}=\-\\frac\{1\}\{M\}\\sum\_\{i=1\}^\{M\}\\Big\[&w\_\{i\}^\{\(r\)\}\\log\\\!\\left\(1\-\\lambda\_\{i\}^\{\(r\)\}\\right\)\\\\ &\+\\left\(1\-w\_\{i\}^\{\(r\)\}\\right\)\\log\\lambda\_\{i\}^\{\(r\)\}\\Big\]\.\\end\{split\}\(2\)
The total Stage 2 objective averages over all non\-final recurrent steps:
ℒstage2=1R−1∑r=1R−1ℒadapt\(r\)\.\\mathcal\{L\}\_\{\\mathrm\{stage2\}\}=\\frac\{1\}\{R\-1\}\\sum\_\{r=1\}^\{R\-1\}\\mathcal\{L\}\_\{\\mathrm\{adapt\}\}^\{\(r\)\}\.This objective penalizes early exits when another recurrent step would substantially reduce the loss, while also penalizing unnecessary continuation when the additional step provides little improvement\. At inference time, the calibrated gate is used with the same cumulative\-probability threshold rule as in Stage 1\.
## Appendix CMano: Knowledge Manipulation
#### Dataset\.
We use the Mano task to evaluate the model’s ability to manipulate knowledge stored in its parameters without explicit intermediate reasoning\. The dataset consists of arithmetic expressions over the finite field𝔽23\\mathbb\{F\}\_\{23\}\. Each instance corresponds to a tree\-structured expression withℓ\\ellbinary operations, whereℓ≤L\\ell\\leq LandLLdenotes the maximum expression length\. During training,ℓ\\ellis uniformly sampled from\[1,L\]\[1,L\]; during evaluation, we use only the hardest expressions withℓ=L\\ell=L\. The expressions are presented in prefix notation\. A length\-3 instance is written as
<bos\> <len\_3\> \- \* a b \+ c d <ans\> ans\.\\texttt\{<bos\> <len\\\_3\> \- \* a b \+ c d <ans\> ans\}\.All operations are performed in𝔽23\\mathbb\{F\}\_\{23\}, and the task uses only the operators\{\+,−,∗\}\\\{\+,\-,\\ast\\\}\. The vocabulary consists of the numbers0,…,220,\\ldots,22, the three operators, the special tokens<bos\>,<ans\>,<eos\>, and<pad\>, and length tokens<len\_i\>fori∈\[0,L\]i\\in\[0,L\]\.
#### Model\.
We replace the Transformer blocks in the looped\-model architecture with Mamba\-2 blocks and study recurrent parameter sharing in this architecture\. A model is denoted byN⊗RN\\otimes R, whereNNis the number of independently parameterized Mamba layers andRRis the number of recurrent steps\. Thus,R=1R=1corresponds to a standard non\-looped Mamba model, whileR\>1R\>1corresponds to a Looped Mamba model\. All models use hidden dimension 1024\. Training and evaluation are performed using only the output from the final recurrent step\. Since our Mamba implementation disables the feed\-forward sublayer, we use two consecutive Mamba\-2 mixer blocks as a proxy for the two\-sub\-layer structure of a Transformer block, namely attention followed by an MLP\. This design matches the number of residual transformations at a coarse level, but should not be interpreted as an exact architectural equivalence\. We therefore report results under explicit effective\-depth, iso\-parameter, and iso\-FLOPs comparisons\.
#### Training details\.
We use the AdamW optimizer withβ1,β2=0\.9,0\.98\\beta\_\{1\},\\beta\_\{2\}=0\.9,0\.98,ϵ=10−6\\epsilon=10^\{\-6\}, weight decay 0\.1, and gradient clipping with maximum norm 1\.0\. We employ 1000 warmup steps followed by a cosine learning\-rate schedule whose minimum learning rate is 0\.1 of the peak learning rate\. Training uses bf16, sequence packing, context length 1024, and global batch size 128\. For the three difficulty levels, we train for 80K, 110K, and 200K steps, respectively\. The corresponding peak learning rates are5⋅10−45\\cdot 10^\{\-4\},1⋅10−41\\cdot 10^\{\-4\}, and5⋅10−55\\cdot 10^\{\-5\}\. Training samples are generated online, and multiple Mano instances are packed into each 1024\-token chunk\. For each configuration, we train the model three times with different random seeds and report the run with the maximum average evaluation accuracy on the test set\. We adopt this best\-of\-three reporting protocol because our goal is to study the expressivity power of the model under a given configuration, rather than to characterize average training stability across random seeds\.
#### Evaluation\.
During evaluation, each sample contains a single expression of lengthℓ=L\\ell=L, without packing\. Since the final answer is a single token in𝔽23\\mathbb\{F\}\_\{23\}, we report exact\-match accuracy at the answer position\. Specifically, the model predicts the single answer token following<ans\>, and the prediction is counted as correct if it matches the ground\-truth answer\. For each difficulty level, we evaluate on 2000 deterministic samples generated with a fixed random seed\.
## Appendix Dp\-hop Induction
#### Dataset\.
We use the p\-hop induction task to evaluate recursive contextual retrieval\. Given a sequencev=\(v1,…,vn\)v=\(v\_\{1\},\\ldots,v\_\{n\}\)of lengthn=256n=256, each token is sampled from an alphabetΣ\\Sigmaof size 4\. A 1\-hop induction step starts from the final position, finds a previous occurrence of the target token, and returns the token immediately following that occurrence\. The p\-hop task recursively applies this operationpptimes\. We evaluate two difficulty levels,p∈\{16,32\}p\\in\\\{16,32\\\}\. To ensure that every instance has a valid answer, our implementation uses a chain\-first construction: it first samples a valid hop chain of lengthpp, fills the remaining positions subject to constraints, and then verifies that the actual hop trace exactly matches the constructed chain\.
#### Model\.
We use the same Looped Mamba architecture as in the Mano experiments\. A model is denoted byN⊗RN\\otimes R, whereNNis the number of independently parameterized Mamba\-2 layers andRRis the number of recurrent steps\. Therefore, the number of parameters scales withNN, while FLOPs scale with the effective depthNRNR\. All p\-hop experiments use hidden dimensiond=256d=256\. The final recurrent step is used for both training loss and evaluation\.
#### Training details\.
We train with AdamW, weight decay 0\.1, and maximum gradient norm 1\.0\. We use 2000 warmup steps followed by a cosine learning\-rate schedule\. The peak learning rate is2×10−42\\times 10^\{\-4\}, and the minimum learning rate is2×10−52\\times 10^\{\-5\}, i\.e\., 0\.1 of the peak value\. Training uses bf16\. The context length is set to 272\. Since each p\-hop sample has length1\+256\+3=2601\+256\+3=260, this setting effectively corresponds to using one p\-hop instance per training sequence rather than packing multiple instances into a longer context\. The loss is applied only to the answer token position\. For each configuration, we train the model three times with different random seeds and report the run with the maximum average evaluation accuracy on the test set\. We adopt this best\-of\-three reporting protocol because our goal is to study the expressivity power of the model under a given configuration, rather than to characterize average training stability across random seeds\.
#### Evaluation\.
The evaluation set contains 10,000 deterministic samples generated with a fixed random seed\. During evaluation, each sample is fed independently into the model, and we compute exact\-match accuracy at the answer position after<ans\>\. Specifically, the prediction is the argmax of the logits at the answer position, and it is counted as correct if it matches the ground\-truth answeryy\. Since the answer is one of four alphabet tokens, random guessing gives approximately 25% accuracy\.
## Appendix ELanguage Model Training Details
### E\.1Common Language Model Pre\-training Details
#### Overall procedure\.
Our language\-model experiments are organized into three steps\. First, we pre\-train fixed\-loop language models without exit gates\. This pre\-training setting is shared by the LLaMA\-2 Transformer, Mamba\-2, and Mamba\-2 Hybrid models described in Sections[E\.2](https://arxiv.org/html/2607.10110#A5.SS2),[E\.3](https://arxiv.org/html/2607.10110#A5.SS3), and[E\.4](https://arxiv.org/html/2607.10110#A5.SS4), respectively\. Second, we evaluate the pretrained fixed\-loop models under the corresponding architecture and loop configurations\. Third, for the exit\-gate experiments, we initialize from a pretrained fixed\-loop checkpoint and continue training with the exit gate enabled, as described in Section[E\.5](https://arxiv.org/html/2607.10110#A5.SS5)\.
#### Pre\-training configuration\.
Unless otherwise specified, the 140M models usedmodel=768d\_\{\\mathrm\{model\}\}=768and 24 layers\. Larger 370M models modify only the architecture\-specific dimensions and depth, as listed in the corresponding architecture tables\. The shared pre\-training hyperparameters are summarized in Table[6](https://arxiv.org/html/2607.10110#A5.T6)\.
ParameterValueContext length2048Vocabulary size32,000OptimizerAdamWAdamW betas\(0\.9,0\.95\)\(0\.9,0\.95\)Weight decay0\.1Learning rate schedulecosineMaximum scheduler LR8\.0×10−48\.0\\times 10^\{\-4\}Minimum scheduler LR8\.0×10−58\.0\\times 10^\{\-5\}Warmup ratio0\.1Gradient clipping1\.0Precisionbfloat16Random seed3407Table 6:Common pre\-training hyperparameters for the fixed\-loop language models\. The LLaMA\-2 Transformer, Mamba\-2, and Mamba\-2 Hybrid models are first pretrained under this shared setting before any exit\-gate continued training\.
### E\.2LLaMA\-2 Transformer Details
#### Architecture\.
For the LLaMA\-2 Transformer experiments, we use the Transformer ablation branch of our looped language model implementation\. The model replaces the Mamba\-2 mixer with LLaMA\-style causal self\-attention while keeping the same looped backbone interface\. Unless otherwise specified, the Transformer models use a model dimension ofdmodel=768d\_\{\\mathrm\{model\}\}=768\. To make the Transformer baseline comparable to the Mamba\-2 model withdmodel=768d\_\{\\mathrm\{model\}\}=768under a similar parameter budget, we set the gated MLP intermediate dimension todintermediate=640d\_\{\\mathrm\{intermediate\}\}=640\. This intermediate dimension is not intended to follow the standard LLaMA MLP expansion ratio; instead, it is chosen as a parameter\-matching configuration for comparison with Mamba\-2\. The configuration is summarized in Table[7](https://arxiv.org/html/2607.10110#A5.T7)\.
ParameterValueArchitectureLLaMA\-2dmodeld\_\{\\mathrm\{model\}\}768dintermediated\_\{\\mathrm\{intermediate\}\}640Layers6, 12, or 24Attention head dimension64Number of attention heads12Position embeddingRoPENormalizationRMSNormFused add\-normyesResidual in FP32yesTied embeddingsnoContext length2048Vocabulary size32,000Table 7:LLaMA\-2 Transformer architecture configuration\.
### E\.3Mamba\-2 Details
#### Architecture\.
For the Mamba\-2 experiments, we use the same looped language model interface as the LLaMA\-2 Transformer baseline, but replace the Transformer self\-attention mixer with a Mamba\-2 mixer\. The 140M Mamba\-2 model usesdmodel=768d\_\{\\mathrm\{model\}\}=768with 24 layers, while the 370M model increases the model dimension todmodel=1024d\_\{\\mathrm\{model\}\}=1024and uses 48 layers\. For both sizes, we set the external MLP intermediate dimension todintermediate=0d\_\{\\mathrm\{intermediate\}\}=0, since the Mamba\-2 block already contains its own sequence mixing and channel projection components\. The configuration is summarized in Table[8](https://arxiv.org/html/2607.10110#A5.T8)\.
Parameter140M370MArchitectureMamba\-2Mamba\-2dmodeld\_\{\\mathrm\{model\}\}7681024Number of layers2448External MLPnonoAttention layersnonoNormalizationRMSNormRMSNormFused add\-normyesyesResidual in FP32yesyesTied embeddingsnonoContext length20482048Vocabulary size32,00032,000Table 8:Mamba\-2 architecture configurations\.
### E\.4Mamba\-2 Hybrid Details
#### Architecture\.
For the Mamba\-2 Hybrid experiments, we use the same looped language model interface as the Mamba\-2 models, but replace a small number of Mamba\-2 layers with causal attention layers\. The attention layers are specified by the 0\-based indices inattn\_layer\_idx; all remaining layers use the Mamba\-2 mixer\. The detailed configurations are summarized in Table[9](https://arxiv.org/html/2607.10110#A5.T9)\. The hybrid models do not use learned position embeddings, while the inserted attention layers use RoPE\(Suet al\.,[2023](https://arxiv.org/html/2607.10110#bib.bib32)\)with rotary dimension 64\.
140M370Mdmodeld\_\{\\mathrm\{model\}\}7681024Layers2448Attention indices\[5,10,15,20\]\[5,10,15,20\]\[6,18,30,42\]\[6,18,30,42\]Attention heads12×6412\\times 6416×6416\\times 64RoPE dim\.6464Table 9:Mamba\-2 Hybrid architecture configurations\. Attention layer indices are 0\-based\. All non\-attention layers are Mamba\-2 layers\. The models use no external MLP and no learned position embedding; RoPE is used only in the inserted attention layers\.
### E\.5Continued Training: Exit\-Gate Stage\-1 Training Details
#### Training procedure\.
For the exit\-gate experiments, we initialize from a pretrained fixed\-loop checkpoint and continue training with the exit gate enabled\. The backbone architecture is kept unchanged, while a lightweight linear gate is added on top of the hidden states to predict whether each token should exit at each recurrent loop\. The exit\-gate bias is initialized to−2\.0\-2\.0, giving a conservative initial exit probability\. The configuration is summarized in Table[10](https://arxiv.org/html/2607.10110#A5.T10)\.
#### β\\betaablation\.
We conduct a stage\-1 ablation on the entropy coefficientβ\\betafor the exit\-gate continuation training\. In this ablation, we focus on the loop\-4 setting with the maximum number of recurrent loopsR=4R=4\. For each token, the exit gate defines an exit distribution over recurrent loops:
pi,r=λi,r∏j<r\(1−λi,j\),λi,r=σ\(g\(hi,r\)\),p\_\{i,r\}=\\lambda\_\{i,r\}\\prod\_\{j<r\}\(1\-\\lambda\_\{i,j\}\),\\qquad\\lambda\_\{i,r\}=\\sigma\(g\(h\_\{i,r\}\)\),forr<Rr<R, and the remaining probability mass is assigned to the final loop:
pi,R=∏j<R\(1−λi,j\)\.p\_\{i,R\}=\\prod\_\{j<R\}\(1\-\\lambda\_\{i,j\}\)\.The exit entropy for tokeniiis computed as
Hi=−∑r=1Rpi,rlog\(pi,r\+ϵ\),H\_\{i\}=\-\\sum\_\{r=1\}^\{R\}p\_\{i,r\}\\log\(p\_\{i,r\}\+\\epsilon\),and the exit entropy is averaged over valid tokens:
H¯exit=1\|𝒱\|∑i∈𝒱Hi,\\bar\{H\}\_\{\\mathrm\{exit\}\}=\\frac\{1\}\{\|\\mathcal\{V\}\|\}\\sum\_\{i\\in\\mathcal\{V\}\}H\_\{i\},where𝒱\\mathcal\{V\}denotes the set of valid non\-padding tokens\. ForR=4R=4, the maximum possible entropy is
Hmax=logR=log4≈1\.386,H\_\{\\max\}=\\log R=\\log 4\\approx 1\.386,which is achieved when the exit distribution is uniform over the four loops\. In practice, however, a fully uniform exit distribution is not ideal because it does not sufficiently favor later, more accurate loops\. We therefore regard an exit entropy around1\.01\.0as a desirable regime: the distribution remains non\-collapsed, while still allocating more probability mass to later recurrent loops\. We also monitor the expected selected exit step\. For a tokenii, this is
si=∑r=1Rrpi,r\.s\_\{i\}=\\sum\_\{r=1\}^\{R\}r\\,p\_\{i,r\}\.The logged metric is the average of this quantity over valid tokens:
s¯=1\|𝒱\|∑i∈𝒱si,\\bar\{s\}=\\frac\{1\}\{\|\\mathcal\{V\}\|\}\\sum\_\{i\\in\\mathcal\{V\}\}s\_\{i\},For the loop\-4 exit\-gate models, we find that an average expected step around3\.23\.2is a good operating point, since it provides adaptive exit\-state selection while favoring later recurrent\-step outputs\. The continuation stage is sensitive to the entropy coefficientβ\\beta: too small a value can lead to a collapsed exit distribution, whereas too large a value can over\-regularize the exit behavior\. Based on the stage\-1 ablation, we useβ=0\.25\\beta=0\.25as the final condition for the exit\-gate continuation experiments\.
ParameterValueExit\-gate bias initialization−2\.0\-2\.0Entropy coefficientβ\\beta0\.25Learning rate schedulecosineMaximum learning rate1\.5×10−41\.5\\times 10^\{\-4\}Minimum learning rate8\.0×10−58\.0\\times 10^\{\-5\}Warmup ratio0\.002OptimizerAdamWAdamW betas\(0\.9,0\.95\)\(0\.9,0\.95\)Weight decay0\.1Gradient clipping0\.6Table 10:Exit\-gate continuation training configuration\. The exit\-gate models are initialized from pretrained fixed\-loop checkpoints and trained with an expected loop loss plus entropy regularization\.Figure 2:Exit entropy during stage\-1 exit\-gate continuation training for different entropy coefficientsβ\\beta\. All curves are truncated to the common training range\. ForR=4R=4, the maximum entropy islog4≈1\.386\\log 4\\approx 1\.386, and we use an entropy around1\.01\.0as a practical reference for a non\-collapsed exit distribution\.Figure 3:Expected selected exit step during stage\-1 exit\-gate continuation training for different entropy coefficientsβ\\beta\. We use this metric as a diagnostic for the prediction depth selected by the exit gate\.
### E\.6Continued Training: Exit\-Gate Stage\-2 Training Details
#### Training procedure\.
In Stage\-2, we further refine the exit gate after the Stage\-1 continuation training\. We initialize the model from the Stage\-1 exit\-gate checkpoint and keep the architecture unchanged\. Unlike Stage\-1, all parameters except the exit gate are frozen, including the token embedding, recurrent backbone, final normalization, and language modeling head\. Thus, Stage\-2 only updates the lightweight exit\-gate parameters\. The configuration is summarized in Table[11](https://arxiv.org/html/2607.10110#A5.T11)\.
ParameterValueImprovement thresholdγ\\gamma0\.005Sharpnesskk50Learning rate schedulecosineMaximum learning rate1\.5×10−41\.5\\times 10^\{\-4\}Minimum learning rate1\.5×10−51\.5\\times 10^\{\-5\}Warmup ratio0\.01OptimizerAdamWAdamW betas\(0\.9,0\.95\)\(0\.9,0\.95\)Weight decay0\.1Gradient clipping1\.0Table 11:Exit\-gate Stage\-2 continued training configuration\. Stage\-2 freezes the pretrained language model and refines only the exit gate using an adaptive continuation objective\.
## Appendix FFixed\-Loop Inference Ablation
Tables[12](https://arxiv.org/html/2607.10110#A6.T12)–[15](https://arxiv.org/html/2607.10110#A6.T15)report the fixed\-loop inference ablation results for the 140M and 370M Mamba\-2 and Mamba\-2 Hybrid looped models\. Each model is trained as a 4\-loop model without an exit gate and evaluated by forcingr∈\{1,2,3,4\}r\\in\\\{1,2,3,4\\\}recurrent loops at inference time\. PPL denotes validation perplexity, and Avg\. denotes the average score over the seven downstream tasks shown in the tables\. Across all models, forcing fewer than four loops yields severely degraded perplexity, indicating that intermediate recurrent steps are not calibrated with the language\-modeling head under the final\-loop pre\-training objective; this motivates the gate\-based adaptive\-depth training described in Appendix[B](https://arxiv.org/html/2607.10110#A2)\.
Table 12:Fixed\-loop inference ablation for the 140M Mamba\-2 looped model\. The model is trained as a 4\-loop model without an exit gate, and evaluated by forcing different numbers of recurrent loops at inference time\. Avg\. denotes the average score over the seven downstream tasks shown in the table\.Forced LoopsPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.146374\.6252\.3461\.9647\.6725\.2025\.4927\.0524\.0337\.68212906\.8651\.4161\.9347\.5925\.0025\.2228\.2423\.7437\.593464\.8252\.3461\.4149\.4925\.6025\.9227\.1324\.9238\.12414\.6852\.3458\.8748\.2227\.0025\.2128\.0724\.8337\.79
Table 13:Fixed\-loop inference ablation for the 140M Mamba\-2 Hybrid looped model\. The model is trained as a 4\-loop model without an exit gate, and evaluated by forcing different numbers of recurrent loops at inference time\. Avg\. denotes the average score over the seven downstream tasks shown in the table\.Forced LoopsPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.124061\.8753\.2137\.8350\.2025\.8026\.0627\.7326\.0135\.26212304\.0052\.1837\.8351\.2225\.8025\.6229\.0126\.1435\.403508\.3852\.6137\.8350\.6725\.8025\.6827\.6525\.8035\.15414\.3751\.3145\.8747\.4326\.8025\.0728\.9226\.0935\.93
Table 14:Fixed\-loop inference ablation for the 370M Mamba\-2 looped model\. The model is trained as a 4\-loop model without an exit gate, and evaluated by forcing different numbers of recurrent loops at inference time\. Avg\. denotes the average score over the seven downstream tasks shown in the table\.Forced LoopsPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.11933\.7050\.9248\.9949\.6426\.0024\.1627\.9026\.1436\.2521906\.4951\.6937\.8349\.4927\.2021\.9728\.4125\.4234\.573832\.5351\.0938\.0450\.4325\.4023\.0728\.7526\.0534\.69412\.0651\.9061\.9648\.1527\.8024\.7229\.1025\.9338\.51
Table 15:Fixed\-loop inference ablation for the 370M Mamba\-2 Hybrid looped model\. The model is trained as a 4\-loop model without an exit gate, and evaluated by forcing different numbers of recurrent loops at inference time\. Avg\. denotes the average score over the seven downstream tasks shown in the table\.Forced LoopsPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.13282\.6051\.4760\.8049\.0126\.0024\.5426\.6225\.1737\.6624692\.7849\.8944\.3150\.0425\.6021\.2727\.3024\.8334\.7531660\.1751\.4137\.7751\.1425\.4023\.2327\.8225\.0434\.54411\.5951\.7461\.2249\.3328\.0025\.9729\.8626\.2638\.91
## Appendix GExit\-Gate Threshold Sweeps
We provide the full exit\-gate threshold sweep results for the four looped models studied in the main text\. Starting from fixed\-loop checkpoints, we attach a lightweight exit gate and continue training the models with the adaptive\-depth objectives described in Appendix[B](https://arxiv.org/html/2607.10110#A2)\. We report results for both Stage 1 and Stage 2 exit\-gate training\. For each stage, we evaluate inference with threshold valuesq∈\{0\.2,0\.4,0\.6,0\.8,1\.0\}q\\in\\\{0\.2,0\.4,0\.6,0\.8,1\.0\\\}, where smaller thresholds select earlier recurrent\-step outputs, whereas larger thresholds select later recurrent\-step outputs\.
Tables[16](https://arxiv.org/html/2607.10110#A7.T16),[17](https://arxiv.org/html/2607.10110#A7.T17),[18](https://arxiv.org/html/2607.10110#A7.T18), and[19](https://arxiv.org/html/2607.10110#A7.T19)report the downstream performance of the 140M Mamba\-2, 140M Mamba\-2 Hybrid, 370M Mamba\-2, and 370M Mamba\-2 Hybrid models, respectively\. These sweeps illustrate how the exit threshold controls the selected\-depth–performance trade\-off after gate training, and show that the best threshold is not universal across model scales, architectures, or downstream tasks\. As in the fixed\-loop ablation, different tasks can prefer different effective recurrent depths, which motivates adaptive exit instead of using a single fixed loop count for all inputs\.
Table 16:Threshold sweep for the 140M Mamba\-2 exit\-gated model\. We report Stage 1 and Stage 2 exit\-gate models withq∈\{0\.2,0\.4,0\.6,0\.8,1\.0\}q\\in\\\{0\.2,0\.4,0\.6,0\.8,1\.0\\\}\. Iter denotes the average exit step selected by the gate\.StageqqIterPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.Stage 11\.004\.0015\.5452\.1861\.5950\.5126\.8024\.6028\.0725\.0038\.39Stage 10\.804\.0015\.5452\.1861\.5950\.0426\.6025\.2727\.9025\.1738\.39Stage 10\.603\.4315\.5451\.4762\.2347\.7526\.8024\.6729\.7825\.0438\.25Stage 10\.402\.8215\.5752\.5062\.2350\.6724\.6026\.1228\.8425\.0438\.57Stage 10\.202\.3816\.2252\.7262\.1747\.0424\.4026\.1428\.5024\.7537\.96Stage 21\.004\.0015\.5452\.1861\.5949\.8826\.6025\.3227\.9025\.1738\.38Stage 20\.803\.0015\.5452\.9962\.2350\.2026\.0026\.0128\.9225\.7238\.87Stage 20\.602\.6315\.7853\.0562\.1749\.0924\.6026\.3628\.3325\.2138\.40Stage 20\.402\.2716\.5052\.6162\.1748\.8624\.2026\.1128\.5824\.7138\.18Stage 20\.201\.9218\.1452\.3962\.1750\.2026\.4026\.3427\.4724\.4138\.48
Table 17:Threshold sweep for the 140M Mamba\-2 Hybrid exit\-gated model\. We report Stage 1 and Stage 2 exit\-gate models withq∈\{0\.2,0\.4,0\.6,0\.8,1\.0\}q\\in\\\{0\.2,0\.4,0\.6,0\.8,1\.0\\\}\. Iter denotes the average exit step selected by the gate\.StageqqIterPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.Stage 11\.004\.0015\.0750\.7147\.5849\.8026\.4025\.3627\.9026\.2636\.29Stage 10\.804\.0015\.0750\.7147\.6149\.3326\.4025\.2827\.9926\.2636\.23Stage 10\.603\.5515\.0550\.7648\.0149\.2527\.0025\.6128\.2426\.5636\.49Stage 10\.402\.8115\.0751\.7454\.7749\.9626\.8025\.2928\.1625\.3437\.44Stage 10\.202\.3515\.6652\.3940\.5550\.2026\.2025\.2928\.5025\.8435\.57Stage 21\.004\.0015\.0750\.6047\.5550\.0426\.4025\.3427\.9926\.3536\.32Stage 20\.803\.2915\.0750\.3854\.8049\.2526\.4025\.3927\.7326\.6837\.23Stage 20\.602\.6415\.2851\.4754\.6547\.9127\.0025\.1927\.9025\.4237\.08Stage 20\.402\.2715\.9552\.7238\.2650\.4326\.0025\.5728\.8425\.6735\.36Stage 20\.201\.9317\.5652\.6137\.9849\.7224\.8025\.6628\.5026\.0135\.04
Table 18:Threshold sweep for the 370M Mamba\-2 exit\-gated model\. We report Stage 1 and Stage 2 exit\-gate models withq∈\{0\.2,0\.4,0\.6,0\.8,1\.0\}q\\in\\\{0\.2,0\.4,0\.6,0\.8,1\.0\\\}\. Iter denotes the average exit step selected by the gate\.StageqqIterPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.Stage 11\.004\.0012\.5552\.5661\.0450\.0426\.4025\.1928\.2426\.3038\.54Stage 10\.804\.0012\.5552\.7261\.1050\.0426\.4025\.1327\.9026\.2238\.50Stage 10\.603\.4012\.5452\.0759\.2449\.8824\.4025\.0926\.9626\.4737\.73Stage 10\.402\.6412\.5850\.7155\.9949\.4124\.0025\.1426\.9626\.3036\.93Stage 10\.201\.9113\.2151\.1457\.8349\.8824\.0024\.8128\.4126\.3037\.48Stage 21\.004\.0012\.5552\.5060\.8050\.0426\.4025\.1427\.3026\.2238\.34Stage 20\.803\.1812\.5552\.3457\.6150\.5924\.8024\.7726\.8826\.0937\.58Stage 20\.602\.5312\.6950\.5456\.2150\.7524\.2025\.4928\.0726\.0537\.33Stage 20\.401\.9413\.1750\.8757\.4949\.4924\.0024\.7627\.3025\.1337\.01Stage 20\.201\.6714\.4150\.6057\.2249\.4924\.6024\.3227\.9025\.5937\.10
Table 19:Threshold sweep for the 370M Mamba\-2 Hybrid exit\-gated model\. We report Stage 1 and Stage 2 exit\-gate models withq∈\{0\.2,0\.4,0\.6,0\.8,1\.0\}q\\in\\\{0\.2,0\.4,0\.6,0\.8,1\.0\\\}\. Iter denotes the average exit step selected by the gate\.StageqqIterPPLPIQABoolQWGOBQAHSARC\-CARC\-EAvg\.Stage 11\.004\.0012\.1352\.0157\.0051\.1427\.8025\.6330\.1225\.4638\.45Stage 10\.804\.0012\.1352\.2956\.8851\.1427\.8025\.7330\.1225\.5938\.51Stage 10\.603\.3412\.1251\.0339\.5450\.6727\.6026\.3630\.6325\.8435\.95Stage 10\.402\.7412\.1551\.0339\.4249\.4128\.8025\.2431\.4825\.3435\.82Stage 10\.202\.2412\.6651\.3142\.0250\.5929\.4025\.3931\.1425\.8836\.53Stage 21\.004\.0012\.1352\.1857\.0051\.3827\.8025\.6330\.1225\.5938\.53Stage 20\.803\.0112\.1251\.0939\.4550\.0427\.4025\.8430\.2926\.0935\.74Stage 20\.602\.5712\.2151\.1439\.5150\.0428\.2025\.3631\.2325\.8035\.90Stage 20\.402\.2212\.6051\.6941\.9648\.8629\.4025\.1531\.0625\.8836\.29Stage 20\.201\.8813\.6251\.2041\.9950\.2829\.2025\.0330\.9725\.8036\.35Similar Articles
Looped World Models
Looped World Models introduce iterative latent state refinement through shared transformer blocks, achieving 100x parameter efficiency while adapting computational depth to prediction complexity.
DiscoLoop: Looping Discrete Embeddings and Continuous Hidden States for Multi-hop Reasoning
DiscoLoop introduces a looping architecture that carries both discrete embedding and continuous hidden-state channels to improve multi-hop reasoning in transformers, achieving near-perfect accuracy on synthetic tasks and stronger performance on real-world language modeling.
Solve the Loop: Attractor Models for Language and Reasoning
This paper introduces Attractor Models, which use fixed-point solving and implicit differentiation for efficient iterative refinement, achieving superior language modeling and reasoning performance with reduced computational costs compared to traditional transformers.
Memory-Efficient Looped Transformer: Decoupling Compute from Memory in Looped Language Models
Proposes Memory-Efficient Looped Transformer (MELT), a novel recurrent LLM architecture that decouples reasoning depth from memory consumption by sharing a single KV cache across loops and using chunk-wise training with interpolated transition and attention-aligned distillation.
Memory-Efficient Meta-Reinforcement Learning for Adaptive Safety-Critical Control in Adversarial Spacecraft Proximity Operations
This paper investigates memory-efficient meta-reinforcement learning architectures for adaptive safety-critical control in adversarial spacecraft proximity operations, finding that state space models like Mamba with PPO achieve superior task completion, safety, and fuel savings compared to LSTM and GRU.