Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation — part 3/6

Cleaned-text code points [30708, 46201). Bibliography omitted here; source.md includes it.

As illustrated in 5(a) , MoR exhibits a distinct compute-optimal scaling behavior compared to baselines under isoFLOPs constraints. The flatter slope of MoR’s optimal path (a line connecting stars) indicates that it benefits more significantly from increases in parameter count (i.e., less data-hungry). This is likely because the performance of the shared parameter block itself becomes important, even more than feeding in additional data. Therefore, the optimal scaling policy for MoR models favors allocating resources to increasing model capacity by using larger models trained for shorter steps.

### 5.2 Routing Analysis

##### The allocation of recursion depth reflects contextual predictability of the subsequent token.

In the Right of Figure 1 , we illustrate each token’s recursion depth, which directly indicates how easily the next token can be predicted within the given context. For example, the second subword part (e.g., “-ensively”) of a word is often straightforward to predict, thus needing fewer steps. As also shown in Table 17 , while the initial generation (opening) of function words like “—”, “(”, “.”, and “,” appears easy for models, predicting their closing parts or the first token immediately following their opening might be more challenging.

##### Expert-choice router with auxiliary loss perfectly separates selected from unselected tokens.

5(b) visualizes one example of the expert-choice router output distribution at each recursion step in a MoR model with N r = 3 N_{r}=3 . For each recursion step, the normalized counts of routing scores are plotted, distinguishing between tokens selected by the expert (blue) and those not selected (orange). In all steps, auxiliary loss achieves a perfect separation in router outputs, with selected tokens sharply concentrated near a routing score of 1.0 and unselected tokens clustering near 0.0. Router output distributions for the two routing strategies and their associated design choices are detailed in Appendix I .

### 5.3 Test-time Scaling Analysis

##### MoR enables test-time scaling via deeper recursion.

We visualize how log-likelihood evolves across recursion steps in MoR models with N r = { 2 , 3 , 4 } N_{r}=\{2,3,4\} in 5(c) . The overlaid bars illustrate the performance of each model when the maximum thinking (recursion) depth of tokens gradually increases. This suggests that deeper recursion not only provides additional compute but also enables each subsequent step to specialize further in refining the token representation or “thought process” at its particular depth, leading to better performance. Thereby, these results support the view that MoR enables test-time scaling: allocating more recursion steps at inference can improve generation quality.

## 6 Related Work

##### Recursive Transformers.

Parameter sharing provides an orthogonal path to efficiency ( Dehghani et al., 2018 ; Lan et al., 2019 ; Xia et al., 2019 ; Jaegle et al., 2021 ; Takase and Kiyono, 2021 ; Tan et al., 2023 ; Ng and Wang, 2024 ; Bae et al., 2024 ) . The Universal Transformer first showed that repeatedly applying a single block can match the representational power of deep, non-shared stacks ( Dehghani et al., 2018 ) . Looped Transformers shown to be effective, can act as programmable computers ( Giannou et al., 2023 ) , learn iterative data-fitting algorithms ( Yang et al., 2023 ) , generalize to much longer inputs on algorithmic tasks ( Fan et al., 2024 ) , and illuminate few-shot learning by mimicking multi-step optimizers ( Gatmiry et al., 2024 ) . Furthermore, Bae et al. (2024) mitigate the accuracy loss often associated with weight tying by adding low-rank adaptation (LoRA) adapters ( Hu et al., 2022 ) in each loop, yielding Relaxed Recursive Transformers . Recent work further demonstrates that Recursive Transformers excel at latent reasoning via recurrent depth ( Geiping et al., 2025 ) . While most prior studies focus on the efficiency gains from weight tying, the recursive architecture itself offers a second level: inspired by early-exiting ( Schuster et al., 2022 ) and compute routing ( Raposo et al., 2024 ) , one can vary the number of recursions per input (e.g., per token), allocating compute only where it is most beneficial during both training and inference.

##### Adaptive computation.

Many works have shown that dynamic compute allocation can markedly reduce the cost of training and inference, from traditional neural networks ( Bengio et al., 2015 ; Huang et al., 2016 ; Teerapittayanon et al., 2016 ; Panda et al., 2016 ; Graves, 2016 ) to large language models ( Hou et al., 2020 ; Elbayad et al., 2020 ; Fedus et al., 2022 ; Bae et al., 2023 ; Elhoushi et al., 2024 ; Raposo et al., 2024 ) . Early exiting methods learn to halt processing for “easy” samples (e.g., tokens or sequences in language modeling) by skipping the remaining layers ( Elbayad et al., 2020 ; Schuster et al., 2022 ; Dehghani et al., 2018 ; Mofakhami et al., 2024 ) . Alternatively, early exits can be combined with speculative decoding techniques ( Chen et al., 2023 ; Leviathan et al., 2023 ) during inference by leveraging lower layers for fast drafting ( Bae et al., 2023 ; Elhoushi et al., 2024 ) . Recently, Mixture-of-Depths (MoD) ( Raposo et al., 2024 ) reframed adaptive depth as a routing problem: a router at each layer selects a subset of tokens to receive the full computation, while the rest bypass the layer, yielding finer-grained conditional compute. This new form of adaptive allocation is well suited to Transformer architectures and has already been extended to other modalities ( Zhang et al., 2024a ; Luo et al., 2024 ) , highlighting a promising paradigm of dynamic compute at token-level granularity. MoR applies this routing idea to recursive Transformers: tokens are dynamically sent through repeated calls of a single, weight - tied block instead of through distinct layers. This shift keeps parameter count constant, allows arbitrarily deep (adaptive) compute beyond the model’s physical depth.

##### Routing mechanism.

LLMs have increasingly employed routers to enable adaptive computation, primarily in sparse Mixture-of-Experts (MoE) frameworks ( Shazeer et al., 2017 ; Lepikhin et al., 2020 ; Dai et al., 2022 ; Zoph et al., 2022 ) , i.e., each token is processed by a subset of expert networks chosen by a learned router, dramatically increasing model capacity without a computational overhead. Early MoE architectures ( Lepikhin et al., 2020 ; Fedus et al., 2022 ; Jiang et al., 2024 ) adopted a token-choice routing strategy, wherein the router selects the top-k experts for each token based on its hidden state. While effective, this approach often leads to load imbalance across experts, necessitating auxiliary balancing losses. To address this, expert-choice routing ( Zhou et al., 2022 ; Guo et al., 2025 ) has been proposed, wherein each expert selects the tokens to serve, ensuring perfect load balancing and improved efficiency. Building on this, a few works employed trainable routers to determine which layers to skip ( Zeng et al., 2023 ; Raposo et al., 2024 ; Gadhikar et al., 2024 ) . Unlike traditional early-exit methods, these expert-choice routing mechanisms enforce a static compute budget by capping the number of tokens processed per layer (or depth).

##### Key-value caching.

Key–value (KV) caching stores the per-token key and value tensors produced at each layer during autoregressive decoding; reusing them eliminates quadratic-time recomputation and boosts throughput ( Shazeer, 2019 ; Ge et al., 2023 ; Liu et al., 2024a ; Xiao et al., 2023 ; Pope et al., 2022 ; Kang et al., 2024 ; Brandon et al., 2024 ) . Unfortunately, retaining these tensors quickly saturates GPU memory, especially for long contexts and large batches ( Chowdhery et al., 2023 ; Brandon et al., 2024 ) . Prior work tackles this issue by quantizing KV activations to lower precision ( Hooper et al., 2024 ; Zhang et al., 2024b ) , discarding entries that contribute little to the final output ( Zhang et al., 2023 ; Liu et al., 2023a ) , and sharing keys and values across attention heads ( Shazeer, 2019 ; Ainslie et al., 2023b ) . Brandon et al. (2024) push this idea further, allowing adjacent layers to share the same key and value tensors and achieving additional memory savings with negligible quality loss. Our Mixture-of-Recursions offer a complementary avenue: KV caches generated in early recursions can be reused in later ones, potentially reducing memory consumption even further. This also provides the advantage of only needing to run the first recursion during prefill phase ( Sun et al., 2024 ) (only with Cycle strategy), promising significant speedups for prompt settings over 1 million tokens. Two caching strategies in MoR can be optimized based on their distinct benefits to suit various deployment settings.

##### Latent reasoning.

An emerging line of work enables LLMs to perform reasoning internally within hidden states rather than through explicit verbalization ( Goyal et al., 2023 ; Pfau et al., 2024 ; Cheng and Van Durme, 2024 ; Tack et al., 2025 ; Kong et al., ) . Many approaches adopt a fixed latent reasoning depth: they insert special tokens or structured prompts (e.g., a learnable “pause” token ( Goyal et al., 2023 ) or filler punctuation ( Pfau et al., 2024 ) ) that allow the model to execute a predetermined number of hidden reasoning passes before producing an answer. Others reuse the model’s hidden states in a closed loop for a fixed number of iterations by feeding final hidden states back as input to simulate chain-of-thought ( Hao et al., 2024 ; Shen et al., 2025 ; Saunshi et al., 2025 ) . Another line of research enhances latent reasoning by augmenting hidden states with intermediate semantic signals ( Zelikman et al., 2024 ; Tack et al., 2025 ) . However, these methods lack the flexibility to allocate computation where it is most needed, leading to unnecessary overhead on easy inputs and insufficient reasoning on complex ones. This motivates leveraging looping mechanisms for more adaptive latent reasoning ( Chen et al., 2025 ; Geiping et al., 2025 ; Saunshi et al., 2025 ; Zeng et al., 2025 ) .

## 7 Conclusion

Mixture-of-Recursions (MoR) presents a unified Transformer architecture that simultaneously leverages parameter sharing, adaptive recursion depth, and efficient KV caching without compromising model quality. By dynamically assigning recursion depth to tokens via lightweight routers and selectively caching key-value states for selected tokens, MoR reduces both quadratic attention computation and redundant memory access costs. Extensive empirical evaluations show that MoR lowers validation perplexity and improves average few-shot accuracy compared to both vanilla and previous recursive baselines, even with higher inference throughput. These results demonstrate that MoR offers an effective path towards achieving large-model capabilities with significantly reduced computational and memory overhead.

### 7.1 Limitations and Future Works

##### Reasoning MoR models.

Recent studies have highlighted the redundancy within reasoning chains and address it by applying token-level adaptive computation, like early-exit mechanisms ( Yang et al., 2025 ; Jiang et al., 2025 ; Dai et al., 2025 ) . Our MoR framework inherently enables latent reasoning by adaptively determining the necessary recursion depth for individual tokens. Therefore, a crucial future work involves exploring how the router can dynamically learn to adjust to the necessity of chain-of-thought (CoT) chains when post-trained on actual reasoning datasets. Developing advanced routing strategies that explicitly align recursion depth with reasoning complexity may enhance reasoning accuracy, computational efficiency, and even interpretability for deliberative reasoning process.

##### Further scaling model family.

Due to current compute constraints, our experiments have been limited to models of up to 1.7 billion parameters. The next step is to train Mixture-of-Recursions (MoR) models at a larger scale (over 3 billion parameters) on substantially larger corpora. To enhance the scalability of the MoR architecture, we can first increase the size of the non-shared block. A potential follow-up work could introduce depth-specific LoRA ( Bae et al., 2024 ) or experts (i.e., incorporating Mixture-of-Experts principles) and leverage expert parallelism ( Rajbhandari et al., 2022 ) to efficiently compute the input at each depth in parallel, which is expected to improve model quality without huge speed bottlenecks. Furthermore, to reduce overall pre-training costs, we could also explore continued pre-training (uptraining), starting from existing pre-trained vanilla LLM checkpoints. As future work, we plan to investigate MoR performance using various initialization strategies for recursive models, as explored in prior work ( Bae et al., 2024 ) .

##### Adaptive capacity control.

Expert-choice routing offers the significant advantage of guaranteeing perfect load balancing through pre-determined capacity factors ( Raposo et al., 2024 ; Zhou et al., 2022 ) . However, a limitation arises when we want to allocate different capacities during inference. Specifically, in our MoR models, we observe that when using an auxiliary loss, the router outputs for selected and unselected tokens are almost perfectly separated. This makes it challenging to adjust top-k values after training. Therefore, a more adaptive model design, which can leverage different capacities during both training and inference phases, is needed to address this limitation.

##### Compatibility with sparse algorithms.

Given MoR’s token-level adaptive recursion, we can further optimize computation by integrating structured sparsity. This approach allows for the selective activation of subnetworks or parameters ( Liu et al., 2023b ) , dynamically pruning unnecessary computations at both the token and layer levels ( Raposo et al., 2024 ; Elhoushi et al., 2024 ) . This investigation into sparse model designs promises significant efficiency improvements. We believe many sparsity-based techniques, such as pruning ( Han et al., 2015 ) or quantization ( Jacob et al., 2018 ) , are highly complementary to our MoR framework. This will provide deeper insights into effective sparse architectures within recursive models, offering promising directions for future research.

##### Expansion to multimodal and non - text domains.

MoR’s recursion block is inherently modality-agnostic, allowing its adaptive depth mechanism to extend beyond text processing. This crucial property enables MoR to readily integrate into vision, speech, and unified multimodal transformer architectures. Applying token-adaptive recursion to long-context video or audio streams holds the potential for even greater memory efficiencies and substantial throughput gains, crucial for real-world applications. By dynamically adjusting the processing depth for each token or segment, MoR could unlock these significant benefits.

### 7.2 Acknowledgements

We thank Jacob Eisenstein for valuable feedback on an earlier version of the paper. We also acknowledge the support and helpful conversations from Seungyeon Kim, Mostafa Elhoushi, Pascal Vincent, Irina Rish, Sangdoo Yun, and Baeseong Park. Finally, we thank the Google Cloud Platform for awarding Google Cloud credits for this project, and Google’s research grant project for its support.