Mixture-of-Recursions: Apprendimento di Profondità Ricorsive Dinamiche per il Calcolo Adattivo a Livello di Token

Abstract

Il ridimensionamento dei modelli linguistici sblocca capacità impressionanti, ma le relative esigenze computazionali e di memoria rendono sia l'addestramento che il dispiegamento costosi. Gli sforzi esistenti per l'efficienza si concentrano tipicamente sulla condivisione dei parametri o sul calcolo adattivo, lasciando aperta la questione di come ottenere entrambi simultaneamente. Introduciamo Mixture-of-Recursions (MoR), un framework unificato che combina i due assi dell'efficienza all'interno di un singolo Transformer ricorsivo. MoR riutilizza una pila condivisa di strati attraverso i passi di ricorsione per ottenere efficienza nei parametri, mentre router leggeri abilitano un pensiero adattivo a livello di token assegnando dinamicamente diverse profondità di ricorsione ai singoli token. Ciò consente a MoR di concentrare il calcolo quadratico dell'attenzione solo tra i token ancora attivi a una data profondità di ricorsione, migliorando ulteriormente l'efficienza dell'accesso alla memoria memorizzando selettivamente solo le loro coppie chiave-valore. Oltre a questi meccanismi principali, proponiamo anche una variante di condivisione KV che riutilizza le coppie KV dalla prima ricorsione, specificamente progettata per ridurre la latenza di prefill e l'impronta di memoria. Su scale di modelli che vanno da 135M a 1.7B parametri, MoR forma una nuova frontiera di Pareto: a pari FLOP di addestramento e dimensioni di modello più piccole, riduce significativamente la perplessità di validazione e migliora l'accuratezza few-shot, offrendo un throughput più elevato rispetto ai baseline ricorsivi esistenti e vanilla. Questi guadagni dimostrano che MoR è un percorso efficace verso la qualità dei modelli di grandi dimensioni senza incorrere nei costi dei modelli di grandi dimensioni.

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: Apprendimento di Profondità Ricorsive Dinamiche per il Calcolo Adattivo a Livello di Token

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

Abstract

Support