Mixture-of-Recursions: Dynamische rekursive Tiefen für adaptive Berechnungen auf Token-Ebene lernen

papers.abstract

Das Skalieren von Sprachmodellen erschließt beeindruckende Fähigkeiten, doch die damit verbundenen Rechen- und Speicheranforderungen machen sowohl das Training als auch den Einsatz kostspielig. Bestehende Bemühungen zur Effizienzsteigerung zielen typischerweise entweder auf Parameter-Sharing oder adaptive Berechnungen ab, wobei die Frage offen bleibt, wie beides gleichzeitig erreicht werden kann. Wir stellen Mixture-of-Recursions (MoR) vor, ein einheitliches Framework, das beide Effizienzachsen innerhalb eines einzigen rekursiven Transformers kombiniert. MoR verwendet einen gemeinsamen Stapel von Schichten über Rekursionsschritte hinweg, um Parameter-Effizienz zu erreichen, während leichte Router eine adaptive Token-Level-Verarbeitung ermöglichen, indem sie unterschiedliche Rekursionstiefen dynamisch einzelnen Tokens zuweisen. Dies ermöglicht es MoR, die quadratische Aufmerksamkeitsberechnung nur auf die Tokens zu beschränken, die in einer bestimmten Rekursionstiefe noch aktiv sind, und verbessert die Speicherzugriffseffizienz weiter, indem selektiv nur deren Key-Value-Paare zwischengespeichert werden. Neben diesen Kernmechanismen schlagen wir auch eine KV-Sharing-Variante vor, die KV-Paare aus der ersten Rekursion wiederverwendet, speziell entwickelt, um die Prefill-Latenz und den Speicherbedarf zu verringern. Über Modellgrößen von 135M bis 1,7B Parametern hinweg bildet MoR eine neue Pareto-Grenze: Bei gleichen Trainings-FLOPs und kleineren Modellgrößen senkt es die Validierungs-Perplexität signifikant und verbessert die Few-Shot-Genauigkeit, während es einen höheren Durchsatz im Vergleich zu herkömmlichen und bestehenden rekursiven Baselines liefert. Diese Gewinne zeigen, dass MoR ein effektiver Weg ist, um die Qualität großer Modelle zu erreichen, ohne die Kosten großer Modelle zu verursachen.

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.

Mixture-of-Recursions: Dynamische rekursive Tiefen für adaptive Berechnungen auf Token-Ebene lernen

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

papers.abstract

Support