勾配スケーリングとは、ディープラーニングの混合精度訓練において、数値のアンダーフローによる勾配消失を防ぐ技術です。損失値や勾配を一時的に大きな定数で倍増させて計算し、重み更新前に元のスケールに戻すことで、FP16などの低精度浮動小数点数でも安定したモデル学習を可能にします。
勾配スケーリングとは
一言でいうと、16ビットなどの低精度データ型でAIモデルを学習させる際に、小さな勾配値がゼロになって計算不能になる現象を防ぐための数値補正技術です。
詳しく解説
ディープラーニングのモデル学習では通常32ビット浮動小数点数が使われますが、計算速度の向上とメモリの節約を目的として16ビット浮動小数点数を併用する混合精度訓練が普及しています。しかし、16ビット精度では表現できる数値の範囲が狭く、小さな数値がアンダーフローを起こして0になり、学習が進まなくなる勾配消失に似た問題が生じます。勾配スケーリングはこの問題を解決するために、順伝播で算出された損失値に大きなスケール因子を掛け合わせて値を拡大し、逆伝播で計算される勾配が小さくなりすぎないように保持します。重みを更新する直前にスケール因子で割り戻すことで、正確な勾配でパラメータを学習させます。
具体例・使われ方
具体的な利用例として、PyTorchなどのフレームワークを用いた大規模言語モデルの学習が挙げられます。自動混合精度訓練機能を使用する際、内部で自動的に勾配スケーリングが実行され、グラフィックボードのメモリ消費を抑えながら高速かつ安定した学習を実現します。
似た用語との違い
勾配クリッピングは、勾配の大きさが極端に大きくなって学習が破綻する勾配爆発を防ぐために、勾配の上限値を制限する技術です。一方、勾配スケーリングは、低精度計算におけるアンダーフローを防ぐために数値を一時的に倍増させる技術であり、対処する課題と処理の方向性が異なります。
注意点
スケール因子が大きすぎると、計算結果が数値の表現可能範囲を超えてオーバーフローを引き起こし、計算結果が非数（NaN）になります。そのため、計算過程でオーバーフローを検知した際にスケール因子を自動で縮小させる動的な制御が必要です。また、より広い数値範囲を表現できるBF16形式を使用する場合は、勾配スケーリングが不要となる場合もあります。