Research map / Architectures and efficiency / Attention and transformers
Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Mixture-of-Recursions (MoR), a weight-tied recursive Transformer with learned per-token routers that choose recursion depth and recursion-wise KV caching, pretrained from scratch at 135M–1.7B scale.
Problem
Recursive Transformers save parameters via layer tying but apply fixed compute per token, and prior dynamic-depth methods need post-hoc routing or extra training; combining parameter efficiency with token-level adaptive computation and efficient KV caching remains open.
Approach
Learn token-specific recursion depths end-to-end during pretraining via expert-choice or token-choice routers over a shared block, with hierarchical filtering so only tokens selected at one depth proceed to the next, and cache KV pairs per recursion depth (or share first-depth KV) to cut memory and attention cost.
Key results
- At 16.5e18 FLOPs (315M-class model), MoR with expert-choice routing and 2 recursions reaches 43.1% average few-shot accuracy vs 42.3% for vanilla, with ~50% fewer parameters.
- MoR-4 achieves up to 2.06x inference throughput speedup over the vanilla baseline with maximum batch size.
- With 20B tokens, MoR (Nr=2) uses 25% fewer training FLOPs and reduces training time by 19% and peak memory by 25% vs vanilla.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use expert-choice adaptive routing with two recursions and recursion-wise caching instead of vanilla Transformer layers | MoR · FineWeb-Edu validation, 16.5e18 FLOPs | NLL: 2.7824 -> 2.7511 | Helped |
| Use expert-choice adaptive routing with two recursions and recursion-wise caching instead of vanilla Transformer layers | MoR · Few-shot average accuracy, 16.5e18 FLOPs | accuracy: 42.3 -> 43.1 | Helped |
| Use MoR expert-choice routing with two recursions instead of vanilla at fixed 20B tokens | MoR · FineWeb-Edu validation, 20B tokens | NLL: 2.7824 -> 2.7749 | Helped |
| Use token-choice routing instead of expert-choice routing (ablation) | MoR, N_r=3, Cache · Few-shot average accuracy, 16.5e18 FLOPs | accuracy: 42.6 -> 40.0 | Hurt |
| Use recursive KV sharing instead of recursion-wise caching (ablation) | MoR expert-choice · FineWeb-Edu validation | NLL | Hurt |
| Use recursive KV sharing instead of recursion-wise caching (ablation) | MoR token-choice · FineWeb-Edu validation | NLL | Helped |
| Use auxiliary loss rather than an auxiliary router (ablation) | Expert-choice MoR · Inference-time routing behavior | few-shot accuracy | Helped |
| Use sigmoid normalization instead of tanh (ablation) | Expert-choice MoR · FineWeb-Edu validation | NLL: 2.9933 -> 2.8667 | Helped |
| Use a Linear router instead of an MLP router with auxiliary loss and sigmoid (ablation) | Expert-choice MoR · Few-shot accuracy | accuracy: 40.0 -> 40.1 | Helped |
| Add auxiliary z-loss (ablation) | Expert-choice MoR with auxiliary loss, sigmoid, Linear router · Few-shot accuracy | accuracy: 40.1 -> 40.0 | Hurt |
| Add auxiliary z-loss (ablation) | Expert-choice MoR with auxiliary loss, sigmoid, Linear router · Dead token ratio | dead token ratio: 0.1 -> 0.0 | Helped |
| Use explicit balancing loss instead of loss-free router bias (ablation) | Token-choice MoR · Router load balance | MaxVio: 0.852 -> 0.266 | Helped |
Surprising findings
- Recursive Transformers with layer tying remain relatively robust to KV sharing, with KV pairs at shared depths showing highly consistent magnitude and direction.
- KV sharing degrades expert-choice MoR but benefits token-choice MoR, where weak routing is complemented by shared context.
Limitations
- Experiments limited to models up to 1.7B parameters; MoR underperforms vanilla at 1.7B scale and smallest 135M size.
- Expert-choice routing with auxiliary loss yields near-perfect separation, making it hard to adjust top-k capacity at inference.
- Token-choice routing suffers load imbalance even with balancing losses.
Experimental setup
Datasets: FineWeb-Edu (deduplicated subset, SmolLM-Corpus)
Benchmarks: FineWeb-Edu validation NLL LAMBADA HellaSwag PIQA WinoGrande ARC MMLU
Metrics: negative log-likelihood few-shot accuracy throughput (tokens/s) MaxVio
Baselines: Vanilla Transformer Recursive Transformer (Middle-Cycle) Relaxed Recursive Transformer
Builds on: Universal Transformer Mixture-of-Depths Relaxed Recursive Transformers expert-choice routing continuous depth-wise batching
Closest papers by meaning
- Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach (NeurIPS 2025 · similarity 0.88)
- Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRA (ICLR 2025 · similarity 0.88)
- R2R: Efficiently Navigating Divergent Reasoning Paths with Small-Large Model Token Routing (NeurIPS 2025 · similarity 0.88)
- SkipGPT: Each Token is One of a Kind (ICML 2025 · similarity 0.87)
- Recursive Inference Scaling: A Winning Path to Scalable Inference in Language and Multimodal Systems (NeurIPS 2025 · similarity 0.87)
- Memory Layers at Scale (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.