SMELT сокращает до 18% вычислений при обучении MoE-трансформеров

Зацикленные трансформеры (Looped Transformers) повышают эффективную глубину сети за счёт повторного прогона одного и того же блока слоёв, но большинство прежних оценок сравнивали такие модели с обычными при одинаковом размере, из-за этого выигрыш архитектуры путался с выигрышем от лишних вычислений, которые даёт повтор. Авторы работы изучили зацикливание на трансформерах со смесью экспертов (Mixture-of-Experts, MoE) и намеренно уравняли три параметра между зацикленной моделью и обычной («базовой»): число операций с плавающей точкой (FLOP) на токен, суммарное число необучаемых-с-нуля (не считая эмбеддингов) параметров и объём KV-кэша, памяти, где хранятся промежуточные представления ключей и значений внимания.

Перебрав несколько вариантов, они остановились на рецепте SMELT (Sparse MoE Transformer, middle layers Loop Twice, «разреженный MoE-трансформер, средние слои прогоняются дважды»): средняя половина слоёв сети прогоняется через себя ещё раз, а остальные параметры архитектуры подгоняются так, чтобы FLOP, число параметров и KV-кэш совпадали с обычной, незацикленной моделью. SMELT протестировали на четырёх размерах моделей, вплоть до 54 млрд параметров без учёта эмбеддингов, и для каждой архитектуры (зацикленной и обычной) отдельно подобрали закон масштабирования в стиле Chinchilla, описывающий, как потери на обучении падают с ростом вычислительного бюджета.

Потери SMELT падают быстрее при росте вычислений, чем у обычной модели: на оптимальной по вычислениям границе (compute-optimal frontier, сочетаниях размера модели и объёма данных, при которых для заданного бюджета достигается наименьшая потеря) SMELT экономит от 6,8 до 18,0% вычислительного бюджета обучения. Преимущество проявляется и на прикладных бенчмарках сильнее, чем можно было бы предсказать по одной только метрике потерь на валидации: оно наиболее заметно в задачах кода и растёт с длиной примера и числом примеров, которые модель видит в контексте (few-shot).

Отдельно авторы разобрали, почему это работает: механистический анализ показывает, что при повторном («втором») проходе через зацикленные слои модель меньше проваливается в эффект стока внимания (attention sink), склонность непропорционально много внимания отдавать первым, часто малоинформативным токенам, и вместо этого перераспределяет внимание на токены, содержательно значимые для задачи. Авторы описывают это как индуктивное смещение (inductive bias), которое может лежать в основе наблюдаемого прироста качества. Общий вывод работы: зацикливание слоёв улучшает трансформеры даже при строгом уравнивании вычислительного бюджета, а не только за счёт скрытого увеличения вычислений, и SMELT даёт практический рецепт, превращающий повторное использование глубины сети в измеримый выигрыш.

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

SMELT прогоняет среднюю половину слоёв MoE-трансформера дважды, при этом FLOP на токен, число параметров и объём KV-кэша уравнены с обычной моделью
Проверено на четырёх размерах моделей вплоть до 54 млрд параметров (без эмбеддингов), для каждой архитектуры отдельно подобран закон масштабирования в стиле Chinchilla
На оптимальной по вычислениям границе SMELT экономит от 6,8 до 18,0% вычислительного бюджета обучения по сравнению с незацикленной базовой моделью
Выигрыш на прикладных бенчмарках больше, чем предсказывает метрика потерь; сильнее всего проявляется в задачах кода и растёт с длиной примеров и числом few-shot-примеров
Механистический анализ: повторный проход снижает эффект стока внимания и перенаправляет внимание на содержательно значимые токены

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

Большинство прежних работ о зацикленных трансформерах сравнивали их с обычными моделями при одинаковом размере, а не при одинаковом вычислительном бюджете, из-за этого было неясно, даёт ли зацикливание архитектурное преимущество или просто использует больше вычислений под видом того же размера модели. Эта работа честно уравнивает FLOP, число параметров и KV-кэш между зацикленной и обычной MoE-моделью и показывает, что выигрыш сохраняется даже при таком строгом сравнении: SMELT снижает потери быстрее при равном бюджете и экономит от 6,8 до 18,0% вычислений на оптимальной по вычислениям границе.

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

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

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

SMELT, это рецепт архитектуры, а не готовый продукт: средняя половина слоёв MoE-трансформера прогоняется дважды, а остальные гиперпараметры подбираются так, чтобы FLOP на токен, число параметров и KV-кэш совпадали с обычной моделью. Рецепт проверен на масштабах вплоть до 54 млрд параметров и приносит наибольшую пользу в задачах кода и на длинных/многопримерных контекстах, это стоит учитывать при выборе, где в первую очередь опробовать зацикливание.

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

Источник, карточка препринта на Hugging Face Papers с полным текстом аннотации; ни автор(ы), ни организация, ни точная дата исследования в самом тексте аннотации не указаны. Имя «Shaowen Wang», числящееся в карточке, принадлежит человеку, разместившему препринт на Hugging Face, а не обязательно автору работы, поэтому в пересказе оно не используется как атрибуция. Сравнение дано только с одной незацикленной базовой моделью, данных о сопоставлении с другими рецептами зацикливания или конкурирующими архитектурами в источнике нет; абсолютные значения FLOP и потерь также не приведены, только относительная экономия и потолок параметров.

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

Все цифры получены на компьютерных экспериментах с моделями до 54 млрд параметров без эмбеддингов, перенос результата на ещё большие модели или на другие семейства архитектур (не MoE) в источнике не проверен. Экономия в 6,8, 18,0%, это диапазон в зависимости от размера/бюджета, а не фиксированное число, и относится именно к точке на оптимальной по вычислениям границе, а не к произвольному режиму обучения. Механистическое объяснение через снижение эффекта стока внимания, это гипотеза авторов о причине эффекта («может лежать в основе»), а не доказанный закон.