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

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

\reportnumber

# Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

###### Abstract

Abstract: Scaling language models unlocks impressive capabilities, but the accompanying computational and memory demands make both training and deployment expensive. Existing efficiency efforts typically target either parameter sharing or adaptive computation, leaving open the question of how to attain both simultaneously. We introduce Mixture-of-Recursions (MoR), a unified framework that combines the two axes of efficiency inside a single Recursive Transformer. MoR reuses a shared stack of layers across recursion steps to achieve parameter efficiency, while lightweight routers enable adaptive token-level thinking by dynamically assigning different recursion depths to individual tokens. This allows MoR to focus quadratic attention computation only among tokens still active at a given recursion depth, further improving memory access efficiency by selectively caching only their key-value pairs. Beyond these core mechanisms, we also propose a KV sharing variant that reuses KV pairs from the first recursion, specifically designed to further decrease memory footprint. Across model scales ranging from 135M to 1.7B parameters, MoR forms a new Pareto frontier: at equal training FLOPs and smaller model sizes, it significantly lowers validation perplexity and improves few-shot accuracy, while delivering higher throughput compared with vanilla and existing recursive baselines. These gains demonstrate that MoR is an effective path towards large-model quality without incurring large-model cost.

## 1 Introduction

Scaling Transformer networks to hundreds of billions of parameters has unlocked impressive few-shot generalization and reasoning abilities ( Brown et al., 2020 ; Chowdhery et al., 2023 ; Llama Team, 2024 ; OpenAI, 2023 ; Gemini Team, 2024 ; DeepSeek-AI, 2024 ; Gemini Team, 2025 ) . However, the accompanying memory footprint and computational requirements make both training and deployment outside hyperscale data centers challenging ( Patterson et al., 2021 ; Momeni et al., 2024 ) . This has motivated researchers to seek alternative “efficient” designs ( Tay et al., 2022 ; Wan et al., 2023 ) . Among the different axes of efficiency, parameter efficiency ( Dehghani et al., 2018 ; Bae et al., 2024 ; Shazeer et al., 2017 ; Fedus et al., 2022 ; LeCun et al., 1989 ) —reducing or sharing model weights—and adaptive computation ( Raposo et al., 2024 ; Schuster et al., 2022 ; Fedus et al., 2022 ; Leviathan et al., 2023 ) —spending more compute only when it is needed—are promising, actively studied research directions.

One proven route to parameter efficiency is layer tying , in which a shared set of weights is reused across multiple layers ( Dehghani et al., 2018 ; Lan et al., 2019 ; Takase and Kiyono, 2021 ; Gholami and Omar, 2023 ; Bae et al., 2024 ) . For adaptive computation, a common approach is early-exiting , which dynamically allocates compute by exiting earlier in the network when predicting simpler tokens ( Elhoushi et al., 2024 ; Schuster et al., 2022 ; Elbayad et al., 2020 ; Bae et al., 2023 ) . Despite the progress achieved along each of these individual efficiency axes, an architecture that effectively unifies both parameter efficiency and adaptive computation is still missing. Recursive Transformers ( Bae et al., 2024 ; Fan et al., 2024 ; Giannou et al., 2023 ; Yang et al., 2023 ; Saunshi et al., 2025 ; Geiping et al., 2025 ; Aleksandrov et al., 2025 ) , models that repeatedly apply the same set of shared layers multiple times, offer a strong foundation due to their built-in weight sharing. However, prior attempts at dynamic recursion have often been constrained by practical hurdles, such as requiring additional specialized training procedures or facing challenges in efficient deployment. This has led most approaches to still default to a simpler fixed-depth recursion, which applies the same amount of computation to every token and is thus incapable of delivering truly adaptive token-level compute allocation.

In this work, we introduce Mixture-of-Recursions (MoR), a unified framework that fully leverages the potential of Recursive Transformers (see Figure 1 ). MoR trains lightweight routers end-to-end to assign token-specific recursion depths: it decides how many times a shared parameter block is applied to each token according to its required depth of “thinking”, thereby directing computation to where it is most needed. This dynamic, token-level recursion inherently facilitates recursion-wise key–value (KV) caching, selectively storing and retrieving key–value pairs corresponding to each token’s assigned recursion depth. This targeted caching strategy reduces memory traffic, thereby improving throughput without relying on post-hoc modifications. Therefore, MoR simultaneously (i) ties weights to cut parameters, (ii) routes tokens to cut redundant FLOPs, and (iii) caches key-values recursion-wise to cut memory traffic—all inside a single architecture.

Conceptually, MoR provides a pre-training framework for latent space reasoning—performing non-verbal thinking by iteratively applying a single parameter block ( Hao et al., 2024 ; Geiping et al., 2025 ; Goyal et al., 2023 ) . However, unlike approaches that deliberate on augmented continuous prompts before generation ( Liu et al., 2024b ; Goyal et al., 2023 ; Hao et al., 2024 ; Shen et al., 2025 ) , MoR enables this latent thinking directly during the decoding of each token ( Zelikman et al., 2024 ) . Furthermore, the routing mechanism facilitates adaptive reasoning along the model’s vertical axis 1 1 1 While this thinking occurs along the depth axis, it is analogous to generating continuous thoughts along the horizontal sequence axis. , moving beyond the uniform, fixed thinking depth common in prior work ( Geiping et al., 2025 ; Tack et al., 2025 ) . In essence, MoR enables models to efficiently adjust their thinking depth on a per-token basis, unifying parameter efficiency with adaptive computation.

##### Contributions.

In summary, our key contributions in this paper are as follows. • Unified framework for efficient language modeling: We present Mixture-of-Recursions (MoR), the first architecture to unify efficiency paradigms—parameter sharing (§ 2.1 ), token-level adaptive thinking depth (§ 2.2.1 ), and memory-efficient KV caching (§ 2.2.2 )—within a single framework.

• Dynamic recursion routing: We introduce a router trained from scratch to assign dynamic per-token recursion depths. This aligns training with inference-time behavior and eliminates the need for costly, performance-degrading post-hoc routing stages used in conventional early-exit methods.

• Extensive empirical validation: Across models from 135M to 1.7B parameters 2 2 2 These are base model sizes, while MoR models have fewer unique parameters due to parameter sharing. under equal compute budgets, MoR establishes a new Pareto frontier by improving validation loss and few-shot accuracy relative to vanilla and recursive baselines (§ 3.1 , § 3.2 ).

• Efficient architecture: MoR dramatically reduces training FLOPs by selectively engaging only essential sequences in attention operations. Simultaneously, reduction in KV cache sizes leads to enhanced inference throughput itself, further boosted by continuous depth-wise batching (§ 3.3 ).

## 2 Method

### 2.1 Preliminary

##### Recursive Transformers.

The standard Transformer ( Vaswani et al., 2017 ) constructs token representations through a stack of L L unique layers, each with a self-attention and a feed-forward network. At time step t t , the hidden state h h evolves as: 𝐡 t ℓ + 1 = f ⁡ ( 𝐡 t ℓ , Φ ℓ ) , \mathbf{h}_{t}^{\ell+1}=f\bigl(\mathbf{h}_{t}^{\ell};\Phi_{\ell}\bigr), where ℓ = 0 , … , L − 1 \ell=0,\dots,L{-}1 and Φ ℓ \Phi_{\ell} represents the parameters of the ℓ \ell -th layer. Recursive Transformers ( Bae et al., 2024 ; Fan et al., 2024 ; Giannou et al., 2023 ; Yang et al., 2023 ; Saunshi et al., 2025 ) aim to reduce parameter count by reusing layers across depth. Instead of having L L distinct sets of weights, they partition the model into N r N_{r} recursion blocks , where each block uses a shared pool of parameters Φ ′ \Phi^{\prime} . This design allows for more computation (by increasing the effective network depth) without increasing parameter size.

##### Parameter-sharing strategies.

We examine four parameter-sharing strategies: Cycle , Sequence , and their variants Middle-Cycle and Middle-Sequence . Table 1 summarizes two main designs, and the full list is provided in Table 5 in the Appendix. In Cycle sharing, recursion blocks are reused cyclically. For example, consider an original non-recursive model with L = 9 L{=}9 layers and its recursive counterpart using N r = 3 N_{r}{=}3 recursions. Under the “Cycle” strategy, the layers are shared and unrolled as [ ( 0 , 1 , 2 ) , ( 0 , 1 , 2 ) , ( 0 , 1 , 2 ) (0,1,2),(0,1,2),(0,1,2) ]. In “Sequence” sharing, each recursion block reuses the same layer consecutively before moving to the next, resulting in [ ( 0 , 0 , 0 ) , ( 1 , 1 , 1 ) , ( 2 , 2 , 2 ) (0,0,0),(1,1,1),(2,2,2) ] for the same configuration. Both have the same effective number of layers when unrolled ( L = 9 L{=}9 ), but with a different order. Furthermore, the “Middle” variants preserve full-capacity parameters at the first and last layers ( Φ 0 \Phi_{0} and Φ L − 1 \Phi_{L-1} ), while sharing weights among the intermediate layers.

##### Enhanced training and inference efficiency in recursive models.

Parameter sharing strategies can reduce the number of unique trainable parameters by a factor of the recursion number, effectively amortizing the memory footprint of the model. From a distributed training perspective, this becomes highly efficient when using Fully Sharded Data Parallel (FSDP) ( Zhao et al., 2023 ) . While a single all-gather operation would only support one iteration previously (i.e., 1 iter/gather), a recursive model reuses the same gathered parameters across all recursive steps (i.e., N r N_{r} iter/gather). Furthermore, recursive architectures enable a novel inference paradigm, continuous depth-wise batching ( Bae et al., 2024 ; Hooper et al., 2023 ) . This technique allows tokens at different stages to be grouped into a single batch, as they all utilize the same block of parameters. This can eliminate the bubbles—idle periods spent waiting for other samples to complete—thereby leading to significant throughput gains.

##### Limitations in prior works.

Although model parameters are tied, the distinct KV caches are typically used for each depth. This design fails to reduce the cache sizes, meaning the high retrieval latency still remains a severe inference bottleneck. Moreover, most existing recursive models simply apply a fixed recursion depth to all tokens, ignoring the varying complexity. While post-hoc methods like early-exiting methods can introduce some adaptivity, they often require separate training phases that can degrade performance ( Schuster et al., 2022 ; Elhoushi et al., 2024 ; Bae et al., 2024 ) . Ideally, the recursion depth should be learned dynamically during pretraining, allowing the model to adapt its computational path to each token’s difficulty in a data-driven manner. However, such dynamic paths introduce a new challenge: exited tokens will have missing KV pairs at subsequent recursion depths. Addressing this would require a parallel decoding mechanism ( Bae et al., 2023 ; Elhoushi et al., 2024 ; Kim et al., 2023b ) to efficiently compute the actual KV pairs, but this requires separate, complex engineering and complicates the system.

### 2.2 Mixture-of-Recursions

We propose Mixture-of-Recursions (MoR)—a framework that dynamically adjusts recursion step for each token during pretraining and inference . The core of MoR lies in two components: a routing mechanism that assigns token-specific recursion steps to adaptively concentrate computation on more challenging tokens, and a KV caching strategy that defines how KV pairs are stored and selectively utilized for attention at each recursive step.

#### 2.2.1 Routing Strategies: Expert-choice vs. Token-choice

##### Expert-choice routing. ( 2(a) )

Inspired by top- k k gating in MoD models ( Raposo et al., 2024 ) , in expert-choice routing, each recursion depth becomes an expert and selects their preferred top- k k tokens (e.g., for N r = 3 N_{r}=3 , we have three experts: Expert 1 applies the first recursion step, Expert 2 applies the second recursion step, and so on). At each recursion step r r , the corresponding router uses the hidden state ℋ t r \mathcal{H}_{t}^{r} (input to the r r -th recursion block) and its routing parameters θ r \theta_{r} to compute a scalar score g t r = 𝒢 ⁡ ( θ r ⊤ ​ ℋ t r CLOSE g_{t}^{r}=\mathcal{G}(\theta_{r}^{\top}\mathcal{H}_{t}^{r} ) for token t t . Here, 𝒢 \mathcal{G} represents an activation function like sigmoid or tanh . Then, the top- k k tokens are selected to pass through the recursion block: ℋ t r + 1 = { g t r ​ f ​ ( ℋ t r , Φ ′ ) + ℋ t r , if ​ g t r > P β ​ ( G r ) ℋ t r , otherwise \displaystyle\mathcal{H}_{t}^{r+1}=\begin{cases}g_{t}^{r}f(\mathcal{H}_{t}^{r},\,\Phi^{\prime})+\mathcal{H}_{t}^{r},&\text{if }g_{t}^{r}>P_{\beta}({G}^{r})\\ \mathcal{H}_{t}^{r},&\text{otherwise}\end{cases} (2.1) where P β ​ ( G r ) P_{\beta}({G}^{r}) is the β \beta -percentile threshold over all scores at recursion step r r .

To ensure coherent progression through steps, we adopt hierarchical filtering : only tokens selected at recursion step r r can be re-evaluated at r + 1 r{+}1 . This simulates early-exit behavior while learning from scratch. As deeper layers tend to encode increasingly abstract and sparse information ( Li et al., 2022 ; Yang et al., 2024 ; Nawrot et al., 2024 ) , this mechanism prioritizes computation for only the most demanding tokens.

##### Token-choice routing. ( 2(b) )

Unlike expert-choice, where token selection is made at each recursion step, token-choice commits each token to a full sequence of recursion blocks from the start. Formally, given the hidden state ℋ t 1 \mathcal{H}_{t}^{1} (in Middle-Cycle strategy, ℋ t 1 = h t 1 \mathcal{H}_{t}^{1}=h_{t}^{1} ), the router computes a non-linear function ( softmax or sigmoid ) over experts: g t = 𝒢 ⁡ ( θ r ⊤ ​ ℋ t 1 ) {g}_{t}=\mathcal{G}(\theta_{r}^{\top}\mathcal{H}_{t}^{1}) , where g t j g_{t}^{j} denotes the routing score for expert j ∈ { 1 , … , N r } j\in\{1,\dots,N_{r}\} . The token is assigned to expert i = arg ⁡ max j ⁡ g t j i=\arg\max_{j}g_{t}^{j} (top-1 gating), which corresponds to sequentially applying the recursion i i times. The hidden state is then updated recursively as: ℋ t r + 1 = { g t r ​ f ​ ( ℋ t r , Φ ′ ) + ℋ t 1 , if ​ r = i g t r ​ f ​ ( ℋ t r , Φ ′ ) , otherwise \displaystyle\mathcal{H}_{t}^{r+1}=\begin{cases}g_{t}^{r}f(\mathcal{H}_{t}^{r},\,\Phi^{\prime})+\mathcal{H}_{t}^{1},&\text{if }r=i\\ g_{t}^{r}f(\mathcal{H}_{t}^{r},\,\Phi^{\prime}),&\text{otherwise}\\ \end{cases} (2.2)