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

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

This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) ([RS-2019-II190075, Artificial Intelligence Graduate School Program (KAIST), 5%], [No. RS-2024-00457882, AI Research Hub Project, 5%], and [No. 2022-0-00871, Development of AI Autonomy and Knowledge Enhancement for AI Agent Collaboration, 90%]). This research was also enabled in part by computational resources, software, and technical assistance provided by Mila , ULaval , Calcul Québec , and Digital Research Alliance of Canada .

## Appendix A Details of Design Choices for Mixture-of-Recursions

In this section, we provide detailed descriptions of the design choices employed in Mixture-of-Recursions, expanding upon the summary provided in the main pages.

### A.1 Parameter-sharing Strategy

Table 5 shows formulation and visualization of four parameter-sharing strategies: Cycle, Middle-Cycle, Sequence, and Middle-Sequence. These strategies determine how a shared pool of blocks Φ ′ \Phi^{\prime} are reused across a total of L L unrolled layers. The optimal strategy for parameter sharing in recursive models remains an open question.

In the Cycle strategy, a fixed set of parameters is reused cyclically across all recursion steps. By forcing the model to re-engage with the input through the same shared block, it encourages a deeper, iterative refinement process, akin to “rethinking” the problem from the ground up at every stage. However, because the same transformations are applied repeatedly regardless of input variation, it may limit the model’s capacity to learn diverse or highly specialized features.

On the other hand, the Sequence strategy assigns distinct parameters to each recursion block in sequential order. A potential drawback is that simply applying similar transformations twice in a row may lead to redundant features with diminishing returns. Nevertheless, the use of a fixed, sequential order of layers may provide a stable and predictable structure.

Building upon these strategies, the Middle sharing variant further refines parameter reuse by preserving unique parameters at the first and last layers while sharing weights only among the intermediate layers. This approach aims to balance the trade-off between parameter efficiency and representational flexibility, maintaining distinct entry and exit transformations while benefiting from reduced parameter redundancy in the middle layers. In line with recent findings ( Kim et al., 2023a ; Geiping et al., 2025 ) , Middle sharing can capture important input and output nuances more effectively than pure Cycle or Sequence sharing, without significantly increasing model size.

### A.2 Routing Strategy

In this section, we provide an in-depth explanation of the two routing strategies employed in Mixture-of-Recursions: Expert-choice and Token-choice routers. Each approach has distinct advantages and inherent limitations, which we first outline before discussing the mitigation techniques we utilized.

##### Expert-choice routing.

The expert-choice router offers several advantages, including a fixed compute budget that simplifies resource management. However, it suffers from a key issue: the top- k k selection operation, which requires information of tokens that appear later in the sequence, violates causality in autoregressive inference. This non-causal dependency (i.e., information leakage) can cause unexpected behavior during inference, potentially reducing model reliability.

To address these challenges, we explore two approaches: the auxiliary router and the auxiliary loss ( Raposo et al., 2024 ) . The auxiliary router is an additional lightweight network trained jointly but used only during inference; it predicts whether a token will be among the top- k k selection. This additional router is trained with a binary cross-entropy loss, where the top-k selections from the main router are defined as the targets. Importantly, its training is isolated from the main objective through gradient blocking, so it does not affect the primary model training. Meanwhile, the auxiliary loss applies the binary cross-entropy loss to the main router itself, enabling it to simultaneously learn to push top- k k tokens towards one and others towards zero during training. This ensures the router can reliably predict which tokens will be selected as top- k k during inference.

##### Token-choice routing.

In contrast, the token-choice router assigns recursion depths on a per-token basis without enforcing a fixed compute budget, thus avoiding leakage of information across tokens and preserving autoregressive properties. However, this introduces load imbalance across experts, which results in uneven token distribution across experts (or recursion depths), potentially causing inefficient compute allocation and unbalanced training.

To mitigate load imbalance, we employ two solutions from existing literature. Balancing Loss ( Lepikhin et al., 2020 ; Fedus et al., 2022 ) regularizes for a more uniform distribution of tokens across experts. For a sequence of length T T , a balancing loss for MoR is calculated as follows: ℒ Balance \displaystyle\mathcal{L}_{\text{Balance}} = α ​ ∑ i = 1 N r f i ​ P i , \displaystyle=\alpha\sum_{i=1}^{N_{r}}f_{i}P_{i}, f i \displaystyle f_{i} = N r T ​ ∑ t = 1 T 𝕀 ⁡ ( Token ​ t ​ selects Expert ​ i ) , \displaystyle=\frac{N_{r}}{T}\sum_{t=1}^{T}\mathbb{I}(\text{Token }t\text{ selects Expert }i), P i \displaystyle P_{i} = 1 T ​ ∑ t = 1 T g t i , \displaystyle=\frac{1}{T}\sum_{t=1}^{T}g_{t}^{i}, where N r N_{r} is the total number of experts (which is also the number of recursion), g t i g_{t}^{i} is the routing score of expert i i for token t t , f i f_{i} represents the fraction of tokens routed to expert i i , P i P_{i} denotes the average routing scores of expert i i , and λ \lambda is a hyperparameter controlling the strength of the auxiliary loss.

Loss-free ( Wang et al., 2024 ) utilizes router biasing without explicit regularization loss. Specifically, this method adjusts per-expert bias terms b i b_{i} to balance token assignments across experts. During each training batch, routing scores are computed, and the number of tokens assigned to each expert ( c i c_{i} ) is counted. The load violation error is calculated as e i = c ¯ i − c i e_{i}=\bar{c}_{i}-c_{i} where c ¯ i \bar{c}_{i} is the average token count for expert i i . Biases are then updated via b i ← b i + u × sign ⁡ ( e i ) b_{i}\leftarrow b_{i}+u\times\operatorname{sign}(e_{i}) , where u u is a bias update rate. The biased routing scores for selecting top- k k expert are calculated as g i t = { g i t , if ​ g i t + b i ∈ topk ​ ( { g j t + b j ∣ 1 ≤ j ≤ N } , k ) 0 , otherwise g_{i}^{t}=\begin{cases}g_{i}^{t},&\text{if }g_{i}^{t}+b_{i}\in\texttt{topk}\left(\{g_{j}^{t}+b_{j}\mid 1\leq j\leq N\},k\right)\\ 0,&\text{otherwise}\end{cases} Note that the expert bias term is only utilized to adjust the routing strategy by influencing the top- k k selection.

### A.3 KV Caching Strategy

This work investigates two principal strategies for key-value (KV) caching to optimize memory usage during Recursive Transformer computations: recursion-wise caching and recursive KV sharing .

##### Recursion-wise caching.

This keeps separate KV caches for each recursion step, ensuring tokens attend only to the KV pairs generated in their current recursion block. This prevents distribution mismatches between recursion steps, helping to maintain model accuracy while reducing memory and computational costs.

##### Recursive KV sharing.

In contrast, recursive sharing reuses KV pairs computed in the first recursion step for all subsequent steps. Although this approach further lowers memory usage and eliminates the need to compute deeper recursion during the prefill phase, it introduces potential mismatches as later recursion steps receive KV representations originally intended for earlier steps. Such mismatch can negatively impact model performance when token routing is precise. Therefore, recursion-wise caching is generally preferred in settings with selective token routing to avoid performance degradation, while recursive KV sharing may be considered when memory efficiency is prioritized and prefill time is main bottleneck in the system.

## Appendix B Experimental Setup

##### Training settings.

We utilized a Llama-based Transformer architecture ( Llama Team, 2024 ) , referring to the configurations of the open-source SmolLM models ( Allal et al., 2024 ) . All models were pretrained on a deduplicated subset of the FineWeb-Edu dataset ( Penedo et al., 2024 ) in SmolLM-Corpus ( Ben Allal et al., 2024 ) , which comprises 220 billion tokens sourced from educational materials. Pretraining was conducted using four H100 or A100 GPUs. In our main and isoFLOPs analysis experiments, we utilized a Trapezoid learning rate scheduler, which consists of warmup (about 5%), stable, and cooldown (20%) phases. This approach allows us to efficiently continue pretraining for scaling laws from intermediate checkpoints, eliminating the need to train all models independently. In contrast, for all other experiments, we used a simple cosine annealing scheduler.

##### Evaluation settings.

To assess model performance, we evaluated few-shot accuracy on six benchmarks using the Language Model Evaluation Harness: LAMBADA (LD), HellaSwag (HS), PIQA (PQ), WinoGrande (WG), ARC (Easy and Challenge), and MMLU. For all few-shot datasets, excluding LAMBADA, WinoGrande, and MMLU, we normalized accuracy by the byte length of the target string. We adhered to the standard number of shots for each dataset, and used the continuation task specifically for MMLU for simplicity. All evaluation performance measurements were conducted using a single H100 or A100 GPU.

##### Model architecture details.

Table 6 summarizes the architectural specifications of the four Vanilla Transformer models used as the base for our recursive models. Each model variant differs in scale, ranging from 135M to 1.7B total parameters (including both non-embedding and embedding components). For consistency and comparability, all models are trained using a vocabulary size of 49K and a maximum input sequence length of 2K tokens.

## Appendix C Expanded Results of Main Experiments

### C.1 Performance Validation at 1.7B Scale

We validate our models by scaling up the size to 1.7B, all under the same FLOPs budget. As shown in Table 7 , the MoR model with expert-choice routing shows strong performance, but the vanilla model performs slightly better. While this could be due to the differences in the optimal number of tokens associated with each model’s optimal scaling policy, it might also signal that the current architecture design of MoR is not suitable for scaling. Nevertheless, we may be able to further improve performance by increasing the number of non-shared blocks (despite reduced efficiency), applying depth-specific LoRA or MoE to each recursion ( Bae et al., 2024 ) , or uptraining with initialization from a well-pretrained checkpoint. Further validation with larger sizes or training on more tokens is required.

### C.2 Increasing Recursion Depth under Fixed Parameters

Table 8 illustrates the performance change when the number of recursions ( N r N_{r} ) is increased while the number of unique parameters is fixed at 118M. As the number of recursions rises from 1 (Vanilla) to 2 and 3 (MoR), the negative log-likelihood decreases on both the train and validation sets, and the average few-shot accuracy improves. This demonstrates that the MoR model can effectively scale performance by merely increasing the recursion steps, even with a fixed model size.

## Appendix D Expanded Results of IsoFLOP Analysis

In the main paper (§ 3.2 ), we compared Vanilla, Recursive and our Mixture - of - Recursions (MoR) models under matched training compute . Four base model capacities were studied—135M, 360M, 730M and 1.7B parameters. For recursive and MoR models, we fix the recursion count to N r = 3 N_{r}\!=\!3 , so the number of unique parameters is roughly one - third of the vanilla counterpart. Each architecture is trained once for the largest compute budget (16.5EB) 6 6 6 1EB = 10 18 \,{=}\,10^{18} floating - point operations. and the resulting checkpoint is re - used to obtain the 5EB and 2EB variants, as detailed below.

##### FLOPs approximated calculation of Transformers.

We follow the approximation for calculating FLOPs as detailed in Kaplan et al. (2020) . Our analysis solely focuses on forward pass FLOPs, since the FLOPs involved in the backward pass are typically just double those of the forward pass. For most operations within Transformers, which primarily consist of linear projections, the forward pass FLOPs are calculated as two times the number of parameters, excluding the attention mechanism.

Regarding attention, we specifically account for the operations from the dot product between queries and keys and the scaling of values with softmax values. We only calculate FLOPs that contribute to the actual loss, excluding redundant computations in the upper triangular portion due to causality masking. Furthermore, we omit any additional computational costs associated with FlashAttention ( Dao et al., 2022 ) , normalization, and non-linearity operations from our overall FLOPs calculation.

As a result, Vanilla and Recursive Transformers have the same FLOPs. For MoR, the FLOPs calculation varies based on the routing and KV caching strategy. Especially, we calculated FLOPs based on the sequence length at each recursion depth, which is determined by the capacity factor and caching mechanism. In the case of token-choice routing, since the actual token allocation changes at every step, we approximated the FLOPs by assuming perfect balancing. Furthermore, we add extra layers to a few MoR models to ensure their effective depth is divisible by the recursion number. For example, in a 135M model with 30 layers, setting a base depth of 10 and applying recursion three times (as in the Middle-Cycle strategy) results in a total of 32 layers. These additional layers introduce extra FLOPs, so we reduce the number of training steps accordingly to maintain our predefined FLOP budget.

##### Trapezoid learning - rate schedule with checkpoint reuse.

To avoid retraining every model from scratch for each FLOPs budget, we employ the trapezoid schedule ( Xing et al., 2018 ) . The rule of this scheduler is as follows: η ⁡ ( t ) = { t w ​ η max , 0 ≤ t < w (warm - up) , η max , w ≤ t < p (plateau) , η max ​ ( 1 − t − p d ) , p ≤ t < p + d (cool - down) , \eta(t)=\begin{cases}\frac{t}{w}\,\eta_{\max},&0\leq t<w\quad\text{(warm\mbox{-}up)},\\[4.0pt] \eta_{\max},&w\leq t<p\quad\text{(plateau)},\\[2.0pt] \eta_{\max}\!\Bigl(1-\tfrac{t-p}{d}\Bigr),&p\leq t<p+d\quad\text{(cool\mbox{-}down)},\end{cases}

where w w denotes the warm - up interval, p − w p-w is the constant - LR plateau, and d d represents the cool-down segment. This stable phase allows us to efficiently manage experiments by saving intermediate checkpoints and then running additional cool-down steps from those points, according to each budget. For the warmup, we allocate 5 % 5\,\% of the total training steps of the smallest budget (2EB), and we set the cool - down steps to 20 % 20\,\% of the total training steps for each corresponding budget.

##### Results at a glance.