Assigning Distinct Roles to Quantized and Low-Rank Matrices Toward Optimal Weight Decomposition

生成日:

Assigning Distinct Roles to Quantized and Low-Rank Matrices Toward Optimal Weight Decomposition

論文の面白いところ

この論文の中心は、量子化と低ランク近似を単に組み合わせるのではなく、それぞれにどの仕事を任せるべきかを明確にする点にある。従来の共同最適化では、量子化を先に行うか、低ランク近似を先に行うかが手順上の違いとして扱われることが多かった。本論文は、その違いを低ランク成分の初期化の違いとして読み替える。すると、初期化の時点で、量子化行列が主表現になるか、低ランク成分が主表現になるかがほぼ決まってしまうことが見えてくる。これは、反復最適化を十分に回せば自然によい分解へ進む、という素朴な見方に対する修正である。著者らは、量子化が苦手とするのは活性化の外れ値に対応する重みであるとみなし、その部分を低ランク成分に先に取らせる。残りを量子化行列に渡すことで、量子化される重みの分布を扱いやすくする。小さな変更に見えるが、圧縮後のモデルでパープレキシティやゼロショット精度が変わるため、実装上の含意はかなり具体的である。

問題設定

大規模言語モデルを軽く動かすには、重みを低いビット幅で表す量子化がよく用いられる。しかし 2 ビット級の強い量子化では、わずかな丸め誤差が層をまたいで増幅し、出力品質を損ないやすい。近年は、重み行列 W を量子化行列 Q と低ランク成分 LR の和、すなわち W ≈ Q + LR として表す方法が使われている。Q は低ビットで大部分を表し、LR は残差や重要な構造を補う。この共同最適化は、Q の更新と LR の更新を交互に行う形で実装されることが多い。既存法には、まず量子化してから低ランクで誤差を直す方式と、先に低ランク近似を置いて残りを量子化する方式がある。論文は、これらの差を反復順序ではなく、LR をゼロで始めるか、重みの低ランク近似で始めるかという初期化の差として整理する。そのうえで、どの初期化が重みを Q と LR に分けるうえで適切かを問う。

提案手法

提案手法は Outlier-Driven Low-Rank Initialization (ODLRI) と呼ばれる。考え方は、活性化の外れ値に強く反応する重みを、量子化ではなく低ランク成分に先に割り当てるというものである。まず、キャリブレーション用の入力から活性化行列 X を取り、その二次統計としてヘッセ行列 H = XX^T を見る。H の対角成分が大きいチャネルは、活性化の大きい、量子化誤差に敏感な方向と解釈される。ODLRI はその上位 k 個のチャネルを選び、外れ値側の活性化だけを表す制限された行列 Ho を作る。次に Ho に基づいて、重み W のうち外れ値に関わる方向を低ランク成分 L0R0 がよく再構成するように初期化する。実装上は、選ばれたヘッセ部分行列にコレスキー分解を施し、変換後の重みに特異値分解を適用する。こうして外れ値に敏感な部分を LR に吸収させたうえで、残りの W - L0R0 を量子化する。その後は CALDERA のような既存の Q + LR 共同最適化に、この初期化を差し込んで反復更新を行う。

結果

実験は Llama2 の 7B、13B、70B、Llama3-8B、Mistral-7B で行われた。量子化行列 Q は主に 2 ビットで表し、低ランク成分 L と R は 4 ビットまたは 16 ビットで評価されている。評価には WikiText-2 と C4 のパープレキシティ、Winogrande、RTE、PiQA、ARC Easy、ARC Challenge のゼロショット精度が用いられた。Llama2 では、4 ビットの LR を用いた場合でも、ODLRI を加えることで多くの設定でパープレキシティが下がった。たとえば Llama2-7B、ランク 256、Q が 2 ビット、LR が 4 ビットの条件では、WikiText-2 のパープレキシティが 6.47 から 6.33 に、C4 が 8.47 から 8.27 に下がっている。Llama2-70B でも、ランク 256 で WikiText-2 が 3.99 から 3.94、C4 が 5.78 から 5.73 へ改善した。Llama3-8B と Mistral-7B でも同じ傾向が見られ、Llama3-8B のランク 256 では WikiText-2 が 8.70 から 8.12 に下がった。著者らはまた、ODLRI が量子化スケールと活性化を考慮した誤差を小さくすることを層ごとの分析で示している。ただし対象は重みのみの量子化であり、活性化や KV キャッシュの量子化は本論文の範囲外である。

具体例

たとえば、Llama2-7B のある注意層にある key projection の重み行列を 2 ビットに圧縮するとする。まず、RedPajama から取った短い文書をいくつかモデルに通し、その層に入る活性化を集める。多くのチャネルは平均的な値を取るが、一部のチャネルだけが非常に大きくなり、対応する重みの丸め誤差を目立たせる。普通に量子化すると、その大きなチャネルに合わせて量子化の幅を広く取る必要があり、残りの普通の重みを粗く表すことになる。ODLRI では、ヘッセ行列の対角成分を見て、そのような外れ値チャネルを先に選ぶ。選ばれたチャネルに関わる重みの構造は、低ランク成分 L0R0 が担うように初期化される。すると量子化行列 Q が扱うのは、外れ値の影響をある程度取り除いた残りの重みになる。期待される出力は、元の W に近い Q + LR であり、推論時にはこれを使って元のモデルに近い振る舞いを保つ。間違えやすい点は、低ランク成分を単なる量子化誤差の後始末と見ることで、この論文ではむしろ量子化が苦手な部分を先に低ランク側へ逃がす役割を与えている。