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

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

Table 9 reports NLL on FineWeb-Edu validation set and few - shot accuracy on six benchmarks. Our findings reveal a clear trend: more compute consistently leads to better models, evidenced by lower NLL and improved accuracy with higher FLOPs. However, weight-sharing alone in recursive models degraded performance compared to Vanilla, a clear trade-off for the reduced parameters. Crucially, token-routed MoR models overcome this shortcoming, catching up to and then surpassing Vanilla models from 360M parameters upward, all while utilizing only one-third of the parameters. This performance advantage persists at 730M and 1.7B parameter scales.

## Appendix E Details of Experimental Settings for Throughput Measurement

We implement a continuous depth-wise batching inference system ( Bae et al., 2024 ; Hooper et al., 2023 ) to evaluate decoding throughput of MoR models. Queries are enqueued and scheduled dynamically during decoding using 1K samples from the FineWeb-Edu validation set. In particular, for MoR, when some queries exit early, the vacant slots in the batch are immediately filled with new queries waiting in the queue, maintaining a fully utilized batch at all times.

We compare the throughput of Vanilla and MoR models (at a 360M parameter scale) for generating a certain length of tokens per sample, where the number is sampled from a normal distribution with a mean of 256, starting without any input prefix. The speeds of the MoR models are normalized against the speed of the Vanilla Transformer. For pretraining MoR-4 models, we add two additional layers (34 layers in total) before applying recursion. This ensures the total effective depth is divisible by the recursion number (specifically for the Middle-Cycle strategy). Consequently, the speed comparison for MoR-4 is made against a modified vanilla model that includes these two extra layers, resulting in a total of 34 layers (32 original layers + 2 added layers).

We use two batching settings: (1) a fixed batch size of 32 and (2) a relative maximum batch size , derived by multiplying 32 by the ratio of the maximum batch sizes of vanilla and MoR models. Specifically, based on the H100 GPU’s VRAM size, we calculated the maximum batch sizes by considering model parameters and their KV cache memory. For simplicity, we omit the memory size from the hidden states at the current position. Under these adaptive conditions, MoR-2 supports a batch size of 42, MoR-3 supports 48, and MoR-4 supports up to 51. By employing recursion-wise KV caching, MoR allows a substantial increase in batch size stemming from its reduced parameter and KV cache memory footprint.

For implementation, we use a queue to enable continuous depth-wise batching and employ FlashAttention 2 ( Dao, 2023 ) to support variable-length KV caches within a batch. We adopt a static -sized cache where each position is updated over time, since this is compatible with torch.compile ( Paszke et al., 2019 ) to further optimize inference speeds. Furthermore, mimicking real-world deployment scenarios ( Kwon et al., 2023 ; Zhong et al., 2024 ) , we decouple the transformer block phase from the rest of the computation by pre-processing the input embeddings or the first non-shared layer before passing into the transformer blocks. Then, we measured the actual time taken during the forward pass. Note that we include a warmup stage by running the model for 100 iterations before actual measurement, in order to obtain stable timing results. For further optimization, tokens that exited early were accumulated up to the maximum batch size before being processed by the last non-shared layer, classifier, and embedding layers (including the non-shared first layer in the case of MoR models). After this, we queue them for sequential batching by following a FIFO (First-In, First-Out) strategy. For implementation convenience, we exclude the time spent on caching and updating for KV pairs, as these aspects can be significantly optimized through various engineering techniques ( Kwon et al., 2023 ) . We leave a more precise speed comparison, which accounts for these considerations, as future work.

## Appendix F Expanded Results of Parameter Sharing Strategy

This section complements the ablation in § 4.1 by providing the full quantitative panorama behind 4(b) . We revisit the four weight - tying schemes— Cycle , Sequence , Middle - Cycle , and Middle - Sequence —on two base model scales (135M and 360M non - embedding parameters) and two different recursion depths ( N r = 2 N_{r}=2 and 3). All models were trained from scratch for 10B tokens under identical optimization hyperparameters. Validation NLL on FineWeb - Edu and averaged few - shot accuracy over six benchmarks are summarized in Table 10 .

##### Middle - Cycle is consistently the safest choice.

For the 360M models, Middle-Cycle achieves the lowest NLL at both depths ( N r = 2 N_{r}=2 and N r = 3 N_{r}=3 ) and also shows the largest improvement in average accuracy compared to vanilla reduced models. For the 135M models, while Cycle is slightly ahead at two recursion setting (3.0071 vs. 3.0330), Middle-Cycle overtakes when recursion depth rises (3.1048 vs. 3.1154) and shows a steadier accuracy profile. Meanwhile, pure Sequence sharing records the worst NLL in all four settings, and its accuracy gap widens with recursion depth. The Middle strategy slightly improves the performance of the Sequence, but it still performs worse than the Cycle-based methodology. We visualized the results in Figure 6 .

For fairer comparison, we also compare the Middle-Cycle and Cycle sharing strategies, fixing the number of unique parameters for both models. For 730M model comparison, both strategies use 3 recursion steps. Especially, Middle-Cycle uses 8 layers for the recursion block and 2 unshared layers (1 + 8 × \times 3 + 1), resulting in 10 unique layers and 26 effective layers. We compare this to the Cycle strategy with 10 layers per recursion block (10 × \times 3), which also has 10 unique layers but a total of 30 effective layers. Despite having fewer effective layers, as shown in Table 11 , the Middle-Cycle model achieves better performance, with a lower negative log-likelihood (2.7552 vs. 2.7573) and higher average few-shot accuracy (42.32 vs. 41.90). While these results are promising, we leave further validation across more scales as future work.

##### Behavior under continued pre - training (up - training).

Table 12 extends the study by “up - training” models—continuing from open-sourced SmolLM ( Allal et al., 2024 ) checkpoints for an additional 5B tokens. Both Middle strategies demonstrate superior performance across all settings, and notably, they significantly outperform the reduced baseline models that are initialized in the same manner but without recursion. The other strategies reach a performance plateau earlier, suggesting that they have limited room for further improvement in capacity.

## Appendix G Expanded Results of Design Choices for Router

### G.1 Details of Design Configurations

We investigate various router design choices to optimize performance and stability. Specifically, we tune the coefficient values controlling the strength of auxiliary or balancing loss terms ( Coeff ), and adjust the scaling factor applied after the router function ( α \alpha ) to modulate routing weights. Moreover, we test different activation functions ( Func ), such as sigmoid or softmax, are evaluated, with architectural variations ( Arch ) of the router network, including linear layer, 2-layer MLP with GELU activation, Wide-MLP that expands the hidden layer size by a factor of four.

We also incorporate several techniques to stabilize training. To improve training stability, we utilize the router z-loss ( Zoph et al., 2022 ) , which penalizes large logits produced by the gating network. Large logits can cause numerical instability and hinder effective training of the router. The z-loss is computed as follows: L z ( x ) = 1 B ∑ i = 1 B ( log ∑ j = 1 N r e x j ( i ) ) 2 , L_{z}(x)=\frac{1}{B}\sum_{i=1}^{B}\left(\log\sum_{j=1}^{N_{r}}e^{x_{j}^{(i)}}\right)^{2}, where B B is the number of tokens in the batch, N r N_{r} is the number of experts, and x ∈ ℝ B × N r x\in\mathbb{R}^{B\times N_{r}} denotes the logits input to the router. This regularization encourages the gating network to produce smaller logits, promoting more stable and reliable routing decisions.

### G.2 Router Performance Evaluation Metrics

##### Expert-choice routing.

We evaluate dead token ratio and sampling accuracy to assess the router’s selection behavior. The dead token ratio measures the proportion of tokens at specific positions within the batch that are consistently unselected during the final recursion step, indicating a positional bias where certain token positions are systematically neglected by the router. The sampling accuracy how well the router used during inference predicts whether a token belongs to the top- k k tokens identified during training, reflecting the router’s ability to consistently select the most relevant tokens. Ideally, high sampling accuracy with a low dead token ratio indicates a router that both identifies important tokens accurately and maintains diversity in token selection.

##### Token-choice routing.

We evaluate the router’s ability to balance token assignments across experts using MaxVio (maximum violation) and entropy metrics. MaxVio ( Wang et al., 2024 ) measures the load imbalance across experts: MaxVio = max i ⁡ Load i − Load ¯ i Load ¯ i , \mathrm{MaxVio}=\frac{\max_{i}\mathrm{Load}_{i}-\overline{\mathrm{Load}}_{i}}{\overline{\mathrm{Load}}_{i}},\addcontentsline{lla}{section}{\hbox to20.74pt{\crtrefnumber{eq:maxvio}\hfil}eq:maxvio} where Load i \mathrm{Load}_{i} denotes the actual number of tokens assigned to the i i -th expert, and Load ¯ i \overline{\mathrm{Load}}_{i} represents the expected load per expert assuming perfect balance.

To measure the diversity of token assignments across experts, we also compute the Entropy of the average selection probabilities for each expert: H = − ∑ i = 1 N r p ¯ i log p ¯ i , H=-\sum_{i=1}^{N_{r}}\overline{p}_{i}\log\overline{p}_{i},\addcontentsline{lla}{section}{\hbox to20.74pt{\crtrefnumber{eq:entropy}\hfil}eq:entropy} where p ¯ i \overline{p}_{i} is the average probability of selecting the i i -th expert over all tokens in the evaluation batch, and N r N_{r} is the total number of experts. A higher entropy indicates a more uniform distribution of tokens among experts, reflecting balanced and diverse routing decisions.

### G.3 Extended Evaluation Results of Router Designs

The results presented in Table 13 indicate that although both the auxiliary router and auxiliary loss methods enhance sampling accuracy, they are also associated with high dead token ratios. In particular, some auxiliary router variants exhibit dead token ratios as high as 66.7%, suggesting that the router always selects tokens from the same positions across inputs, reflecting a positional bias. Notably, employing a linear router architecture in conjunction with auxiliary loss effectively reduces the dead token ratio without compromising sampling accuracy.

Results from Table 14 reveal that applying an explicit balancing loss significantly reduces MaxVio and increases entropy, leading to improved load balance without sacrificing overall model performance. Loss-free approaches, while simpler, tend to show higher MaxVio and lower entropy, indicating less balanced token routing. Architectures such as MLP and Linear routers perform comparably under balancing loss, with z-loss often contributing to improved routing stability and model accuracy. Nevertheless, it still struggles to achieve balance during quite long initial stage. The heterogeneity among the experts, stemming from the use of computation blocks with varying recursion depths as experts, likely complicates load balancing.

## Appendix H Expanded Results of KV Cache Sharing Mechanism

### H.1 Key Value Representation Trends in Recursive Transformers

Sharing KV caches across model depths has emerged as a promising approach to improve inference throughput in Vanilla Transformers ( Brandon et al., 2024 ) . This technique can reduce the memory footprint required for KV caches, enabling larger inference batch sizes. Significant speedups can be also achieved by skipping the KV projection and even prefill operations at shared depths, especially with Cycle strategy ( Sun et al., 2024 ) . Due to the high degree of freedom in Vanilla models—where trainable parameters can be well optimized for shared caches—these models exhibit only marginal performance drops when KV caches are shared between adjacent layers. In contrast, Recursive Transformers have far fewer parameters available for being optimized to tied KV states. Nevertheless, we hypothesize that similar patterns may emerge between shared blocks. To investigate this, we decomposed the KV states from pretrained Recursive Transformers into magnitude and directional components.

As shown in Figure 7 , the sharing of key and value projection layers across recursion depths leads to clear recursive patterns in the magnitude values. Although the magnitudes of hidden states tend to increase, the projection layers appear to be trained to produce similar signal sizes at corresponding depths within each recursion.

When we measured the cosine similarity in Figure 8 , distinct diagonal patterns emerge, suggesting that shared projection layers generate highly similar key and value representations. While sharing value states across recursions appears to be more challenging than sharing key states, these findings suggest that the performance drop from KV cache sharing can be marginal even in Recursive Transformers.

### H.2 Performance Comparison of KV Sharing Strategy

##### Experimental results of key-value cache sharing.

In Table 15 , we present the performance results when applying KV cache sharing to Vanilla, Recursive, and MoR models. Especially, we tested various strategies for KV caches, including Cycle or Sequence strategies that share the same concepts as parameter sharing (see § A.1 for details). Interestingly, KV cache sharing even improves the performance of vanilla models, where sharing acts as a regularization technique. In case of recursive models, we align the sharing strategy for parameters and KV caches. Despite some variations in the results after applying KV sharing, the Middle-Cycle strategy (the best parameter sharing strategy) showed a slight perplexity drop, albeit not substantial.

When moving to MoR models, they still introduced a small amount of degradation in our best settings (expert-choice router). However, considering the reduced parameter sizes and cache sizes, we believe this minor drop is acceptable. Furthermore, we explored an alternative sharing strategy (indicated by † \dagger ) that utilized shared caches for inactive (unselected) positions while updating active positions through actual computation. This method is analogous to a recursive caching scheme but initializes inactive positions with key-value pairs from the first recursive iteration. Although it did not provide additional benefits, it is still worth exploring combinations of KV sharing and actual updates.

##### Relaxation for key-value sharing constraints.

We also investigated relaxing the constraints on KV sharing in Table 16 , similar to the relaxation approach in Bae et al. (2024) for parameter sharing constraints. Specifically, we first re-examined four relaxation techniques for standard Recursive Transformers with very small ranks ( Hu et al., 2022 ; Liu et al., 2024c ) or prefix lengths ( Liu et al., 2021 ) . We also experimented with the position encoding ( Chen et al., 2025 ) , where trainable embeddings are element-wise multiplied with the output of each recursion block.