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

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

To compare routing strategies under equal compute, we align the token allocation budgets of expert-choice with that of token-choice. Specifically, we calibrate token capacity (i.e., top- k ) of expert-choice to match the expected token distribution of token-choice routing with perfect load balancing. In perfectly balanced token-choice, each token is assigned to recursion depth i ∈ { 1 , … , N r } i\in\{1,\dots,N_{r}\} with equal probability 1 / N r 1/N_{r} . Thus, recursion step j j processes a fraction ( N r − j + 1 ) / N r (N_{r}-j+1)/N_{r} of the tokens. For example, when N r = 3 N_{r}=3 , recursion steps 1, 2, and 3 handle { 3 / 3 , 2 / 3 , 1 / 3 } \{3/3,2/3,1/3\} of tokens, respectively. Therefore, we apply this same fractional allocation in the top- k k selection of the expert-choice routing (i.e., k k is sequenced like N r / N r N_{r}/N_{r} , ⋯ \cdots , 1 / N r 1/N_{r} over N r N_{r} recursion steps).

##### Strengths and limitations. ( Table 2 – Left )

Although expert-choice routing guarantees perfect load balancing with static top- k k selection, it suffers from information leakage ( Zhou et al., 2022 ; Wang et al., 2024 ; Raposo et al., 2024 ) . This violation of causality during training forces to exploit an auxiliary router or a regularization loss ( Zhou et al., 2022 ; Raposo et al., 2024 ) , aiming to precisely detect top- k k tokens at inference without access to future token information. Meanwhile, token-choice is free from such leakage, but typically requires a balancing loss or loss-free algorithms ( Wang et al., 2024 ; Fedus et al., 2022 ; Zoph et al., 2022 ) due to its inherent load balancing challenges. We explore each of these components for MoR in further detail (§ 4.2 ).

#### 2.2.2 KV Caching Strategies: Recursion-wise Caching vs. Recursive Sharing

Dynamic-depth models often struggle with KV cache consistency during autoregressive decoding. When a token exits early, its corresponding keys and values in deeper layers will be missing, which can be crucial information for subsequent tokens. Some methods attempt to reuse stale entries ( Schuster et al., 2022 ) or run parallel decoding ( Bae et al., 2023 ) , but these solutions still introduce overhead and complexity. To this end, we design and explore two KV cache strategies tailored to MoR models: recursion-wise caching and recursive sharing .

##### Recursion-wise KV caching. ( 2(c) – Top )

Inspired by Raposo et al. (2024) , we cache KV pairs selectively: only tokens routed to a given recursion step store their key–value entries at that level. Thereby, the KV cache size at each recursion depth is determined exactly by the capacity factor in expert-choice, or according to actual balancing ratios in token-choice. Attention is then restricted to those locally cached tokens. This design promotes block-local computation, which improves memory efficiency and reduces IO demands.

##### Recursive KV sharing. ( 2(c) – Bottom )

A key design choice for our MoR model is that all tokens traverse at least the first recursion block 3 3 3 Though this is not a strict requirement of the MoR framework itself. . We leverage this by caching KV pairs exclusively at this initial step and reusing them across all subsequent recursions. Therefore, the query length might get shorter at each recursion depth based on the selection capacity, but the key and value lengths will consistently maintain the full sequence. This ensures that all tokens can access to past context without recomputation, despite any distribution mismatch.

##### Strengths and limitations. ( Table 2 – Right )

Recursion-wise caching cuts KV memory and IO to approximately ( N r + 1 ) / 2 ​ N r (N_{r}+1)/2N_{r} times across the entire model (when assuming capacity factors follow a sequence like N r / N r N_{r}/N_{r} , ⋯ \cdots , 1 / N r 1/N_{r} over N r N_{r} recursion steps). It also reduces per-layer attention FLOPs to a factor of ( k / N ctx ) 2 (k/N_{\text{ctx}})^{2} of those in vanilla models, resulting in substantially improved efficiency for both training and inference phases. Meanwhile, recursive sharing can yield maximal memory savings by globally reusing context. Speedups can be further achieved by skipping KV projection and prefill at shared depths (compatible only with Cycle strategy ( Sun et al., 2024 ) ). However, attention FLOPs only decrease by k / N ctx k/N_{\text{ctx}} , and high volume of KV IO still leads to a decoding bottleneck.

## 3 Experiments

We pretrain our models from scratch using a Llama-based Transformer architecture 4 4 4 Experiments on Llama are conducted without direction or involvement from Google advisors. ( Llama Team, 2024 ) , referring to the configurations of SmolLM open-source models ( Allal et al., 2024 ) , on a deduplicated subset of the FineWeb-Edu dataset ( Penedo et al., 2024 ) in SmolLM-Corpus ( Ben Allal et al., 2024 ) . We evaluate the models on validation set of FineWeb-edu and six few-shot benchmarks ( Gao et al., 2024 ) . Detailed training and evaluation procedures, as well as throughput measurement protocols, are described in Appendix B .

### 3.1 Main Results

##### MoR outperforms baselines with fewer parameters under equal train compute.

Under an equal training budget of 16.5e18 FLOPs, we compared our Mixture-of-Recursions (MoR) model against both Vanilla and Recursive Transformers. As shown in Table 3 , the MoR model, using an expert-choice router and two recursions, achieves a lower validation loss and surpasses the vanilla baseline in average few-shot accuracy (43.1% vs. 42.3%). Remarkably, this superior performance is achieved despite using nearly 50% fewer parameters. This is attributed to MoR’s higher computational efficiency, which allows it to process more training tokens within the same FLOPs budget. Furthermore, as N r N_{r} increases to 3 or 4, MoR maintains its competitive accuracy, consistently outperforming the recursive baselines while remaining within a tight margin of the full-capacity vanilla model.

##### MoR outperforms baselines with less compute at equal data.

To isolate architectural differences, we analyze performance under a fixed number of training tokens (20B). Specifically, our MoR model with N r = 2 N_{r}=2 outperforms both vanilla and recursive baselines—achieving lower validation loss and higher accuracy—despite using 25% fewer training FLOPs. This theoretical efficiency translates into significant practical gains: compared to the vanilla baseline, our model reduces training time by 19% and cuts peak memory usage by 25%. These improvements stem from our hierarchical filtering and recursion-wise attention mechanism, which shortens sequence lengths to achieve a superior compute-accuracy trade-off, even during pretraining.

##### MoR performance varies with routing and caching strategies.

We also evaluate a few design variants within MoR framework, specifically with N r = 3 N_{r}=3 that is lightweight and still comparable with Vanilla. In this case, using token-choice routing yields lower performance (40.0%) compared to expert-choice routing (42.6%), indicating that routing granularity plays a pivotal role in model performance. Additionally, applying KV cache sharing slightly reduces performance compared to independent caching, while providing improved memory efficiency. This trade-off remains favorable for practical deployment when memory usage is a key concern.

### 3.2 IsoFLOP Analysis

A core criterion for evaluating a new model architectural design is whether performance continues to improve as model and compute scales grow ( Kaplan et al., 2020 ) . Therefore, we evaluate MoR against both Vanilla and Recursive Transformers across a wide range of model sizes and computational budgets to show that it maintains competitive or superior predictive performance as the scale increases.

##### Experimental Setup.

We experiment with four scales—135M, 360M, 730M, and 1.7B parameters—fixing the number of recursions to three for both Recursive and MoR configurations, resulting in roughly one-third the number of unique parameters. Each model is pretrained under three FLOPs budgets: 2e18, 5e18, and 16.5e18.

##### MoR is a scalable and parameter-efficient architecture.

As shown in Figure 3 , MoR consistently outperforms recursive baselines across all model sizes and compute budgets. While it underperforms the vanilla model at the smallest model size (135M)—likely due to a recursive capacity bottleneck—this gap closes rapidly at scale. For >360M parameters, MoR not only matches but often exceeds the Vanilla Transformer, particularly under low and mid-range budgets. Overall, these results highlight that MoR is a scalable and efficient alternative to standard Transformers. It achieves strong validation performance with significantly lower parameter counts, making it a strong candidate for both pretraining and large - scale deployment. Further details are presented in Appendix D .

### 3.3 Inference Throughput Evaluation

As a parameter-shared architecture, MoR can leverage continuous depth-wise batching ( Bae et al., 2024 ) to dramatically boost inference throughput compared to Vanilla Transformers. This maintains high and consistent GPU utilization by immediately replacing completed sequences with incoming tokens during decoding. The early-exiting mechanism in MoR further eliminates bubbles in the computational batch.

##### Experimental Setup.

We measure throughput for 360M scale-based MoR models with recursion depths of 2, 3, and 4, trained under a 16.5e18 FLOPs budget. Throughput (tokens/second) is measured based on the generation time for tokens per sample, where the number of tokens is sampled from a normal distribution with a mean of 256, starting without any input prefix. We examine two batching configurations: a fixed batch size of 32 and a (relative) maximum batch size derived by multiplying 32 by the ratio of the maximum batch sizes of vanilla and MoR models. Further details on the experimental setup are provided in Appendix E .

##### MoR boosts inference throughput with continuous depth-wise batching.

In 4(a) , across both batch settings, all MoR variants outperform the vanilla baseline, which even leverages continuous sequence-wise batching ( Yu et al., 2022 ; Kwon et al., 2023 ) . Increasing recursion depth leads to more tokens exiting early and a further reduction in KV cache usage. This, in turn, boosts throughput significantly (e.g., MoR-4 achieves up to a 2.06 × \times speedup with B = Max B=\text{Max} ). While there’s a slight performance degradation, it can be a favorable trade-off given the substantial throughput gain. These results support that the integration of the depth-wise batching paradigm with early-exiting can significantly accelerate MoR’s actual deployment throughput.

## 4 Ablation Studies

### 4.1 Parameter Sharing Strategies

##### Middle-Cycle is the most effective parameter sharing strategy.

As discussed in § 2.1 , parameter sharing is a key component of Recursive Transformers and MoR. To identify the most effective sharing configuration, we empirically compare four aforementioned strategies: Cycle, Sequence, Middle-Cycle, and Middle-Sequence. We evaluate each strategy on Recursive Transformers based on 135M and 360M model sizes. As shown in 4(b) , “Middle-Cycle” consistently achieves the lowest validation loss, and its superiority is further confirmed by the detailed results in Appendix F . Based on these findings, we adopt the “Middle-Cycle” configuration for all subsequent MoR and Recursive Transformers presented in this paper.

### 4.2 Routing Strategies

We conduct an extensive ablation study to understand the impact of various design choices for expert-choice and token-choice routing in our MoR framework. Detailed results are summarized in Appendix G .

##### In expert-choice routing, auxiliary loss and linear router yield the best performance.

For the expert-choice routing setup ( Left of Table 4 ), we evaluate several design aspects: solution to mitigate causality violation (auxiliary router vs. auxiliary loss), normalization functions ( sigmoid vs. tanh ), router architectures (MLP, Linear, or Wide-MLP), and the impact of an auxiliary z-loss ( Zoph et al., 2022 ) . To assess how well the router performs dynamic allocation, we measure the proportion of “dead” tokens—those never selected by the final recursion in a batch—on the validation dataset. Our key findings are as follows: First, using an auxiliary loss is more effective for inference-time behavior than training a separate auxiliary router. Second, a sigmoid normalization function and a simple linear router architecture yield the best performance. Finally, the auxiliary z-loss has a negligible impact on accuracy, though it does slightly reduce the proportion of dead tokens.

##### In token-choice routing, balancing loss yields stable and accurate routing.

For token-choice routing ( Right of Table 4 ), we follow common MoE practices and enable z-loss by default. We compare two balancing strategies: using a balancing loss and training in a loss-free manner using router bias. While both approaches achieve similar log-probability and few-shot accuracy, the explicit balancing loss yields a significantly lower MaxVio ( Wang et al., 2024 ) in our MoR architectures, making it the preferable choice for stable routing. However, despite this, the model often struggles to balance loads among its heterogeneous experts, even for nearly half of the training steps. Softmax activation with an MLP router performs best, and removing z-loss—though we add back with a very small coefficient in the final design—results in higher performance and routing stability.

### 4.3 KV Caching Strategies

##### KV sharing robustly works even in parameter-shared architectures.

In 4(c) , we first investigate the effect of KV sharing in Vanilla and Recursive Transformers. As consistent with prior works ( Brandon et al., 2024 ; Wu and Tu, 2024 ; Sun et al., 2024 ) , if we pretrain models from the scratch, KV sharing does not often compromise performance due to the greater parameter flexibility. Surprisingly, the Recursive Transformer remains relatively robust to KV sharing, despite its reduced degrees of freedom. We found evidence for this by decomposing the KV pairs at each recursion depth into their magnitude and direction. As detailed in Appendix H , depths that share parameters exhibit highly consistent magnitude patterns and high cosine similarity, providing a clear justification for why KV sharing results in only a slight performance drop.

##### KV sharing degrades expert-choice but benefits token-choice routing in MoR.

We compare recursion-wise KV caching and recursive KV sharing mechanisms in our MoR framework. We observe that while recursive KV sharing offers the advantages of reduced memory footprint and overall FLOPs 5 5 5 Although attention FLOPs increase by N ctx / k N_{\text{ctx}}/k than recursion-wise KV caching, reduced KV projection FLOPs lead to an overall reduction. , it leads to quite large performance degradation in expert-choice routing under a fixed token setting. This suggest that exclusively updating and attending to the tokens active in that recursion depth may be more beneficial. Conversely, MoR with token-choice routing could benefit from KV sharing, where its weaker, inaccurate routing decisions can be complemented by the additional contextual information provided by shared KV pairs.

## 5 Analysis

### 5.1 Compute-optimal Scaling Analysis

##### MoR scaling favors model size over training length under isoFLOPs.