Mixture-of-Recursions: Aprendendo Profundidades Recursivas Dinâmicas para Computação Adaptativa em Nível de Token

Resumo

A escalabilidade de modelos de linguagem desbloqueia capacidades impressionantes, mas as demandas computacionais e de memória associadas tornam tanto o treinamento quanto a implantação caros. Os esforços existentes de eficiência geralmente visam o compartilhamento de parâmetros ou a computação adaptativa, deixando em aberto a questão de como alcançar ambos simultaneamente. Apresentamos o Mixture-of-Recursions (MoR), uma estrutura unificada que combina os dois eixos de eficiência dentro de um único Transformer Recursivo. O MoR reutiliza uma pilha compartilhada de camadas em etapas de recursão para alcançar eficiência de parâmetros, enquanto roteadores leves permitem pensamento adaptativo em nível de token ao atribuir dinamicamente diferentes profundidades de recursão a tokens individuais. Isso permite que o MoR concentre a computação quadrática de atenção apenas entre os tokens ainda ativos em uma determinada profundidade de recursão, melhorando ainda mais a eficiência de acesso à memória ao armazenar em cache seletivamente apenas seus pares chave-valor. Além desses mecanismos principais, também propomos uma variante de compartilhamento de KV que reutiliza pares KV da primeira recursão, projetada especificamente para diminuir a latência de preenchimento e a pegada de memória. Em escalas de modelos que variam de 135M a 1,7B de parâmetros, o MoR forma uma nova fronteira de Pareto: com FLOPs de treinamento iguais e tamanhos de modelo menores, ele reduz significativamente a perplexidade de validação e melhora a precisão em poucos exemplos, enquanto oferece maior taxa de transferência em comparação com baselines recursivos existentes e vanilla. Esses ganhos demonstram que o MoR é um caminho eficaz para a qualidade de modelos grandes sem incorrer no custo de modelos grandes.

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: Aprendendo Profundidades Recursivas Dinâmicas para Computação Adaptativa em Nível de Token

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

Resumo

Support