Видели, кстати, что творится:
https://arxiv.org/abs/2610.31642
Исследователи (вполне вероятно, вооружённые современными ИИ-инструментами) нашли и починили два бага — в софтмаксе и кросс-энтропии, а именно:
—
Softmax Collapse — случается, когда один логит сильно больше остальных. Из-за округления градиент по этому логиту может незаметно обнулиться (а по остальным — нет).
—
Softmax Cancellation — если при вычислении кросс-энтропии распределения близки (например, при дистилляции, когда студент инициализируется учителем), градиенты могут заметно отличаться от нуля.
Причём был в реализациях Softmax и Cross-Entropy как в Jax, так и PyTorch. Починка помогает стабилизировать обучение моделей.
Интересно, что косяк был, что называется, «у всех навиду».
А вообще, мой всем совет, используйте, где это возможно, вероятностное округление вместо обычного — в мире low precision-обучения это мастхэв