混合精度と損失スケーリング

モデルが実際にどう学習するか — 普通の勾配降下法からAdamまで

混合精度訓練は、速度とメモリのために、より小さな数値形式を使う。すべての計算を全精度(標準的な32ビット浮動小数点数)で保持する代わりに、多くの演算にfloat16やbfloat16を使う。これらは16ビットの形式で、メモリを半分に抑える代わりに精度が落ち、float16の場合は表現できる大きさの範囲も狭くなる。

リスクとなるのは数値の範囲だ。ごく小さな勾配もある。小さな数がゼロに丸められてしまうと、オプティマイザーは情報を失う。損失スケーリングは、逆伝播の前に損失を掛け算し、そのあとで勾配を割り戻すことで、その小さな勾配を守る。

1グラム単位に丸めるキッチンスケールは、ひとつまみのごく少量のスパイスを見落としてしまう。同じひとつまみを10回分まとめて量れば、スケールは合計を見ることができる。それを10で割れば、ひとつまみ分を取り戻せる。損失スケーリングも同じ手口を使う。小さな値を表現しやすくしてから、あとでスケールを戻すのだ。下の図は、何が懸かっているかを思い出させてくれる。降下がうまくいくのは、各ステップの勾配が演算を生き延びた場合だけだ。精度はループそのものを変えはしないが、最小値の近くにあるごくわずかな傾きが、そのループにまだ見えているかどうかを左右する。

機械学習における位置づけ混合精度は、大規模なニューラルネットワークが現代のハードウェア上で高速に訓練できる理由のひとつだ。オプティマイザーに必要な概念自体は変わらないが、数値のスケールが訓練のレシピの一部になる。
▶ 混合精度と損失スケーリング
← 交互最小二乗法バッチサイズのスケーリング →