Учёные научились предсказывать скорость обучения MoE-моделей на триллионах токенов без дорогого перебора

Авторы работы предлагают двухэтапный метод переноса гиперпараметров, который позволяет заранее подобрать оптимальную скорость обучения для больших MoE-моделей (архитектур со «смесью экспертов», где на каждый токен активируется лишь часть параметров сети), не запуская дорогой полный перебор настроек на целевом масштабе.

Первый шаг, адаптация Maximal Update Parameterization (μP, метод параметризации сети, при котором оптимальные гиперпараметры не меняются при изменении ширины модели) под архитектуру MoE. Авторы применили её вместе с механизмом внимания Multi-head Latent Attention (MLA) и оптимизатором Muon и показали, что при таком сочетании оптимальная скорость обучения устойчиво переносится между моделями разной ширины, то есть подобранное на узкой модели значение остаётся оптимальным и для более широкой.

Второй шаг, перенос уже не по ширине сети, а по числу токенов обучения. Авторы обучили несколько маленьких прокси-моделей на ограниченных бюджетах токенов, зафиксировали для каждой оптимальную скорость обучения и с помощью линейной регрессии построили закон масштабирования, экстраполирующий это значение на гораздо большие горизонты обучения, например, на 10 триллионов токенов. Точность такой экстраполяции оказалась высокой: R² = 0,95.

Из этого авторы делают вывод, что обучения маленьких прокси-моделей достаточно, чтобы заранее определить оптимальную скорость обучения для полномасштабного обучения большой MoE-модели. Чтобы проверить метод на практике, они с нуля предобучили собственную базовую MoE-модель на 155 млрд параметров, из которых 17 млрд активны на каждом токене. Обучение прошло стабильно, а результаты оценки подтвердили, что оптимальные настройки для модели полного масштаба можно точно предсказать заранее, с минимальными затратами на дополнительные проверочные эксперименты (абляции), вместо того чтобы подбирать их прямым перебором на целевом масштабе.

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

Предложен двухэтапный метод переноса гиперпараметров для MoE-моделей: сначала, перенос оптимальной скорости обучения между моделями разной ширины через адаптацию Maximal Update Parameterization (μP) для MoE с механизмом внимания MLA и оптимизатором Muon; затем, перенос по числу токенов обучения через закон масштабирования.
Закон масштабирования по токенам построен линейной регрессией по маленьким прокси-моделям и экстраполирует оптимальную скорость обучения на горизонты вплоть до 10 триллионов токенов с точностью R² = 0,95.
Метод проверен на практике: авторы предобучили с нуля собственную MoE-модель на 155 млрд параметров, из которых 17 млрд активны на каждом токене.
Обучение этой модели прошло стабильно, а оценка подтвердила, что оптимальные гиперпараметры для полного масштаба можно предсказать заранее, без дорогого перебора настроек на целевом масштабе.
Вывод авторов: обучения маленьких прокси-моделей достаточно, чтобы заранее определить оптимальную скорость обучения для полномасштабного обучения крупных MoE-моделей.

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

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

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

В первую очередь, командам и лабораториям, которые сами предобучают крупные MoE-модели с нуля: ML-инженерам и исследователям, отвечающим за подбор гиперпараметров при масштабировании, а также тем, кто планирует вычислительный бюджет обучения (GPU-часы) и заинтересован не тратить его на дорогие пробные прогоны с разными скоростями обучения прямо на целевом масштабе.

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

Метод описан как последовательность из двух шагов. Сначала, адаптировать Maximal Update Parameterization (μP) под архитектуру MoE, добавив механизм внимания Multi-head Latent Attention (MLA) и оптимизатор Muon: это даёт параметризацию, при которой оптимальная скорость обучения остаётся одной и той же при изменении ширины модели, то есть подобранное на узкой модели значение можно перенести на более широкую. Затем, обучить несколько маленьких прокси-моделей на ограниченных бюджетах токенов, зафиксировать для каждой оптимальную скорость обучения и построить по этим точкам закон масштабирования линейной регрессией; это позволяет экстраполировать искомое значение на целевой, гораздо больший бюджет токенов, в работе показано на примере вплоть до 10 триллионов, вместо того чтобы подбирать его прямым перебором на этом бюджете.

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

Это единичная научная работа: в доступном тексте источника не указаны ни имена авторов, ни организация, ни дата публикации, поэтому независимой проверки или воспроизведения результата пока нет. При этом заявление подкреплено не только теорией, авторы довели дело до реального большого прогона: предобучили с нуля собственную MoE-модель на 155 млрд параметров (17 млрд активных) и показали, что обучение с предсказанными гиперпараметрами прошло стабильно, а оценка подтвердила качество итоговой модели. Экстраполяция скорости обучения на большие бюджеты токенов также подкреплена количественно, точность R² = 0,95. Это весомее, чем просто теоретическая гипотеза, но результат стоит воспринимать как заявление одной исследовательской группы до появления независимых повторений на других MoE-архитектурах.

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

В тексте не указано, на скольких разных конфигурациях и вариантах MoE-роутинга проверялся перенос гиперпараметров, устойчивость подхода за пределами протестированных авторами настроек не подтверждена. Точность R² = 0,95 у экстраполяции по токенам получена линейной регрессией по ограниченному набору маленьких прокси-моделей, это не гарантия, что зависимость останется линейной для любых бюджетов токенов и любых MoE-архитектур за пределами протестированных. В источнике нет прямого численного сравнения с альтернативными методами переноса гиперпараметров или поиска по сетке, поэтому по одной этой работе нельзя оценить, насколько именно предложенный способ дешевле конкурирующих. Не указано и то, публикуют ли авторы код, веса модели или датасет, независимо воспроизвести результат по одной статье может быть сложно.