Отключение слоёв сокращает затраты на обучение LLM и ускоряет вывод

Глубину LLM можно сделать гибкой: модель учится работать как со всеми блоками, так и с сокращённым набором. Работа Cerebras Systems и MBZUAI показала экономию до 25% вычислительных операций при обучении; поскольку это нерецензированный препринт, числа получили сами авторы. Для разработчиков собственных моделей это связывает две задачи: сократить бюджет предобучения и заранее подготовить несколько режимов вывода с разной задержкой.

Слои нужно отключать целиком и возвращать к концу обучения

Случайное отключение слоёв (layer dropout) временно пропускает часть трансформерных блоков для отдельных последовательностей в пакете. Пропущенный блок не выполняет вычисления, поэтому экономия складывается из структурно удалённых операций, а не из разреженности отдельных весов, которую оборудование не всегда умеет ускорять.

Авторы сравнили два уровня отключения. В первом внимание и полносвязная часть пропускались независимо, во втором — вместе как единый трансформерный блок. Второй вариант давал меньшую ошибку: связанные части слоя лучше оставлять или убирать одновременно.

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

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

Лучший режим сочетает распределение по глубине и расписание по времени. Нижние блоки отключаются реже, верхние — чаще; по мере обучения вероятность пропуска снижается. Модель начинает работу в сокращённом виде, постепенно наращивает эффективную глубину и завершает обучение близко к полной архитектуре.

Одна схема сокращает вычисления и готовит три режима вывода

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

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

В крупном опыте модель на 3,9 млрд параметров ускорила такое декодирование в 1,54 раза, тогда как плотная версия — лишь в 1,02 раза. При пропуске чередующихся слоёв ошибка выросла до 2,129 у обученной с отключением модели и до 6,446 у плотной. Значит, случайный пропуск во время предобучения не просто сокращает работу, а учит сеть сохранять результат при другой глубине.

Рецепт стоит закладывать до предобучения, а не после него

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

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

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

Границы проверки достаточно широки для решения о пилоте, но не для автоматического переноса на любую архитектуру. Проведено более 2400 запусков декодерных трансформеров размером от 271 млн до 8,2 млрд параметров на корпусах естественного языка и кода объёмом до 160 млрд токенов; предобучение выполнялось на Cerebras CS-3. В работу не вошли модели с экспертными слоями и прямое сравнение с обучаемыми механизмами выбора глубины вроде Mixture-of-Depths.

Для нового предобучения разумный следующий шаг — небольшой контрольный запуск с плотной базовой линией и полным рекомендуемым расписанием. Если задача состоит только в ускорении уже выбранной модели через API, результаты не меняют архитектурный план: метод требует доступа к обучению и к выполнению отдельных блоков.

Источники

• Don't Drop Dropout: Optimizing Layer Sparsity for Efficient LLM Training and Inference

Иллюстрация: рисунок из статьи «Don't Drop Dropout: Optimizing Layer Sparsity for Efficient LLM Training and Inference», Mostafa Elhoushi, Alex Pretko, Nolan Dey и др., CC BY 4.0