Зацикленный трансформер на 1B параметров обошёлся в 1,4 раза дешевле по FLOP

Зацикленные трансформеры в этом году стали модной темой: статьи выходят почти каждый день. Идея простая: вместо того чтобы добавлять новые слои, модель несколько раз прогоняет через одни и те же. Параметров столько же, а «думает» она глубже. Новая работа на 149 страниц от исследователей, среди которых Майк Льюис и Артидоро Паньони, наконец отвечает на практичный вопрос: как именно зацикливать, чтобы это окупалось.

В чём проблема зацикливания

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

Что выяснили авторы

Команда предобучила плотные декодеры до 1B параметров и до 32 циклов и сравнила два режима. Full-Stack Looping повторяет все слои. Middle-Stack Looping повторяет только средние, а первые и последние остаются как есть.

Вывод зависит от того, что у вас в дефиците. Если жёстко ограничено число параметров, лучше повторять все слои. Если ограничены вычисления (FLOP), выигрывают средние слои. Для инженера это сразу превращается в правило выбора архитектуры.

Вторая часть работы про параллельное декодирование. В схеме Parallel Loop Transformers состояния циклов сдвигаются к соседним токенам, так что разные циклы считаются одновременно. Авторы заметили, что при нескольких циклах качество деградирует, потому что в одну сумму смешиваются эмбеддинг текущего токена и состояние предыдущего. Их фикс, Normalized Shifted PLT, нормализует эти два входа по отдельности через RMSNorm перед смешиванием. Кроме того, они добавили общие низкоранговые адаптеры для циклов (Loop Matrix), по сути LoRA-подобную поправку, чтобы каждый проход чуть отличался от предыдущего.

Цифры такие. На датасете DCLM при 200M параметров четырёхцикловая версия снижает биты на байт до ~1,060, у оригинального PLT это ~1,087, у модели без циклов ~1,088. При 1B параметров двухцикловая модель даёт ~0,86 против ~0,90 при 80 токенах на параметр. Точность на восьми задачах составила ~47,5% против ~46,1% при 160 токенах на параметр. При долгом обучении двухцикловые модели требуют до 1,4 раза меньше FLOP, чем модели без циклов.

Что с этим делать

Если вы не обучаете модели с нуля, прямо сейчас делать нечего. Но если вы планируете собственную компактную модель для устройства или сервера с жёстким лимитом памяти, зацикливание становится реальным рычагом. Памяти нужно столько же, а качество растёт. Схема выбора простая: упираетесь в параметры, берите полный цикл, упираетесь в FLOP, зацикливайте середину. Если важна скорость генерации, смотрите на нормализованный сдвиг: он как раз делает параллельный декодинг рабочим.

Ограничения

Есть три важные оговорки. Во-первых, эксперименты остановились на 1B параметров, и неизвестно, сохранятся ли зависимости на моделях в десятки миллиардов. Во-вторых, реальное время генерации авторы не замеряли. Экономия FLOP ещё не означает, что ответ придёт быстрее: на практике многое решают память и ядра. В-третьих, выигрыш в точности небольшой, около 1,4 процентного пункта, так что это скорее инженерная оптимизация, чем прорыв.

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