PaperlayerView exact revision on arXiv

Source-linked AI summary

Mixture-of-Recursions: 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

arXiv:2507.10524v3cs.CLcs.LG

TL;DR

Large language models are costly to train and deploy, while existing efficiency methods generally address parameter sharing or adaptive computation separately. MoR combines shared recursive layers, token-level routing, and selective KV caching in one architecture; under equal training compute, it reports lower validation loss and higher few-shot accuracy with fewer parameters, alongside higher throughput.

  • Problem

    The paper addresses how to combine parameter efficiency and adaptive computation to reduce the computational and memory costs of scaling language models.

  • Method

    MoR reuses a shared recursion block, routes each token to a learned recursion depth, and uses recursion-wise or recursive KV caching for attention.

  • Results

    Under an equal training budget of 16.5e18 FLOPs, MoR achieved 43.1% average few-shot accuracy versus 42.3% for Vanilla while using nearly 50% fewer parameters, and outperformed recursive baselines.

  • Takeaways & Limitations

    Across models from 135M to 1.7B parameters, MoR lowers validation perplexity, improves average few-shot accuracy, and delivers higher inference throughput than vanilla and recursive baselines.

  • Takeaways & Limitations

    Experiments were limited by current compute constraints to models of up to 1.7 billion parameters; larger-scale validation remains future work.

Abstract

from arXiv · show

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

1. Introduction

Mixture-of-Recursions (MoR) unifies parameter sharing, token-level adaptive recursion, and memory-efficient KV caching in one Recursive Transformer. It addresses the cost of scaling language models by routing harder tokens through more shared recursion steps while selectively reducing computation and memory traffic.

  • Scaling language models creates computational and memory demands that make training and deployment challenging outside hyperscale data centers.
  • Parameter sharing reuses weights across depth, while adaptive computation allocates more compute only when needed, but prior approaches largely study these axes separately.
  • MoR trains lightweight routers to assign token-specific recursion depths by repeatedly applying a shared parameter block according to each token’s required thinking depth.
  • MoR combines weight tying, token routing, and recursion-wise KV caching to reduce parameters, redundant FLOPs, and memory traffic within one architecture.
  • The framework is presented as a unified efficiency architecture combining parameter sharing, adaptive thinking depth, and memory-efficient KV caching.

2. Method

MoR combines recursive parameter sharing with token-adaptive routing and KV caching to vary computation by token difficulty while reducing memory and IO demands.

  • Recursive Transformers: Recursive Transformers reuse shared parameter pools across recursion blocks, increasing effective depth without increasing parameter size.
  • MoR framework: MoR dynamically adjusts each token’s recursion depth during pretraining and inference through lightweight routing.The framework assigns token-specific recursion steps to concentrate computation on more challenging tokens.
  • Routing strategies: Expert-choice routing selects top-k tokens at each recursion step, while hierarchical filtering lets only selected tokens continue to deeper steps.The active token set therefore narrows progressively with recursion depth.
  • Routing strategies: Token-choice routing assigns each token a complete recursion path at the outset, with top-1 gating selecting how many recursive blocks it traverses.
  • KV caching strategies: Recursion-wise KV caching stores entries only for tokens routed to each depth, restricting attention to locally cached tokens and reducing memory and IO demands.
  • KV caching strategies: Recursive KV sharing caches all tokens’ KV pairs at the first recursion and reuses them across later recursions, preserving access to past context.All tokens traverse the first recursion block, while later query lengths may shrink as tokens are selected.

3. Experiments

Experiments evaluate MoR across model scales, training budgets, and deployment settings, finding strong accuracy and throughput with substantially fewer parameters or lower compute.

  • Main Results: MoR achieves 43.1% average few-shot accuracy versus 42.3% for Vanilla under equal 16.5e18-FLOP training, despite nearly 50% fewer parameters.With three or four recursions, MoR remains competitive with the full-capacity vanilla model and outperforms recursive baselines.
  • Main Results: At fixed 20B training tokens, MoR uses 25% fewer training FLOPs while achieving lower validation loss and higher accuracy than vanilla and recursive baselines.It also reduces training time by 19% and peak memory usage by 25% versus vanilla.
  • Scaling: For model sizes above 360M parameters, MoR matches or often exceeds Vanilla Transformer validation performance across compute budgets.At 135M parameters, MoR underperforms vanilla, likely because of a recursive capacity bottleneck.
  • Inference Throughput: MoR variants outperform the vanilla baseline in inference throughput across fixed and maximum batching, with MoR-4 reaching up to a 2.06× speedup at maximum batch size.Greater recursion depth increases early exits and reduces KV-cache usage, with a slight performance degradation.

4. Ablation Studies

Ablations identify effective parameter sharing and routing choices, while showing that KV sharing trades memory savings for routing-dependent performance changes.

  • Parameter Sharing Strategies: Middle-Cycle consistently achieves the lowest validation loss among the four evaluated parameter-sharing strategies.The authors adopt Middle-Cycle for subsequent MoR and Recursive Transformer experiments.
  • Routing Strategies: Auxiliary loss with a linear router provides the best expert-choice routing performance and improves inference-time behavior over a separate auxiliary router.The expert-choice ablation also evaluates normalization, router architecture, and auxiliary z-loss choices.
  • Routing Strategies: Balancing loss yields similar log-probability and few-shot accuracy to router bias but significantly lower MaxVio for token-choice routing.Despite this stability advantage, heterogeneous expert loads often remain difficult to balance.
  • KV Sharing: KV sharing is relatively robust in Recursive Transformers pretrained from scratch, despite their reduced degrees of freedom.The authors investigate KV sharing in Vanilla, Recursive, and MoR architectures.
  • KV Sharing: In MoR, KV sharing degrades expert-choice performance but benefits token-choice routing while reducing memory footprint and overall FLOPs.The degradation is observed under a fixed token setting.

5. Analysis

Analysis examines MoR’s scaling behavior, routing decisions, and deeper-recursion effects, linking adaptive depth to token predictability and additional computation.

  • Compute-Optimal Scaling: Under isoFLOPs, MoR benefits more from increasing parameter count than from extending training length.Its compute-optimal path has a flatter slope than the baselines, suggesting that the shared block’s performance is especially important.
  • Learned Routing: Router-assigned recursion depth reflects the contextual predictability of the subsequent token.Some token continuations are easier to predict, while tokens following openings can be more challenging.
  • KV Caching: Recursive KV sharing increases attention FLOPs by N_ctx/k relative to recursion-wise KV caching but can reduce total FLOPs through lower KV-projection cost.This comparison concerns the memory-saving KV-sharing mechanism.
  • Compute-Optimal Scaling: The Figure 5a analysis fits polynomial functions across architectures and FLOPs budgets to derive compute-optimal model sizes.Stars mark the optimal model size for each compute budget.
  • Learned Routing: With auxiliary loss, the expert-choice router perfectly separates selected and unselected tokens across recursion steps in the illustrated 3-recursion model.Figure 5b plots normalized routing-score distributions for selected and unselected tokens.
  • Test-Time Scaling: Increasing the maximum recursion depth improves log-likelihood across MoR models with two, three, and four recursions.The analysis interprets deeper recursion as additional computation that can refine token representations or thought processes.

6. Related Work

Related work spans recursive parameter sharing, adaptive computation, learned routing, KV caching, and latent reasoning, providing the conceptual foundations for MoR.

  • Recursive Transformers: Recursive Transformer research reuses layers across depth to reduce parameter count while increasing effective computation.The Universal Transformer demonstrated that repeatedly applying a shared block can match the representational power of a deep non-shared stack.
  • Adaptive Computation: Adaptive-computation methods dynamically reduce cost by allocating less processing to easier samples or tokens.Early exiting is a common approach for halting computation earlier in the network.
  • Routing Mechanism: Routing mechanisms use learned routers to select subsets of expert networks for each token in sparse Mixture-of-Experts models.This increases model capacity without proportional computational overhead.
  • Key-Value Caching: KV caching reuses per-token key and value tensors during autoregressive decoding, avoiding quadratic recomputation but increasing GPU-memory pressure.The memory problem is especially pronounced for long contexts and large batches.
  • Latent Reasoning: Latent-reasoning methods perform reasoning in hidden states, often using fixed reasoning depths introduced through special tokens or structured prompts.MoR is positioned within this emerging line of work.

7. Conclusion

MoR unifies parameter sharing, token-level adaptive recursion, and efficient KV caching in one Transformer architecture. The paper also identifies scaling, routing flexibility, reasoning alignment, sparsity, and multimodal extension as open directions or boundaries.

  • Conclusion: MoR combines parameter sharing, adaptive recursion depth, and efficient KV caching within a single Transformer architecture.Lightweight routers assign token-specific recursion depths, while selective KV caching reduces redundant computation and memory access.
  • Future directions: Future work may align recursion depth with reasoning complexity and extend MoR to structured sparsity, vision, speech, and multimodal Transformers.These directions are presented as potential ways to improve reasoning, efficiency, interpretability, or applicability beyond text.
  • Limitations and future work: Experiments are limited to models of at most 1.7B parameters, leaving larger-scale validation as future work.Proposed extensions include larger non-shared blocks, depth-specific LoRA or experts, expert parallelism, and continued pre-training.
  • Routing strategies: MoR uses expert-choice and token-choice routing, trading fixed compute budgets and load balancing against autoregressive causality and adaptive capacity.Expert-choice routing can leak information from later tokens, whereas token-choice routing avoids leakage but may create load imbalance.
  • KV caching: Recursive KV sharing lowers memory use and prefill computation but can mismatch later recursion steps with earlier KV representations.The paper therefore generally prefers recursion-wise caching when routing is selective, while recursive sharing suits memory-constrained settings where prefill time dominates.

B. Experimental Setup

The experiments use Llama-based models trained on FineWeb-Edu and evaluate language modeling and few-shot accuracy across model scales and recursive configurations. The setup includes matched-compute comparisons and fixed-parameter recursion studies.

  • Training settings: Models use a Llama-based Transformer configuration derived from SmolLM and train on a 220-billion-token deduplicated FineWeb-Edu corpus.Pretraining used four H100 or A100 GPUs.
  • Evaluation: Evaluation measures few-shot accuracy on LAMBADA, HellaSwag, PIQA, WinoGrande, ARC, and MMLU using standard shot settings.Accuracy is byte-length normalized for most tasks except LAMBADA, WinoGrande, and MMLU.
  • Model architecture: Four model variants range from 135M to 1.7B total parameters, with a 49K vocabulary and maximum sequence length of 2K tokens.The size definition includes both non-embedding and embedding parameters; smaller models use Grouped-Query Attention.
  • Matched-compute comparisons: At 1.7B scale and fixed 68.5e18 FLOPs, Table 7 compares MoR, Recursive, and Vanilla Transformers using few-shot accuracy.IsoFLOP rows vary training-token counts according to model efficiency, and recursive models use fixed rather than adaptive recursion depths.
  • Recursion scaling: With 118M unique parameters, increasing recursion from 1 to 3 decreases train and validation NLL while improving average few-shot accuracy.This study uses 10B training tokens and Middle-Cycle sharing for MoR models.

D. Expanded Results of IsoFLOP Analysis

The isoFLOP analysis compares Vanilla, Recursive, and MoR models across four sizes and three compute budgets. Token-routed MoR overcomes the performance degradation of weight sharing and surpasses Vanilla from 360M parameters upward while using one-third as many parameters.

  • Experimental design: The analysis studies 135M, 360M, 730M, and 1.7B model capacities under matched training compute, with three recursion steps for Recursive and MoR models.Each architecture is trained at 16.5EB and reused to obtain 5EB and 2EB checkpoints.
  • FLOPs accounting: The FLOPs accounting includes forward-pass projections and causal attention contributions while omitting backward-pass, FlashAttention, normalization, and non-linearity costs.The approximation follows the stated Transformer FLOPs convention and excludes redundant causally masked operations.
  • FLOPs accounting: MoR FLOPs depend on routing and KV caching, with attention cost estimated from the sequence length active at each recursion depth.For token-choice routing, the calculation assumes perfect token balancing.
  • Training schedule: The trapezoid learning-rate schedule uses intermediate checkpoints to evaluate multiple FLOPs budgets without retraining every model from scratch.Warm-up uses 5% of the smallest-budget steps, and cool-down uses 20% of the corresponding budget’s steps.
  • Results at a glance: Token-routed MoR catches up to and surpasses Vanilla from 360M parameters upward while using only one-third of the parameters.The advantage persists at the 730M and 1.7B scales, whereas weight sharing alone degraded performance relative to Vanilla.

E. Details of Experimental Settings for Throughput Measurement

Throughput is measured with continuous depth-wise batching, dynamic query scheduling, and variable-length KV caches. The evaluation compares normalized MoR and Vanilla decoding speeds under fixed and memory-adjusted batch settings.

  • Batching system: Continuous depth-wise batching immediately replaces early-exiting queries with queued queries, keeping the inference batch fully utilized.The system evaluates decoding throughput on 1K FineWeb-Edu validation samples.
  • Throughput comparison: Throughput compares 360M-parameter Vanilla and MoR models generating lengths sampled from a normal distribution with mean 256.MoR speeds are normalized against Vanilla speed, and MoR-4 adds two layers before recursion.
  • Batching settings: The evaluation uses fixed batch size 32 and a relative maximum batch size scaled by the Vanilla-to-MoR maximum-batch ratio.Maximum batch sizes account for model parameters and KV-cache memory on an H100 GPU.
  • Implementation: Implementation uses a queue, FlashAttention 2 for variable-length KV caches, and a static cache compatible with torch.compile.The design decouples transformer-block processing from the remaining inference stages to support deployment-oriented evaluation.

F. Expanded Results of Parameter Sharing Strategy

Across model scales and recursion depths, Middle-Cycle is the strongest parameter-sharing strategy, improving validation NLL and few-shot accuracy relative to alternatives. Continued pre-training preserves this advantage, while router evaluations compare routing stability and balance.

  • Parameter-sharing comparison: For 360M models, Middle-Cycle achieves the lowest NLL at both recursion depths and the largest average-accuracy improvement over vanilla reduced models.At 135M, Cycle is slightly ahead at two recursions, but Middle-Cycle overtakes as recursion depth rises and has a steadier accuracy profile.
  • Parameter-sharing comparison: Pure Sequence sharing records the worst NLL across all four settings, while its accuracy gap widens with recursion depth.Middle sharing improves Sequence modestly but remains below Cycle-based methods.
  • Parameter-sharing comparison: Middle-Cycle consistently outperforms other strategies in validation NLL and average few-shot accuracy, especially at higher recursion depth.The comparison covers 135M and 360M models, two recursion depths, 10B-token pretraining, FineWeb-Edu NLL, and six few-shot tasks.
  • Continued pre-training: Both Middle strategies outperform similarly initialized reduced baselines during continued pre-training, while other strategies plateau earlier.The up-training evaluation continues open-sourced SmolLM checkpoints for 5B additional tokens and measures train NLL and six-benchmark few-shot accuracy.
  • Evaluation dimensions: Router evaluation examines auxiliary and balancing losses, scaling factors, activations, architectures, dead-token ratios, sampling accuracy, MaxVio, and entropy.These measurements assess expert-choice selection behavior and token-choice load balance.

G.3. Extended Evaluation Results of Router Designs

Router-design experiments show a trade-off between selection accuracy and positional dead tokens for expert-choice routing, while balancing losses improve token-choice load distribution. KV-sharing analyses find similar projected representations across recursion depths, supporting potentially small performance costs.

  • Expert-choice routing: Auxiliary routers and auxiliary losses improve sampling accuracy but can produce dead-token ratios as high as 66.7%.A linear router with auxiliary loss reduces the dead-token ratio without compromising sampling accuracy.
  • Token-choice routing: Explicit balancing loss significantly reduces MaxVio and increases entropy, improving token assignment balance without sacrificing overall model performance.Loss-free approaches are simpler but tend to have higher MaxVio and lower entropy.
  • KV representation analysis: Shared key and value projection layers produce recursive magnitude patterns, with hidden-state magnitudes increasing while corresponding projected signals remain similar.The observation comes from a 360M Middle-Cycled Recursive Transformer with three recursion steps.
  • KV representation analysis: Cosine-similarity matrices show highly similar key and value representations across shared recursion depths, although value-state sharing is more challenging than key-state sharing.The results suggest KV-sharing performance drops can be marginal in Recursive Transformers.

H.2. Performance Comparison of KV Sharing Strategy

KV sharing reduces cache requirements but introduces a performance trade-off in MoR. The paper finds small degradation in best MoR settings and no substantial gains from tested constraint-relaxation techniques.

  • MoR trade-off: MoR incurs a small degradation in its best expert-choice-router settings after KV sharing, which the authors consider acceptable given reduced parameter and cache sizes.The comparison concerns the trade-off between quality and memory footprint.
  • Hybrid caching: Hybrid caching of inactive positions while recomputing active positions does not provide additional benefits.The strategy initializes inactive positions from the first recursive iteration and updates active positions through actual computation.
  • KV-sharing comparison: KV cache sharing is evaluated across Vanilla, Recursive, and MoR Transformers using FineWeb-Edu NLL and few-shot accuracy.The comparison includes Cycle, Sequence, recursion-wise, and hybrid cache strategies.
  • Constraint relaxation: Relaxation techniques for KV-sharing constraints do not substantially improve performance when relaxed models are pretrained from scratch.The experiments use limited additional parameters, and more sophisticated relaxation remains an open direction.
  • Adaptive computation paths: MoR assigns different recursion depths to subword tokens, with some exiting early and others receiving deeper processing.The visualization uses three recursion steps and recursion-wise KV caching.

I.2. Analysis on Router Weights

Router-weight analyses show that expert-choice routing with auxiliary loss can achieve nearly binary selected-versus-unselected decisions and strong load balancing. Token-choice strategies generally balance well but often miss optimal balance consistently.

  • Expert-choice routing: Expert-choice routers with auxiliary loss achieve nearly perfect load balancing and almost binary weights for selected and unselected tokens.The analysis visualizes router outputs across recursion steps under four routing configurations.
  • Router behavior: The auxiliary router distinguishes selected and unselected tokens while allowing different capacity factors during training and inference.The auxiliary router is described as capturing intra-token differences rather than only extreme points.
  • Token-choice routing: Other token-choice strategies show good balancing properties with reasonable router values, but most fail to converge to optimal load balancing consistently.The finding highlights challenges in heterogeneous expert settings.
Loading 2507.10524v3…