ByteDance Seed предложила расписание шага обучения AIR: итоговая ошибка на 0,01 ниже, чем при cosine

Исследователи ByteDance Seed (Qifan Yu, Xinyu Ma, Xingyan Bin, Di He) предложили расписание скорости обучения AIR. На предобучении модели Mixture-of-Experts (MoE) на 128 млрд токенов оно снизило итоговую ошибку на обучающих данных (training loss) примерно на 0,01 по сравнению с косинусным спадом шага (cosine). Сравнение вели при одинаковой сумме шагов за всё обучение. Под скоростью обучения (learning rate, LR) здесь понимается размер шага, которым оптимизатор меняет веса модели. Её обычно уменьшают по ходу тренировки, чаще всего по косинусной кривой.

Почему вообще заговорили про нормировку весов

Для контроля размера весов обычно используют weight decay (WD): на каждом шаге он слегка тянет веса к нулю. Недавно в качестве альтернативы стали рассматривать weight normalization (WN), обучение с нормированными весами. В нескольких работах сообщали о заметном выигрыше по loss на предобучении.

Авторы проверили, откуда он берётся, и пришли к другому объяснению. При WN номинальный LR превращается в другой реальный размер обновления, чем при WD. Получается своя траектория эффективного шага обучения (effective learning rate, ELR). Исследователи подогнали ELR у обоих методов так, чтобы траектории совпали. Кривые loss у WN и WD при этом почти слились. Значит, выигрыш давала форма ELR, а не нормировка сама по себе.

Отсюда практический вывод: расписание надо строить под ELR, а не только под номинальный LR.

Как устроен AIR

Название расшифровывается как Affine Inverse-square-Root decay, то есть аффинный спад по обратному квадратному корню. Расписание состоит из двух частей.

• База. Шаг падает пропорционально 1/√t, где t — номер шага обучения. Эту форму авторы вывели теоретически. • Поправка. Линейное (аффинное) преобразование. Обучение длится конечное время, и чистый спад 1/√t не даёт нужных значений в начале и в конце. Поправка позволяет заранее задать пиковый и итоговый LR.

Пользователь, как и в случае с cosine, указывает две крайние точки: пик и финал.

Есть и наблюдение про позднюю фазу. Если у запусков одинаковый хвост расписания ELR, он ослабляет влияние более ранних выборов LR. Авторы считают позднее планирование ELR одним из главных источников выигрыша.

Что показал эксперимент

Проверку провели на предобучении MoE-модели на 128 млрд токенов. В MoE на каждый токен работает только часть экспертных блоков. AIR сравнивали с cosine и другими расписаниями-кандидатами. Результаты такие:

• Training loss ниже у AIR. Итоговый loss меньше примерно на 0,01 относительно cosine при одинаковой сумме LR за всё обучение. • AIR лучше и по качеству на прикладных задачах (downstream). • Расписание устойчивее к выбору гиперпараметров.

Отдельный опыт: ту же траекторию ELR применили к обычному weight decay. Динамика обучения воспроизвелась. Авторы делают вывод, что польза не ограничена обучением с фиксированной нормой весов, и называют ELR тем объектом, который стоит планировать.

Что из этого следует для практики

Разница в 0,01 по loss небольшая, но получена она сменой одного расписания шага. Работа даёт три практических ориентира:

• проверить, как ваша схема весов и оптимизатор превращают номинальный LR в эффективный; • попробовать расписание AIR, где задаются только пиковый и итоговый LR; • уделить больше внимания концу обучения, потому что общий хвост расписания ослабляет влияние более ранних настроек.

Эксперименты в аннотации описаны на MoE-предобучении с 128 млрд токенов. О других размерах моделей и объёмах данных она не сообщает.

Статья опубликована на arXiv под номером 2610.40362 и попала в подборку Hugging Face Daily Papers.