Дублирование слоёв при обучении подняло GSM8K с 12% до 19% на 3B-модели

Исследователи Ferdinand Kapl, Kaitlin Maile, Stefan Bauer и Johannes von Oswald сравнили три способа построить языковую модель до 3 млрд параметров. Первый — обычная плотная сеть. Второй — сеть, которая во время обучения постепенно становится глубже за счёт копирования средних слоёв (метод LIDAS). Третий — «зацикленная» сеть, где одни и те же слои прогоняют несколько раз за один проход. Смотрели на in-context learning (ICL): умение модели понять задачу по паре примеров в промпте без дообучения. Оказалось, что наращивание и зацикливание выучивают ICL заметно лучше плотной модели. На 3B-масштабе LIDAS набрал около 19% точного совпадения на GSM8K с 8 примерами в промпте, плотная модель — около 12%.

Два способа переиспользовать глубину

У обычной модели число слоёв задано заранее, и каждый слой обучается с нуля. У LIDAS всё иначе: обучение начинается с неглубокой сети, а потом средние слои дублируются, и сеть растёт прямо в процессе. Поначалу копии почти одинаковы, затем расходятся. Ранний этап обучения идёт на маленькой модели, поэтому он дешевле. По данным авторов, на 3B LIDAS тратит на обучение на 26% меньше FLOPs (операций с плавающей точкой) и при этом обходит плотную модель.

Зацикливание устроено проще. Берётся небольшой набор слоёв, и данные проходят через него несколько раз за один прямой проход. Уникальных параметров получается меньше, а эффективная глубина больше.

Эксперименты шли на моделях от 360M до 3B параметров. Код построен на библиотеке nanotron от Hugging Face. Собственного кода и весов авторы не публикуют.

Где модели выигрывают

GSM8K — набор школьных текстовых задач по математике. Восемь решённых примеров кладут в промпт и смотрят, совпадёт ли ответ модели с эталоном. Результаты 3B-моделей:

• LIDAS: около 19% точных совпадений;
• плотная модель: около 12%.

Второй тест — RegBench: новые формальные языки, которые модель раньше не видела. По нескольким примерам она должна уловить правила и предсказать следующий символ. Качество измеряют расстоянием полной вариации — это отличие предсказанных вероятностей от настоящих. Чем меньше, тем лучше:

• наращивание глубины: ~0.119;
• зацикливание: ~0.125;
• плотная модель: ~0.137.

Где это именно ICL: в тесте нечего вспомнить из весов, остаётся только читать контекст.

Размноженные индукционные головы и цена выигрыша

Авторы объясняют результат через индукционные головы. Это головки внимания, которые находят в контексте прошлое вхождение текущего токена и подсказывают то, что шло за ним. Они считаются одним из основных механизмов ICL.

В 3B-модели с LIDAS 10 из 12 сильнейших индукционных голов происходят от двух предков из ранней четырёхслойной сети. Схема, найденная на маленькой модели, при росте копируется по всей глубине. Выглядит так, будто она не вытесняется потом знаниями, которые сеть запоминает в весах.

Ещё одна проверка — отключение. Авторы выключали топ-20% таких голов и смотрели, как падает средний F1 на 19 задачах с несколькими примерами:

• плотная модель: падение примерно на 0.08;
• LIDAS: примерно на 0.18;
• зацикленная модель: примерно на 0.28.

Значит, наращенные и зацикленные модели сильнее держатся на этих схемах, и вся способность учиться по примерам у них сосредоточена в них.

Что учесть инженерам

Для тех, кто предобучает модели, вывод практичный. Если нужна экономия вычислений и хороший ICL, постепенное наращивание глубины выглядит разумным вариантом. На 3B оно дало и лучшие результаты, и 26% экономии FLOPs.

У зацикливания цена выше. Уникальных параметров в такой модели меньше, и хуже получается извлекать знания, записанные в весах. Выигрыш в ICL за это приходится оплачивать: модель лучше читает контекст, но хуже помнит факты.

Есть и оговорка по надёжности. Каждую конфигурацию от 360M до 3B обучили только с одним сидом, то есть с одним начальным состоянием генератора случайных чисел. Разброс между запусками не измерен. Семь процентных пунктов на GSM8K выглядят убедительно, но без повторов нельзя сказать, насколько они устойчивы. Размер до 3B тоже мал по меркам современных моделей, и перенос выводов на десятки миллиардов параметров нужно проверять отдельно.

Тем не менее результат неожиданный. Плотная модель, обученная привычным способом, оказалась хуже именно там, где её считали сильной: в способности учиться на примерах из контекста. Объяснение авторов — раннее копирование схем внимания, и его можно проверять на других размерах.