紹介する論文

タイトル: DuQuant: Distributing Outliers via Dual Transformation Makes Stronger Quantized LLMs
学会: NeurIPS 2024 (Oral)

DuQuantはLLMのPTQです。外れ値対策をして4bit量子化でSOTAを達成しています。

LLMではアクティベーションに外れ値が発生し量子化で劣化する課題を、Smooth変換、回転変換、ジグザグ順列による変換でMassiveな外れ値を散らすことによって解決しています。

SmoothQuantによる重み行列の変換

LLMにおけるLinear層 $\mathbf{Y}=\mathbf{X}\cdot\mathbf{W}$を考えます。この入力$\mathbf{X}$の外れ値を抑制したいです。

まず、SmoothQuantと同様の考えにより、入力のアクティベーションと重みをスケーリングします。

\mathbf{Y}=\mathbf{X}\cdot\mathbf{W}=(\mathbf{X}\cdot\mathbf{\Lambda}^{-1})(\mathbf{\Lambda}\cdot\mathbf{W})

ここで、前のLinear層の重みを$\mathbf{W}'\to\mathbf{W}'\mathbf{\Lambda}^{-1}$に置き換えておけば、入力は$\mathbf{X}\to\mathbf{X}\mathbf{\Lambda}^{-1}$になります。さらに今考えているLinear層の重みを$\mathbf{W}\to\mathbf{\Lambda}\mathbf{W}$に置き換えれば、今考えているLinear層のアクティベーションは従来通り$\mathbf{Y}$のままとなり、変換前後で整合性が取れることになります。
また、$\Lambda$はSmoothQuantに記載されている方法で選びます。

これによって外れ値がある程度抑制されますが、低bit量子化では外れ値課題が残ります。DuQuantではさらに外れ値対策を行っています。

回転行列による対策

入力のアクティベーションに回転行列$\mathbf{R}$ ($\mathbf{R}\mathbf{R}^T=\mathbf{I}, |\mathbf{R}|=\pm 1$)を適用することを考えます。うまく回転させれば、アクティベーションの値の分布を小さくすることができます。

例えば、2次元平面で$(0, 1)$のベクトルを45度回転させると$(1/ \sqrt{2}, 1/ \sqrt{2})$になりますが、値の分布は小さくなっています。つまり、うまく座標系を回転させれば、値の分布を小さくして外れ値を他の次元に散らすことができていると見ることもできます。回転行列を使っているのは、このようなアイデアだと解釈しました。

回転行列の構築方法は論文中の(2)式で記載されています。外れ値のチャンネルを1列目に移動させて($\mathbf{E}_{d^{(1)}}$)、回転行列$\tilde{\mathbf{R}}$を右からかけることによって1列目の外れ値を軽減させます。その後1列目を保持しながら2列目以降も同様に回転行列を適用していきます($\mathbf{Q}$)。

実用上には、計算コストを小さくするために、行列全体をブロック分割し、各ブロック内で回転行列を構築します((3)式)。

ジグザグ順列変換による対策

ジグザグ順列変換では、分割されたブロックをまたぐ入れ替えによって外れ値対策します。

論文Figure 2の2の変換のように、各列の要素の最大値をジグザグ順に並び変えます。これによりブロック内の値の分布が均一になっていきます。

論文では、ジグザグ変換の後、再度回転行列を適用することによって、さらに外れ値を抑制しています((5)式における$\hat{\mathbf{R}_{(2)}}$のこと)。

以上の対策をまとめると(5)式になります。SmoothQuantの考え方同様、前のLinear層の重みを$\mathbf{W}'\to\mathbf{W}'\mathbf{G}$に置き換え、今考えているLinear層の重みを$\mathbf{W}\to\mathbf{G}^{-1}\mathbf{W}$に置き換えれば、今考えているLinear層のアクティベーションは従来通り$\mathbf{Y}$のままとなり、変換前後で整合性が取れることになります。

参考にしたURL

【DL輪読会】DuQuant: Distributing Outliers via Dual Transformation Makes Stronger Quantized LLMs