混合精度学習

Mixed Precision Training

混合精度学習とは

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

表:混合精度学習の要点まとめ
ひとことで言うと計算の種類によって16ビットと32ビットを使い分け、学習を速くする手法。
何がうれしいか学習時間とGPUメモリを大きく削減でき、より大きなモデルを扱える。
注意点16ビットでは値が小さすぎると0になる。損失スケーリングで対策する。

混合精度のメカニズム

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

図:混合精度学習のしくみ
1順伝播は16bit高速・省メモリで計算
▶
2損失を拡大小さな勾配が消えないよう
▶
3勾配を計算16bitのまま逆伝播
▶
4更新は32bit重みは精度を保つ

損失スケーリング

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

BFloat16と最新の動向

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