損失スケーリングは、混合精度訓練（FP16）において小数の勾配がアンダーフローして消失するのを防ぐ技術です。損失値に一定のスケール因子を掛けて大きな値に変換してから逆伝播を行い、勾配計算後に元のスケールに戻すことで、数値精度を保ちつつメモリ消費を抑え、計算を高速化します。
損失スケーリングとは
詳しく解説
ディープラーニングのモデルが大規模化するにつれて、16ビット浮動小数点数（FP16）を用いた半精度演算（混合精度訓練）が広く採用されるようになりました。しかし、FP16は表現できる最小値が小さいため、逆伝播の過程で勾配値が小さくなりすぎるとゼロとして処理されるアンダーフローが発生し、モデルの学習が正しく進まなくなる問題が生じます。損失スケーリングは、順伝播で算出された損失に大きなスケール因子を掛け算し、全体的に勾配の値を大きくした状態で逆伝播を行います。その後、重みの更新直前にスケール因子で割り戻すことで、本来の勾配の大きさを保ったまま、数値表現の範囲内に勾配を収めます。手動で固定値を設定する静的スケーリングと、訓練中に動的にスケール因子を調整する動的スケーリングが存在します。
具体例・使われ方
例えば、PyTorchなどのディープラーニングフレームワークでは、自動混合精度訓練の機能とともに損失スケーリングが標準で組み込まれています。大型言語モデルや画像生成AIの学習において、GPUのメモリ消費を抑えて計算速度を向上させつつ、モデルの精度劣化を防ぐために適用されます。
似た用語との違い
混合精度訓練そのものはFP16やFP32といった異なる精度表現を組み合わせて計算を効率化する手法全体を指します。これに対し、損失スケーリングは混合精度訓練の中で発生するアンダーフローを防ぐための具体的な要素技術です。また、勾配クリッピングは勾配爆発を防ぐために大きすぎる勾配を制限する手法であり、小さな勾配の消失を防ぐ損失スケーリングとは逆の目的を持っています。
注意点
スケール因子を大きく設定しすぎると、今度は数値が表現可能範囲を超えてオーバーフローを引き起こし、Inf（無限大）やNaN（非数）が発生する可能性があります。動的損失スケーリングではオーバーフローを検知した際にそのステップの重み更新をスキップし、スケール因子を小さく修正するため、学習が一時的に停滞する場合があります。なお、数値範囲の広いBF16形式を使用できる最新環境では、損失スケーリングが不要になることもあります。