混合精度演算

Mixed Precision Training

混合精度演算(Mixed Precision Training)とは、ディープラーニングの学習において、FP32(32ビット)とFP16またはBFloat16(16ビット)など、異なる精度の数値フォーマットを組み合わせて使用する手法です。学習速度を大幅に向上させながら、モデルの精度をほぼ維持できるため、現在の大規模モデル学習では事実上の標準となっています。

混合精度演算の仕組み

混合精度学習では、順伝播と逆伝播の行列演算をFP16/BFloat16で実行し、パラメータのマスターコピーはFP32で保持します。FP16の演算はTensor Coreなどの専用ハードウェアで高速に処理され、FP32のマスターコピーにより累積的な丸め誤差を防ぎます。損失スケーリング(Loss Scaling)によりFP16のアンダーフローを回避します。

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

性能上のメリット

混合精度演算により、FP32のみの学習と比較して2〜3倍の高速化が実現できます。メモリ使用量も大幅に削減されるため、より大きなバッチサイズやモデルサイズでの学習が可能になります。NVIDIA GPUのTensor Coreは混合精度演算に最適化されており、最大限の性能を引き出せます。

図:混合精度演算の流れ
1順伝播はFP16高速・省メモリ
▶
2損失を拡大小さな勾配を守る
▶
3勾配を計算16ビットのまま
▶
4重み更新はFP32精度を維持

実装と対応フレームワーク

PyTorchのtorch.cuda.amp、TensorFlowのtf.keras.mixed_precision、NVIDIA Apexなどにより簡単に混合精度学習を導入できます。BFloat16はGPUやTPUでの学習で広く使われており、損失スケーリングなしでも安定した学習が可能なため、実装がさらにシンプルです。