Mixture-of-Recursions: Dynamische Recursiedieptes Leren voor Adaptieve Token-Level Berekenen
Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation
July 14, 2025
Auteurs: Sangmin Bae, Yujin Kim, Reza Bayat, Sungnyun Kim, Jiyoun Ha, Tal Schuster, Adam Fisch, Hrayr Harutyunyan, Ziwei Ji, Aaron Courville, Se-Young Yun
cs.AI
Samenvatting
Het opschalen van taalmodelen ontgrendelt indrukwekkende mogelijkheden, maar de bijbehorende reken- en geheugeneisen maken zowel training als implementatie kostbaar. Bestaande efficiëntie-inspanningen richten zich doorgaans op het delen van parameters of adaptieve berekening, waardoor de vraag open blijft hoe beide tegelijkertijd kunnen worden bereikt. Wij introduceren Mixture-of-Recursions (MoR), een uniform raamwerk dat de twee assen van efficiëntie combineert binnen een enkele Recursive Transformer. MoR hergebruikt een gedeelde stapel lagen over recursiestappen om parameter-efficiëntie te bereiken, terwijl lichtgewicht routers adaptief denken op tokenniveau mogelijk maken door dynamisch verschillende recursiediepten toe te wijzen aan individuele tokens. Hierdoor kan MoR de kwadratische aandachtberekening concentreren op alleen de tokens die nog actief zijn op een gegeven recursiediepte, wat de efficiëntie van geheugentoegang verder verbetert door selectief alleen hun sleutel-waardeparen in cache op te slaan. Naast deze kernmechanismen stellen we ook een KV-delingvariant voor die KV-paren van de eerste recursie hergebruikt, specifiek ontworpen om de prefill-latentie en geheugenvoetafdruk te verkleinen. Over modelschalen variërend van 135M tot 1,7B parameters vormt MoR een nieuwe Pareto-grens: bij gelijke trainings-FLOPs en kleinere modelgroottes verlaagt het significant de validatieperplexiteit en verbetert het few-shot-nauwkeurigheid, terwijl het een hogere doorvoersnelheid biedt in vergelijking met standaard- en bestaande recursieve basislijnen. Deze winsten tonen aan dat MoR een effectieve weg is naar kwaliteit van grote modellen zonder de kosten van grote modellen te maken.
English
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 decrease prefill latency and 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.