混合精度学習

Mixed Precision Training

混合精度学習とは

混合精度学習(Mixed Precision Training)は、ニューラルネットワークの学習において、計算の種類に応じて異なる数値精度(FP32とFP16/BF16)を使い分ける手法です。精度を維持しながら学習速度を高速化し、GPUメモリの使用量を削減することが可能です。NVIDIAのTensorコアなど、低精度演算に特化したハードウェアの性能を最大限に活用します。

混合精度のメカニズム

混合精度学習では、順伝播と逆伝播の大部分をFP16(半精度浮動小数点数)で実行し、マスターウェイト(重みの正確なコピー)をFP32で保持します。パラメータの更新はFP32のマスターウェイトに対して行われ、更新後にFP16に変換して次の順伝播に使用します。これにより、FP16の速度メリットとFP32の精度を両立します。

損失スケーリング

FP16は表現できる数値の範囲が狭いため、勾配の値がアンダーフロー(0に丸められる)するリスクがあります。損失スケーリング(Loss Scaling)は、損失関数の値を大きな係数で掛けて勾配の値を増幅し、アンダーフローを防ぐ手法です。動的損失スケーリング(Dynamic Loss Scaling)はスケーリング係数を自動調整し、オーバーフローとアンダーフローの両方を回避します。

BFloat16と最新の動向

BFloat16(BF16)はGoogleのTPUで導入された16ビット浮動小数点フォーマットで、FP32と同じ指数部(8ビット)を持つため、表現できる数値の範囲が広く損失スケーリングが不要な場合が多いです。最新のGPU(A100、H100)ではFP8演算もサポートされ、さらなる高速化が可能になっています。混合精度学習は大規模モデルの学習において標準的な技術となっています。