紹介する論文
タイトル：OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models
学会：ICLR2024
OmniQuantは学習ベースの量子化方法です。Learnable Weight Clipping (LWC)とLearnable Equivalent Transformation (LET)の2つの工夫を持ち、学習でパラメータを決めます。
画像・数式は論文から引用しています。
Learnable Equivalent Transformation (LET)
LETの考え方はSmoothQuantに似ています。つまり、線形層においてチャンネル毎にアクティベーションを$1/s$倍、重みを$s$倍して、アクティベーションの外れ値を重みに移そうというものです。
SmoothQuantとの違いを挙げます。
(1) SmoothQuantでは線形層がバイアス無しの場合で定義されていましたが、LETではバイアス有りの場合でも定義されています。
(2) SmoothQuantは、SelfAttentionの最初のLayerNorm〜線形層とFFNの最初のLayerNorm〜線形層間でのみ変換を適用していましたが、LETではFFNの2個目の線形層を除く全ての線形層の直前に変換を適用します。
この記事ではLayerNorm〜線形層で変換を適用できることを説明しましたが、線形層〜線形層でも$\mathrm{diag}(\mathbf{s})^{-1}$を直前の線形層の重みに適用することによって変換を適用できます。
バイアスが無い場合で確認しましょう。重み$\mathbf{W}$の線形層への入力が$\mathbf{X}$、出力が$\mathbf{Y}$であるとします。この1つ前の線形層の重みが$\mathbf{W}’$、入力が$\mathbf{X}’$とします。この時
\mathbf{Y}=\mathbf{X}\mathbf{W},
\mathbf{X}=\mathbf{X}’\mathbf{W}’
が成り立ちますので、
\mathbf{Y}=\mathbf{X}\mathrm{diag}(\mathbf{s})^{-1}\mathrm{diag}(\mathbf{s})\mathbf{W}=\hat{\mathbf{X}}\hat{\mathbf{W}},
\hat{\mathbf{X}}=\mathbf{X}\mathrm{diag}(\mathbf{s})^{-1}=\mathbf{X}’\mathbf{W}’\mathrm{diag}(\mathbf{s})^{-1}=\mathbf{X}’\hat{\mathbf{W}}’
です。ここで、$\hat{\mathbf{X}}=\mathbf{X}\mathrm{diag}(\mathbf{s})^{-1}, \hat{\mathbf{W}}=\mathrm{diag}(\mathbf{s})\mathbf{W}, \hat{\mathbf{W}}’=\mathbf{W}’\mathrm{diag}(\mathbf{s})^{-1}$と置きました。
従って、線形層の重みに左からスケーリングパラメータ$\mathrm{diag}(\mathbf{s})$を掛けるのと同時に、一つ前の線形層の重みに右から$\mathrm{diag}(\mathbf{s})^{-1}$を掛けておけば良いということです。図で書くと下図のようになります。
加えて、OmniQuantでは$\mathbf{Q}\mathbf{K}^T$の計算でも同様の変換を適用します。
(3) SmoothQuantでは、スケーリングパラメータ$\mathbf{s}$は人が決めていましたが、LETでは学習で決めます。具体的な方法は後で述べます。
Learnable Weight Clipping (LWC)
LWCの考え方はPACTに似ています。非常に大きな外れ値がある場合、大きな外れ値を正確に量子化することを諦めて、それ以外をなるべく正確に量子化しようというものです。
例えば、[0, 128, 255, 255000]を(a)のように全て量子化しようとすると、0, 128, 255がつぶれて0になってしまいますが、(b)のように大きな255000は正しい量子化を諦める(clippingと言います)と、それ以外の値は正確に量子化できます。
LWCではclip幅$\beta, \gamma$を学習で決めます。重み$\mathbf{W}$の最小値・最大値に対する比$\beta, \gamma (0\le\beta, \gamma\le 1)$を学習で決め、$\beta\min(\mathbf{W})$以下と$\gamma\max(\mathbf{W})$以上はclipします。
式で書くと(2)式のようになります。
Block-wise Quantization Error Minimization
OmniQuantでは、以上に述べた学習パラメータを学習で決めるのですが、そもそも学習に大きな計算量がかかるLLMに対して全てを同時に学習するのは大変です。
そこで、全ての学習パラメータを同時に決めるのではなく、Transformerブロック毎に先頭から最適化していくことによって学習パラメータを決めることを提案しています。
ここで、$\Theta_1, \Theta_2$はそれぞれLWCとLETの学習パラメータの集合、$Q_w(\cdot), Q_a(\cdot)$はそれぞれ重みとアクティベーションの量子化です。
実験では、SmoothQuantで得られたパラメータを初期値に、WikiText2からキャリブレーション用データを128個ランダムに選択して学習しています。
実験
重みのみ量子化方法ではGPTQやAWQと、重み・アクティベーション量子化方法ではSmoothQuantなどと比較し、性能悪化が小さいことを示しています。