Ошибка первая: обвал softmax
Когда модель очень уверена, оценка правильного класса сильно превосходит остальные. Знаменатель softmax становится единицей плюс крошечная сумма, и в арифметике с плавающей точкой эта сумма исчезает при округлении. Градиент по правильному классу считается как вероятность около единицы минус единица — и превращается в ровный ноль.
Градиенты по остальным классам при этом остаются положительными, и сумма градиентов перестаёт быть нулевой, как должна. Порог, за которым это происходит, зависит от точности вычислений.
Данные авторов. В bf16, популярном при обучении больших моделей, порог самый низкий.
Ошибка вторая: взаимное сокращение
При дистилляции ученик учится повторять распределение учителя. Стандартная реализация отдельно считает обе вероятности и вычитает их. Когда ответы ученика почти совпадают с учителем, разница теряется при округлении — и градиент получается ненулевым даже там, где ученик и учитель полностью совпадают.
Как исправили
- Для обычного обучения. Сумму экспонент по неправильным классам считают отдельно и аккуратно объединяют с правильным. Градиент правильного класса строится как минус сумма остальных, а не вычитанием единицы из округлённой вероятности.
- Для дистилляции. Разницу считают напрямую из разности логитов функциями, которые сохраняют точность около нуля. При совпадении ученика и учителя градиент становится ровно нулевым.
Что это даёт
На задачах модульной арифметики исчезают «рогатки» — повторяющиеся всплески функции потерь на поздних этапах обучения с Adam. Без регуляризации весов эффект резче: со стандартной кросс-энтропией точность на тесте падает с 61 до 23%, с исправленной — продолжает расти до 85%.
Данные авторов, задача — квадрат суммы по модулю 97.
При дистилляции, где ученик стартует с весов учителя, исправленная версия держит расхождение на нуле и точность 100%. Стандартная в bf16 уводит ученика от учителя. Любопытная деталь: ошибку нашла и исправила модель Gemini, а большую часть экспериментов тоже провёл ИИ.
Что это значит на практике
Если вы дистиллируете модели в bf16 или обучаете до очень высокой уверенности, стоит проверить реализацию функции потерь — исправление подставляется вместо стандартной. Авторы честно оговаривают: такие режимы редки в обычном крупном обучении, а проверки шли на небольших задачах; замеров скорости в работе нет.