Смесь рекурсий: обучение динамической глубины рекурсии для адаптивных вычислений на уровне токенов

Аннотация

Масштабирование языковых моделей открывает впечатляющие возможности, но сопутствующие вычислительные и ресурсные затраты делают как обучение, так и развертывание дорогостоящими. Существующие подходы к повышению эффективности обычно сосредоточены либо на совместном использовании параметров, либо на адаптивных вычислениях, оставляя открытым вопрос о том, как достичь обеих целей одновременно. Мы представляем Mixture-of-Recursions (MoR) — унифицированную структуру, которая объединяет оба направления повышения эффективности в рамках единого Рекурсивного Трансформера. MoR повторно использует общий набор слоев на каждом шаге рекурсии для достижения параметрической эффективности, в то время как легковесные маршрутизаторы обеспечивают адаптивное мышление на уровне токенов, динамически назначая различную глубину рекурсии отдельным токенам. Это позволяет MoR сосредоточить квадратичные вычисления внимания только на токенах, активных на текущей глубине рекурсии, дополнительно повышая эффективность доступа к памяти за счет избирательного кэширования только их пар ключ-значение. Помимо этих основных механизмов, мы также предлагаем вариант с совместным использованием пар ключ-значение (KV), который повторно использует KV-пары из первой рекурсии, специально разработанный для уменьшения задержки предварительного заполнения и объема занимаемой памяти. На масштабах моделей от 135 млн до 1,7 млрд параметров MoR формирует новый Парето-фронт: при равных затратах на обучение в FLOPs и меньших размерах моделей он значительно снижает перплексию на валидации и повышает точность в условиях малого количества примеров, обеспечивая при этом более высокую пропускную способность по сравнению с базовыми и существующими рекурсивными моделями. Эти достижения демонстрируют, что MoR является эффективным путем к качеству крупных моделей без затрат, связанных с их использованием.

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: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation

Аннотация

Support