Mixture-of-Recursions : Apprentissage de profondeurs récursives dynamiques pour un calcul adaptatif au niveau des tokens

papers.abstract

La mise à l'échelle des modèles de langage débloque des capacités impressionnantes, mais les exigences computationnelles et en mémoire qui l'accompagnent rendent l'entraînement et le déploiement coûteux. Les efforts existants en matière d'efficacité ciblent généralement soit le partage de paramètres, soit le calcul adaptatif, laissant ouverte la question de savoir comment atteindre les deux simultanément. Nous introduisons Mixture-of-Recursions (MoR), un cadre unifié qui combine ces deux axes d'efficacité au sein d'un seul Transformer récursif. MoR réutilise une pile partagée de couches à travers les étapes de récursion pour atteindre l'efficacité des paramètres, tandis que des routeurs légers permettent une réflexion adaptative au niveau des tokens en assignant dynamiquement différentes profondeurs de récursion à des tokens individuels. Cela permet à MoR de concentrer le calcul d'attention quadratique uniquement parmi les tokens encore actifs à une profondeur de récursion donnée, améliorant ainsi l'efficacité d'accès à la mémoire en mettant en cache de manière sélective uniquement leurs paires clé-valeur. Au-delà de ces mécanismes de base, nous proposons également une variante de partage KV qui réutilise les paires KV de la première récursion, spécifiquement conçue pour réduire la latence de préremplissage et l'empreinte mémoire. À travers des échelles de modèles allant de 135M à 1,7B de paramètres, MoR forme une nouvelle frontière de Pareto : à égalité de FLOPs d'entraînement et avec des tailles de modèles plus petites, il réduit significativement la perplexité de validation et améliore la précision en few-shot, tout en offrant un débit plus élevé par rapport aux baselines récursives existantes et classiques. Ces gains démontrent que MoR est une voie efficace pour atteindre la qualité des grands modèles sans en supporter le coût.

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 : Apprentissage de profondeurs récursives dynamiques pour un calcul adaptatif au niveau des tokens

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

papers.abstract

Support