Recurrent Looped Transformer: архитектура с обратной связью лучше обобщает на длинные последовательности

Авторы статьи отталкиваются от ограничения Transformer: отслеживание состояния требует обновления на каждом входе, а глубина, которую Transformer применяет к каждому токену, фиксирована и не зависит от длины последовательности.

Их решение, Recurrent Looped Transformer (RLT). Слои делятся между параллельным каузальным кодировщиком и рекуррентным декодером. На каждом токене декодер объединяет выход кодировщика с финальным состоянием декодера на предыдущем токене. Поэтому путь вычислений растёт вместе с длиной последовательности, а стоимость на один токен остаётся фиксированной.

Эксперименты проведены на шести алгоритмических задачах. Пять вариантов разделения восьми слоёв между кодировщиком и декодером сравнивали с восьмислойным Transformer на трёх случайных инициализациях (seed).

Результаты из аннотации. Три задачи описаны подробно. Чётность (parity): при обучении на последовательностях не длиннее 40 бит два варианта RLT обобщили задачу на 256 бит со 100% точностью при каждой инициализации, а Transformer остался на уровне случайного угадывания. Отслеживание перестановок S_5 на основе обменов (swap-based S_5) при длине в восемь раз больше обучающей: RLT показал 97% точности итогового состояния, Transformer, меньше 1%; точность растёт с глубиной декодера. Модульная арифметика за пределами обучающих длин: RLT достигает до 93% против 33% у Transformer.

Абляции показывают, что выигрыш зависит от обратной связи: если её убрать, чётность и swap-based S_5 падают до уровня случайного угадывания при любом разделении слоёв. Если обновлять обратную связь один раз на блок из четырёх токенов, известные токены блока можно обрабатывать параллельно, а 64-битная чётность остаётся на 99%. Но отслеживание перестановок зависит от обратной связи на каждом токене: при блоках точность swap-based S_5 на длине 64 падает со 100% до 20%.

Ключевые факты

RLT делит слои между параллельным каузальным кодировщиком и рекуррентным декодером, который на каждом токене учитывает финальное состояние декодера с предыдущего токена.
Чётность: после обучения не более чем на 40 битах два варианта RLT дали 100% точности на 256 битах при каждой инициализации, Transformer остался на уровне случайного угадывания.
Отслеживание перестановок S_5 при длине в восемь раз больше обучающей: 97% у RLT против менее 1% у Transformer; модульная арифметика, до 93% против 33%.
Без обратной связи чётность и swap-based S_5 падают до случайного уровня при любом разделении слоёв.
Обновление обратной связи раз в четыре токена сохраняет 99% на 64-битной чётности, но снижает swap-based S_5 на длине 64 со 100% до 20%.

Почему это важно

Авторы указывают на конкретную слабость Transformer: на каждый токен тратится фиксированная глубина вычислений, хотя отслеживание состояния требует обновления на каждом входе. RLT пытается решить это, давая вычислениям расти вместе с длиной последовательности при фиксированной стоимости на токен. В аннотации это даёт заметный отрыв на задачах, где обычный Transformer не выходит за уровень случайного угадывания, например на чётности и перестановках S_5 за пределами обучающих длин.

Кому это важно

Исследователям архитектур нейросетей, обобщения на длины, больше обучающих, и рассуждений по шагам. Практикам эта работа пока скорее как ориентир: в аннотации описаны только алгоритмические задачи.

Как это применить

Прямого применения в аннотации нет. Идею можно рассматривать как принцип проектирования: часть слоёв остаётся параллельной, а часть образует рекуррентный контур с обратной связью от предыдущего токена. Из абляций следует, что для отслеживания перестановок обратная связь нужна на каждом токене, а для чётности её можно обновлять раз в четыре токена и сэкономить на параллелизме.

Можно ли доверять

Это препринт в листинге Hugging Face Papers, и мы опираемся на аннотацию. Результаты усреднены по трём инициализациям, а для чётности 100% заявлено при каждой из них. Но эксперименты ограничены шестью алгоритмическими задачами и сравнением с одним восьмислойным Transformer. Для модульной арифметики не указана длина, на которой получены 93% и 33%. Для 97% на S_5 не приведён разброс по инициализациям. Не указано, какие именно два из пяти разделений слоёв дали 100% на чётности.

Риски и подводные камни

Успех на синтетических задачах не гарантирует выигрыша на естественном языке или реальных данных: в аннотации таких результатов нет. Рекуррентная обратная связь по токенам ограничивает параллелизм: вариант с блоками по четыре токена ускоряет обработку, но для перестановок S_5 качество резко падает. Кроме того, сравнение идёт только с восьмислойным Transformer, а о размерах моделей и затратах на обучение в аннотации не сказано.

«Отслеживание состояния требует обновления на каждом входе, но глубина, которую Transformer применяет к каждому токену, фиксирована независимо от длины последовательности.»

— Из аннотации статьи