紹介する論文

タイトル：SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models
学会：ICML2023

画像は論文から引用しています。

量子化

ディープラーニングにおける量子化とは、浮動小数点数をよりビット数の少ない整数に変換して表現する技術です。
量子化によって、ディープラーニングの重みデータのサイズを小さくでき、メモリ量削減や高速化を行うことができます。
スマートフォン等の機器でディープラーニングを動かすのに使われるほか、近年ではLLMを一般的なパソコンで動かす際によく使われます。

例えば、[0, 256, 510]というベクトルを8ビット量子化してみましょう。8ビット量子化では、値を0～255の整数で表します。

[0, 256, 510]は[0, 128, 255]に量子化されました。量子化前の値を$X$ 、量子化後の値を$X_q$と書くと、$X_q=X/2$です。

より一般的には、ビット数を$b$、ステップ幅を$\Delta$、ゼロ点を$z$と書くと、

X_q = \mathrm{clamp}\left(\left\lfloor\frac{X}{\Delta}\right\rceil+z, 0, 2^b-1\right),

\Delta = \frac{\max(X)-\min(X)}{2^b-1}, 
z = -\left\lfloor\frac{\min(X)}{\Delta}\right\rceil

です。¹ $\lfloor\cdot\rceil$は最近接への丸めを表します。

では、[0, 256, 510]というベクトルと[0, 256, 255000]というベクトルを8ビット量子化することを考えます。どちらの方が量子化しやすいでしょうか？

後者のベクトルは$X_q=X/1000$ですので、2つ目の要素がつぶれて1つ目の要素と同じ値になってしまいました。
このように、値の分布の中に外れ値があると、量子化で情報が落ちてしまいます。
LLMでは、レイヤの出力(アクティベーション)の中に外れ値が発生しやすいことが知られており、量子化による性能悪化が課題となっています。
ここで紹介する論文は、量子化によるLLMの性能悪化を防止する方法を紹介しています。

SmoothQuant

量子化では、重み$\mathbf{W}$とアクティベーション$\mathbf{X}$の両方を量子化します。
先程も述べたように、LLMではアクティベーション$\mathbf{X}$に外れ値が発生することがあると知られています(Figure 2)。これを抑えることを考えましょう。

SmoothQuantでは、SelfAttentionの入口のLayerNorm～Q, K, VのProjectionの線形層と、FFNの入口のLayerNorm～最初の線形層の間に以下の変換を適用します(Figure 6)。

いま、重み$\mathbf{W}$の線形層への入力が$\mathbf{X}$、出力が$\mathbf{Y}$であるとします。このとき

\mathbf{Y}=\mathbf{X}\mathbf{W}

が成り立ちます。ここで、トークン数を$L$、入力特徴次元を$D_i$、出力特徴次元を$D_o$と書くと、
$\mathbf{X}\in\mathbf{R}^{L\times D_i}, \mathbf{W}\in\mathbf{R}^{D_i\times D_o}, \mathbf{Y}\in\mathbf{R}^{L\times D_o}$です。

Figure 2(a)のように、$\mathbf{W}$の方はあまり外れ値が発生していないとしましょう。このとき。$\mathbf{X}$は外れ値があるため量子化が難しいですが、$\mathbf{W}$は外れ値が無いため量子化は非常に容易です。

ここで、$\mathbf{X}$の第$i$列を$1/s_i$倍して$\mathbf{W}$の第$i$行を$s_i$倍したとし、それぞれ$\hat{\mathbf{X}}, \hat{\mathbf{W}}$と書きます。すると$\hat{\mathbf{X}}$の外れ値は(b)のように$\mathbf{X}$に比べてマシになります。$\hat{\mathbf{W}}$は少し外れ値ができてしまいますが、$\hat{\mathbf{X}}, \hat{\mathbf{W}}$とも外れ値の大きさはマシなため、量子化が容易です。

具体例で見てみましょう。

Figure 5のOriginalでは、アクティベーション$\mathbf{X}$の最大値は16となっており、外れ値になっています。そこで$\mathbf{X}$の2列目を1/4し、$\mathbf{W}$の2行目を4倍します。このように変換しても$\mathbf{X}\mathbf{W}$の掛け算の結果は変わりませんが$\mathbf{X}, \mathbf{W}$それぞれの外れ値は抑制され、量子化しやすくなっています。

この変換を式で書きます。

\mathbf{Y}=\mathbf{X}\mathbf{W}=\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}, \hat{\mathbf{W}}=\mathrm{diag}(\mathbf{s})\mathbf{W}$と置きました。また、$\mathbf{s}\in\mathbf{R}^{D_i}$は予め決めるスケーリングパラメータです。

$\hat{\mathbf{W}}$は推論の度に計算する必要はありません。$\hat{\mathbf{W}}=\mathrm{diag}(\mathbf{s})\mathbf{W}$を予め計算しておき、線形層の重み$\mathbf{W}$を$\hat{\mathbf{W}}$に置きかえておけば良いです。

また、$\hat{\mathbf{X}}=\mathbf{X}\mathrm{diag}(\mathbf{s})^{-1}$も推論の度に計算する必要はありません。
$\mathbf{X}\mathrm{diag}(\mathbf{s})^{-1}$は$\mathbf{X}$に対してチャネル毎($D_i$毎)に$s_i$倍する変換ですが、線形層の直前のLayerNormでも、チャネル毎($D_i$毎)の重み倍する変換があります。
従って、LayerNormの$i$番目の重み$\gamma_i, \beta_i$に予め$1/s_i$倍を掛け算しておけば、LayerNormの処理で$\mathrm{diag}(\mathbf{s})^{-1}$倍の変換ができます。

論文では、スケーリングパラメータ$\mathbf{s}$は、学習データからキャリブレーションデータを取ってきてLLMに流したときの$s_i = \max(|\mathbf{X}_i|)^\alpha / \max(|\mathbf{W}_i|)^{1-\alpha}, 1\le i\le D_i$で決めることを提案しています。

図で表すとこのようになります。

以上の議論は、LlamaのようにLayerNormの代わりにRMSNormが採用されているLLMでも同様です。

公式コードを確認しましょう。(図は公式コードからの引用です)

69行目でLayerNormの重みに対してスケーリングパラメータの逆数倍をし、70~71行目で線形層の重みに対してスケーリングパラメータ倍しています。

また下記の部分から、SmoothQuantの適用範囲は、QKVのProjectionの前、及びFFNの最初の線形層の前であることが確認できます。

Table 3 はFloat計算や先行研究との比較です。SmoothQuantはper-tensor量子化方法ですが、一般的にはper-tensor量子化より精度が良くなるFloat計算やper-channel量子化方法と同等の精度となっています。

https://arxiv.org/abs/2406.01721 ↩