В мире больших языковых моделей (LLM) доминирующей архитектурой уже много лет остается трансформер с механизмом внимания, работающий исключительно в режиме декодирования. Однако у этой архитектуры есть фундаментальное ограничение, о котором часто забывают в погоне за производительностью: информация, вычисленная на последнем слое токена t, не передается напрямую на первый слой токена t+1. Позиции общаются друг с другом только через механизм внимания, опираясь на кэшированные ключи и значения (KV-cache). Это создает своего рода «амнезию» на границе между токенами, где контекст должен быть восстановлен заново, а не продолжен.
Исследователь из Принстонского университета, Ифань Чжан (Yifan Zhang), в своем техническом отчете предлагает радикально изменить этот подход. Его новая архитектура, названная Recurrent Looped Transformer (RLT) (Рекуррентный зацикленный трансформер), предлагает «замкнуть эту петлю». Идея заключается в том, чтобы финальное скрытое состояние декодера и его кэш внимания с скользящим окном (Sliding-Window Attention, SWA) переносились в следующий токен. Важно отметить, что этот перенос происходит как во время обработки промпта, так и во время генерации ответа, без какого-либо сброса состояния на границе между ними.
Это не просто теоретическая выкладка. RLT — это детальная спецификация дизайна, которая определяет архитектуру, расписание выполнения и контракт для переигровки в обучении с подкреплением (RL). Однако важно сразу обозначить границы текущего исследования: автор явно сообщает, что измеренных результатов эффективности, качества рассуждений или масштабирования пока нет. Это архитектурная концепция, требующая валидации, но открывающая новые горизонты для понимания того, как модели могут обрабатывать временную глубину.
01Как устроен Recurrent Looped Transformer
Архитектура RLT представляет собой гибрид, сочетающий в себе преимущества параллельной обработки и рекуррентности. Она состоит из двух основных частей: причинного энкодера и рекуррентного декодера. Энкодер обрабатывает токены параллельно, накладывая причинную маску (causal mask), и производит представления et. Из этих представлений проецируется память M≤t. Гибкость системы позволяет группам памяти быть общими для всех слоев декодера (G = 1) или оставаться специфичными для каждого слоя (G = LD).
Сердцем системы является декодер, который удерживает рекуррентное состояние. Его полное состояние H определяется как пара (s, C), где s — это финальный выход декодера, а C содержит сохраненные ключи и значения (KV) для скользящего окна внимания на каждом слое декодера. Для каждого нового токена происходит процесс «гибридизации»:
- Входные представления et от энкодера объединяются с предыдущим выходом декодера st-1 с помощью вентильного механизма (gated merge).
- Каждый блок декодера выполняет причинное скользящее окно внимания (SWA) над активациями декодера.
- Выполняется кросс-внимание (cross-attention) к памяти энкодера.
- Применяется полносвязный блок (FFN - Feed-Forward Network).
Размер окна W включает текущий токен, что означает, что на каждом слое сохраняется не более W – 1 исторических записей. Распределение вероятностей следующего токена считывается из st. Инициализация происходит один раз перед токеном начала последовательности (BOS) с использованием обучаемого начального состояния s* и пустого кэша.
В эталонной конфигурации используется 48 слоев энкодера и 48 слоев декодера, с совместным использованием весов внимания и FFN между ними. Это означает, что каждый токен выполняет 96 логических блоков. Однако, поскольку блоки декодера добавляют кросс-внимание, вычислительная нагрузка (FLOPs) на блок не является равной. Чжан называет это переиспользованием параметров, а не копированием активаций, что является важным семантическим различием в оптимизации памяти.
02Три принципа дизайна RLT
Архитектура RLT опирается на три фундаментальных принципа, которые отличают ее от предыдущих попыток внедрения рекуррентности в трансформеры.
1. Латентные рассуждения с неограниченной временной глубиной
После обработки t токенов путь состояния от st проходит через t·LD блоков декодера. В эталонной конфигурации с 48 слоями декодера это означает путь длиной в 48t блоков. При этом работа на один токен остается фиксированной, но структурная глубина пути растет вместе с длиной последовательности. Это создает эффект «глубокой памяти» без необходимости увеличивать размер окна внимания.
Однако автор предупреждает: наличие структурной глубины не гарантирует качества рассуждений. Вентили (gates) и механизмы сжатия могут подавлять длинные пути, если они не настроены должным образом. Это вызов для будущих исследований: как заставить модель эффективно использовать эту глубину, не теряя важные детали.
2. Совместный дизайн модели и оборудования
RLT учитывает особенности современных GPU и TPU. Вычисления для энкодера и проекции памяти для известных токенов используют токено-параллельные ядра (token-parallel kernels). Переходы декодера остаются последовательными внутри одной последовательности, но обновления, готовые от независимых последовательностей, могут быть объединены в один батчевый kernel.
В отчете прямо заявлено: не предполагается никакого параллельного сканирования (parallel scan) для нелинейного декодера. Также не заявляется о ускорении этапа префилла (prefill). Стандартный параллельный проход декодера с SWA не эквивалентен рекуррентному подходу. Батчинг, слияние ядер (kernel fusion) и чекпоинтинг указаны как цели реализации, а не как уже завершенные ядра. Это честное признание того, что оптимизация под железо — это отдельная большая задача.
3. Совместный дизайн модели и алгоритма RL
Это, пожалуй, самый инновационный аспект. Предобучение, обучение с учителем (SFT), сэмплирование и переигровка в RL (RL replay) используют одно и то же переходное состояние. Для обучения с подкреплением сэмплер записывает логарифмическую вероятность поведения каждого действия в соответствии с его фактическим распределением сэмплирования, включая температуру и усечение.
Тренер перестраивает память энкодера, рекуррентный вывод и каждый кэш SWA с начала последовательности, используя текущие параметры, перед оценкой каждого действия. Старые состояния роутинга (rollout) никогда не переиспользуются. Это гарантирует, что градиенты вычисляются корректно для текущей версии модели. Предложение 3.1 формализует выигрыш: перемещение границы между промптом и ответом не меняет условное распределение для фиксированной истории токенов. Это критически важно для стабильности обучения.
03Обучение и обслуживание (Serving)
Предобучение RLT осуществляется путем предсказания следующего токена для полной последовательности с полной обратной передачей через время (full backpropagation through time). Это вычислительно затратно, но необходимо для сохранения целостности рекуррентного состояния.
Для обслуживания в многошаговых сценариях (multi-turn serving) требуется точный снимок префикса. Он включает:
- Кэш энкодера и память.
- Полное состояние декодера.
- Метаданные позиций.
- Конвенцию окна (window convention).
- Версию модели.
Снимок с фиксированными весами можно переиспользовать, так как состояние независимо от разделения обслуживания. Однако обновления весов аннулируют старые состояния. Редактирование префикса требует пересчета от более раннего чекпоинта. Внешние токены в многошаговом RL обновляют состояние, но не получают факторов отношения важности (importance-ratio factors), что упрощает, но и ограничивает некоторые аспекты оптимизации.
04Связь с предыдущими работами
RLT не возникла в вакууме. Она строится на плечах гигантов индустрии и академии:
- Память, производная от энкодера: Следует подходу YOCO, который кэширует KV один раз для кросс-декодера, и DeepSeek-V4.1-Flash, который проецирует глобальный KV декодера из финальных состояний энкодера. RLT сохраняет эту память, но отказывается от пропуска декодера по всему промпту, который был в этих моделях.
- Временная обратная связь: Строится на Feedback Transformer и Recurrent Transformer. Однако RLT подает предыдущий финальный вывод декодера в следующий вход декодера и запускает рекуррентность также по самому промпту, а не только по ответу.
- Переиспользование по глубине: Связано с Universal Transformers и рекуррентной латентной рассуждающей способностью. Аргумент переигровки расширяет заметку Чжана о несоответствии ядер префилла и декодирования (prefill-decode kernel mismatch).
05Что это значит на практике
Для разработчиков и исследователей RLT представляет собой интересный экспериментальный полигон. Вот ключевые выводы, которые стоит учитывать:
- Непрерывность состояния: RLT переносит полное состояние декодера (финальный вывод плюс кэш SWA на каждом слое) через каждый токен промпта и ответа без сброса границы. Это может улучшить контекстную целостность в длинных диалогах.
- Вычислительная сложность: Эталонная конфигурация использует 48 связанных слоев энкодера и декодера, что дает 96 логических блоков на токен. Путь состояния после t токенов составляет 48t блоков. Это означает линейный рост вычислительных затрат на последовательность, что может быть проблемой для очень длинных контекстов.
- Возможности для оборудования: Энкодер может использовать параллелизм, а батчинг возможен между последовательностями. Однако автор честно признает, что параллельного сканирования или ускорения префилла не заявлено. Это значит, что на текущем этапе RLT может быть медленнее стандартных трансформеров на этапе инициализации.
- RL Replay: Переигровка в RL перестраивает все состояния под текущими параметрами, сохраняя записанные логарифмические вероятности поведения как знаменатели отношений. Это обеспечивает стабильность обучения, но требует больших вычислительных ресурсов.
- Отсутствие измеренных результатов: Качество рассуждений, эффективность и масштабирование RL остаются открытыми целями для валидации. RLT — это спецификация, а не готовый продукт.
Если вы хотите углубиться в технические детали, вы можете ознакомиться с Техническим отчетом, репозиторием на GitHub и страницей проекта. Все заслуги принадлежат исследователю этого проекта. Следите за обновлениями в области AI, так как подобные архитектурные инновации могут в будущем изменить то, как мы строим модели с долгой памятью.
Источник: MarkTechPost ↗