Mezcla de Recursiones: Aprendizaje de Profundidades Recursivas Dinámicas para Cálculo Adaptativo a Nivel de Tokens

Resumen

El escalado de modelos de lenguaje desbloquea capacidades impresionantes, pero las demandas computacionales y de memoria asociadas hacen que tanto el entrenamiento como la implementación sean costosos. Los esfuerzos existentes en eficiencia suelen enfocarse en el uso compartido de parámetros o en el cálculo adaptativo, dejando abierta la cuestión de cómo lograr ambos simultáneamente. Presentamos Mixture-of-Recursions (MoR), un marco unificado que combina estos dos ejes de eficiencia dentro de un único Transformer Recursivo. MoR reutiliza una pila compartida de capas a través de los pasos de recursión para lograr eficiencia en parámetros, mientras que enrutadores ligeros permiten un pensamiento adaptativo a nivel de token al asignar dinámicamente diferentes profundidades de recursión a tokens individuales. Esto permite que MoR enfoque el cálculo de atención cuadrática solo entre los tokens aún activos en una profundidad de recursión dada, mejorando aún más la eficiencia de acceso a la memoria al almacenar selectivamente solo sus pares clave-valor. Más allá de estos mecanismos centrales, también proponemos una variante de uso compartido de KV que reutiliza los pares KV de la primera recursión, diseñada específicamente para reducir la latencia de prellenado y la huella de memoria. En escalas de modelos que van desde 135M hasta 1.7B de parámetros, MoR forma una nueva frontera de Pareto: con igual cantidad de FLOPs de entrenamiento y tamaños de modelo más pequeños, reduce significativamente la perplejidad de validación y mejora la precisión en pocos ejemplos, al tiempo que ofrece un mayor rendimiento en comparación con líneas base recursivas existentes y modelos convencionales. Estas ganancias demuestran que MoR es un camino efectivo hacia la calidad de modelos grandes sin incurrir en el costo 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.

Mezcla de Recursiones: Aprendizaje de Profundidades Recursivas Dinámicas para Cálculo Adaptativo a Nivel de Tokens

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

Resumen

Support