Google DeepMindの研究者、softmaxとクロスエントロピーの標準実装にある数値誤差2件と修正版を公開

論文「A Fix for Softmax and Cross-Entropy」が指摘する1つ目のバグはSoftmax Collapseで、1つのロジットが他より極端に大きいと、その勾配が丸めで0になる一方、他の勾配は残り、勾配の総和が0にならなくなる。2つ目のSoftmax Cancellationは、蒸留のように生徒と教師のロジットが非常に近い場合に精度が落ち、両者が等しくても0でない勾配が出る問題。修正版はsoftplus、expm1、log1pを使って対数空間で計算し直すもので、内部で高精度型にキャストせずにbf16、fp32、fp64で正確な勾配を得られるとしている。論文にはJAXの実装が掲載されている。これによりgrokking時の「スリングショット」と呼ばれる損失の急上昇がなくなり、蒸留で生徒が教師から離れていく現象も防げたという。著者らは、極端なケースは大規模学習では多くないとしつつ、速度低下はほぼなく置き換えられると述べている。