FP4量子化学習のシミュレーション

概要

Microsoft Research Asiaによる、FP4量子化学習に関する提案手法。FP4の量子化は特にアクティベーションの量子化が難しく、ナイーブに適用すると学習が発散するが、提案手法はBF16と遜色ない性能で安定して学習させることに成功した。

背景

FP4は[-6, 6]の範囲の15階層の値を表現する。メモリや計算効率のメリットがある一方で、量子化誤差の影響をいかに回避するかが鍵となる。

ナイーブな量子化関数の適用は、勾配が至る所で0になるため、学習時の勾配推定値の誤差が深刻となる。また、これを回避するために提案された線形の量子化関数も勾配推定の誤差は依然として存在する。これに対し、提案手法は微分可能な量子化関数を提案する。

また、アクティベーションの量子化では外れ値が問題となる。外れ値をcropすることである程度問題を回避できるが、彼らはさらに外れ値とcropした値の誤差をSparseな余剰行列で保持することで外れ値を可能な限り保持するアプローチをとった。

例えば、0.999の閾値を外れ値と定義すると、余剰行列は0.2%を除いて値がゼロになるため、追加で導入されるオーバーヘッドはちいさい。閾値を下げれば表現できる値の精度は上がるが、余剰行列のゼロでない要素の値が増えるため、トレードオフが存在する。

Limitation

なお、この手法はFP4精度に対応したNvidiaのBlackberryシリーズのGPUで使うことを想定して開発されているが、論文では実機が手元になかったため、速度改善の実験は行なっていない。