Mixture-of-Recursions: Learning Dynamic Recursive Depths for Efficient Language Modeling
Sangmin Bae, Yujin Kim, Reza Bayat, Sungnyun Kim, Tal Schuster, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Jiyoun Ha, Aaron Courville, Se-Young Yun
Scaling Transformer 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 a lightweight token-level router dynamically assigns recursion depth to tokens, thereby focusing quadratic attention computation only where it is most useful. Further enhancing its efficiency, MoR incorporates a recursion-wise key-value caching mechanism that eliminates redundant memory access across recursion steps by selectively storing only the key-value caches for designated tokens. Across pretraining runs at model scales ranging from 135M to 1.7B parameters, MoR forms a new Pareto frontier: at equal training FLOPs, it significantly lowers validation perplexity and improves few-shot accuracy, while delivering up to 2.18x higher inference 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.