Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
Authors: Sangmin Bae
Yujin Kim
Reza Bayat
Sungnyun Kim
Jiyoun Ha
Tal Schuster
Adam Fisch
Hrayr Harutyunyan
Ziwei Ji
Aaron Courville
Se-Young Yun
Venue: NeurIPS 2025 poster ()
Keywords: Adaptive Computation
Dynamic Routing
Early-exiting
Parameter Sharing
Language Model
Efficiency
Primary area: deep_learning

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.

Mixture-of-Recursion Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
Sangmin Bae!*
Yujin Kim
Reza Bayat
Sungnyun Kim Jiyoun Ha Tal Schuster
Adam Fisch
Hrayr Harutyunyan" Ziwei Ji
Aaron Courville Se-Young Yun
KAIST AI
Mila
Google Cloud
+Google DeepMind
*Google Research
Université de Montréal
https:///gtthub.comm/ com http:p///gggh.i..ooo////ooo/o/aa
Abstract
Scaling language models unlocks impressive capabilities, but the accompanying computational and memory demands make both training and deployment expen- sive. Existing efficiency efforts typically target either parameter sharing or adaptive computation, leaving open the question of how attain both simultaneously. We in- troduce Mixture-of- Recursiesss (MoR), a unified framework that combines the two axes of efficiency inside 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 given recursion depth, further im- proving memory access efficiency by selectively caching only their key-value pairs. Beyond these core mechanisms, we also propose 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 new Pareto frontier: at equal training FLOPs and smaller model sizes, it significantly lowers validation perplexity and improves few-shot accuracy, while de- livering higher throughput compared with vanilla and existing recursive baselines.
Introduction
Scaling Transformer networks hundreds of billions of parameters has unlocked impressive few- shot generalization and reasoning abilities [10, 14, 19, 30, 31, 65, 71J. However, the accompanying memory footprint and computational requirements make both training and deployment outside hyperscale data centers challenging [68, 74]. This has motivated researchers to seek alternative "efficient" designs [89, 92]. Among the different axes of efficiency, parameter efficiency [6, 20, 24, 55, 83] 83]-reducing or sharing model weights and adaptive computation [24, 57, 79, 81] spenddnng more compute only when it is needed-are promising, actively studied research directions.
One proven route to parameter efficiency is layer lying, in which a shared set of weights is reused across multiple layers [6, 20, 32, 54, 87]. For adaptive computation, common approach is early- exiting, which dynamically allocates compute by exiting earlier in the network when predicting simpler tokens [5,21, 21, 22, 81]. Despite the progress achieved along each these individual efficiency axes, an architecture that effectively unifies both parameter efficiency and adaptive computation still missing. Recursive Transformers [3, 6, 23, 29, 33, 80, 100], models that repeatedly apply the same set of shared layers multiple times, offer strong foundation due to their built-in weight sharing. However, prior attempts at dynamic recursion have often been constrained by practical hurdles, such requiring additional specialized training procedures or facing challenges in efficient
"qqual contribution. fCoreespondngg authors. The Google co co-authors performed only advisory role.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Layer L-1
Computed
Skipped
Layer L-2
Layer
11
Layer
Recursion Block
Layer
0.7
Sequence Position
Router
People who feel comfortable defending their views--de def ensively confident- may also eventually change those views corresponding behaviors."" Drugs August 26,
Layer O
Figure I: Overview of Mixture-of-Reeursions (MoR). ). (Left) Each recursion step consists of fixed stack layers and a router that determines whether each token should pass through or exit. This recursion block corresponds to the gray box in the middle. (Middle) The full model structure, where the shared recursion step applied up to N, times for each token depending on the router decision. (Right) An example routing pattern showing token-wiss wise recursion depth, where darker cells indicate active computation through the recursion block. Below shows the number recursion steps that each subword token undergoes to predict the nex/ token, shown in colors:
2 and
deployment. This has led most approaches to still default to 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 the potential of Recursive Transformers (see Figure 1). MoR trains lightweight routers end-to-end to assign coken-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 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 '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 cut parameters, (ii) routes tokens cut redundant FLOPs, and (iii) caches key-values recursion- wise to cut memory traffic-all inside a single architecture.
(MoR), a unified framework that fully leverages
Conceptually, MoR provides a pre-training framework for latent space reasoning-performiinn non-verbal thinking by iteratively applying a single parameter block [29, 34, 38]. However, unlike approaches that deliberate on augmented continuous prompts before generation [34, 38, 60, 84], MoR enables this latent thinking directly during the decoding of each token [102]. Furthermore, the routing mechanism facilitates adaptive reasoning along the model's vertical axis moving beyond the uniform, fixed thinking depth common in prior work [29, 86). In essence, MoR enables models to efficiently ad- just their thinking depth on per-token basis, unifying parameter efficiency with adaptive computation.
Contributions. In summary, our key contributions in this paper follows.
Unified framework for efficient language modeling:
We present Mixture-of Reccrrsionns (MoR), sharing ($2.1), token-level adaptive single framework.
the first architecture to unify efficiency
thinking depth ($2.2.1), and memory-efficient KV caching ($2.2.2)-
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-degradig post-hoc routing stages used in conventional early-exit methods.
Extensive empirical validation: Across models from 135M to 1.7B parameters 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 FLOPPs by selectively engaging only essential sequences attention operations. Simultaneously, reduction KV cache sizes leads to enhanced inference throughput itself, further boosted by continuous depth-wist wise batching ($3.3).

---

Table 1: Parameter-sharing strategies in Recursive Transformers. This table shows Cycle and Middle- Cycle schemes with cyclic layer reuse, where Middle- Cycle retains unique first and last layers.
Cycle Strategy
Myddle-Cyele Strategy
Layers
Equation
Figure
Equation
Figure
Last
(bb 1 1
Laver
Recursion
(ut 4/ mod (2,28.))
Laysr
bi: D (e-1 mod (LE- ((E-2)/N)44)
Loyer
Layar D
Layer
First
5(62:
Lager
2 Method
2.1 Preliminary
Recursive Transformers. The standard Transformer [91] constructs token representations through stack of L. unique layers, each with a self-attention and feed-forward network. At time step t, the hidden state evolves as: h f [b': 42), where f = O..... L-1 1 and D represents the parameters of the f-th layer. Recursive Transformers [6. 23, 33, 80, 100] aim to reduce parameter count by reusing layers across depth. Instead of having I, distinct sets of weights, they partition the model into N, recursion blocks, where each block uses a shared pool of parameters '' 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-Cycla and Middle-Sequence. Table summarizes two main designs, and the full list is provided in Table 5 in the Appendix. In Cycle sharing, recursion blocks reused cyclically. For example, consider an original non-recursive model with L-9 layers and its recursive counterpart using N,=3 recursions. Under the "Cycle" strategy, the layers are shared and unrolled as [(0, 1, 2), (0, 2). (0, 1, 2)]. In "Sequence" sharing, each recursion block reuses the same layer consecutively before moving the next, resulting in [(0, 0), (1, (2, 2)] for the same configuration. Both have the same effective number of layers when unrolled (L=9), but with a different order. Furthermore, the "Middle" variants preserve full-capacity parameters at the first and last layers (Lo and 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) [108]. While single all-gather operation would only support one iteration previously (i.e., iter/gather), a recursive model reuses the same gathered parameters across all recursive steps (i.e., N, iter/gather). Furthermore, recursive architectures enable a novel inference paradigm, continuous depth-wise batching [6, 39]. This technique allows tokens at different stages to be grouped into single batch, they all utilize the same block of parameters. This can climinate 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 reduce the cache sizes, meaning the high retrieval latency still remains 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 carly-exiting methods can introduce some adaptivity, they often require separate training phases that can degrade performance [6, 22, 81]. Ideally, the recursion depth should be learned dynamically during pretraining, allowing the model to adapt its computational path to each token's 's difficulty in 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 [5, 22, 51] to efficiently compute the actual KV pairs, but this requires separate, complex engineering and complicates the system.
3

---

Loop
Loop 2
Loop 2
Loop
Loop 1
Loog
Itauter
Loop
Router
Routerr
O MOOEEOOO
(a) Expert-choice routing
(b) Token-choice routing
(c) Caching mechanism
Figure 2: Architectural components of Conture-of-Recuusiion (MoR). (a) Expert-choice routing: At each recursion step, a router selects top-k tokens to continue, progressively narrowing the set of active tokens with depth. (b) Token-choie routing: Each token is assigned fixed recursion step at the outset via a single routing decision, defining its complete compute path through the model. (c) KV caching strategies: Each square in the matrix represents whether token (row) attends to another token's ' cached key (column). "recursion- wise KV caching" (Top). only the keys of currently selected (non moon-rropped) tokens at each recursion step cached blue ), and attention is restricted only to these entries. In "recursive KV sharing" (Bottom), all keys of previous tokens are cached at the first recursion step purple and shared across subsequent recursion steps for attention operations.
2.2 Mfixture-of-Recursie
We propose Mfxttre--ofRRecusiinn (MoR)-4 framework that dynamically adjusts recursion step for each token during pretraining and inference. The core of MoR lies in components: routing mechanism that assigns token-specific recursion steps to adaptively concentrate computation on more challenging tokens, and 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 Token-choice
Expert-choice routing. (Figure 2a) Inspired by top-k gating in MoD models [79], in expert-choice routing, each recursion depth becomes an expert and selects their preferred top-k tokens (e.g.. for N, 3, have three experts: Expert applies the first recursion step, Expert 2 applies the second recursion step, and so on). At each recursion step r, the corresponding router uses the hidden state H. (input to the r-th recursion block) and its routing parameters 0, compute scalar score of (0, HD) for token t. Here, G represents activation function like sigmoid tanh. Then, the top-k tokens are selected to pass through the recursion block:
[6 [HH 4¹) + Hi, if of Ps(G") otherwise
HC+
(H,
where Ps(G") is the 8-percentile threshold over all scores at recursion step r.
To ensure coherent progression through steps, we adopt hierarchical filtering: only tokens selected at recursion step r can re-evaluated at r+1. This simulates early-exit behavior while learning from scratch. As deeper layers tend to encode increasingly abstract and sparse information [58, 69, 99], this mechanism prioritizes computation for only the most demanding tokens.
Token-choice routing, (Figure 2b) Unlike expert-choice, where token selection made at each recursion step, token-choice commits each token full sequence of recursion blocks from the start. Formally, given the hidden state Hi (in Middle-Cycle strategy. 71 = h ), the router computes non-linear function (softmax or sigm sigmoid) over experts: 9: = G( H.) ). where gf denotes the routing score for expert j E (I....,N,, J. The token is assigned to expert gating), which corresponds to sequentially applying the recursion times. The hidden state is then updated recursively as:
arg max, gf (top-1
SHH Φ') + H3. if [HHE, Φ'), otherwise
H
(2)

---

Table 2: Comparison of routing strategies and key-value caching strategies. (Left) Summary of two routing strategies: expert-choice and token-choice, highlighting their pros, cons, and mitigating solutions from previous works [79, 93, 111]. (Right) Relative cost efficiency of caching strategies against a vanilla Transformer (normalized to 1). Here, N, denotes the number of recursions, and k (k < Netx) denotes the number of selected tokens per layer. We compare only the recursion blocks, excluding the non-shared layers. KV cache memory and IO are measured across the entire model, whereas attention FLOPs are reported per layer.
Expert-choice
Token-choice
Recursion-wist Caching
Recursive Sharing
Pros
Static compute budget
leakae
KV Memory
(N, + 1),2N,
I/N,
Cons
Causality violation
Load imbalance
KV Cache IO
(N. + 1)/2N,
1 k/Nax
LSol
Aux Rout, Aux Loss
Boss. Loss-free
Atto FLOPs
12/N2x
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 bal- ancing. In perfectly balanced token-choice, each token is assigned to recursion depth i E (1...., NN,...,,,)) with equal probability 1/N,. Thus, recursion step j processes fraction (N, j 1),N, of the tokens. For example, when N, = 3, recursion steps 1, 2. and 3 handle [3/3, 2/3, 1/3} of tokens, respectively. Therefore, we apply this same fractional allocation in the top-k selection of the expert-choice routing (i.e., k is sequenced like N,,N,. ... I/N, over N, recursion steps).
Strengths and limitations. (Table 2-Leff) Although expert-choice routing guarantces perfect load balancing with static top-k selection, it suffers from information leakage [79, 93, 110]. This violation of causality during training forces exploit an auxiliary router or regularization loss [79, 110], aiming to precisely detect top-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 [24, 93, 111] 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 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 [81] or run parallel decoding [5], 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. (Figure 2c-Top) Inspired by Raposo et al. [79], 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. (Figure 2c-Bottom) A key design choice for our MoR model is that all tokens traverse at least the first recursion block?. 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-wist caching cuts KV memory and IO to approximately (N, + 1)/2N, times across the entire model (when assuming capacity factors follow a sequence like ...., I/N, over N, recursion steps). It also reduces per-layer attention FLOPs to a factor of (k/Nex)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 [85]). However, attention FL.OPs only decrease by k/ /NNax. and high volume of KV I0 still leads to a decoding bottleneck.
TThough this is not a strict requirement of the MoR framework itself.
5

---

Table 3: Comparison of MoR, Recursive, and Vanilla Transformers under both fixed FLOPs (16.5e18) and token (20B) settings. All models are trained on FineWeb- and evaluated by validation negative log- bogliikelhhod (NLL) and few-shot accuracy. For the isoFLOP rows, the number of training tokens (Niok) varies by model efficiency. For the fixed-token rows, we report the effective FLOPs consumed. For the model sizes, we report non-embedding parameter counts. For the KV mechanisms, distinguish between Cache (recursion-wist eaching) and Share (recursive sharing). "In recursive models, all tokens go through fixed recursion depths (N,), instead of adaptive depths.
MoR
Recursion
Pretrain
NLL
Few-shot Accuracy t
Models
Type
KV
Share
N,
Param
FLOPs
Nist
FineWeb
LD
IIS
PQ
WG
ARC
MMLU
Avg
Vanilla
315M
16.5
2.7824
32.0
65.6
50.5
39.6
28.0
42.3
M-Cye
2
167M
16.5
20B
2.8079
31.0
37.1
66.7
52.3
40.8
27.5
42.6
Rocursive!
Cyc
3
118M
16.5
20B
2.8466
29.8
35.9
65.0
52.3
39.0
27.2
41.5
M-Cyc
4
98M
16.5
19B
2.8781
28.2
35.4
65.5
52.5
38.0
26.8
41.0
Expert
Cache
M.Cyc
2
167M
16.5
27B
2.7511
34.4
39.3
65.7
51.2
39.6
28.1
43.1
Export
Cache
M-Cyc
3
118M
16.5
30B
2.7925
33.1
66.9
52.1
38,3
27.4
42.6
Expert
Cache
M-Cyc
4
98M
16.5
30B
2.8204
30.1
37.3
65.0
51.1
38.9
27.4
41.6
Experi
Cache
M. MyCyc
2
167M
12.3
2013
2.7749
33.2
38.3
65.2
52.6
40.1
28.1
42,9
MoR (ours)
Expert
Cache
M-Cyc
188
11.0
20B
2.8246
31.9
37.0
65.7
50.5
38.3
27.4
41.8
Expert
Cache
M. Cyc
98M
11.0
201
2.8519
30.2
36.5
64.3
52.3
38.6
27.2
41.5
Token
Cache
M Cyc
118M
2.9163
27.6
63.8
37.4
26.8
|40.0
Expert
Share
M-Cyc
3
118M
16.5
31B
2.7983
31.7
37.2
65.1
51.0
39.0
27.1
41.9
3 Experiments
We pretrain our models from scratch using Llama-based Transformer architecture? [65], referring to the configurations of SmolLM open-source models [4], on a deduplicated subset of the WebeWeb-Edu dataset [75] in SmolLMMCorpus [7]. We evaluate the models on validation of Fine Web-edu and six few-shot benchmarks [26]. Detailed training and evaluation procedures, as well as throughput measurement protocols, are described in Appendix D.
3.1 Main Results
MoR outperforms baselines with fewer parameters under equal train compute. Under an equal training budget of 16.5e18 FL.Oss, we compared our Nfxture-of-RRccusiions (MoR) model against both Vanilla and Recursive Transformers. As shown Table 3, the MoR model, using an expert- choice router and two recursions, achieves 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 effi- ciency, which allows to process more training tokens within the same FLOPs budget. Furthermore, as N, 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, = 2 outperforms both vanilla and recursive baselines maselines-achieving lower validation loss and higher accuracy-deppite using 25% fewer training FL.OPs. 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-win 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, that is lightweight and still comparable with Vanilla. In this case, using token-choice routing yields lower performance (40.0%) compared to expert-choics 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 eaching, while providing improved memory efficiency. This trade-off remains favorable for practical deployment when memory usage is a key concern.
Experiments Llama conducted without direction involvement from Google advisors.

---

Vanilla
Recursive
MoR (ours)
3.5
3.2
3.5
Validation
3.0
3.1
2.8
2.0 5.0
16.5
2.0 5.0
16.5
2.0
5.0
16.5
2.0
5.0
16.5
Compute Budget (x10'))
Compute Budget (x10"))
Compute Budget (x10
Compute Budgel (x10')
(a) 335M-baseh model
(b) 360M-basco model
(c) 730M-basee model
(d) 1.7B-based model
Figure 3: Validation loss across different compute budgets across four model sizes: 135M, 360M, 730M, and 1.7B parameters. For MoR models, we use expert-choice routing and recursion-wis caching. MoR consistently outperforms recursive baselines and matches or exceeds the standard Transformers larger scales, despite using significantly fewer parameters (approximately one-third due to layer tying with NR = 3).
3.2 IsoFLOP Analysis
core criterion for evaluating new model architectural design is whether performance continues to improve as model and compute scales grow [49]. 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 eters sters-fixing the number of recursions 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.
360M, 730M, and 1.7B param-
MoR is a scalable and parameter-efficies 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) due to recursive capacity bottle- neck- 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 scalable and efficient alternative 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 Appendix
3.3 Inference Throughput Evaluation
As parameter-shared architecture, MoR can leverage continuous depth-wise wise batching [6] to dramat- ically boost inference throughput compared to Vanilla Transformers. This maintains high and con- sistent 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 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 normal distribution with 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 G.
MoR boosts inference throughput with continuous depth-wise batching. In Figure 4a, across both batch settings, all MoR variants outperform the vanilla baseline, which even leverages continuous sequence-wis -iiee batching [53, 101]. Increasing recursion depth leads 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 2.06x speedup with = Max). While there's s a slight performance degradation, it can be 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 s actual deployment throughput.
7

---

3.3
3.0
Sag M-Seq
KV Sharing
Log Likellhood
-2.55
x1.28
x1.60
MoR-2
Vanilla
Max
Neg. Log Likelhoodd
3.1
0.2
Oyc
M-Oye
Neg. Log Likelihood
2.9
KV Sharing
MoR-3
MoR-4
3.0
-2.80
x1.33
x1.95
x1.42
x2.06
1.0
1.5
2.0
Throughput (x)
Vanilla
MoR-T
(a) Parcto frontier of throughput
(b) Sharing strategy
(c) KV cache sharing
Figure 4: (a) Pareto frontier of inference throughput and log-likehood for MoR and Vanilla Transformer under fixed and maximum batching scenarios. Setting details are in Appendix G. (b) Negative log-likelihood (NLL) of Recursive Transformers with N,- 3 across four different parameter-sharing strategies. We pretrained the models 10 billion tokens. The dashed red and black lines denote the full-size Vanilla Transformer and (approximately one- third scales), respectively. (c) NLL performance comparison across four different architectures with KV sharing. For MoR, green (disabled) and blue (enabled) refer to recursion-we KV caching and recursive KV sharing strategies. MoR-E and MoR-T denotes expert-choice and token-choice MoR, respectively. All models are based on 360M scale and trained on 10 billion tokens.
vanilla models
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,
empirically compare four aforementioned strategies: Cycle, Sequence, Middle-Cycle, and Middle Selddee-Seuunce. We evaluate each strategy on Recursive Transformers based on 135M and 360M model sizes. As shown in Figure 4b, "Middle-Cycle" consistently achieves the lowest validation loss, and its superiority is further confirmed by the detailed results in Appendix H. Based on these findings, adopt the "Middle-Cycle" configuration for all subsequent MoR and Recursive Transformers presented in this paper.
4.2 Routing Strategies
conduct an extensive ablation study to understand the impact of various design choices for routing schemes in our MoR framework. Detailed results are summarized in Appendix I.
In expert-choice routing, auxiliary loss and linear router yield the best performance. For the expert-choice routing setup (Left of Table 4), 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 auxiliary z-loss [1 [1]1]. To assess how well the router performs dynamic allocation, we measure the proportion of "dead" tokens- those never selected by the final recursion in 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, sigmoid normalization function and a simple linear router architecture yield the best performance. Finally, the auxiliary z-loss has negligible impact on accuracy, though 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 proobabbbyyytyy and few-shot accuracy, the explicit balancing loss yields a significantly lower Max Vio [93] 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 MI.P router performs best, and removing z-loss- though we add back with a very small coefficient in the final results in higher performance and routing stability.

---

Table 4: Ablation results for expert-choice (Left) and token-choice (Right) routers with various design choices. We use MoR models that apply three recursions to a 360M model with recursion- wise caching. Model performance is measured by NLL and average few-shot accuracy. We evaluate router metrics-dead token ratio (for expert-choice) and Max Vio (for token-choice)-on the validation set. The dead token ratio denotes the proportion of tokens that remain unselected during the last recursion step, measured on 500 samples, each with 2K sequence length. The selected design choice highlighted in gray.
Expert-choice Router
Performance (L L/1/1)
Token-choice Router
Performance (1/ (1///1)
Sampling Func Arch
z-loss
Dead
NLL
Few-shot
Balancing
Func
Arch z-loss
M-Vio
NLL
Few-shot
Aux Rout
MLP
0.0
2.8893
39.4
Loss (0.1)
soft
MLP /
0.200
3.0239
38.5
Aux Rout tanh
MLP
66.7
2.8720
36.2
Loss (0.01)
soft
MLP
0.682
2.9118
39.4
Aux Loss Aux Loss
MLP
0.0
2.8816
40.0
Loss-free Loss-free
soft
ML.
/
0.852
2.9081
39.4
tanh
MLP
0.0
2.9933
38.8
MLP
/
1,281
3.0188
37.6
Aux Loss
Linear
0.1
2.8667
40.1
(0.1)
soft Linear /
0.492
2,9974
38.4
Aux Loss
W-MLP
x
0.4
2.8716
39.4
Loss (0.1)
soft W-MLP /
0.384
3.0293
38.8
Aux Loss
Linear
0.0
228822
40.0
Loss (0.1)
soft Linear
x
0.266
[.9358
39.1
4.3 KV Caching Strategies
KV sharing robustly works even in parameter-shared architectures. In Figure 4c, we first investigate the effect of KV sharing in Vanilla and Recursive Transformers. As consistent with prior works [9, 85, 94], 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 J, 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-wist 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 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 isoFL.Ps. As illustrated in Figure 5a, MoR exhibits a distinct compute-optimal scaling behavior compared to baselines under isoFL.OPs constraints. The flatter slope of MoR's '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 's recursion depth, which directly indicates how easily the next token can be predicted within the given context. For example, the second subword
AAthhouhh attention FLOPs increase by Nax/k than recursion KV caching, reduced KV projection FLOPs lead to overall reduction.

---

Selected
Unselected
Vanilla Feccrsive MoR (ours)
1.0
Validalion Loss
3.3
Normalized Counts
0.8
Log Likelhood
-2.8
3.1
-2.9
0.4
-3.0
2.9
0.2
Loop
Loop
Loop
2.7
100
1000
0.0
0.5
0.0
1.0
0.5
-3.2
MoR-2 MoR-3 MoR-4
Rec
Rec #2
Rec
Model Size (M)
Models (Max Recursion)
(a) Compute-optimal scaling
(b) Learned router scores
(c) Test-time scaling
Figure 5: (a) Compute-optimal scaling analysis for three model architectures, Each star indicates the optimal model size for a given compute budget. We visualize the results in $3.2 by fitting polynomial functions for each architecture and FLOPs budget, and derive the optimal points from these fits. (b) Distribution router outputs for selected and unselected tokens at each recursion step. As an example, 360M size- bizee MoR model with N, = 3, expert-choice router with auxiliary loss, and recursion- wise caching, used. (c) Test-time scaling analysis illustrating the cumulative aig-likelhhood improvement with increasing recursion depth, measured over 500 samples. As we increase N, based on a 360M model size, the number of unique parameters in MoR decreases, resulting in gradual decline in overall performance (i.e., decrease in log-likelihood), All models are trained by an expert-choice router with auxiliary loss and recursion-wis caching mechanism.
part (e.g.. "-ensively"" of 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 the first token immediately following their opening might be more challenging.
Expert-choice router with auxiliary loss perfectly separates selected from unselected tokens. Figure 5b visualizes one example of the expert-choice router output distribution at each recursion step in MoR model with N, 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 perfect separation router outputs, with selected tokens sharply concentrated near a routing score of 1.0 and unselected tokens clustering near 0.0. Router output distri- butions for the two routing strategies and their associated design choices are detailed Appendix K.
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, {2, 3, 4} in Figure 5c. 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 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 Conclusion
Mixture-off oaiee-ff--RRurionnnn (MoR) presents a unified Transformer architecture that simultancously lever- ages 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 putation 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.
10

---

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.
This work was supported by Institute of Information & communications Technology Planning Evaluation (IITP) grant funded by the Korea government (MSIT)
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.
References
[1] Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: training generalized multi-query transformer models from multi- head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pages 4895-4901, 2023. doi: 10.18653/V1/2023 htt.:/2222...22PPPPE
[2] Joshua Ainslie, James Lee-Thorp. Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13. 13245, 2023.
[3] Preslav Aleksandrov, Meghdad Kurmanji, Fernando Garcia Redondo, David O' O'Shea, William Shen, Alex Iacob, Lorenzo Sani, Xinchi Qiu, Nicola Cancedda, and Nicholas D Lane. Abbie: Autoregressive block-based iterative encoder for efficient sequence modeling. arXiv preprint arXiv:2507.08500 085677 2025.
[4] Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Leandro von Werra, and Thomas Wolf. Smollm blazingly fast and remarkably powerful, 2024.
[5] Sangmin Bac, Jongwoo Ko, Hwanjun Song, and Se- Young Yun. Fast and robust carly-exiting framework for autoregressive language models with synchronized parallel decoding. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 5910-5924, Singapore, dec 2023. Association for Computational Linguistics. doi: 10. 11611881120223333 http://22...mmmmmmaaa....... main.362. URL https://aclantlaee
org 2rg/2023. emnlp-me mttppmaaan...................... 362.
[6] Sangmin Bae, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Seungyeon Kim, and Tal Schuster. Relaxed recursive transformers: Effective parameter sharing with layer-wise lora. arXiv preprini
2024.
[7] Loubna Ben Allal, Anton Lozhkov, Guilherme Penedo, Thomas Wolf, and Leandro von Werra. Smollm-corpus, 2024. URL https://hugginge co daattastte/uuggigggeeee.. smollm-com corpus.
[8] Emmanuel Bengio, Pierre-Lue Bacon, Joelle Pineau, and Doina Precup. Conditional computa- tion in neural networks for faster models. arXiv preprint 1oXiv:15/1.0622 2015.
[9] William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Kelln-Kelley. Reducing transformer key-value cache size with cross-layer attention. Neural Information Processing Systems, 2024. doi: (0.48550/arXivv24405 12981.
[10] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877- 1901, 2020.

---

[11] Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptisn Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling, arXiv preprint arXiv: 2300.01316 2023.
[12] Yilong Chen, Junyuan Shang, Zhenyu Zhang, Yanxi Xie, Jiawei Sheng, Tingwen Liu, Shuo- huan Wang. Yu Sun, Hua Wu, and Haifeng Wang. Inner thinking transformer: Leveraging dynamic depth scaling to foster adaptive internal thinking. arXiv preprint arXiv:2502. 13842, 2025.
[13] Jeffrey Cheng and Benjamin Van Durme. Compressed chain of thought: Efficient reasoning through dense representations. arXiv preprint 1rXiv:24/2.13171, 2024.
[14] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):
1-113, 2023.
[15] Damai Dai, Dong. Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. Stablemoe: Stable routing strategy for mixture of experts. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7085-7095, 2022.
[16] Muzhi Dai, Chenxu Yang, and Qingyi Si. S-grpo: Early exit via reinforcement learning in reasoning models. arXiv preprint 2Xii:2505.076860 2025.
[17] Tri Dao. Flashattention-22 2: Faster attention with better parallelism and work partitioning. arXiv preprint (rxXiv:307708669, 2023.
[18] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- memorr-efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344-16359, 2022.
[19] DeepSeek-AI. Deepseek-v3 technical report. arXiv preprint arXiv: 2024.
[20] Mostafa Dehghani, Stephan Gouws, Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. International Conference on Learning Representations, 2018.
[21] Maha Elbayad, Jiatao Gu, Edouard Grave, and Michael Auli. Depth-adaptive transformer. In International Conference on Learning Representations, 2020. URL https://openrevie net/fforum?id=SIggHNYY.
[22] Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich, Basil Hosmer, Bram Wasti, Liangzhen Lai, Anas Mahmoud, Bilge Acun, Saurabh Agarwal, Ahmed Roman, et al. Layerskip: Enabling early exit inference and self selfuspeculative decoding. arXiv preprint arXiv:2404. 16710, 2024.
[23] Ying Fan, Yilun Du, Kannan Ramchandran, and Kangwook Lee. Looped transformers for length generalization. arXiv preprint UrXiv:2409.156// 2024.
[24] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-555 5270, 2022.
[25] Advait Gadhikar, Souptik Kumar Majumdar, Niclas Popp, Piyapat Saranrittichi,, Martin Rapp. and Lukas Schott. Attention is all you need for mixture-of-deptte routing. arXiv preprint arXiv:2412.20875, 2024.
[26] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024. URL https: /ttp://zz/oo/////////////e/o//// www.eoooooo//eeooooooooeooooooo. org//ecords/ 12608602.
12

---

[27] Khashayar Gatmiry, Nikunj Saunshi, Sashank Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn implement multi-step gradient descent for in-context learning? arXiv preprint arXiv:2410.08290 2024.
[28] Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for Ilms. International Conference on Learning Representations, 2023. doi: (0.48550/arXX..310..001
[29] Jonas Geiping, Sean McLeish, Neel Jain, John Kirchenbauer, Siddharth Singh, Brian R Bartoldson, Bhavya Kailkhura, Abhinav Bhatele, and Tom Goldstein. Scaling up test-time compute with latent reasoning: A recurrent depth approach. arXiv preprint arXiv:2502.05177, 2025.
[30] Google Gemini Team. Gemini 1.5:
Unlocking multimodal understanding across millions of tokens of context. arXiv preprint arXiv: 2403.05530, 2024.
[31] Google Gemini Team. Gemini 2.5: Pushing the frontier with advanced reasoning, mul- timodality, long context, and next generation agentic capabilities, 2025, URL https: //arxiv.org abs/ 0tt./257..0000 06261.
[32] Sia Gholami and Marwan Omar. Do generative large language models need billions of parameters? arXiv preprint arXiv: 2309.06589 2023.
[33] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. In International Conference on Machine Learning, pages 1398-11442. PMLR, 2023.
[34] Sachin Goyal, Ziwei Ji, Ankit Singh Rawat, Aditya Krishna Menon, Sanjiv Kumar, and Vaishnavh Nagarajan. Think before you speak: Training language models with pause tokens. arXiv preprint 1rXiv:2310.02226, 2023.
[35] Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint arXiv:1603.0mm 08993, 2016.
[36] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song. Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepscek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501. 1244924,0. 2025.
[37] Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
[38] Shibo Hao, Sainbayar Sukhbaatar, DiJia Su, Xian Li, Zhiting Hu, Jason Weston, and Yuandong Tian. Training large language models to reason in a continuous latent space. arXiv preprint arXiv:2412.06779, 2024.
[39] Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Hasan Gene, Kurt Keutzer, Amir Gholami, and Sophia Shao. Speed: Speculative pipelined execution for efficient decoding. arXiv preprint arXiv:2310.1200 12072, 2023.
[40] Coleman Hooper, Schoon Kim, Hiva Mohammadzadeh. Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. arXiv preprint arXiv: 2401. 2024.
[41] Lu Hou, Zhiqi Huang, Lifeng Shang, Xin Jiang, Xiao Chen, and Qun Liu. Dynabert: Dynamic bert with adaptive width and depth. Advances in Neural Information Prucessing Systems, 33: 9782-9793, 2020.
[42] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, al. Lora: Low-rank adaptation of large language models. ICLR, (2):3, 2022.
13

---

[43] Gao Huang, Yu Sun, Zhuang Liu. Daniel Sedra, and Kilian Q. Weinberger. Deep networks with stochastic depth. In Bastian Leibe, Jiri Matas, Nicu Sebe, and Max Welling, editors, Computer Vision ECCV 2016- 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV, volume 9908 Lecture Notes in Computer Science, pages 646-661. Springer, 2016. doi: 00.1007/978-3331944469999 //doi.org/10.11 10..0.1007997000 3-1---------444--0--.399
URL https:
[44] Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer-aritmmeticonne inference. In Proceedings of the IEEE conference on computer vision and pattern recognition. pages 2704-2713, 2018
[45] Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning, pages 4651-4666 PMLR, 2021.
[46] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:240I /04088, 2024.
[47] Guochao Jiang, Guofeng Quan, Zepeng Ding. Ziqin Luo, Dixuan Wang, and Zheng Hu. Flashthink: An early exit method for efficient reasoning. arXiv preprint arXiv:2505.1394 13949, 2025.
[48] Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong. Zaoxing Liu, Tushar Krishna, and Tuo Zhao Gear: An efficient kv cache compression recipe for near-lossless generative inference of Ilm. arXiv preprint arXiv: 2403.05527, 2024.
[49] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0883 2020.
[50] Dahyun Kim, Chanjun Park, Sanghoon Kim, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeon- woo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, et al. Solar 10.7 b: Scaling large language models with simple yet effective depth up-scaling. arXiv preprint 1rX1v:2312.15166 2023.
[51] Sehoon Kim, Karttikeya Mangalam, Suhong Moon, Jitendra Malik, Michael W Mahoney, Amir Gholami, and Kurt Keutzer. Speculative decoding with big little decoder. Advances in Neural Information Processing Systems, 06-39236-39256, 2023.
[52] Deqian Kong, Minglu Zhao, Dehong Xu, Bo Pang, Shu Wang, Edouardo Honig, Zhangzhang Si, Chuan Li, Jianwen Xie, Sirui Xie, al. Latent thought models with variational bayes inference-time computation. In Forty-second International Conference on Machine Learning.
[53] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng. Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and lon Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th symposium operating systems principles, pages 611-66666 1-6666, 2023.
[54] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. International Conference on Learning Representations, 2019.
[55] Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in neural information processing systems, 2, 1989.
[56] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang. Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
[57] Yaniy Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
14

---

[58] Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank J Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, et al. The lazy neuron phenomenon: On emergence of activation sparsity in transformers. arXiv preprint arXiv:2210006333. 2022.
[59] Akide Liu, Jing Liu, Zizheng Pan, Yefei He, Reza Haffari, and Bohan Zhuang. Minicache: Kv cache compression in depth dimension for large language models. Advances in Neural Information Processing Systems, 37:139997-140031 2024.
[60] Luyang Liu, Jonas Pfeiffer, Jiaxing Wu, Jun Xie, and Arthur Szlam. Deliberation in latent space via differentiable cache augmentation. arXiv preprint arXhv:2412.17747, 2024.
[61] Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang-Ting Cheng, and Min-Hung Chen. Dora: Weight-decomposed low-rank adaptation. In Forty-first International Conference on Machine Learning, 2024.
[62] Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Lam Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang. P-tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks. arXiv preprint 2rXiv:21 2021.
[63] Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang. Victor Xie, Zhaozhuo Xu, Anas- tasios Kyrillidis, and Anshumali Shrivastava. Scissorhands: Exploiting the persistence of importance hypothesis for IIm kv cache compression test time. Advances in Neural Informa- tion Processing Systems, 36:52342-52366 2023.
[64] Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shri- vastava, Ce Zhang, Yuandong Tian, Christopher Re, al. Deja vu: Contextual sparsity for efficient lims at inference time. In International Conference on Machine Learning, pages 22137-22176. PMLR, 2023.
[65] AI @ Meta Llama Team. The llama 3 herd of models. arXiv preprint arXiv: 2407.21783. 2024.
[66] Yaxin Luo, Gen Luo, Jiayi Ji, Yiyi Zhou, Xiaoshuai Sun, Zhiqiang Shen, and Rongrong Ji. m-mod: Exploring -ixture-of-depth adaptation for multimodal large language models. arXiv preprint arXiv:2410.13859, 2024.
[67] Mehmnaz Mofakhami, Reza Bayat, Ioannis Mitliagkas, Joao Monteiro, and Valentina Zant- edeschi. Performance control in early exiting to deploy large models at the same cost of smaller ones. arXiv preprint arXiv:2412.19325, 2024.
[68] Ali Momeni, Babak Rahmani, Benjamin Scellier, Logan G. Wright, Peter L. McMahon, Clara C. Wanjura, Yuhang Li, Anas Skalli, Natalia G. Berloff, Tatsuhiro Onodera, Ilker Oguz, Francesco Morichetti, Philipp del Hougne, Manuel Le Gallo, Abu Sebastian, Azalia Mirhoseini, Cheng Zhang, Danijela Markovic, Daniel Brunner, Christophe Moser, Sylvain Gigan, Florian Marquardt, Aydogan Ozcan, Julie Grollier, Andrea J. Liu, Demetri Psaltis, Andrea Alù, and Romain Fleury. Training of physical neural networks. arXiv preprint arXiv: 0332,0337, 2024.
[69] Piotr Nawrot, Adrian Lańcucki, Marcin Chochowski, David Tarjan, and Edoardo M Ponti. Dynamic memory compression: Retrofitting Ilms for accelerated inference. arXiv preprint arXiv:2403. 09666,, 2024.
[70] Kei-Sing Ng and Qingchen Wang. Loop neural networks for parameter sharing. arXiv preprint 14999 2024.
[71] OpenAL. Gpt-4 technical report. PREPRINT, 2023.
[72] Priyadarshini Panda, Abhronil Sengupta, and Kaushik Roy. Conditional deep learning for energy-efficient and enhanced pattern recognition. In 2016 design, automation & test in europe conference & exhibition (DATE), pages 475-480, IEEE, 2016.
[73] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
15

---

[74] David Patterson, Joseph Gonzalez, Quoe Le, Chen Liang, Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. Carbon emissions and large neural network training. arXiv preprint arXiv: 2/04. 10350, 2021.
[75] Guilherme Penedo, Hynek Kydliček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. URL https:/ /ttp:////opeeeeeeeeeeee net pet/forumnidd-nSCcnn22
[76] Jacob Pfau, William Merrill, and Samuel R Bowman. Let's think dot by dot: Hidden computa- tion in transformer language models. arXiv preprint arXiv:2404. 15750. 2024.
[77] Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. arXiv preprint arXiv: 2211005102, 2022.
[78] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-moe: Advancing mixture-of- experts inference and training to power next-generation ai scale. In International conference on machine learning, pages 18332-18346. PMLR, 2022.
[79] David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. Mixture-of language models. arXiv preprint arXiv:2404.02258, 2024.
Dynamically allocating compute in transformer-baa
[80] Nikunj Saunshi, Nishanth Dikkala, Zhiyuan Li, Sanjiv Kumar, and Sashank J Reddi. Reasoning with latent thoughts: On the power of looped transformers. arXiv preprint arXiv:2502.17416 2025.
[81] Tal Schuster, Adam Fisch, Jai Gupta, Mostafa Dehghami, Dara Bahri, Vinh Q. Tran, Yi Tay, and Donald Metzler. Confident adaptive language modeling. arXiv preprint arXiv: 2207.07000 200.10706, 2022.
[82] Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint urXiv:/911.02150, 2019.
[83] Noam M. Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gasd gated mixture-of- experts layer. International Conference on Learning Representations, 2017.
[84] Zhenyi Shen, Hanqi Yan, Linhai Zhang, Zhanghao Hu, Yali Du, and Yulan He. Codi: Compressing chain-of-thought into continuous space via self-distillation. arXiv preprint arXiv:2502.21074, 2025.
[85] Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. Advances in Neural Information Processing Systemx, 37:7339-7361, 2024.
[86] Jihoon Tack, Jack Lanchantin, Jane Yu, Andrew Cohen, Ilia Kulikov, Janice Lan, Shibo Hao, Yuandong Tian, Jason Weston, and Xian Li. Llm pretraining with continuous concepts. arXiv preprint arXiv:2502.08853 2025.
[87] Sho Takase and Shun Kiyono. Lessons on parameter sharing across layers in transformers. arXiv preprint arXiv:2 06022, 2021.
[88] Shawn Tan, Yikang Shen, Zhenfang Chen, Aaron Courville, and Chuang Gan. Sparse universal transformer. arXiv preprint arXiv:2310.07096, 2023.
[89] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28. 2022.
[90] Surat Teerapittayanon, Bradley McDanel, and Hsiang-Tsung Kung. Branchynet: Fast inference via early exiting from deep neural networks. 2016 23rd international conference on pattern
recognition (ICPR), pages 2464-2469. IEEE, 2016.
16

---

[91] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[92] Zhongwei Wan, Xin Wang, Che Liu, Samiul Alam, Yu Zheng, Jiachen Liu, Zhongnan Qu, Shen Yan, Yi Zhu, Quanlu Zhang, Mosharaf Chowdhury, and Mi Zhang. Efficient large language models: A survey. arXiv preprint arXiv: 2312.03866 2023.
[93] Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Auxiliary-loss-free load balancing strategy for mixture-of- expertsss. arXiv preprint 0rXiv:2408.15664 2024.
[94] Haoyi Wu and Kewei Tu. Layer-condensed kv cache for efficient inference of large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1 11175- 188, 2024.
[95] Yingce Xia, Tianyu He, Xu Tan, Fei Tian, Di He, and Tao Qin. Tied transformers: Neural machine translation with shared encoder and decoder. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 5466-5473, 2019.
[96] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, arXiv preprint arXiv: 2309. 3309.174533 17453, 2023.
[97] Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd. arXiv preprint arXiv: 080770070 2018.
[98] Chenxu Yang, Qingyi Si, Yongjie Duan, Zheliang Zhu, Chenyu Zhu, Qiaowei Li, Zheng Lin, Li Cao, and Weiping Wang. Dynamic early exit in reasoning models. arXiv preprint arXiv:2504.1589 2025.
[99] Dongjie Yang. XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramid- infer: Pyramid kv cache compression for high-throughput Ilm inference. arXiv preprint arXiv:2405 12532, 2024.
[100] Liu Yang, Kangwook Lee, Robert Nowak, and Dimitris Papailiopoolos.. Looped transformers are better at learning learning algorithms. arXiv preprint arXiv:2311.12422 2023.
[101] Gyeong-In Yu, Joo Scong Jeong, Geon- Woo Kim, Soojeong Kim, and Byung- Syung-Gon Chun. Orca: A distributed serving system for (Transformer-Basee generative models. In 16th USENIX Symposium Operating Systems Design and Implementation (OSDI 22), pages 521-538, 2022.
[102] Eric Zelikman, Georges Harik, Yijia Shao, Varuna Jayasiri, Nick Haber, and Noah D Goodman. Quiet-star: Language models can teach themselves to think before speaking. arXiv preprint arXiv:2403.09629, 2024.
[103) Boyi Zeng, Shixiang Song, Siyuan Huang, Yixuan Wang, He Li, Ziwei He, Xinbing Wang, Zhiyu Li, and Zhouhan Lin. Pretraining language models to ponder in continuous space. arXiv preprint 2rXiv:2505.206:
2025.
[104] Dewen Zeng. Nan Du Tao Wang, Yuanzhong Xu, Tao Lei, Zhifeng Chen, and Claire Cui. Learning to skip for language modeling. arXiv preprint arXiv:2311.1544 2023.
[105] Jun Zhang, Desen Meng, Qi, Zhenpeng Huang, Tao Wu, and Limin Wang. p-mod: Building mixture-of-depths milms via progressive ratio decay. arXiv preprint 2411v:2412.04449, 2024.
[106] Tianyi Zhang, Jonah Yi, Zhaozhuo Xu, and Anshumali Shrivastava. Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization. Advances in Neural Information Processing Systems, 37:3304-3333 3331, 2024.
[107] Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song. Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661-347100 2023.
17

---

[108] Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri. Myle Ott, Sam Shleifer, et Pytorch fsdp: experiences on scaling fully sharded data parallel. arXiv preprint arXiv:2304. 1277, 2023.
[109] Yinmin Zhong, Shengyu Liu, Junda Chen, Jianbo Hu, Yibo Zhu, Xuanzhe Liu, Xin Jin, and Hao Zhang. (DistServe): Disaggregating prefill and decoding for goodput-optimized large language model serving. In 18th USENIX Symposium Operating Systems Design and Implementation (OSDI 24), pages 193-210, 2024.
[110] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoe V Le, James Laudon, et al. Mixuure-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114, 2022.
[111] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint 2709::222.00906, 2022.
18

---

Contents
Introduction
2
Method
2.1
Preliminary
2.2
Mixture-of ofeeeerrsooennns
2.2.1
Routing Strategies: Expert-choice vs. Token-choice
2.2.2
KV Caching Strategies: Recursion-wise Caching vs. Recursive Sharing
3
Experiments
3.1 Main Results
6
3.2
IsoFLOP Analysis
7
3.3
Inference Throughput Evaluation
Ablation Studies
4.1
Parameter Sharing Strategies
8
4.2
Routing Strategies
8
4.3
KV Caching Strategies
5
Analysis
5.1
Compute-optimal Scaling Analysis
5.2
Routing Analysis
5.3
Test-time Scaling Analysis
10
Conclusion
10
Related Work
21
B Limitations and Future Works
22
Details of Design Choices for Mixture-of-Reeursion
23
C.1 Parameter-sharing Strategy
23
C.2 Routing Strategy
................................
24
C.3 KV Caching Strategy
25
Experimental Setup
25
Expanded Results of Main Experiments
26
E.1
Performance Validation at 1.7B Scale
................................
26
E.2
Increasing Recursion Depth under Fixed Parameters
................................
26
Expanded Results of IsoFLOP Analysis
27
Details of Experimental Settings for Throughput Measurement
29
H Expanded Results of Parameter Sharing Strategy
30
19

---

I Expanded Results of Design Choices for Router
33
LI Details of Design Configurations
33
I.2 Router Performance Evaluation Metrics
I.3 Extended Evaluation Results of Router Designs
33
J Expanded Results of KV Cache Sharing Mechanism
35
J.1 Key Value Representation Trends in Recursive Transformers
35
J.2 Performance Comparison of KV Sharing Strategy
36
K Expanded Qualitative Results
37
K.1 Analysis Adaptive Computation Paths
37
K.2 Analysis on Router Weights...
39
20

---

A Related Work
Recursive Transformers. Parameter sharing provides an orthogonal path to effi- ciency [6, 20, 45, 54, 70, 87, 88, 95]. The Universal Transformer first showed that repeatedly applying a single block can match the representational power of deep, non-shared stacks [20]. Looped Transformers shown to be effective, can act as programmable computers [33], learn iterative data- fitting algorithms [100], generalize to much longer inputs on algorithmic tasks [23], and illuminate few-shot learning by mimicking multi-step optimizers [27]. Furthermore, Bae et al. [6] mitigate the accuracy loss often associated with weight tying by adding low-rank adaptation (LoRA) adapters [42] in each loop, yielding Relaxed Recursive Transformers. Recent work further demonstrates that Recur- sive Transformers excel at latent reasoning via recurrent depth [29]. While most prior studies focus on the efficiency gains from weight tying, the recursive architecture itself offers a second level: inspired by carly-exiting [81] and compute routing [79], 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 [8, 35, 43, 72, 90] to large language models [5, 21, 22, 24, 41, 79]. Early exiting methods learn to halt processing for "easy" samples (e.g., tokens sequences in language modeling) by skipping the remaining layers [20, 21, 67, 81]. Alternatively, early exits can be combined with speculative decoding techniques [11, 57] during inference by leveraging lower layers for fast drafting [5, 22]. Recently, Mixture-of-Deplle (MoD) [79] reframed adaptive depth as a routing problem: a router at each layer selects a subsel 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 [66, 105], 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 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-Fxperre (MoE) frameworks [15, 56, 83, 111], i.e., each token is processed by a subset of expert networks chosen by learned router, dramatically increasing model capacity without computational overhead. Early MoE architectures [24, 46, 56] adopted 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 [36, 110] 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 [25, 79, 104]. Unlike traditional early-exit methods. these expert-choice routing mechanisms enforce 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 [9, 28, 48, 59, 77, 82, 96]. Unfortunately, retaining these tensors quickly saturates GPU memory, especially for long contexts and large batches [9, 14]. Prior work tackles this issue by quantizing KV activations lower precision [40, 106], discarding entries that contribute little to the final output [63, 107], and sharing keys and values across attention heads [2, 82). Brandon al. [9] 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-Recursee 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 [85] (only with Cycle strategy), promising significant speedups for prompt settings over million tokens. Two caching strategies in MoR can be optimized based on their distinct benefits to suit various deployment settings.
21

---

Latent reasoning. An emerging line of work enables LLMs to perform reasoning internally within hidden states rather than through explicit verbalization [13, 34. 52, 76, 86]. Many approaches adopt a fixed latent reasoning depth: they insert special tokens or structured prompts (e.g.. a learnable "pause" token [34] or filler punctuation [76]) that allow the model to execute a predetermined number of hidden reasoning passes before producing an answer. Others reuse the model's i hidden states in closed loop for a fixed number of iterations by feeding final hidden states back as input to simulate shain-of-thougee [38, 80, 84]. Another line of research enhances latent reasoning by augmenting hidden states with intermediate semantic signals [86, 102]. However, these methods lack the flexibility to allocate computation where it is most needed, leading to unnecessary overhead easy inputs and insufficient reasoning on complex ones. This motivates leveraging looping mechanisms for more adaptive latent reasoning [12, 29, 80, 103].
B 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 [16, 47, 98]. Our MoR framework inherently enables latent reasoning by adaptively determining the necessary recursion depth for individual tokens. Therefore, crucial future work involves exploring how the router can dynamically learn to adjust to the necessity of chain-of-thouge 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.
(CoT) chains when
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 larger scale (over 3 billion parameters) 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 [6] or experts (i.e., incorporating Mixture-of-ENae -f-eExs principles) and leverage expert parallelism [78] to efficiently compute the input 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 [6].
Adaptive capacity control. Expert-choice routing offers the significant advantage of guaranteeing perfect load balancing through pre-determined capacity factors [79, 110]. 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 [64], dynamically pruning unnecessary computations at both the token and layer levels [22, 79]. This investigation into sparse model designs promises significant efficiency improvements. We believe many sparsity-based techniques, such as pruning [37] or quantization [44], are highly complementary 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 moken-adaptin recursion to long-context video or audio streams holds the potential for even greater memory efficiencies and substantial throughput gains, crucial for world applications. By dynamically adjusting the processing depth for each token or segment, MoR could unlock these significant benefits.
22

---

C Details of Design Choices for Nixture-of-Recuriions
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.
C.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 t' are reused across total of L unrolled layers. The optimal strategy for parameter sharing in recursive models remains an open question.
In the Cycle strategy, fixed set of parameters is reused eyclically across all recursion steps. By forcing the model to re-engage with the input through the same shared block, encourages a deeper, iterative refinement process, akin "rethinking" the problem from the ground up at every stage. However, because the same transformations are applied repeatedly regardless of input variation, may limit the model's capacity to learn diverse highly specialized features.
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 row may lead to redundant features with diminishing returns. Nevertheless, the use of a fixed, sequential order of layers may provide stable and predictable structure.
Building upon these strategies, the Middle sharing variant further refines parameter reuse by preserv- ing unique parameters at the first and last layers while sharing weights only among the intermediate layers. This approach aims balance the trade-off between parameter efficiency and representational flexibility, maintaining distinct entry and exit transformations while benefiting from reduced parame- ter redundancy the middle layers. In line with recent findings [29, 50], Middle sharing can capture important input and output nuances more effectively than pure Cycle or Sequence sharing, without significantly increasing model size.
Table 5: Parameter-sharing strategies in Recursive Transformers. This table shows Cycle, Middle- Cycle, Sequence, and Middle-SSquence schemes with layer reuse, where Middle retains unique first and last layers.
Cycle Strategy
Middle-Cyck Strategy
Layers Last
Equation
Figure
Equation
Figure
f(h 001))
Recursion
(bf: o mod (1.1N)))
Layer 1
h1t o 1 mad (ce ((2-2)/M)3-1
Layer
Layer
Layer 1
First
f(b; 40)
Sequence Strategy
Middle-Sequence Strategy
Layers
Equation
Figure
Equation
Figure
Last
f (b (2-2))
Recursion
Abb ((eNN))
Layer 1
bt: o ((-1)x.1-1)
Layer
First
1(1::
0
23

---

C.2 Routing Strategy
In this section, we provide an in-depth explanation of the two routing strategies employed in Mixture- of Recursesssn::: 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 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, explore two approaches: the auxiliary router and the auxiliary loss [79]. The auxiliary router is an additional lightweight network trained jointly but used only during inference; it predicts whether token will be among the top-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 tokens towards one and others towards zero during training. This ensures the router can reliably predict which tokens will be selected top-k during inference.
Token-choice routing. In contrast, the token-choice router assigns recursion depths on a per-token basis without enforcing 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 [24, 56] regularizes for more uniform distribution of tokens across experts. For sequence of length T, a balancing loss for MoR is calculated as follows:
Lalannce
f.P
N fi
T
I(Token selects Expert i),
T
T
1 P. = T
9
where N,, is the total number of experts (which is also the number of recursion), of is the routing score of expert for token 4, Si represents the fraction of tokens routed to expert i, P. denotes the average routing scores of expert and À is hyperparameter controlling the strength of the auxiliary loss.
Loss-free [93] utilizes router biasing without explicit regularization loss. Specifically, this method adjusts per-expert bias terms 6, to balance token assignments across experts. During each training batch, routing scores are computed, and the number of tokens assigned each expert (c:) is counted. The load violation eтror is calculated as Biases are then updated via b; + b, + x sign(e,). where is a bias update rate. The biased routing scores for selecting top-k expert are calculated as
G where è; is the average token count for expert i.
of 0,
if of + b, topk (1g) + by 1 SSSSNN &),K) otherwise
Note that the expert bias term only utilized to adjust the routing strategy by influencing the top-k selection.
24

---

C.3 KV Caching Strategy
This work investigates two principal strategies for key-value (KV) caching to optimize memory usage during Recursive Transformer computations: recursion-wis caching and recursive KV sharing.
Recursion- wisee caching. This keeps separate KV caches for each recursion step, ensuring tokens attend only to the KV pairs generated their current recursion block. This prevents distribution mismatches between recursion steps, helping maintain model aceuracy while reducing memory and computational costs.
Recursive KV sharing. In contrast, recursive sharing reuses KV pairs computed in the first re- cursion step for all subsequent steps. Although this approach further lowers memory usage and eliminates the need 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-wer caching is generally preferred in settings with selective token routing avoid per- formance degradation, while recursive KV sharing may be considered when memory efficiency is prioritized and prefill time is main bottleneck in the system.
D Experimental Setup
Training settings. We utilized Llama-based Transformer architecture [65], referring the configurations of the open-source SmolLM models [4]. All models were pretrained on a deduplicated subset of the FineWeb-Edu dataset [75] in SmolLM-Corpus [7], which comprises 220 billion tokens sourced from educational materials. Pretraining was conducted using four H100 or A100 GPUs. In utilized Trapezoid learning rate scheduler, which our main and isoFLOPs analysis experiments,
consists of warmup (about 5%). stable, and cooldown (20%) phases. This approach allows 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 simple cosine annealing scheduler.
Evaluation settings. To assess model performance, we evaluated few-shot accuracy on six bench- marks 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, exclud- ing LAMBADA, WinoGrande, and MMLU, 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 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 maximum input sequence length of 2K tokens.
Table 6: Key parameters of four model size variants. A model's size is defined by the total number of its non-embedding and embedding parameters. Three small models utilize Grouped-Quee Attention [1], reducing the number of key-value heads. We refer the base configurations of the open-sourced SmolLM models [4].
Base Configuration
Attention & Feed-Forwarr
Input
Models
N-emb
Emb
N₂
dmodet
Nhead
Ngv
dhead
dexter
Vocab
Lete
Vanilla 135M
106M
28M
30
576
9
3
64
1536
49K
2K
Vanilla 360M
315M
47M
32
960
15
5
64
2560
49K
2K
Vanilla 730M
654M
7SM
26
1536
24
8
64
4096
49K
2K
Vanilla L7B
1.61B
101M
24
2048
32
32
64
8192
49K
2K
25

---

E Expanded Results of Main Experiments
E.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, 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 MoE to each recursion [6], or uptraining with initialization from well-pretrained checkpoint. Further validation with larger sizes or training on more tokens is required.
Table Comparison of MoR, Recursive, and Vanilla Transformers 1.7B scale under both fixed FLOPs (68.5e18) settings. All models are trained on FineWWebEEdt and evaluated by few-shot accuracy. For the isoFLOP rows, the number of training tokens (Nisk) varies by model efficiency. For the model sizes, report non-embedding parameter counts. 'n recursive models, all tokens go through fixed recursion depths (N,,), instead adaptive depths.
MoR
Recursion
Pretrain
Few-shot Accuracy
Models
Type
KV
Share
N.
Param
FLOPs
Nish
LD
HS
PQ
WG
ARC
MMLU
Avg
Vanilla
161BB
68.5
20B
40.8
49.4
70.6
54.8
47.4
30.2
48,9
Recursive
M-Cye M-Cye
0.87B
68.5
18B
37.3
46.5
68.9
52.6
44,2
29.6
46.5
0.67B
68.5
20B
36.4
45.3
69.5
52.7
43.9
29.1
46.2
Expert Export
Cache
M-Cye M-Cyc
0.87B
68.5
26B
41.1
47.5
70.0
55.6
46.0
30.1
48.4
MoR (ours)
Cache
0.67B
68.5
27B
37.2
46.6
53.7
44.1
29.7
46.7
M-Cye
0.67B
30B
29.0
45.4
E.2
Increasing Recursion Depth under Fixed Parameters
Table 8 illustrates the performance change when the number of recursions (N,,) is increased while the number of unique parameters is fixed at 1 18M. As the number of recursions rises from (Vanilla) to 2 and 3 (MoR), the negative log- likeeikhehood 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.
Table 8: Comparison of MoR with different numbers recursions under the same number of parameters. Vanilla equals MoR with recursion 1. All models are trained on Fine with 10B tokens, and we apply the Middle-Cycle parameter sharing for MoR models. We report negative log-likelihood (NIL)) on the train and validation sets and few-shot accuracy across six benchmarks.
Pretrain
Recursion
NLL
Few-shol Accuracy
Models
N-Emb
Nz.
Nick
N;
Share
Train
Valid
L.D
HS
PQ
WG
ARC
MMLU
Avg
Vanilla
118M
12
10B
2.9182
2.9678
25.81
33.08
61.81
51.46
38.10
26.64
39.48
MoR
118M
1+10+1
10B
M-Cye
2.8767
2.9205
27.09
33.84
62.13
53.04
36.77
26.88
39.96
MoR
1+10+1
10B
M-Cye
2.8667
2.9111
34.60
63.17
51.46
37.15
26.83
40.10
26

---

F Expanded Results of IsoFLOP Analysis
In the main paper ($3.2), we compared Vanilla, Recursive and our models under matched training compute. Four base model capacities were studied- -155M, 360M, 730M and 1.7B parameters. For recursive and MoR models, we fix the recursion count to =3 so the number of unique parameters is roughly one-third of the vanilla counterpart. Each architecture trained once for the largest compute budget (16.5EB)² and the resulting checkpoint re-used to obtain the SEB and 2EB variants, as detailed below.
(MoR)
FLOPs approximated calculation of Transformers. We follow the approximation for calculating FLOPs as detailed in Kaplan et al. [49]. 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 [18], 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- ooken-choice routing, since the actual token allocation changes at every step, approximated the FLOPs by assuming perfect balancing. Furthermore, layers to a few MoR models to ensure their effective depth is divisible by the recursion number. For example, in 135M model with 30 layers, setting a base depth of 10 and applying recursion three times (as in the Middle-Cycle Cycle strategy) results in a total of 32 layers. These additional layers introduce extra FLOPPs, budget.
add extra
reduce the number of training steps accordingly to maintain our predefined FLOP
Trapezoid learning-rate schedule with checkpoint reuse. To avoid retraining every model from scratch for each FLOPs budget, we employ the trapezoid schedule [97]. The rule of this scheduler is as follows:
to 7max:
S t w (warm-up),
(()
Tmax
w S t < ps < + (cool-down),
(plateau).
11 L-p
1mmax
where denotes the warm-up interval, p is the constant-LR plateau, and 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 % of the total training steps of the smallest budget (2EB), and we set the cool-down steps to 20 % of the total training steps for each corresponding budget.
jEBB 10t floating- point operations.
27

---

Results at a glance. Table 9 reports NLL on Fine 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 NL.. and improved accuracy with higher FI.OPs. 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.
Table 9: Detailed results of isoFLOP analysis across three compute budgets. We evaluate negative log-likelihood (NLL) on the FineWeb-FFl Web-Edu validation set and few-shot accuracy on six downstream tasks for four base model sizes (1 35M, 360M, 730M, 1.7B). Each model was initially trained up to 16,5EB and sliced back to SEB and 2EB via mid mid-training checkpoints using a trapezoid learning-rate schedule. All models used three recursion steps. For MoR models, we use expert-choice routing and recursion-wise caching mechanisms. We highlight the best-performing model in each setting in gray.
Pretrain
Recursion
NLL
Few-shot Accuracy
Models
Base
N-Fmb
N,.
FLOPps
Nick
Share
Loop
FineWeb
L.D
PQ
ARC
MMLUU|
Avg
Vanilla
135M
106M
30
10
6.5B
3.0922
22.80
30.93
62.35
51,14
36,28
26.29
38.30
Recursive
135M
42M
1+10+1
2.0++18
6.1B
Cy
3
3.2058
19.79
29.32
60.17
50.59
34,83
25.40
36.68
MoR
135M
42M
1+10+1
2.0++18
9:28
M-Cye
3
3.1077
21.13
31,00
59.79
49.09
34.87
25.63
36.92
Vanilla
135M
106M
30
5.0c+18
16.18
2,9464
26.88
33.69
63.98
51.46
37.08
27.07
40.03
Recursive
135M
42M
I+10+1
5.0c+18
18.18
M-Cye
3
3.0534
24:51
31.57
62.40
50.83
35.78
25.94
38.51
MoR
135M
42M
1++0+
5.0e+18 2B
CyCye
3
3.0192
22.01
32.53
61.75
49.88
35.39
26.19
37.96
Vanilla
135M
106M
16.5e+18
2,8432
30.16
80
42.15
Recursive
135M
42M
1+10+1
16.Se+18
50.08
M-Cye M-Cyo
3
2,9552
25.98
33.36
6.3.98
51.78
36.96
26.68
39.79
MoR
135M
42M
1+10+1
16.5c+18
76.28
3
2,9490
22.61
33.99
47.83
26.36
38.11
Vanilla
360M
315M
32
2.0e+18
2.4B
3.3785
17.27
27.90
59.36
51.38
32.10
25.49
35.58
Recursive
360M
118M
1+10+1
2.0c+18
2.4B
M-Cye M-Cye
3
3,4864
10.34
26.66
58.00
51.54
30,94
24.94
33.74
MoR
360M
118M
1+10+1
20
3.68
3
3.1026
24.14
30.53
61.86
50.99
34.74
25.50
37.96
Vanilla
360M
315M
32
5.0c+18
6.0B
3.0097
25.17
32.10
63.22
48.62
36.01
26.69
38.63
Recursive
360M
118M
1+10+1
5.0e+188
6.0B
Cye
3
3.0722
23.29
31.19
62.62
51.30
35.85
25. 95.99
38.37
MoR
360M
118M
1+10+1
5.0c+18
9:0B
M-Cye
3
2.9161
28.33
34.53
63.22
51.07
36.70
26.98
40.14
Vanilla
360M
315M
32
16.5c+18
19.88
2.7824
31.94
37.92
66.10
51.30
39.70
27.95
42.49
Recursive
360M
118M
[+10+1
16.5e+18
19.88
M-Cyo
2.8466
29:75
35.92
64.91
51.46
39.12
27.
41.39
MoR
360M
118M
1++1
16.5e+18
29.78
MyCye
2,7924
33.15
37.94
66.97
52.09
38.46
27. .9
42.68
Vanilla
730M
654M
26
2.0c+18
1.28
3.7164
07.74
26.58
57.62
$1.
29.
24.46
32,88
Recursive
730M
1+8+1
2.0e+18
1.2B
M.Cyc
3.8136
05.53
55.77
50.59
29.88
32.11
MoR
730M
252M
1+8+1
2.0e+18
188B
M Cyc
3
3,3300
17.93
28.74
59.30
51.46
33,14
25.37
35.99
Vanilla
730M
654M
26
5.0e+18
3.1B
3.0821
22.05
31.99
62.68
50.67
35,88
26.12
38.23
Recursive
730M
252M
1+8+1
5.0c+18
3.18
M-Cye
3
3.1640
18:51
30.72
62.13
47.83
35,97
84
36.83
MoR
730M
252M
1+8+1
5.0e+18
4.58
M-Cye
3
3.0067
26.18
32,76
62.46
50.91
36,93
26.37
39.27
Vanilla
730M
654M
26
16.5c
10.18
2.7048
34.50
40.29
66.81
49.49
40.82
28. 66
43.43
Recursive
730M
252M
1+8+1
16.5e+18
10.1B
Cye
3
2.7886
30.76
37.84
65.51
52.41
39.26
27.51
42.21
MoR
730M
252M
1+8+1
16.5e+18
14.98
M-Cye
3
2.7438
32,93
39.55
66.32
54.38
40.00
28.09
43.55
Vanilla
1.61B
1.
0.6B
5.1349
00.00
51.00
25.75
23.07
29.37
Recursive
177
0.67B
1+8+1
2.0c+18
0.5B
M-Cye
3
5,3277
00.00
25.27
51.36
48.62
26.52
22.98
29.13
MoR
0.67B
1+8+1
20c+188
0.88
M-Cyo
4.1175
25.80 80
53.97
49.64
27.56
24.08
Vanilla
177
1661B
24
5.0c+18
1.58
3.6926
33
26.84
57.
$1.30
29.72
24.51
33.00
Recursive
0.67B
1+8+1
5.0e+18
M-Cyc
3
3.8876
03.14
26.57
29.01
24.49
31.19
MoR
L7B
0.67B
1+8+1
5.0e+18
2.0B
M Cyc
3
3,2905
17.62
28.32
59.03
49,80
32,14
25.28
35.37
Vanilla
1.7
1.61B
24
16.5e+18
4.8B
2.8658
26.94
35.61
64.74
50.59
38.55
26.81
40.54
Recursive
L7B
0.67B
1+8+1
16.5c+18
4.58
M-Cyc M-Cye
3
3.0042
23.25
32.09
62,95
50.75
37.64
26.53
38.87
MoR
178
0.678
1+8+1
16.5e+18
6.58
2.8316
36.18
64.64
50.99
38.68
27.25
40.97
28

---

G Details of Experimental Settings for Throughput Measurement
We implement continuous depth-wise batching inference system [6. 39] to evaluate decoding throughput of MoR models. Queries are enqueued and scheduled dynamically during decoding using 1K samples from the Feb-Eu 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 certain length of tokens per sample, where the number is sampled from a normal distribution with 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 4 is made against a modified vanilla model that includes these two extra layers, resulting in 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 batch size of 42, MoR-3 supports 48, and MoR-4 supports up 51. By employing recursion-wis 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 FlashAt- tention 2 [17] support variable-length KV caches within a batch. We adopt static-sized cache where each position is updated over time, since this is compatible with torch. compile [73] to fur- ther optimize inference speeds. Furthermore, mimicking real-world deployment scenarios [53, 109], 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-sharee 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 (Fi 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 [53]. We leave a more precise speed comparison, which accounts for these considerations, future work.
29

---

H Expanded Results of Parameter Sharing Strategy
This section complements the ablation in $4.1 by providing the full quantitative panorama be- hind Figure 4b. We revisit the four weight-tying schemes schemes-yycle, Sequence, Siidle-Cyycle and Selidde-SSquence--aa on two base model scales (135M and 360M non-embedding parameters) and two different recursion depths (N, 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, = 2 and N, = 3) and also shows the largest improvement in average accuracy compared 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.1 3.11 and shows 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 still performs worse than the Cycle-based methodology. We visualized the results in Figure 6.
Table 10: Comparison of parameter-sharing strategies (Cycle, Sequence, Middle-Cycle, Middle- Sequence) across two model scales (135M and 360M) and two recursion depths (Nn and NR 3). All models are pretrained from scratch on 10B tokens. We report validation negative log-likelihood (NLL) on FineWeb-Edd and few-shot accuracy across six tasks. Middle-Cyck consistently outperforms other strategies in both NLL and average task accuracy, especially at higher recursion depth. We highlight the optimal strategy for each setting in gray.
Preirain
Recursion
NLL
Few-shot Accuracy 1
Base Model
N-Emb
N₁.
Nick
Share
Loop
FineWeh
L.D
HS
PQ
WG
ARC
MMIJU
Avg
Vanilla 135M
106M
30
10B
3.0323
24.14
31.12
61.15
52.01
34.74
25,95
18.19
Vanilla 135M
53M
15
1OB
3,0818
23.64
30.10
60.94
50.99
35.38
25.93
37.83
Vanilla 135M
35M
10
10B
3.1582
21.46
29.30
60.01
52.01
34.40
25.53
37.12
Vanilla 135M
53M
15
10B
Cyc
2
3,0071
25.52
31.25
61.10
50.99
36.08
26.
38.51
Vanilla 135M
S3M
15
10B
Seg
2
3.1093
22.39
29.60
61.10
50.12
34.46
25.72
37.23
Vanilla 135M
57M
1+14+1
10B
M-Cyc
2
3,0330
23.40
31.20
61.59
50.59
35.44
25.54
37,96
Vanilla 135M
57M
1+14+1
10B
M-Seg
2
3,0991
21.70
30.06
60.45
49.41
35.20
25.74
37,09
Vanilla 135M
35M
10
10B
Cye
3.1154
21.42
30.14
60.61
49.72
34.15
25.57
36.94
Vanilla 135M
35M
10
108
Seg
3
3.1637
19.99
29.39
59.25
51.62
33.79
25.32
36.56
Vanilla 135M
39M
1+9+1
108
M-Cye
3
3.1048
22.41
30.35
61.04
49.01
34,80
25.91
37.26
Vanilla 135M
39M
1+9+1
10B
M-Sey
3
3.1602
20.69
29.35
61.43
51.30
34.40
25.51
37.11
Vanilla 360M
315M
32
10B
2.8471
27.27
34.78
64.20
52.80
38.29
26.72
40.68
Maiilla600M
157M
16
10B
2.8908
27.01
33.49
64.42
52.09
37.40
26.54
40,16
Vanilla 360M
98M
10
10B
2.9449
26.41
32.93
63.38
50.36
37.15
26.48
39.45
Vanilla 360M
157M
16
10B
Cyc
2
2.8487
28.47
34.79
63.06
49.96
37.38
26.81
40.08
Vanilla 360M
157M
16
10B
Seg
2
2.9467
26.33
32.49
62.89
52.41
36.37
26.24
39.46
Vanilla 360M
167M
1+15+1
10B
M-Cyc
2
2.8295
28.59
34.98
64.53
50.51
39. 668
27.20
40.91
Vanilla 360M
167M
1+15+1
10B
M-Sey
2
2.9303
26.14
32.71
62.79
51.38
36.31
25.73
39.18
Vanilla 360M
98M
10
10B
Cyc
3
2.9363
25.87
32.98
62.89
50.28
36.35
26.54
39.11
Vanilla 360M
98M
10
Seq
3
3.0245
24.55
31.48
63.11
49.25
35.65
25.73
38.30
Vanilla 360M
118M
1+10+1
10B
M M-Cyc
3
2.8760
28.51
34.89
64.31
50.51
39.51
27.20
40.82
Vanilla 360M
118M
1+10+1
10B
M-Scq
2.9753 3
24.18
31.89
62.08
49.72
36.47
26.27
38.44
30

---

3.2
Seg M-Seq
Seg M-Seq
Neg. Log Likalihood
3.1
Cyc
Neg. Log Likellhood
3.2
Cyc
M-Cyc
M-Cyc
3.1
3.0
2.9
2.9
2.8
2.8
135M
360M
135M
360M
(a) Recursion two (NR 2)
(b) Recursion three (Nn=
Figure 6: Validation negative log-likelihood (lower is better)
FineWeb-Edu for four parameter-sharing strategies. Bars are grouped by model capacity (135M vs. 360M parameters). Middle- Cycle consistently attains the lowest NLL, with its margin widening either model size depth increases. The horizontal dashed lines mark the untied (non-sharing) baselines: the lower red line represents the full capacity model, while the upper black line represents parameter-matched reduced model with footprint equal to the unique trainable parameter sizes of the recursive model.
For fairer comparison, 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 recursion steps. Especially, Middle-Cycle uses 8 layers for the recursion block and 2 unshared layers (1 + 8x3 + 1), resulting in 10 unique layers and 26 effective layers. We compare this to the Cycle strategy with 10 layers per recursion block (10x3), which also has 10 unique layers but total of 30 effective layers. Despite having fewer effective layers, as shown in Table 1. the Middle-Cycl model achieves better performance, with lower negative ligeliiel (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 future work.
Table 11: Comparison of Middle-Cycle and Cycle sharing strategies under the same number of unique parameters. All models 730M scales (non-embedding parameters) the base model with 3 recursion steps. The models are trained from scratch on 10B tokens using the Fineweb-Eda dataset.
Recursion
Pretrain
NLL |
Few-shot Accuracy
Base Model
Share
Loop
N-Emb
No
Nox
Fineweb
LD
PQ
WG
ARC
MMLU
Avg
Vanilla 730M
Cycle
252M
10B
2.7573
28.72
37.31
65.56
52.01
39.94
27.85
41.90
Vanilla 730M
M-Cycle
252M
1+8+1
10B
2.7552
29.75
37.73
66.00
52.49
40.49
27.50
42.32
31

---

Behavior under continued pre-training (up-training). Table 12 extends the study by "up-training models--cntinnen from open-sourcee SmolLM [4] checkpoints for 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.
Table 12:
Uptraining results across four parameter sharing strategies. Models are trained on 5B tokens from FineWeb-Edu and evaluated by train NLL and few-shot accuracy across six benchmarks. ARC denotes average of ARC-Easy and ARC- CRC-Challenge tasks, MMLU denotes the MMLU-Cont task. We highlight the optimal strategy for each setting in gray.
Pretrain
Recursion
NLL
Few-shot Accuracy
Base Model
N N-EEmh
N1.
Nov
Share
Init
Loop
FineWeb
1.1>
HS
PQ
WG
ARC
MMLU
Avg
Vanilla 360M
315M
32
5B
2.4825
41.67
50.63
70.35
55.09
46.99
30.82
49.26
Vanilla 360M
157M
16
5B
Step
2.7168
31.85
37.59
64.74
53.20
41.06
27.34
42.63
Vanilla 360M
157M
5B
Cye
Avg
2,8603
22.14
30.36
60.07
48.22
34,99
25.56
36.89
Vanilla 360M
157M
16
5B
Seq
Avg
2.7919
25.40
32.35
62.30
50.12
35.88
26.19
38.71
Vanilla 360M
98M
10
5B
Step
2,8915
26.63
35.03
64.42
52.09
38.75
26.86
40.63
Vanilla 360M
98M
Cye
Avg
3.0512
23.19
31.27
62.30
51.22
36.71
26.29
38.50
Vanilla 360M
98M
5B
Seq
Avg
2.9915
25.67
32.21
62.30
51.38
36.68
26.56
39.14
Vanilla 360M
157M
Cyc
Step
2
2.7165
31.30
37.68
64.91
52.17
39.29
27.53
42.15
Vanilla 360M
157M
5B
Cyc
Avg
2
2.8263
23.21
30.52
60.55
50.28
36.01
25.50
37.68
Vanilla 360M
157M
16
58
Cyc
Lower
2
2.8024
27.67
34.71
63.49
49.88
38.12
26.87
40.13
Vanilla 360M
157M
16
5B
Cyc
Upper
2
2.7915
18.26
34.88
63.06
$1.85
39.27
26.88
39.03
Vanilla 360M
157M
5B
Cye
2.7575
25.29
34.78
61.64
52.01
38.09
26.62
39.74
Vanilla 360M
157M
Seq
Step
2
2.6862
34.14
42.49
67,90
53.35
43.24
28.80
44.99
Vanilla 360M
157M
5B
Soq
Avg
2
2.7508
34.13
63.60
52.09
36.31
26.58
40.29
Vanilla 360M
157M
16
5B
Seq
Lower
2
2,8300
27.50
33.28
63.38
51.38
37.44
26.32
39.88
Vanilla 360M
157M
16
5B
Seq Upper
2
2.7498
30.49
40.07
65.61
52.25
40.28
28.17
42.81
Vanilla 360M
157M
Seq
Rand
2.7153 2
32,00
41.31
66.10
53.35
42.13
28.52
43.90
Vanilla 360M
167M
1+15+1
5B
MyC Step
2
2,6800
35.47
42.39
67.19
50.99
42.54
28.79
44.56
Vanilla 360M
167M
1+15+1
M-Cyc Avg
2.7314
33.81
40.42
66.87
51.78
41.68
28.17
43.79
Vanilla 360M
167M
1+15+1
5B
M-Cyc Lower
2
2.7449
30.76
50
66.16
50.99
41.12
28.07
42.77
Vanilla 360M
167M
1+15+1
5B
M M-Cyc Upper
2
2.6605
34.41
43.74
67.46
53.20
43.75
28. 93
45.25
Vanilla 360M
167M
1+15+1
5B
M-Cyc Rand
2.6730 2
35.65
43.04
67.74
52.17
42.60
28.62
44,97
Vanilla 360M
167M
1+15+1
58
M M-Seq
Step
2
2.6627
35.09
43.34
67.57
$1.22
43.66
28. 91
44,97
Vanilla 360M
167M
1+15+1
M-Seq
Avg
2
2.7143
33.92
40.93
66.49
51.70
40.72
28.24
43.66
Vanilla 360M
167M
1+15+1
5B
M-Soq
Lower
2.7696
38.36
65.94
51.78
41.27
27.73
42.64
Vanilla 360M
167M
1+15+1
5B
M-Seq Upper
2
2.6931
32.66
42.35
67.14
52.80
42.83
28.49
44.38
Vanilla 360M
167M
1+15+1
5B
M-Seq
Rand
2.6908 2
35.07
42.03
66.32
53.75
43.11
28.42
44.78
Vanilla 360M
98M
10
5B
Cyc
Step
3
2.8901
27.46
35.26
63.82
51.54
39.35
27.44
40.81
Vanilla 360M
98M
10
5B
Seq
Step
3
2.8258
30.43
37.37
63.76
52.33
40.55
27.65
42.02
Vanilla 360M
118M
1410+1
M-Cyc Step
2.7735
31.38
39.31
65.51
50.51
40.70
27.65
42.51
Vanilla 360M
118M
1+10+1
5B
M-Seg
Step
3
2.7678
31.67
39.23
65.89
52.09
40.65
27.90
42.91
Vanilla 360M
118M
1+10+1
5B
M-Cyc
Upper
2,8100 3
28.86
37.61
65.51
52.57
41.44
28.17
42.36
32

---

I Expanded Results of Design Choices for Router
L11 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 (a) 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 [111], II, 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:
2
(x(x)
log
where B is the number of tokens in the batch, N, the number of experts, and x E BxN,, denotes the logits input to the router. This regularization encourages the gating network produce smaller logits, promoting more stable and reliable routing decisions.
1.2 Router Performance Evaluation Metrics
Expert-choice routing. We evaluate dead token ratio and sampling accuracy assess the router's selection behavior. The dead token ratio measures the proportion of tokens at specific positions within the batch that 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 token belongs the top-k tokens identified during training, reflecting the router's 's ability to consistently select the most relevant tokens. Idcally, 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 using Max Vio (maximum violation) and entropy metrics. MaxVio [93] measures the load imbalance across experts:
experts
Max Vio max; Load, Load;
Load;
where Load, denotes the actual number of tokens assigned to the i-th expert, and Load, 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:
N. Σ D. log
=
where p. is the average probability of selecting the i-th expert over all tokens in the evaluation batch, and N, is the total number of experts. A higher entropy indicates a more uniform distribution of tokens among experts, reflecting balanced and diverse routing decisions.
I.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 positional bias. Notably, employing linear router architecture in conjunction with auxiliary loss effectively reduces the dead token ratio without compromising sampling accuracy.
33

---

Results from Table 14 reveal that applying an explicit balancing loss significantly reduces Max Vio and increases entropy, leading to improved load balance without sacrificing overall model performance. Loss-free approaches, while simpler, tend to show higher Max Vio and lower entropy, indicating less balanced token routing. Architectures such as MI.P 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.
Table 13: Ablation results on using expert-choice router with different routing configurations, We use the recursion-wis KV caching strategy by default. Coeff denotes coefficient values for auxiliary loss term, and a denotes scaling term after router function. Dead token ratio are measured within evaluation batch size of 500. Warmup refers to gradually decreasing the capacity from 1.0 to desired value over warmup steps. The last highlighted row represents chosen final strategy, and intermediate best-performing designs (based on performance and routing metrics) are highlighted to illustrate how it was derived.
Expert-choice Configurations
Router Metrics
NLL.
Few-shot Accuracy 1
Sampling
Coeff' Func
Arch
Warmup 2-loss
Dead
Samp-Ace
FineWeb
LD
HS
PQ
WG
ARC MMLU
Avg
rand
X
x
0.0
2,9335
26.0
33.1
61.6
52.3
35.8
26.2
39.1
Aux Router
MLP
X
X
66.7
50.0
NaN
0.0
25.04
49,5
49.6
23,9
23.0
28.5
Aux Roater
0.1
MI.P
X
X
0.0
89.2
2,8993
26.1
33.8
62.0
51.5
36.6
26.4
39.4
Aux Router
1.0
MLP
K
66.7
50.0
2. 2.8867
26.4
33.6
63.0
52.4
37.0
24.1
39.8
Aux Router
tanh 0.1
MI.P
X
x
66.7
98.6
2,8720
13.9
31.8
60.7
49.3
35.8
25.8
36.2
Aux Router
tanh
1.0
MLP
66.7
97.0
3.0624
18.26
29.7
60.1
50.9
34.6
25.5
36.5
Aux Loss
0.01
ML.P
X
X
66.7
50.0
NaNl
0.0
25.04
49.5
49.6
23,9
23.0
28.5
Aux Loss
0.01
0.1
MLP
0.0
99.6
2. 2.6977
24.8
33.6
63.3
50.3
36.6
26.6
39.2
Aux Loss
0.01
1.0
MLP
X
X
65.9
100.0
2.9189
12.0
31.6
59.4
51.5
33.2
25.3
35.5
Aux AuxLoss Loss
0.01
tanh
0.1
MLP
32.8
99.7
2.9426
23.5
32.4
62.4
49.8
35.6
26.0
38.3
Aux Loss
0.01
tanh
1.0
MLP
X
0.0
98.8
3.2743
16.4
28.14
58.8
52.2
31.6
24.8
35.3
Aux Loss
0.1
0.1
MLP
0.0
99.8
3.0416
21.5
31.0
61.8
50.3
35.0
26.0
37.6
Aux Loss
0.001
0,1
MLP
X
x
0.0
99.1
2. 281816
27.6
34.3
63.0
51.6
36.7
26.5
40.0
Aux Loss
0.001 tanh
0.1
MI.P
x
x
0.0
56.4
2,9933
25.0
32.3
61.5
51.5
36.6
26.0
38.8
Aux Loss
0.001
a
0.1
Linear
x
0.1
99.2
2.8667
27.4
34.6
63.2
51.5
37.2
26.5
40.1
Aux Loss
0.001
0.1
W-MLP
0.4
99.2
2.8716
27.8
33.9
62.4
49.9
36.3
26.3
39.4
Aux Loss
0.001
0.1
Linear
4.9
99.1
2.8744
26.0
33.9
62.0
51.2
36.1
26.1
39.2
Aux Loss
0.001
0.1
Linear
/
0.0
99.3
2.8824
26.9
34.0
63.8
52.3
36.8
26.4
40.0
Table 14: Ablation results on token-choice router under different routing configurations. We use the recursion-wise KV caching strategy by default. Coeff denotes coefficient for balancing loss term, or updating coefficient (u) for loss-free algorithm. The last highlighted row represents chosen final strategy, but we added z-loss back in with a small coefficient of le-3 since it often stabilizes load balancing. The intermediate best pers-rrefoominggg designs (based on performance and routing metrics) are highlighted to illustrate how it was derived.
Token-choice Configurations
Router Metries
NLL
Few-shot Accuracy
Balancing Coeff
Func
Arch
Z-loss
MaxVio
Entropy
FineWeb
LD
HS
PQ
WG
ARC
MMLU
Avg
rand
0.007
1.099
3,0268
24.8
32.0
61.4
52.2
35.5
26.1
38.7
Loss
0.1
soft
1.0
MLP
0.200
1.076
3.0239
24.2
31.9
61.4
51.5
35.7
26.2
38.5
Loss
0.01
soft
1.0
MI.P
0.682
0.921
2,9118
28.0
33.3
62.8
49.7
36.4
26.2
39.4
Loss-free
0.01
soft
1.0
MLP
1.788
0.297
2.9078
25.5
32.5
61.3
52.3
36,1
26.0
38.9
Loss-free
0.01
a
0.1
ML.P
0.956
0.646
3.1144
21.8
29.8
60.3
51.6
34.0
25.7
37.2
Loss-free
0.01
1.0
MLP
0.918
0.749
3.0188
23.4
31.3
59.9
50.0
35.2
25.8
37.6
Loss free
0.001
soft
1.0
MLP
0.852
0.915
2,9081
25.8
33.6
62.8
50.6
37.5
26.7
39.5
Loss-free
0.001
a
0.1
MLP
1.281
0.551
2.9165
23.9
33.1
61.2
51.6
37.3
26.2
38.9
Loss-
0.001
1.0
MLP
0.542
0.941
3.0188
24.9
32.0
61.9
51.4
35.5
25,9
38.6
Loss
0.1
soft
1.0
Linear
0.492
0.960
2,9974
23.7
31.3
62.2
50.3
36.7
26.0
38.4
Loss
0.1
soft
1.0
W-MLP
0.384
1.037
3.0293
25.3
31.5
62.2
$1.2
36.4
26.3
38.8
Loss
0.1
soft
1.0
Linear
0.266
1.056
2.9358
25.7
32.6
61.9
51.7
36.4
26.5
39.1
34

---

Expanded Results of KV Cache Sharing Mechanism
J.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 [9]. This technique can reduce the memory footprint required 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 [85]. Due to the high degree of freedom Vanilla models-where trainable parameters can be well optimized for shared taches-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 increase, the projection layers appear to be trained to produce similar signal sizes corresponding depths within each recursion.
500
16
12
Average Norm
400
14
300
200
10
100
24
16 24
32
32
Layer Index
Layer Index
Layer Index
(a) Hidden states
(b) Key states
(c) Value states
Figure 7: Average L2 norm magnitude of (a) hidden states, (b) key states, and (c) value states across layers in Middle-Cycled Recursive Transformer 360M with 3 recursion steps. Note that the last hidden states correspond to the final hidden states after the last layer normalization.
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.
1.0
8
Layer Index
16
16
0.5
24
32
31
0.0
16
24
16
31
16
31
Layer Index
Layer Index
Layer Index
(a) Hidden states
(b) Key states
(c) Value states
Figure 8: Cosine similarity matrices showing the layer-wise similarity of (a) hidden states, (b) key states, and (c) value states Recursive Transformer with Middle-Cycle strategy and recursion depth
3. Results are from a 360M parameter model with 32 layers. The hidden states matrix includes the final hidden states after the last layer normalization.
35

---

J.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 $C.1 for details). Interestingly, KV cache sharing even improves the performance of vanilla models, where sharing acts as 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 slight perplexity drop, albeit not substantial.
When moving MoR models, they still introduced 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 i) 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, still worth exploring combinations of KV sharing and actual updates.
Table 15:
Comparison of KV cache sharing strategies across Vanilla, Recursive, and MoR Transformers. Models are pretrained on 10B tokens of FineWeb-Edu, and evaluated using negative log-likelihood (NLL) on train set and few-shot accuracy across benchmarks. KV sharing denotes use of recursive KV sharing mechanism. If MoR is mentioned without further specification of KV sharing strategy, it implies the use of the recursion-we wise caching strategy. 'It indicates training with hybrid KV sharing that leverages shared caches for inactive positions while updating active ones through actual computation.
Pretrain
Recursion
MoR
KV Sharing
NLL
Few-shot Accuracy 1
Models
N-Emb
Ni.
Share
Loop
Type
Share
Loop
FineWeh
L.D.
HS
PQ
WG
ARC
MMLU
Avg
Vanilla
315M
32
28847
27.3
34,8
64.2
52.8
38.3
26.7
40.7
Vanilla
315M
Seq
2.7848
30.0
36.5
64.6
50.7
39.4
26.9
41.3
Vanilla
315M
Cyc
2.7650
30.0
36.7
65.2
51.1
39.6
27.5
41.7
Vanilla
295M
2,8069
29.1
35.6
65.1
50.4
38.5
27.3
41.0
Vanilla
295M
30
Seq
2.7879
28.3
36.4
64.3
52.7
39.4
27.3
41.4
Vanilla
Cye
2.7890
39.0
Recursive
157M
Seq
2
2.9467
26.3
32.5
62.9
52.4
36.4
26.2
39.5
Recursive
157M
16
Soq
2
Seg
2,8904
26.4
33.4
64.0
51.0
37.0
26.9
39.8
Recursive
157M
16
Cyc
2
2.8487
28.5
34.8
63.1
50.0
37.4
28.8
40.1
Recursive
Cyc
Cye
37.3
Recursive
167M
1+15+1
M-Cyc
2
2,8295
28.6
35.0
64.5
50.5
39.7
27.2
40,9
Recursive
167M
1+15+1
M-Cyc
M-Cye
2.8451
27.3
34.7
63.7
50.5
37.8
27.0
40.2
Recursive
98M
Soq
3
3.0245
24.6
31.5
63.1
49.3
35.7
25.7
38.3
Recursive
98M
Seg
2.9554
24.2
32.3
62.5
52.7
36.6
26.2
39.1
R Recursive
98M
10
Cyc
2.9363
25.9
33.0
62.9
50.3
36.4
26.5
39.2
Recursive
98M
Cyc
Cyc
2.9155
24.1
32.9
62.4
51.2
37.4
26.7
39.1
Recursive
118M
1+10+1
M-Cyc
2,8760
28.5
34.9
64.3
50.5
39.5
27,2
40.8
Recussive
118M
1+10+1
M-Cye
M-Cye
2.8854
27.3
33.8
63.3
52.3
37.5
26.8
[0.22
MoR
118M
1+10+1
M-Cye
3
Expert
2,8667
27.4
34.6
63.2
51.5
37.2
26.5
40.1
MoR
1+10+1
M-Cyc
Expert
M-Cye
2.8995
34.0
61.6
50.2
26.0
36.5
27.0
39,2
MoR
118M
1+10+1
M-Cyc
Expert
M-Cye!
2.8653
24.8
34.3
62.0
36.7
26.7
39.1
MoR
118M
1+10+1
M-Cyc
Token Token
2.9358
25.7
32.6
61.9
36.4
39.1
MoR
118M
1+10+1
M M-Cye
M Cye
2.9155
25.7
32.6
61.8
49.4
36.2
26.0
38.6
36

---

Relaxation for key-value sharing constraints. We also investigated relaxing the constraints on KV sharing in Table 16, similar the relaxation approach in Bae et al. [6] for parameter sharing constraints. Specifically, we first re-examined four relaxation techniques for standard Recursive Transformers with very small ranks [42, 61] or prefix lengths [62]. We also experimented with the position encoding [12), where trainable embeddings are element-wise multiplied with the output of each recursion block.
Our results show that these techniques do not provide substantial performance improvements when pretraining relaxed models from scratch, consistent with prior studies, as they introduce only limited number of additional parameters. Although we hypothesize that incorporating prefix- based approaches (such adding trainable prefixes attention) into KV sharing might lead to greater benefits, our experiments did not reveal substantial differences in this regard. Further exploration of more sophisticated techniques for efficiently relaxing KV cache sharing constraints remains an open direction for future research.
Table 16: Experimental results of relaxing parameter sharing and KV cache sharing constraints in Recursive Transformers, All models are trained on FineWeb-Edu with 10B tokens, and we apply the Middle-Cycle parameter sharing for 360M models with 3 recursion depths. We evaluate them based on training NLL and few-shot accuracy across six benchmarks. Relaxation types include encoding trainable embeddings on recursion outputs via element- wise multiplication (Enc), applying LoRA and DoRA to query and value weight matrices, and adaptation prompt tuning (Adapt-P).
Pretrain
Relaxation
KV Sharing
NLI.
Few-shot Accuracy
Models
N-Emb
N,
Type
Len
Share
Loop
FineWeb
PQ
WG
MMLU|
Avg
Recursive
118M
1+10+1
2.8854
27.3
33.8
63.3
52.3
37.5
26,8
40.2
Recursive
118M
1+10-1
Enc
2,8604
27.3
34.6
63.9
534
38.6
26.7
40.2
Recursive
124M
1+10+1
LoRA
64
2,8599
27.3
34.6
64.3
50.9
38.0
26.9
39.7
Recursive
124M
1+10+1
DoRA
64
2.8945
26.4
33.6
64.4
50.6
37.4
26.5
39.2
Recursive
126M
1+10+1
Adapt-P
256
2.8626
64.0
51.9
26.8
39.7
Recursive
1+10+1
M-Cye M-Cye
2.8854
52.3
26.8
Recursive
126M
1+10-1
Adapt-P
256
3
2,9030
24.5
33.1
63.0
52.2
26.7
37.6
39.5
K
Expanded Qualitative Results
K.1 Analysis Adaptive Computation Paths
Table 17 illustrates a qualitative analysis of the recursion depth assigned to each subword token. This visualization provides detailed insight into how tokens within each sample exhibit varying levels of recursive processing, showcasing the adaptive computation mechanism within the MoR framework. Notably, some tokens exit early (purple), while others require deeper processing (blue and red), reflecting the model's 's ability to focus more compute on challenging parts of the input.
37

---

Table 17: Visualization of the recursion depth for each subword token, with colors representing the number of recursion steps: 1. 2 and This depth specifically indicates how many times recursion is applied at that token's 's position to predict the subsequent token. Each row corresponds to single sample, offering a clear illustration
the token-level recursion distribution in practice. We use an MoR model with N, = 3, auxiliary loss, and recursion -uiss KV caching. This model is built on a 360M parameter base and trained on 30B tokens.
Sample
Text
Sample #1
People who feel comfortable defending their views - def ensively confident - may also eventually change those views and corresponding behaviors. National Election Studies surveys showed that defensive confidence predicted def ection in the 2006 U.S. House elections, above and beyond the impact of various demographic and political variables. Moreover, defensive confidence was also associated with political knowledge and attention to polities and government affairs, but not
Sample #2
2014. 7:57 AM Space X Falcon 9-R Rocket Suff ers M alf unction. Self Dest ructs During Test Flight August 23, 20 , 9:36 AM Texas Ch osen as Site for SpaceX &#x27;s First Commercial Launch pad August 5, 2014, 4, 1:44 PM South Carolina Prison Find C ras hed Dr one Car rying Drugs, ones August 1. 2014, 14 2 2:49 PM NASA & Mars 2 2020 Rover G ains Seven New Instruments for Exploration August 1, 2012 4. 1:30 PM NASA
Sample #3
9 PR New sw ire. All rights reserved A report released Thursday on the Slide Fire Oak Creek Canyon doesn & hold back in laying out the danger facing the burned out area. The U.
S. Forest Service issued the report Thursday afternoon stating the biggest concern is the st eeper slopes of Oak Creek Canyon is especially subject to debris flows, rocks I ides, flash flooding and erosion that could become concentrated flow or landslide. The report said in smaller streams
Sample #4
bo to accompany the dwar ves to fight the enemy. He says, "S ar uman believes it is only great power that can hold evil in check, but that is not what I have found. found is the small everyday deeds ordinary folk that keep the darkness at bay. Small acts of kindness and love." That s what Jesus teaches us as well. Warning us that we would live in dark times, He reminded us that hecause of Him we are "the light of the world" ( Matt. 5:14)
Sample #5
mitting Plants After four years of research, most of it in total darkness, a Stanford University plant biologist has discovered that some plants have a system fiber optics that can transmit light up and down their tissues in a way similar to the method the telephone company uses to transmit many of its telephone calls. Dr. D Mand oli, seeking to find the light- transmit ting properties of plants, had to rely on her sense of touch as she placed tissues of oat, m ung
Sample #6
an impression of market power abuses or other market failures. In some cases, , however, prices may spike to higher than usual levels and cause public concern and the need more public information. address such events, the proposed amendment includes an event trigger that would require the public release of entity- specific information on a much quicker timeframe The proposed amendment requires that, when the trigger is exceeded, the portion every
Sample #7
and spaceflight and spacecraft covers elementary astronomy. Newton ian mechanics, the Sun and related physics and spaceflight. Also included are Spanish translation. 46 lesson plans, short but complete math course (al gebra + trig ). teachers &#x27: guides, glossary. timelines. 445 questions (uurrent tally) by users and their answers over 100 problems to solve. and more. Learning Design Eng ines as Remote Control to Learning Support Environments. Context
Sample #8
well established role in scientific computing and recent increased presence in desktop computing. it almost certain that contemporary information professionals will encounter Unix based systems in their work. This workshop is an intermediate level look at the Unix operating system as compared to the Unix introduction that our S 5401 I students receive. We will have short review but will be essential to have the S401- UN IX material internalized as much as
Sumple #9
code and will serve as a good introduction to the syntux necessary creating shell programs. We will continue by discussing the importance quoting now that we have used more powerful met ach ar acters (variable and command subs itution) ) and need control their expansion interpretation) ). Finally. , we will end our day revisiting the Unix find utility and looking at some advanced uses for this indispensable tool. SI ides for Day 5 (PDF) A program running in Unix
38

---

K.2 Analysis on Router Weights
To gain insight into the router output distributions, we visualized the results in Figure 9, Our analysis reveals that various routing mechanisms are optimized to balance expert loads according to the desired capacity. Notably, expert-choice routers achieved nearly perfect load balancing with the auxiliary loss, resulting in almost binary values (1 or 0) for selected and unselected tokens, respectively. For the auxiliary router, it was able to distinguish between selected and unselected tokens to some extent, but still showed overlapping. Since this strategy allows for different capacity factors during training and inference", further research into methodologies that can more distinctly separate inter-cluster variations seems necessary.
Other token-choice strategies also exhibited good balancing properties with reasonable router values, which are used to refine the outputs of the corresponding recursion computation blocks. However, most cases failed converge to optimal load balancing (i.e., these were edge cases where they achieved their own optimal load balancing), highlighting the challenges of achieving consistent performance heterogeneous expert settings.
Selected
Unselected
Selected
Unaelected
1.0
1.0
Normalized Counts
Counts
0.8
0.8
0.2
0.2
0.0
0.0
0.0
0.5
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
0.0
0.5
1.0
0.5
1.0
Rec
Rec 42
Rec #3
Rec
Rec #2
Rec #3
(a) Expert-choice (Auxiliary loss)
(b) Expert-choice choice (Auxiliary router)
Normalized Counts
Normalizzd Counts
0.4
0.2-
0.2
0.0
0.0
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
0.5
0.6
0.5
0.6
0.5
0.6
Rec #1
Rec #2
Rec #3
Rec #1
Rec #2
Rec #3
(c) Token-choice (Balancing loss)
(d) Token-choice (Loss-free)
Figure 9: Distribution of router weights for selected and unselected tokens each recursion step in expert-choice and token-choice MoR (N,=3). All models use the recursion-wist KV caching strategy. The subplots show results for (a) expert-choice routing with auxiliary loss, (b) auxiliary router, (c) token-choice routing with balancing loss, and (d) loss-free algorithm. Each subplot uses the best hyperparameter settings identified in Table 4.
'The auxiliary router learns to capture intra-token differences within selected or unselected groups, rather than being biased towards extreme points.
39

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Abstract and Introduction sufficiently contain our main contributions: present- ing Mixture-of-RReursiii that unifies parameter sharing, token-level adaptive depth, and memory-efficient KV caching, within single framework.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
In Appendix B, we provide detailed discussion of the limitations of our work.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
40

---

3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification:
We do not include any theoretical results in this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provided detailed instructions of our new Mixture-of- (MoR) architecture in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution
specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
41

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [No]
Justification: We rely accessible, well-known, and publicly available model architecture
and datasets. Code for both training and evaluation will be released upon publication.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ /nips. wwtpppplc/ggccieecccceooooocccce guides/ /CodeSummddiebiicc cimmisssionn wwwssssinolooooooooeo seiiinnoiccy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Appendix D contain all the experimental settings and details. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Our work involves pretraining models from scratch using a fixed data order, which minimizes sources of randomness. Therefore, there is little no stochastic variation in our experiments that would necessitate reporting error bars or statistical significance measures.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
42

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Appendix D, we elaborate the details on compute resource we have used for the experiments and report the elapsed time.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //werri/////eeeeeeeeeeeeeess htpppli//tthissuuiiiittee
Answer: [Yes]
Justification: We have reviewed the NeurIPS Code of Ethies and couldn't identify any violation of the code of ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., - if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This study does invoke considerable societal concerns.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
43

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: This paper does not contain such risks for data/model misuse. Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification: We follow the conventional LLM evaluation protocols for every experiment and dataset, which we assume are properly cited.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
44

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not provide new assets in this research.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do not deal with crowdsourcing nor human subjects in this research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
We do deal with crowdsourcing nor human subjects in this research.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
45

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
We do involve LLMs for the core method development in this research.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: //tp/r/is....//eeee. cc/Conferences/20022 /LLM) for what should or should be described.
46


Reviews:
[{"review_id": "dLmwsM0EPf", "reviewer": "Reviewer_cDNG", "summary": "This paper proposes a new method named mixture-of-recursive, which considers both parameter and computation efficiency. It combines the (expert-level or token-level) router with recursive layers, and also KV cache sharing strategies to improve model performance, can achieve similar or even better performance than vanilla transformer with 1/3 parameters when training from scratch.", "questions": "could you give some intuitions about the example in Fig.1 Right bottom, why words like \"---\" or \"and\" \"those\" will be pass twise while words like \"comfortable\" \"views\" only once?", "limitations": "authors mention their limitations in appendix A", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The ideas are clear and effective, and the writing is excellent and very easy to follow\n2. The experimental results are comprehensive and impressive. Utilizing recursive layers could save the number of parameters but will always cause a little drop of performance. While here utilizing token router can let model learn better on important tokens and improve performance. This paper considers different scales of models and computations and different combinations of strategies, and show consistent advantage compared with vanilla recursive layers.\n3. Doing many ablations studies to show how to combine different routers and archs, etc.\n\nWeakness:\n1. maybe giving the calculation process of Tab2 (right) in appendix and refer it in Tab 2 would be helpful for readers, because they are dependent on some specific capability factor sequence (N_r/N_r, ..., 1/N_r)\n2. It seems that as increasing compute budget, the performance gain from MoR is decreasing. And on small models (like 135M), it perform worse than vanilla models.", "was_revised": true, "final_justification": "experiment results are very solid and the architecture is interesting, I think definitely can be accepted", "rebuttal": "We thank the reviewer for their careful review and thoughtful comments. We appreciate that the reviewer found our idea to be effective and clear, and our experimental results to be comprehensive. We provide responses to the reviewer’s specific questions and concerns (quoted) below.\n\n---\n > **[W1]** Maybe giving the **calculation process of Tab2 (right)** in appendix and refer it in Tab 2 would be helpful for readers, because they are dependent on some specific capability factor sequence (N_r/N_r, ..., 1/N_r)\n\n**[A1]** Thank you for your valuable suggestion. We will include a detailed explanation of the cost efficiency numbers in the appendix of the revised version and refer to it from the table caption. Below is a summary of the calculation process for **Table 2 (right)**. Please note, these values are calculated based on the *Cycle* sharing strategy. The *Middle-Cycle* strategy differs slightly, as it includes single layers at both the beginning and the end, which may lead to some variations.\n\n**[KV Cache Memory and IO Comparison]** \n\nWe consider the cost of processing a full sequence at each recursion step as 1 unit. **Table 2 (right)** presents the calculated values across the *entire model* for each strategy, normalized against the Vanilla strategy.\n\n- **Vanilla:** Every token is activated over $N_r$​ steps (i.e., all layers). This means it requires $N_r$​​ times the KV memory and corresponding IO (loading time) to fetch the data.\n\n- **Recursion-wise KV Caching:** The number of tokens selected at each recursion step ($r$) decreases linearly. Specifically, a fraction of $(N_r + 1 - r)/N_r$ of tokens are activated (e.g., $N_r / N_r$ for the first step, then $(N_r - 1) / N_r$, and so on, down to $1/N_r$). Summing these fractions across all $N_r$ steps gives us: \n$$\\frac{N_r}{N_r} + \\frac{N_r-1}{N_r} + \\cdots + \\frac{1}{N_r} = \\frac{N_r(N_r+1)}{2N_r} = \\frac{N_r + 1}{2}$$\nThis calculation shows that both the KV memory and IO requirements for this strategy are $(N_r + 1)/2$ in total.\n\n- **Recursive Sharing:** In this approach, the KV cache generated for the full sequence during the very first recursion step is continuously reused. Consequently, the memory cost is just 1 unit. However, fetching operations are still required at all $N_r$ steps, making the IO cost $N_r$​.\n\n**[Attention FLOPs Comparison]**\n\nFor the Attention FLOPs *per layer*, we report values normalized by the Vanilla case.\n\n- **Vanilla:** The attention mechanism exhibits quadratic complexity with respect to the sequence length ($N_{\\mathrm{ctx}}$).\n\n- **Recursion-wise KV Caching:**  In this strategy, we activate only the top-$k$ tokens at each recursion step ($r$), where $k = \\frac{(N_r + 1 - r)}{N_r}N_{ctx}$. Consequently, the attention mechanism has a complexity of $k^2$.\n\n- **Recursive Sharing:** When using sharing, the query length is reduced to top-$k$, but the key states still utilize the full sequence length of $N_{\\mathrm{ctx}}$. This results in a complexity of $k N_{\\mathrm{ctx}}$.\n\n---\n> **[W2]** It seems that as **increasing compute budget, the performance gain from MoR is decreasing**. And on small models (like 135M), it performs worse than vanilla models.\n\n**[A2]** In **Figure 3**, Recursive Transformer and MoR models are constructed by applying recursion to a base Vanilla model. Therefore, the trainable parameters for Recursive Transformer and MoR (with a recursion step of 3) are approximately *three times fewer* than those of Vanilla models. Consequently, a direct comparison based solely on FLOPs, as shown in **Figure 3**, can be misleading, as it doesn't account for the differing parameter counts that impact total training steps and convergence speed. \n\nTherefore, to accurately assess the models' performance in relation to their actual parameter size and computational budget, we refer to **Figure 5a**, which provides a more straightforward comparison by also considering the number of trainable parameters.\n**Figure 5a** visualizes the scaling law results from **Figure 3** by fitting them to axes of model size and validation loss. Recursive (blue) and MoR (red) are positioned to the left of the Vanilla curve (green) because they use a recursion step of 3. **Figure 5a** shows several key insights:\n\n- **As the compute budget increases, the MoR model demonstrates strong scalability.** The gap between the optimal points of MoR (indicated by stars) does not shrink with increasing FLOPs, confirming that MoR is a highly scalable architecture.\n\n- **MoR consistently outperforms the Recursive baseline.** While recursive baselines with fixed recursion depth might appear to converge better at larger compute budgets, the similar differences between optimal points within each MoR and Recursive at higher budgets suggest that MoR's performance advantage will likely persist. It's crucial to note that vanilla recursive models require an additional training phase for early-exiting [1], which can be sensitive to hyperparameters and introduce performance drops. \n\n- **MoR demonstrates comparable performance to Vanilla Transformers given the same training budget.** Across three different budgets, the scaling curve just shifts to the left due to the reduction in model size with recursion. We further validated the performance for the 1.7B model with a 6.85e19 FLOPs budget in **Table 1R**, which also showed comparable performance. \n\n**[Table 1R]** Comparison of MoR, Recursive, and Vanilla 1.7B models under 6.85e19 FLOPs budget. We used expert-choice router and recursion-wise KV caching for MoR, and $N_r = 2$ for MoR and Recursive models.\n\n| Model       | Share | $N_r$ | Param  | FLOPs   | $N_{\\mathrm{tok}}$ | LD    | HS    | PQ    | WG    | ARC   | MMLU  | Average |\n| :---------- | :---- | :---- | :----- | :------ | :----------------- | :---- | :---- | :---- | :---- | :---- | :---- | :------ |\n| **Vanilla** | -     | -     | 1.61B  | 6.85e19 | 20B                | 40.7% | 49.2% | 71.3% | 53.2% | 30.8% | 47.3% | 48.8%   |\n| **Recursive** | M-Cyc | 2     | 0.87B  | 6.85e19 | 20B                | 37.3% | 46.5% | 68.9% | 52.6% | 29.6% | 44.2% | 46.5%   |\n| **MoR** | M-Cyc | 2     | 0.87B  | 6.85e19 | 26B                | 40.2% | 48.0% | 69.6% | 52.8% | 30.6% | 47.0% | 48.0%   |\n\nLooking at the optimal scaling policy in **Figure 5a**, MoR performance can be lower when the shared block size is very small (using an adequately larger model is better at the same FLOPs). Notably, a 135M-based MoR model has roughly 45M unique parameters—nearly three times fewer than vanilla models. However, we would like to emphasize this is simply a **performance-efficiency trade-off**; MoR achieves much faster training speeds (wall-clock time) and nearly double the throughput gain. To ensure a fairer and more accurate comparison, we should consider *“speed”* in addition to performance.\n\n---\n> **[Q1]** could you give some **intuitions about the example in Fig.1 Right bottom**, why words like \"---\" or \"and\" \"those\" will be pass twice while words like \"comfortable\" \"views\" only once?\n\n**[A3]** Thanks for highlighting these valuable points. In **Figure 1 (right)** and **Figure 10** (in the Appendix), we show the recursion depth performed for each subword token. To clarify, this depth indicates how easily the next token can be predicted within its given context. Essentially, the allocation of recursion depth reflects the *contextual predictability* of the subsequent token.\n\nThrough several examples, we made the following observations:\n\n- The second (or subsequent) subword part of a word is often straightforward to predict, thus requiring fewer steps. For instance, in a word like \"defensively,\" only a single recursion is performed to predict \"-ensively,\" (i.e., “def” is highlighted in purple) which aligns with our intuition.\n\n- In many cases, the initial generation (opening) of function words like `---  ( { . ,` appears easy for models to predict. However, predicting their closing parts or the first token immediately following their opening might be more challenging, often requiring 2 (blue) or 3 (red) recursions.\n\n- Generally, more content-rich or less common words are typically assigned deeper recursion, as expected.\n\nWe will make sure to clarify these points in the revised version of the paper. For full transparency, we avoided cherry-picking samples; the examples discussed were simply the first 10 samples from our validation set.\n\n---\n**References**\n\n[1] Bae et al. “Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA”. ICLR 2025."}, {"review_id": "6tta7UXcxR", "reviewer": "Reviewer_MjJH", "summary": "In this work, the authors present Mixture-of-Recursions (MoR), a unified framework for improving both the training and inference efficiency of Transformer models. MoR combines two key ideas: parameter sharing through recursive reuse of a shared stack of layers, and adaptive computation over the tokens (i.e, not all tokens proceed to full depth/receive full compute). This allows the model to allocate computation proportionally to importance, reducing redundant computation and significantly lowering overall cost compared to conventional baselines. MoR is shown to dominate the Pareto frontier in iso-FLOP settings, consistently achieving lower validation perplexity with better efficiency. Additionally, this conditional token computation delivers strong inference speedup on several model sizes. These results highlight MoR as a compelling direction for scaling model quality without incurring full-scale computational costs.", "questions": "The main question I have is whether the authors could explicitly contrast MoR with a conceptual combination of MoE and MoD. Since the layer router could be viewed as a generalization of expert selection, and the token-router can be viewed similarly as depth allocation, it would be useful to understand what architectural or algorithmic elements of MoR are fundamentally new, beyond being a combination of existing approaches (refer to weaknesses).", "limitations": "yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. Clarity and Presentation: This paper is well written, with diagrams and figures that clearly convey the core ideas. In particular, Figure 1 effectively illustrates how computation is allocated per output token during the decoding stage, providing intuitive insight into the mechanism of Mixture-of-Recursions (MoR).\n\n2. KV Caching Design: The approach to key-value caching is both novel and well-motivated. The authors show thoughtful design by selectively storing KV entries only for tokens that pass through routed blocks, minimizing unnecessary memory use. Meanwhile, the layer-sharing component ensures full-context awareness during generation, demonstrating a strong understanding of the importance of maintaining contextual integrity in autoregressive decoding.\n\n3. Strong Empirical Results: The experimental results are compelling. MoR consistently dominates the compute-performance Pareto frontier across model scales and benchmarks. The ablations on cache strategy, throughput, and router configuration provide valuable insights into the contribution of each component and the rationale behind the final design choices. The detailed router analysis further supports the effectiveness of the method.\n\nWeaknesses:\nWhile the paper is overall well-executed and well-motivated, I believe the methodology could benefit from deeper discussion in relation to prior work on conditional depth computation. In particular, \"Mixture-of-Depths: Dynamically Allocating Compute in Transformer-Based Language Models\" (arXiv 2024) is notably relevant. That paper also explores token-level conditional routing using depth-based computation and introduces a similar routing heatmap to Figure 1 in this work.\n\nGiven these similarities, it would be valuable for the authors to clarify how MoR goes beyond a naive combination of Mixture-of-Experts (MoE) and Mixture-of-Depths (MoD). Without this, one might interpret MoR as simply integrating known ideas (layer skipping + routing) under a recursive framing. An in-depth comparison or positioning in the related work would help distinguish the unique contributions of this paper.", "was_revised": true, "final_justification": "I have reviewed the rebuttal and maintain my score. The paper is technically solid, and the authors' agreement to revise the text with the distinctions satisfies my previous concerns.", "rebuttal": "We appreciate the reviewer's thorough review and positive feedback. We're glad that the reviewer found our core idea clear and well-motivated, and our experimental results strong. Below, we address the reviewer’s specific questions and concerns (quoted). \n\n---\n> **[W1]** While the paper is overall well-executed and well-motivated, I believe the methodology could benefit from **deeper discussion in relation to prior work on conditional depth computation, such as MoD and MoE**. An in-depth comparison or positioning in the related work would help distinguish the unique contributions of this paper.\n\n**[A1]** Thank you for suggesting us to clarify the distinctive contribution of MoR compared to MoR and relevant works, MoD and MoE. We will add this discussion to the revised script.\n\n**[Comparison with MoD]**\n\n While MoR and MoD explore **token-level routing over network depth**, MoR extends to be a unified framework by incorporating parameter sharing, early-exiting, and various KV caching mechanisms:\n\n- **Parameter sharing:**  MoR uses **recursive blocks as the routing primitive**, enabling substantial parameter efficiency through weight sharing. This contrasts with MoD, which does not leverage parameter sharing and thus cannot achieve the same reductions in model size without sacrificing performance. \n\n- **Early-exiting and throughput gains:** In MoR, our **hierarchical filtering** scheme does not allow for “deactivated” tokens to be reactivated in subsequent blocks. Once a token is \"early-exited\" by the router at a certain recursion depth, its computation is considered complete, and it does not participate in further recursive blocks. Thus, early-exiting and parameter-sharing structure enables **continuous depth-wise batching** during inference, which significantly boosts throughput by reducing idle time and maximizing resource utilization. \nIn contrast, MoD's throughput gains are potentially limited. While MoD also reduces per-block computation based on the number of activated tokens, its mechanism does allow for deactivated tokens to be reactivated later. This *reactivation mechanism* prevents the same level of idle time reduction achievable with MoR's early-exiting mechanism and depth-wise batching.\n\n- **KV caching mechanism:** Both MoR and MoD improve memory and throughput efficiency by using attention mechanisms only on activated tokens. We further propose a *KV sharing* method to optimize memory even more by leveraging the recursive structure of our MoR blocks and reusing the KV cache from the first recursion step.\n\n**[Comparison with MoE]**\n\nMoE and MoR share a high-level similarity in their pursuit of **token-level conditional computation** via a routing mechanism. This enables us to explore various techniques developed for MoE, such as expert-choice and token-choice strategies and efforts to mitigate their inherent problems. \n\nHowever, a fundamental difference lies in their routing objectives: MoE's router directs tokens to different FFN-based experts, allowing each token to find its optimal parameter path, whereas MoR's router learns to decide whether to apply further recursion to a token. \n\nImportantly, MoR and MoE are *orthogonal*, and they can be integrated easily. It's possible to replace the FFN within a MoR's recursion block with a conventional MoE layer. This combined approach would offer two axes of adaptive computation for activated tokens: not only an adaptive path based on recursion depth (MoR) but also selection among FFN experts (MoE), potentially leading to even greater model capacity and specialization."}, {"review_id": "3UmYtKkSxo", "reviewer": "Reviewer_eKta", "summary": "The paper aims at improving the efficiency of recursive transformers, where the inputs is looped over the model for several recursive steps. The efficiency gains are obtained across three main axis. First, a routing mechanism allocates tokens to different loops. Two variants are proposed, by sequentially dropping tokens at each loop, or by allocating the tokens to different number of loops prior to start each of them. Second, by proposing two alternatives for caring the activations, specifically tailored for the recursive computation performed by the model. Third, different ways of grouping layers. The results show that recursive transformers with these modifications become as compute efficient as vanilla transformers despite using much less parameters. Several ablations compare all the different proposed alternatives in the setup.", "questions": "- **Effective FLOPs-parameter trade-off**: For fixed FLOPs, MoR with $N_r = 2$ strikes an optimal balance, while deeper recursion degrades performance as parameter count is reduced disproportionately. Can the authors expand on this point? Is the model saturating its capacity with too many recursions until becoming inefficient?\n\n- Does the model quickly decode because the prediction is easier, i.e. the model is more confident about predicting? What if the a token was forced to be decoded despite the router allocating it to the next step?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "### Strengths\n\n- **Strong improvements over baselines**: The proposed MoR architecture significantly outperforms other recursive baselines across compute budgets, achieving lower validation loss and higher few-shot accuracy (e.g., Table 3) while maintaining efficiency through dynamic recursion. It is impressive that it is as FLOPs efficient as a vanilla transformer despite using much less parameters. \n- **Efficient use of compute**: The model demonstrates compute optimality consistent with scaling laws (e.g., Chinchilla scaling). The models and baselines seem well tuned. I verified with quick calculations that with the FLOPs budget of $16.5 * 10^{18}$ the vanilla model with the 730M parameter should be the optimal, which indeed appears to be true from Figure 3.\n- **Advances in recurrent efficiency**: The focus on improving efficiency in recurrent architectures, which are typically FLOP-intensive, addresses an important practical bottleneck.\n- **Pareto improvements in throughput**: MoR delivers substantial throughput gains, positioning it as a practical solution for high-speed inference, which is also a remarkable achievement. \n\n### Weaknesses\n\n- **More recursion is detrimental?**: Across the results in the paper, such as in iso-flops comparisons, fixed token setting and throughout plots, it seems that more recursive steps give worse performance. This seems to indicate some issues with the experimental set up. I would be expecting performance to saturate, but not to degrade, with more recursions. In this context, trainability issues do not seem to be addressed. Or could it be that the combination of routing mechanism + high number of recursions causes signal degeneracies? Some of these issues were addressed in related works, see e.g. \"Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach\".\n- **Test-time scaling interpretation**: The conclusion that “allocating more recursion steps at inference can improve generation quality” is overstated. The results suggest that across models, fewer recursion steps are generally better, and additional recursion helps only within a given model. \n- **Shallow semantic analysis**: The discussion of how the router aligns recursion depth with semantic importance is somewhat superficial. The mechanism by which token context determines exit decisions could be better examined. See also questions on this point.\n- **Vanilla still outperforms at iso-parameter**: In iso-parameter settings, the vanilla Transformer continues to outperform recurrent variants, including MoR, as seen in some configurations in Figure 4. While not a major flaw given the paper’s focus on efficiency, this limitation could be acknowledged more explicitly.\n- Minor: **Decoding logic clarity**: The description of the decoding process—particularly how tokens that exit early are dynamically substituted, and how masking and caching are handled—is insufficiently detailed. This is important for interpreting the throughput results of Figure 4a and understanding practical deployment implications.\n- Minor: **Caching mechanism explanation**: The prose explaining the caching mechanism (e.g., lines 143–149) could be clearer, especially in detailing the logic behind the claimed memory savings and consistency of KV caches. Right now, it feels a bit too compressed. \n\nOverall, I am in favor of acceptance, although I will consider increasing my score based on the author's rebuttal, especially on the issue of why more recursions does not improve performance.", "was_revised": false, "final_justification": "", "rebuttal": "We thank  the reviewer for highlighting MoR’s superior performance and compute efficiency to other baselines.\n\nWe’d like to clarify the **misunderstanding regarding MoR performance degradation as recursion number increase**; put simply, **higher recursion number of model means applying more parameter sharing and thus fewer unique trainable parameters**. When applying recursion number as 2, 3, or 4 to a 360M base model using the Middle-Cycle strategy, the unique parameters are reduced to 167M, 118M, and 98M, respectively. Please refer to the “Param” column in **Table 3** for verification.\n\nThis particular setup aims to highlight a **different kind of trade-off than what's explored in [1]**. Specifically, more recursion with fewer parameters can lead to better throughput through our early-exiting concepts (**Figure 4a**), while introducing some performance drop. Conversely, [1] aims to show that more recursion with fixed parameter sizes leads to better performance, which will have lower throughput.\n\nBelow, we address the reviewer’s specific questions and concerns (quoted).\n\n---\n> **[W1]** **More recursion is detrimental?**: It seems that more recursive steps give worse performance. I would be expecting performance to saturate, but not to degrade, with more recursions. \n\n**[A1]** We would like to clarify that the **observed performance degradation when increasing the number of recursions is due to the reduction in the number of unique parameters**, not the wrong experimental setup or an inherent issue with the recursion and routing mechanism itself. By increasing the recursion number of model configuration, MoR proportionally reduces the number of unique parameterized layers, thereby decreasing the overall model capacity.\n\nAt the same time, to understand how the number of recursions affects performance **while keeping the number of unique parameters fixed**, we conducted an experiment as follows.\n\n**[Table 2R]** Comparison of MoR with different numbers of recursions under the same number of parameters. Vanilla equals MoR with recursion 1. All models are trained on 10B tokens.\n\n| Model      | $N_r$ | Param | Train loss | Val loss | LD     | HS     | PQ     | WG     | ARC    | MMLU   | Average |\n|:----------:|:-----:|:-----:|:----------:|:----------:|:------:|:------:|:------:|:------:|:------:|:------:|:-------:|\n| **Vanilla**|   1   | 118M  |  2.9182    |  2.9678    | 25.81% | 33.08% | 61.81% | 51.46% | 38.10% | 26.64% | 39.48%  |\n| **MoR**    |   2   | 118M  |  2.8767    |  2.9205    | 27.09% | 33.84% | 62.13% | 53.04% | 36.77% | 26.88% | 39.96%  |\n| **MoR**    |   3   | 118M  |  2.8667    |  2.9111    | 27.38% | 34.60% | 63.17% | 51.46% | 37.15% | 26.83% | 40.10%  |\n\nThis aligns perfectly with the reviewer's insight and [1], demonstrating that the MoR model can scale performance by increasing recursion steps with fixed model sizes.\n\n---\n> **[W2]** **Test-time scaling interpretation:** The results suggest that across models, fewer recursion steps are generally better, and additional recursion helps only within a given model.\n\n**[A2]** Similar to the previous point, as we increase the recursion step $N_r$  in MoR, **the number of unique parameters per layer decreases proportionally**. This reduction in model capacity is the main reason for the gradual decline in overall performance as $N_r$ increases—not an inherent limitation of MoR itself.\n\nWe also want to emphasize that in **Figure 5c**, we show how the **MoR model scales up performance by allowing more recursion steps within a given model like in [1]**. This demonstrates that, within a single MoR model, performance increases as the model forwards more recursion blocks, and saturates as it reaches the maximum recursion step. \n\n---\n> **[W3, Q2]** **Shallow semantic analysis:** The discussion of how the router aligns recursion depth with semantic importance is somewhat superficial.\n\n**[A3]** Thank you for raising this point. **The allocation of recursion depth reflects contextual predictability of the subsequent token.** If the predicting next token is relatively easy, the model could skip the last recursion steps. \n\nFor example, in **Figure 1 (right)**, the latter subword parts (e.g., “-ensively” in “defensively”) within a word are usually easier to predict and often require only one recursion (i.e., “def” goes one recursion). Also, opening tokens of function words (e.g., `( { . ,` typically require fewer steps, while their corresponding closing parts or immediately following tokens may need deeper recursion. Please refer to **detailed response in [A3] to reviewer cDNG** regarding intuitions behind recursion.\n\nIn **Figure 5c**, we limited the maximum recursion steps within a given MoR model. For example, \"Loop 1 in MoR-3 performance\" denotes that we forced tokens to be decoded at the first loop, even though the router had allocated them to the second loop. While their performance was quite robust, it was lower than when allowing more recursions based on the router's decision.\n\n---\n> **[W4]** **Vanilla still outperforms at iso-parameter:** In iso-parameter settings, the vanilla Transformer continues to outperform recurrent variants, including MoR, as seen in Figure 4.\n\n**[A4]** We’d like to **clarify again that it's not iso-parameter settings**. All of our experiments comparing vanilla, recursive and MoR are not iso-parameter settings. In our experiments, recursive and MoR models actually have “*fewer unique parameters than the vanilla Transformer*” due to parameter sharing through recursion. \n\nThe main point of our work is to demonstrate that **MoR can closely match the performance of a larger vanilla model despite using \"fewer parameters”, while achieving faster speed and greater memory efficiency**. We will state this distinction more clearly in the revised version.\n\n---\n> **[W5]** **Minor: Decoding logic clarity:** The description of the decoding process—particularly how tokens that exit early are dynamically substituted, and how masking and caching are handled—is insufficiently detailed.\n\n**[A5]** We appreciate the reviewer’s feedback and will provide a more thorough explanation in the updated version of our paper.\n\nMoR models can leverage a **continuous depth-wise batching** inference system [2], thanks to their recurrent architecture and early-exiting mechanisms. This enables batch inference even when tokens are at different depths (recursion steps), because they go through the same set of parameters.\n\nOnce queries are enqueued, the router first decides whether to apply more recursion or exit to proceed the next token. Queries that are directed to the next token then pass through the output layer and token embedding. Subsequently, they are queried back to the batch, and efficiently computed together with samples that move to the next recursion step. After this computation of a recursion block, each token obtains a new decision from routers, and this process is repeated iteratively.\n\nAt this point, since the context length will be variable depending on the recursion depth, we need to carefully handle padding and masking for KV caches. In our experiments for **Figure 4a**, we utilized Static Cache, FlashAttention-2 variable-length decoding, and torch.compile. However, there's significant room for further engineering optimization here. \nUltimately, early-exiting reduces redundant recursion depths and the number of activated tokens in deeper recursion steps, leading to improved speed and throughput.\n\n---\n> **[W6]** **Minor: Caching mechanism explanation:** The prose explaining the caching mechanism could be clearer, especially in detailing the logic behind the claimed memory savings and consistency of KV caches.\n\n**[A6]** We acknowledge that the explanation of the caching mechanism was too compressed. Please refer to our **response to reviewer cDNG [A1]** for a more detailed explanation.\n\nTo summarize, **recursion-wise KV caching** selectively cache KV pairs: only tokens routed to a specific recursion step will store their key–value entries at that level. This design means that a linearly decreasing fraction of tokens are active at each recursion step (specifically, $\\frac{N_r +1- r}{N_r}$, where $r$ denotes the recursion step starting from 1). As a result, this method achieves a reduction in total memory and I/O costs by a factor of $\\frac{N_r + 1}{2N_r}$ compared to a vanilla approach.\n\nIn **recursive sharing**, we cache KV pairs exclusively at the initial recursion step and reuse them across all subsequent recursions. This can further lower memory usage, though it still incurs the same I/O to vanilla at each recursion.\n\nBoth strategies also reduce Attention FLOPs relative to vanilla ($N_{ctx}^2$), either by operating attention on fewer active tokens (recursion-wise; $k^2$ per layer) or by shortening the query length while keeping the full key length (recursive sharing; $kN_{ctx}$ per layer).\n\n---\n> **[Q1]** **Effective FLOPs-parameter trade-off:** For fixed FLOPs, MoR $N_r = 2$ with strikes an optimal balance, while deeper recursion degrades performance as parameter count is reduced disproportionately.\n\n**[A7]** We believe this point has already been addressed in our previous responses. To directly answer the additional aspect of your question: for a fixed FLOPs budget, as recursion increases, the effective capacity per step is reduced due to parameter sharing (as indicated by “Param” column in **Table 3**). Thus, the observed degradation can be fair as the performance change is proportional to this parameter allocation. At the same time, we want to emphasize that MoR with more recursions offers significant efficiency benefits: smaller parameter sizes, lower KV memory footprint, and higher throughput.\n\n---\n**References**\n\n[1] Geiping et al. “Scaling up test-time compute with latent reasoning: A recurrent depth approach”. arXiv preprint 2025.\n\n[2] Bae et al. “Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA”. ICLR 2025."}, {"review_id": "fHyPIGlm4O", "reviewer": "Reviewer_kjG6", "summary": "This paper introduces Mixture-of-Recursions, a  that aims to integrate two mainstream model efficiency techniques—parameter sharing and dynamic computation (i.e., adaptive depth) into a single, unified architecture. The core idea is to reuse a shared block of layers within a recursive Transformer to achieve parameter efficiency, while simultaneously employing a lightweight, token-level router to dynamically assign different recursion depths to each token.", "questions": "Please refer to my weakness part.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n-  The idea of combining parameter sharing and dynamic computation within a single recursive framework is highly novel and forward-looking. Developing more efficient large language models without sacrificing performance is a core challenge in the field, and this paper's research direction is well-aligned with this goal, making it highly significant.\n\n- The design of the MoR framework is well-considered, encompassing not only dynamic routing but also a corresponding KV cache strategy to form a complete solution. The paper's explanation of the methodology is clear.\n\n- The authors evaluate their model across multiple scales and computational budgets and conduct several ablation studies (e.g., on parameter sharing and routing strategies), which demonstrates the effort invested in validating their method.\n\n**Weaknesses**\n\n- A core claim of the paper is that MoR is more efficient to train, but this is primarily justified using theoretical FLOPs rather than actual wall-clock time. The dynamic routing and conditional computation logic introduced by MoR will incur additional overhead, which may offset or even negate the FLOPs savings on real hardware. The advantage of a truly efficient architecture should ultimately be reflected in end-to-end training time. The lack of a comparison based on actual training time undermines the claim of efficiency.\n\n-  The paper's main results rely on the expert-choice routing strategy, but the ablation study shows that this method requires an auxiliary loss to function correctly, which adds complexity and potential instability to training. Meanwhile, the simpler \"Token-choice\" strategy performs significantly worse. This raises doubts about whether the success of the MoR framework is overly dependent on a specific, finicky routing mechanism rather than a clean, robust design.\n\n- The comparison between the M-Cyc and Cyc strategies may fail to properly control for variables. M-Cyc has independent first and last layers, meaning its shared recursive module has fewer parameters than that of Cyc. Therefore, it is unclear whether M-Cyc's superiority stems from its structural advantage or simply from its smaller, easier-to-optimize recursive module. A more rigorous ablation study is needed to disentangle these factors.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for highlighting MoR's strengths: unified parameter sharing with dynamic computation, dynamic routing paired with a KV-cache, its clear methodology, and comprehensive multi-scale ablations. Below, we address the reviewer’s specific questions and concerns (quoted).\n\n---\n> **[W1]** A core claim of the paper is that MoR is more efficient to train, but this is primarily justified using **theoretical FLOPs rather than actual wall-clock time**. The dynamic routing and conditional computation logic introduced by MoR will incur additional overhead, which may offset or even negate the FLOPs savings on real hardware. The advantage of a truly efficient architecture should ultimately be reflected in end-to-end training time. The lack of a comparison based on actual training time undermines the claim of efficiency.\n\n**[A1]** Thank you for raising this interesting point. We want to highlight **the theoretical efficiency of MoR directly translates to a reduction in actual wall-clock training time**. In MoR, each expert (recursion step) selects only top-$k$ tokens among a total of $N_{\\mathrm{ctx}}$ input tokens. As shown in **Table 2 (right)**, MoR significantly reduces the theoretical computational cost compared to the vanilla baseline.  For a detailed derivation of these expressions, please refer to our **first response [A1] to reviewer cDNG**.\n\nAs shown in **Table 3R**, MoR achieves a substantial reduction in wall-clock training time compared to the vanilla Transformer under identical training steps. This confirms that the theoretical efficiency of our method directly translates to practical speedup in real-world training scenarios.\n\n**[Table 3R]** Comparison of wall-clock training time for MoR and Vanilla 360M Transformers, measured under the same number of training tokens (20B). Note that recursive baselines have almost the same wall-clock time to Vanilla models. We trained models with FlashAttention-2 and HuggingFace Accelerator on 4 H100 GPUs.\n\n|             | Vanilla | MoR-2 | MoR-3 |\n|:-----------:|:-------:|:-----:|:-----:|\n| Training Time (hr) | 30.2  | 25.1 (17%↓)  | 23.6 (22%↓) |\n\n\nAlthough MoR introduces dynamic routing, the associated overhead (routers or indexing) is minimal in practice: the router consists of a single linear layer, and each token passes through it for only recursion times (e.g., 2, 3, or 4). As a result, the efficiency gains predicted by our theoretical analysis are realized in faster wall-clock performance.\n\nWe would also like to note that the throughput gain reported in **Figure 4a** is measured in wall-clock time rather than theoretical calculation.\n\n---\n> **[W2]** The paper's main results rely on the expert-choice routing strategy, but the ablation study shows that this method requires an auxiliary loss to function correctly, which adds complexity and potential instability to training. Meanwhile, the simpler \"Token-choice\" strategy performs significantly worse. This raises doubts about **whether the success of the MoR framework is overly dependent on a specific, finicky routing mechanism rather than a clean, robust design**.\n\n\n**[A2]** Thank you for raising this point. We'd like to clarify that we **didn't do extensive engineering or pick a finicky router for the MoR framework**. Instead, we found that a **much simpler and more intuitive expert-choice routing consistently demonstrated robust performance**, which is why we adopted it.\n\n**[Auxiliary Loss]** The expert-choice routing strategy **does not introduce significant complexity or instability** in practice. There are typically two solutions to address information leakage in expert-choice: auxiliary loss and auxiliary router. **The auxiliary loss is more straightforward and easy to implement**—for example, it can be expressed with only a few lines of pseudocode as below. \n\n```python\nrouter_weights = Router(x)\nselected_tokens = TopK(router_weights, k=top_k)\ntargets = zeros_like(router_weights)\ntargets[selected_tokens] = 1\nsampling_loss = BCE(router_weights, targets) / (batch_size * seq_len)\n```\nIn contrast, the auxiliary router has the drawback of requiring a different router for training and inference, and both must be used during training. We believe our finding that the auxiliary loss is more robust and stable to use across all experiments is therefore very significant (refer to **Table 4 (left)** and **Table 11**). \n\n**[Expert-Choie vs. Token-Choice]**\nWe would like to emphasize that **expert-choice routing is much simpler than token-choice routing** and is better suited to recursion steps.\n\n- Expert-choice selects a fixed number of tokens per expert, enabling consistent capacity allocation (see **Section 2.2.1** and **Table 2 (left)**), which leads to more stable training.\n- **Token-choice for MoR shows unstable and sensitive performance to hyperparameter settings.** We tested various configurations—such as different load balancing strategies, router activation functions, and z-loss for stability—but as shown in **Table 4 (right)** and **Table 12**, it was very difficult to balance expert allocation. It actually results in high MaxVio values (i.e., the router failed to balance the load) and higher loss across training phases. We believe the heterogeneous experts in token-choice make training more difficult to stabilize.\n- Implementation complexity is even higher for token-choice. Because it does not guarantee static capacity, the number of tokens processed by each expert can vary within a batch, requiring more engineering like additional padding for efficient batching.\n\nWe conducted thorough ablations of these design choices (**Table 4**, **Table 11** and **Table 12**) and concluded that the *simple* expert-choice routing with an auxiliary loss is better suited to MoR’s recursive nature and robustly shows good performance across all settings.\n\n---\n> **[W3]** The **comparison between the M-Cyc and Cyc strategies may fail to properly control for variables**. M-Cyc has independent first and last layers, meaning its shared recursive module has fewer parameters than that of Cyc. Therefore, it is unclear whether M-Cyc's superiority stems from its structural advantage or simply from its smaller, easier-to-optimize recursive module. A more rigorous ablation study is needed to disentangle these factors.\n\n**[A3]** Thank you for pointing out this. **Achieving a perfectly rigorous ablation between Middle-Cycle and Cycle strategies is inherently difficult**. Since the middle-cycle has two unshared layers, it is not trivial to exactly match both the number of unique parameters and effective depth for a given recursion number. For example, in **Table 9**, M-Cyc often has one or two more layers than Cyc cases. This is the main reason why M-Cyc has much better performance than Cyc (not because it has a smaller recursive module as larger scales generally have better performance in language modeling).\n\nAnother reason why we opt for M-Cyc is from its uptraining results. In **Table 10**, despite a parameter difference of only about 10M (i.e., one or two layers), the M-Cyc strategy demonstrated significantly better performance. When uptrained from a pretrained checkpoint, the non-shared first and last layers likely had a substantial impact. As a result, the robust performance of M-Cyc across diverse settings led us to choose it. \n\nHowever, since this is not fair comparison, we compared the Cycle and Middle-Cycle strategies while fixing the number of unique layers, in a setting that *favors the Cycle strategy*. Specifically, with 730M model size, M-Cyc with a recursion number of 3 uses 8 layers for the recursion block and 2 unshared layers (1 + 8 x 3 + 1), resulting in 10 unique layers and *26 effective layers*. We compared this to the Cycle strategy with 10 layers per recursion block (10 x 3), which also has 10 unique layers but a total of *30 effective layers*.\n\n**[Table 4R]** Comparison of M-Cyc and Cyc sharing strategies under the same number of unique parameters. All models use 730M scales (non-embedding parameters) as the base model with 3 recursions. The models are trained from scratch on 10B tokens using the Fineweb-Edu dataset.\n\n| Model |  $N_r$ | Param| Train loss | LD | HS | PQ | WG | ARC | MMLU | Average | \n|:----------------:|:------:|:-----:|:-----:|:--:|:--:|:--:|:--:|:---:|:-----:|:-----:|\n| **Cyc** | 3 | 252M | 2.7573 |  28.72% | 37.31% | 65.56% | 52.01% | 39.94% | 27.85% | 41.90% |\n| **M-Cyc** | 3 | 252M | 2.7552 | 29.75% | 37.73% | 66.00% | 52.49% | 40.49% | 27.50% | 42.32% |\n\nThe M-Cyc has better performance than the Cycle despite having much fewer effective layers (26 vs. 30). However, we need to validate these across more scales to draw a better, trustable conclusion. We appreciate the valuable comments."}]

Author rebuttal:

Meta-review: The manuscript combines effective transformer models in a clever way that improves the computational efficiency of the model.  This "Mixture of Recursions (MoR)" model is a transformer based architecture that shares parameter weights across layers and uses a heuristic to dynamically determine the depth for different aspects of the network.  Both early exit and parameter sharing a proven successful methods, and it isn't a surprise that combining them in an intelligent way is also effective.  

The reviewers are unanimous in their support of the manuscript with numerous strengths listed.  The summary from the author's final rebuttal lists the strengths as:
Novel and significant MoR framework combining parameter sharing with dynamic computation (@cDNG, @MjJH, @kjG6).
Strong empirical results outperforming baselines and dominating the compute-performance frontier (@cDNG, @MjJH, @eKta, @kjG6).
Remarkable compute and throughput efficiency for practical high-speed inference (@MjJH, @eKta).
Extensive ablation studies providing valuable insights into design choices (@cDNG, @MjJH, @kjG6).

These are accurately listed and reflect well the reviewers remarks.  

The main weakness of the manuscript is that there isn't much surprising here, nor is it clear that other researchers developing ML models would gain significant insight from the manuscript.  That said, it can be expected that practitioners would appreciate and benefit from the results.
Decision: Accept (poster)