Google DeepMind нашла две ошибки в softmax и перекрёстной энтропии

Исследователи Google DeepMind Jonathan Barron, Dan Gnanapragasam и Michael Tschannen показали, что стандартные реализации softmax и перекрёстной энтропии — в JAX, PyTorch и учебниках — содержат две численные ошибки.

Первая, Softmax Collapse, возникает при очень уверенном предсказании: градиент верного класса округляется до нуля, а у остальных — нет, и сумма градиентов перестаёт быть нулевой. Вторая, Softmax Cancellation, бьёт по дистилляции: когда логиты ученика почти равны логитам учителя, градиент остаётся ненулевым даже при полном совпадении.

Авторы переписали функции в логарифмическом пространстве через softplus, expm1 и log1p без заметного замедления. Это убирает скачки потерь «slingshot» при гроккинге и не даёт ученику уплывать от учителя...

Источники:
https://arxiv.org/abs/2610.31642
https://jonbarron.info/softmaxfix/

#AI #DeepLearning #GoogleDeepMind #PyTorch #JAX #Исследования