Google DeepMind нашла дефект softmax: градиент обнуляется при отрыве 16,6

22 октября 2026 года Джонатан Баррон, Дэн Гнанапрагасам и Майкл Тшаннен из Google DeepMind выложили препринт A Fix for Softmax and Cross-Entropy. Авторы разобрали два численных дефекта в стандартных реализациях softmax и cross-entropy, которые используют PyTorch и JAX. Первый, Softmax Collapse, обнуляет градиент правильного класса в float32, когда отрыв от конкурентов превышает примерно 16,6. Второй, Softmax Cancellation, добавляет ненулевой градиент в дистилляции, даже если ученик в точности повторил учителя. Оба дефекта лечатся перестановкой членов формулы без изменения математики.

Где float32 теряет маленькое число

Стандартная схема уже защищена от переполнения: из логитов вычитают максимум, и экспонента не взрывается. Но защита от overflow не гарантирует точность малых разностей. Возьмём логиты [20, 0] в float32. Вероятность второго класса получается 2,06·10⁻⁹, и это число прекрасно хранится. Вероятность первого округляется до ровно 1, потому что прибавка в 10⁻⁹ к единице не помещается в семь значащих цифр. После вычитания единицы градиент правильного класса оказывается нулём, хотя должен быть около −2,06·10⁻⁹. Градиенты конкурентов при этом остаются, и сумма производных перестаёт быть нулевой.

Если обозначить отрыв правильного ответа через m, а r = e^(−m), обычная формула считает 1/(1+r) − 1. Исправленная считает −r/(1+r). Это та же функция, но малая величина стоит в числителе и не округляется вместе с единицей. Проверка в float32: при m = 17 обычная схема даёт 0, исправленная даёт −4,14·10⁻⁸. При m = 20 получается −2,06·10⁻⁹, при m = 30 — −9,36·10⁻¹⁴.

Пороги зависят от числа битов дробной части в формате вычисления:

• bfloat16, 7 бит: порог отрыва ≈ 5,54;
• float16, 10 бит: ≈ 7,62;
• float32, 23 бита: ≈ 16,63;
• float64, 52 бита: ≈ 36,74.

Это пороги арифметики конкретной операции. Веса в bfloat16 не означают, что softmax и loss тоже считаются в bfloat16: при смешанной точности чувствительные операции часто выполняются в float32.

Ученик совпал с учителем, а градиент остался

Второй дефект проявляется в дистилляции. Ученик учится повторять вероятности фиксированного учителя, и при полном совпадении градиент по логитам должен быть нулевым. Cross-entropy при этом остаётся равной энтропии учителя, а нулём становятся градиент и KL-дивергенция.

На практике вероятности учителя и ученика часто считаются разными путями: одна через деление экспонент, другая через экспоненту от log_softmax. Для логитов [2, 1, 0] в float32 разница составляет около 1,49·10⁻⁸ по двум компонентам, норма шума 2,1·10⁻⁸. Логиты идентичны, а градиент нет.

Авторы предлагают сначала сравнить логиты, а потом переходить к вероятностям. Считается разность логитов ученика и учителя, из неё вычитается среднее, взвешенное вероятностями учителя, затем малая поправка вычисляется через expm1 и log1p. Эти функции сохраняют точность при малом аргументе. При совпавших логитах получается строгий ноль: expm1(0) возвращает 0.

Почему 10⁻⁸ доходит до обучения

Градиент порядка 10⁻⁹ кажется безобидным, но Adam нормирует шаг на оценку масштаба градиента, поэтому абсолютная малость сигнала не защищает от последствий. При Softmax Collapse градиент целевого логита исчез, а у конкурентов остался. Искажённый сигнал проходит через сеть и попадает в адаптивный оптимизатор. Авторы связывают эту асимметрию с эффектом Slingshot: после периода низкого loss возникают резкие всплески.

Эксперименты в препринте проводились в float32 на пяти запусках с разными seed. В первом сложение по модулю 31 с weight decay λ = 1 и Adam ε = 10⁻⁸. Во втором (a² + b²) mod 97 с λ = 0 и ε = 10⁻¹⁶. В обоих случаях исправление численной схемы убирает всплески и меняет дальнейшую динамику обучения. Это результат на конкретных моделях и задачах. Он не доказывает, что любой скачок loss у LLM вызван softmax.

Для устойчивого вычисления hard-label loss авторы используют logsumexp по конкурентам и softplus вместо прямого сложения с единицей. Учебный NumPy-пример в разборе работает только для случая, когда правильный класс лидирует, и не заменяет autograd.

Что проверить в своём коде

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

• Подайте в ядро loss логиты с отрывом 17–30 в float32 и посмотрите, обнуляется ли градиент правильного класса.
• Убедитесь, что в вашем режиме смешанной точности softmax и loss действительно считаются в float32, а не в bfloat16 или float16.
• В дистилляции подайте ученику логиты учителя. Градиент должен быть ровно нулевым, а не порядка 10⁻⁸.
• Сравните поведение при разных ε в Adam: от этого параметра зависит, насколько шум превращается в шаг.
• Если в обучении после долгого плато видны всплески loss, проверьте численную схему до того, как менять гиперпараметры.

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