Немецкая исследовательская группа предложила способ, при котором Transformer-модели сами выбирают, сколько раз повторять вычисления для решения задачи. В сочетании с дополнительной памятью такой подход заметно превосходит более крупные модели в математических тестах.
Работу выполнили специалисты из Lamarr Institute, Fraunhofer IAIS и Боннского университета. Они изучали, можно ли устранить компромисс между looped transformers, которые экономят параметры за счёт повторного использования одного и того же вычислительного блока, и потерей ёмкости для хранения знаний.
Обычные языковые модели могут решать задачи по шагам с помощью chain-of-thought prompting, но каждый промежуточный шаг требует дополнительных токенов. В looped transformers один и тот же блок вычислений многократно применяется к внутренним представлениям модели без вывода промежуточных шагов в текст.
Архитектура исследователей объединяет два механизма: адаптивные циклы, при которых каждый слой Transformer через обучаемый механизм остановки решает, сколько раз повторять вычисления, и обучаемые банки памяти, которые дают модели дополнительную ёмкость для хранения знаний.
Базовая версия — decoder-only Transformer с 12 слоями и примерно 200 млн параметров. Модель обучали на 14 млрд токенов из дедуплицированного датасета FineWeb Edu. Варианты с циклами позволяли каждому слою выполнять до 3, 5 или 7 повторов.
Банки памяти включают 1024 локальных слота на слой и 512 глобальных общих слотов. По данным исследования, это добавляет около 10 млн параметров.
Результаты показали, что возможность повторять вычисления до трёх раз заметно улучшает математические способности модели. Версия с циклами набрала на 22% больше, чем базовая модель без циклов.
Наибольший прирост зафиксирован в более трудных разделах математики: Precalculus дал улучшение на 31%, а Intermediate Algebra — на 26%. Для задач на повседневные знания, включая вопросы о социальных ситуациях или физической интуиции, циклы почти не помогли. При большем числе повторов качество даже немного снижалось.
Для сравнения исследователи сопоставили 12-слойную модель с тройными циклами и обычную 36-слойную модель без циклов, но с тем же вычислительным бюджетом. Несмотря на втрое меньшее число слоёв, модель с циклами показала результат на 6,4% выше в математических бенчмарках.
Авторы пишут, что для математического рассуждения циклы эффективнее, чем простое увеличение числа слоёв. При этом банки памяти решают другую задачу: повседневные знания нельзя получить одним только повторным вычислением, их нужно где-то хранить.
Дополнительная память частично закрывает этот разрыв. По данным исследования, вместе с банками памяти модель получила ещё 4,2% прироста в математических задачах и 2% в заданиях на повседневные знания по сравнению с вариантом без памяти.
Исследователи также заметили самопроизвольную специализацию слоёв. Ранние слои повторяют вычисления минимально и почти не обращаются к памяти, тогда как поздние делают это заметно чаще.
Это согласуется с прошлыми работами, где ранние слои Transformer связывали с локальными синтаксическими шаблонами, а поздние — с более сложными семантическими и логическими операциями. Простые вычисления мало выигрывают от повторов, а более глубокие — наоборот.
Во время обучения выявилась и чёткая точка перехода. На ранних этапах модели почти не используют циклы, хотя такая возможность у них есть. Повторные вычисления начинают активно применяться лишь после того, как модель достаточно хорошо осваивает понимание и предсказание языка.
По словам авторов, этот порог возникает почти в одной и той же точке для всех конфигураций циклов. Сначала модель должна набрать базовые языковые навыки, и только потом она начинает получать пользу от многократного повторения вычислений.
Исследователи считают, что результаты указывают на разделение функций внутри Transformer. Feed-forward слои работают как память для фактических связей, а attention-слои направляют и обрабатывают информацию. Циклы улучшают этот процесс, но не компенсируют нехватку места для хранения знаний.
Тот факт, что слои с большим числом повторов чаще обращаются к памяти, поддерживает эту интерпретацию. Циклы и память дополняют друг друга: больше вычислений требует большего запаса фактов.
Авторы отдельно отмечают ограничения работы. Эксперименты проводились в сравнительно небольшом масштабе — около 200 млн параметров и 14 млрд обучающих токенов. Сохранятся ли эти результаты для моделей на несколько миллиардов параметров, у которых уже есть значительная встроенная ёмкость, пока неясно.
Источник: The Decoder, исследование Lamarr Institute, Fraunhofer IAIS и Боннского университета.
