Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

arxiv:2507.10524

Original source · Paper JSON

LLMs: architectures, context, training · 2025-07-14

Cites (5) · Cited by (12)

Complete cleaned Markdown · Source Markdown including bibliography

Read in parts

1. Part 1
2. Part 2
3. Part 3
4. Part 4
5. Part 5
6. Part 6

Contents

• Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
• Abstract
• 1 Introduction
• Contributions.
• 2 Method
• 2.1 Preliminary
• Recursive Transformers.
• Parameter-sharing strategies.
• Enhanced training and inference efficiency in recursive models.
• Limitations in prior works.
• 2.2 Mixture-of-Recursions
• 2.2.1 Routing Strategies: Expert-choice vs. Token-choice
• Expert-choice routing. ( 2(a) )
• Token-choice routing. ( 2(b) )
• Strengths and limitations. ( Table 2 – Left )
• 2.2.2 KV Caching Strategies: Recursion-wise Caching vs. Recursive Sharing
• Recursion-wise KV caching. ( 2(c) – Top )
• Recursive KV sharing. ( 2(c) – Bottom )
• Strengths and limitations. ( Table 2 – Right )
• 3 Experiments
• 3.1 Main Results
• MoR outperforms baselines with fewer parameters under equal train compute.
• MoR outperforms baselines with less compute at equal data.
• MoR performance varies with routing and caching strategies.
• 3.2 IsoFLOP Analysis
• Experimental Setup.
• MoR is a scalable and parameter-efficient architecture.
• 3.3 Inference Throughput Evaluation
• Experimental Setup.
• MoR boosts inference throughput with continuous depth-wise batching.
• 4 Ablation Studies
• 4.1 Parameter Sharing Strategies
• Middle-Cycle is the most effective parameter sharing strategy.
• 4.2 Routing Strategies
• In expert-choice routing, auxiliary loss and linear router yield the best performance.
• In token-choice routing, balancing loss yields stable and accurate routing.
• 4.3 KV Caching Strategies
• KV sharing robustly works even in parameter-shared architectures.
• KV sharing degrades expert-choice but benefits token-choice routing in MoR.
• 5 Analysis
• 5.1 Compute-optimal Scaling Analysis
• MoR scaling favors model size over training length under isoFLOPs.
• 5.2 Routing Analysis
• The allocation of recursion depth reflects contextual predictability of the subsequent token.
• Expert-choice router with auxiliary loss perfectly separates selected from unselected tokens.
• 5.3 Test-time Scaling Analysis
• MoR enables test-time scaling via deeper recursion.
• 6 Related Work
• Recursive Transformers.
• Adaptive computation.
• Routing mechanism.
• Key-value caching.
• Latent reasoning.
• 7 Conclusion
• 7.1 Limitations and Future Works
• Reasoning MoR models.
• Further scaling model family.
• Adaptive capacity control.
• Compatibility with sparse algorithms.
• Expansion to multimodal and non - text domains.
• 7.2 Acknowledgements
• Appendix A Details of Design Choices for Mixture-of-Recursions
• A.1 Parameter-sharing Strategy
• A.2 Routing Strategy
• Expert-choice routing.
• Token-choice routing.
• A.3 KV Caching Strategy
• Recursion-wise caching.
• Recursive KV sharing.
• Appendix B Experimental Setup
• Training settings.
• Evaluation settings.
• Model architecture details.
• Appendix C Expanded Results of Main Experiments
• C.1 Performance Validation at 1.7B Scale
• C.2 Increasing Recursion Depth under Fixed Parameters
• Appendix D Expanded Results of IsoFLOP Analysis
• FLOPs approximated calculation of Transformers.
• Trapezoid learning - rate schedule with checkpoint reuse.
• Results at a glance.
• Appendix E Details of Experimental Settings for Throughput Measurement
• Appendix F Expanded Results of Parameter Sharing Strategy
• Middle - Cycle is consistently the safest choice.
• Behavior under continued pre - training (up - training).
• Appendix G Expanded Results of Design Choices for Router
• G.1 Details of Design Configurations
• G.2 Router Performance Evaluation Metrics
• Expert-choice routing.
• Token-choice routing.
• G.3 Extended Evaluation Results of Router Designs
• Appendix H Expanded Results of KV Cache Sharing Mechanism
• H.1 Key Value Representation Trends in Recursive Transformers
• H.2 Performance Comparison of KV Sharing Strategy
• Experimental results of key-value cache sharing.
• Relaxation for key-value sharing constraints.
• Appendix I Expanded Qualitative Results
• I.1 Analysis on Adaptive Computation Paths
• I.2 Analysis on Router Weights
• Instructions for reporting errors