【DL輪読会】DiJiang: Efficient Large Language Models through Compact Kernelization

DEEP LEARNING JP [DL Papers] DiJiang: Efficient Large Language Models through Compact Kernelization Yuta Oshima, Matsuo Lab http://deeplearning.jp/

• http://deeplearning.jp/

DiJiang: Efficient Large Language Models through Compact Kernelization 書誌情報 著者 • ICML2024 (Oral) • Hanting Chen, Zhicheng Liu, et al. (Huawei) 概要 • Transformerでは系列長の二乗の計算量を要する • 訓練済みのTransformerを少ない計算コストで線形計算量に近似する，新しい周波 数領域カーネル化アテンション (FKA) を提案 • 大規模モデル(ex. LLaMA2-7B)で実験 2

背景 • LLMの発展により，Transformerの有効性がますます明確化 • しかし，TransformerのAttentionは系列長の二乗の計算量を要するため，計算資 源が制限される状況下での使用が困難 • Transformerの軽量化手法として，Attentionの計算量を系列長の一乗オーダーで 抑える研究が盛ん（ex. Linear Attention[Katharopoulos et al., 2020] ）だが，全体を 一から学習することが想定されており，計算コストは相変わらず甚大 • 本研究では，準モンテカルロ法とFFTを用いて，誤差が小さいながら計算コスト の低いAttentionの近似手法を提案し，その有効性をLLMで検証 3

準備：Attentionとカーネル • Self Attentionは，以下のような処理である • 𝑛: 系列長, 𝑑: 𝐴𝑡𝑡𝑒𝑛𝑡𝑖𝑜𝑛の次元数 • これは類似度を計算する正定値カーネル関数𝐾( ∙ , ∙ )を用いて下のように書ける 4

準備：Attentionとカーネル • さらに，正定値カーネル関数𝐾( ∙ , ∙ )をカーネルトリックにより線形分解するこ とで，以下のように書き直せる • ただし，𝜙 ∙ ∶ ℝ𝑑 → ℝ𝑚 • すると，Ο(𝑛2 𝑑)のAttentionの計算量を， Ο(𝑛𝑚𝑑)にできる 5

準備：Attentionとカーネル • LLMの文脈では，一から学習するのではコストが甚大なので，元のAttentionを甚 大な再学習なしに，等価に置き換えるカーネルが望ましい • Positive Random Features (PRF) [Choromanski et al., 2020] が有力な候補の一つ • ただし，𝜔 ∈ ℝ𝑚×𝑑 • 以下の関係式が理論的に示されているため，射映後の次元数𝒎が大きいとき， PRFによる計算が元のAttentionと同等になる(モンテカルロ近似) 6

準備：Attentionとカーネル • しかし，裏を返せば，このままでは𝒎がかなり大きくなければ，Attentionと同 等の性能を発揮できない • それでは無視できない計算量が必要 • この問題に対処すべく，Attentionを近似するためのカーネル法に対して理論分析 を行う 7

Bochnerの定理 [Feller, 1966]とモンテカルロ近似 • 並行移動不変なカーネル𝐾(𝑥, 𝑧)に対し，ℝ𝑑 上の確率密度関数𝑝(𝜔)が一対一対応 し，以下のように書ける • よってモンテカルロ法により，以下のように近似可能 • ただし， • ここで，𝜔𝑖 はモンテカルロ法における，𝑝(𝜔)からのサンプリング 8

Bochnerの定理 [Feller, 1966]とモンテカルロ近似 • なお， query, keyは正規化されることから，Attentionは並行移動不変なカーネル の一つであるガウシアンカーネルと等価と見做せるため，この定理の応用範囲内 • 裏を返せば，ガウシアンカーネルに対して，モンテカルロ法より効率的な近似手 法を見つけることで，Attentionの近似も効率化できる 9

ガウスカーネルの準モンテカルロ法による近似 • 以下のように定式化されるPositive Fixed Features (PFF)を導入する • ただし，𝑉 = 𝑣1 , 𝑣2 , ⋯ , 𝑣𝑚 ∈ 𝕊𝑑×𝑚 は漸近的に一様に分布し，𝑡𝑖 ~𝑈(0,1) • このとき𝜙𝑃𝐹𝐹 (𝑥)の内積は，ガウシアンカーネルの不偏推定量 • さらに，一様列を用いた準モンテカルロ法により，積分誤差の収束オーダーを 𝑂(1/𝑚)にできる • モンテカルロ法では， 𝑂(1/ 𝑚) 10

ガウスカーネルの準モンテカルロ法による近似 • さらに，学習可能な重み𝐷を導入することで，PFFよりも更に近似性能が高い Weighted Positive Fixed Features (WPFF) が以下のように定式化される • WPFF法による目的関数の積分推定誤差の上限は、PFF法による目的関数の積 分推定誤差の上限を超えないことが示されている • さらに，学習の高速化のために，高速フーリエ変換(FFT)と離散コサイン変換 (DCT)を用いて，カーネルの計算を𝑂(𝑚)から𝑂(log(𝑚))に削減(Weighted Discrete Cosine Features (WDCF)) 11

周波数領域カーネル化アテンション (FKA) • 以下のように，周波数領域カーネル化アテンション (FKA) を定式化 12

手法のまとめ • 改善の流れとしては，以下のようになる 1. Attentionを，カーネルトリックによりに分解 (PRF) 2. モンテカルロ近似を準モンテカルロ近似に置き換えることで，カーネル数𝑚 に対する収束性を向上 (PFF) 3. 重み付きの準モンテカルロ近似を導入し，収束性能をよりタイトにする (WPFF) 4. FFTとDCTにより，カーネルの計算を𝑂(𝑚)から𝑂(log(𝑚))に削減 (WDCF) 13

実験 • 様々なサイズの言語モデルをベースモデルにし，AttentionをFKAに置き換え，fine-tune • 少ない訓練時間で同等の性能を発揮しつつ，推論速度を速めることに成功した 14

実験 • 更に公開のLLMであるLamma2-7Bにおいても，提案手法が有効性を示した 15

実験 • これまでの線形Attentionモデルに対しても有効性を確認 • 既存のモデルを置き換えて，fine-tuneするという設定は変わらず • 提案手法はカーネル法を用いつつ，効率良くAttentionを近似できている • カーネル法を用いないRetNetはかなり性能が低い • カーネル法を用いたモデルの中でも，提案手法が最も高い性能 16

実験 • これまでの線形Attentionモデルに対しても有効性を確認 • 提案手法のカーネルであるWDCFカーネルは，元のAttentionをよく近似できて いる 17

実験 • Token sizeと計算コスト 18

結論 • 周波数領域カーネル化Attention (FKA) を提案 • これによりAttentionを置き換え，低コストなfine-tuneを行うことで，同等の性能 を保ちながら言語モデルを効率化 • スケーラビリティの検証もできている 19

感想 興味深かった点 • 学習済みの言語モデルのAttentionを低コストで置き換える，という発想自体が画 期的に感じた • 色々な軽量化や収束効率化の工夫が詰め込まれており，学ぶところが多かった 考える余地のある点 • モデルサイズが大きくなると，Attentionの次元数𝐷が系列長𝑁より効いて来そう なので，𝑂(𝑁 2 𝐷)の方が𝑂(𝑁𝐷(log(𝐷) + 𝐷))より良くなることはないのか 20