はじめに

こんにちは！私はGxPの張です。
この記事はグロースエクスパートナーズ Advent Calendar 2023 の22日目の記事です。

最近、会社のプロジェクトでOCRモデルのトレーニングが必要になりました。トレーニングの効率を上げて計算負担を減らすため、計算の複雑さを軽減する方法を探求しました。その過程で、Transformerモデルの計算複雑度を下げる新しい方法について述べた論文に出会いました。特に、Transformer-VQ技術です。この方法をプロジェクトに適用したところ、驚くほどの成果を得ました。この経験を踏まえて、より多くの人にTransformer-VQのメリットを共有したいと思います。

Efficient Transformerのアプローチ

Efficient Transformerの主な目的は、Transformerの計算複雑度¹を削減することです。当初は、attentionメカニズムの改善に焦点を合わせていましたが、その後フーリエ変換や線形RNNなど、より多様な手法を統合するようになりました。これらの理論は魅力的ですが、実用面では、特に大規模な言語モデル（LLM）と比較して、目立った成果が見られないこともあります。

最近、特に注目されているのが「Transformer-VQ」という研究です。この研究では、標準的なattentionメカニズムのKeyベクトルに対してベクトル量子化（VQ）を施し、複雑度を線形時間²にまで減少させつつ、従来のattention機能を維持しています。しかし、Efficient Transformerの多くは効率と性能の間のバランスを取ることが難しく、理論上の複雑度低減が必ずしも性能の向上に直結するわけではありません。また、Flash Attentionのような最新技術は、標準的なTransformerにも改善の余地が残っていることを示唆しています。

Transformer-VQの紹介

Transformer-VQモデルの主要な特徴は、attentionメカニズム内のKeyベクトルへの革新的なアプローチです。ここでは、クラスターの中心点を利用して元のベクトルを置換し、その結果、計算の複雑さを線形レベルまで削減します。このアプローチでは、主にKeyの構造に変更を加えつつ、他のコンポーネントはそのまま保持されます。この手法によって、attentionメカニズムは効率的な線形形式へと変換され、精度の変化も計測できます。

行列$Q, K, V$はそれぞれ異なる次元で設定されます。従来のattentionメカニズムでは、式は$QK^{\top}V$として表現されます。一方、Transformer-VQでは、ベクトル量子化³を施した$\tilde{K}$を使用し、$Q\tilde{K}^{\top}V$という形式になります。ここで$\tilde{K}$は、$K$の量子化された形態を指します。

Transformer-VQの核心構造

Transformer-VQモデルの中心的な構造は、$softmax(Q\tilde{K}^{\top})V$という形式で表され、ここで$\tilde{K} = \mathcal{VQ}(K, C)$となります。もし直接実装する場合、このプロセスの計算複雑度は二乗に達する可能性があります。しかし、$\tilde{K}$がセット$C$の中のベクトルであることを利用することで、まず$exp(QC^{\top})$を計算し、その結果から$exp(Q\tilde{K}^{\top})$を選び出すことが可能です。このアプローチにより、プロセス全体の複雑度は線形に抑えられます。加えて、$C$のサイズは固定されているため、$QC^{\top}$の計算も線形時間内で完了できます。

Bi-Directional Attention

Bi-Directional Attentionを備えたエンコーダーは、次のように表現されます：

\begin{equation}
softmax(QK^{\top})V = \frac{exp(QK^{\top})V}{exp(QK^{\top})1_{n\times 1}}
\end{equation}

ここで、$1_{n\times 1}$は全ての要素が1の$n\times 1$サイズの行列を意味します。分母は分子の特別な形式で、焦点は$exp(Q\tilde{K}^{\top})$に置かれています。$\tilde{K}$が集合$C$の中のベクトルであるため、one hot行列$\Delta \in {0,1}^{n\times c}$を用いて$\tilde{K}=\Delta C$と表現します。

このため、Transformer-VQでは次のように計算されます：

\begin{equation}
exp(Q\tilde{K}^{\top})V = exp(QC^{\top}\Delta^{\top})V = exp(QC^{\top})\Delta^{\top}V = exp(QC^{\top})(\Delta^{\top}V)
\end{equation}

この式の中で特に重要なのは二番目の等式です。one hot行列$\Delta$を使うことで、$exp$部分を分離し、行列の結合法則を利用して最初に$\Delta^{\top}$と$V$を乗算します。これにより、$c\times d_v$の行列が得られます。$\exp(QC^{\top})$は$n\times c$の行列で、これを$\Delta^{\top}V$と乗算することで最終的な$n\times d_v$の行列が得られます。このプロセス全体の理論上の複雑度は$\mathscr{O}(ncd_k + ncd_v + ncd_v) = \mathscr{O}(n)$となります。

最終的に、$softmax(QK^{\top})V = \frac{exp(QK^{\top})V}{exp(QK^{\top})1_{n\times 1}}$の形式に従って、$exp(Q\tilde{K}^{\top})V$の結果を代入することで、最終的なattentionの結果を得ることができます。このプロセスには、オーバーフローを避けるためのいくつかの詳細な調整が必要な場合がありますが、全体のステップは線形の複雑度内で完了することが可能です。

Decoder Attention

デコーダーは、大規模な言語モデル（LLM）の中核部分を構成し、attentionメカニズムを採用しています。エンコーダーの基本原則に基づいているため、比較的に理解しやすい構造を持っています。

簡単に説明すると、デコーダーのattentionは特定のベクトル$Q, \tilde{K}, V$を使用して計算されます。ここで、$Q_i, \tilde{K}_j$はそれぞれキーと値のベクトル、$V_j$は出力ベクトルを意味します。これらのベクトルを使用して、デコーダーは次のように情報を処理します：

\begin{equation}\begin{aligned} 
O_{[i]} =&\, \sum_{j\leq i}\exp\left(Q_i\tilde{K}_j^{\top}\right)V_j
\end{aligned}\end{equation}

この式では、$i$までの各$j$に対して、キーとクエリの積を計算し、対応する値ベクトルを加算しています。

さらに、特定の状況下では、新しいベクトル$U_i$を導入して計算を簡素化することが可能です。これにより、データは次のように処理されます：

\begin{equation}O_{[i]} = \exp\left(Q_i C^{\top}\right)U_i
\end{equation}

推論時にはこの逐次計算が効果的ですが、トレーニング時には計算速度を高めるためにブロック単位の処理を採用します。これにより、データは小さなブロックに分割され、それぞれのブロックが個別に計算されます。このアプローチは計算効率を高めるだけでなく、ハードウェアの並列処理能力を最大限に活用することも可能です。

このようなブロック単位の処理を通じて、最終的なattentionの結果を得ることができます。このプロセスは比較的単純で、データをより効率的に扱うことができます。

Transformer-VQは外見上Kernelized Attentionの一種であるPerformerに似ているかもしれませんが、実際には重要な違いがあります。特に、シーケンスの長さ$n$がエンコード表のサイズ$c$に比べて非常に大きい場合、一部のエンコーダーベクトルが繰り返し出現する可能性があります。すべてのエンコーダーベクトルがシーケンス内で均等に分布していると仮定すると、近隣のトークンに対するattentionが遠方のトークンと同じになることがあり、これはKernelized Attentionの一般的な低ランクの問題につながります。

言語モデルにおいては、通常、近隣のトークンが遠方のトークンよりも重要視されます。したがって、効果的な言語モデルのアーキテクチャは、距離に応じて区別を行う能力が必要です。このため、Transformer-VQでは$Q\tilde{K}$の計算後にスライディングウィンドウ形状のAttention Bias$（B）$を加え、近隣のトークンに重みを付けます。

ウィンドウサイズをブロックサイズ$l$に設定し、$i<j$または$i−j≤l$の場合$B_{i,j}=0$とします。これにより、ブロック計算時に行列$B$は最も近い2つのブロックにのみ影響を与え、それ以外の遠方のブロックは線形化の「pick out」⁴によって処理されます。

推論のためには以下の式を考慮します：

\begin{equation}\begin{aligned} 
O_{[i]} =&\, \exp\left(Q_{[i]}\tilde{K}{}_{[i]}^{\top} + B_{[i,i]}\right)V_{[i]} + \exp\left(Q_{[i]}\tilde{K}{}_{[i-1]}^{\top} + B_{[i,i-1]}\right)V_{[i-1]} + \sum_{j\lt i-1}\exp\left(Q_{[i]}\tilde{K}{}_{[j]}^{\top}\right)V_{[j]} 
\end{aligned}\end{equation}

したがって、以下のように表されます$（(V_{[-1]},U_{[-1]},U_{[-2]})$はすべてゼロ行列とします）：

\begin{equation}\begin{aligned} 
O_{[i]} =&\, \exp\left(Q_{[i]}\tilde{K}{}_{[i]}^{\top} + B_{[i,i]}\right)V_{[i]} + \exp\left(Q_{[i]}\tilde{K}{}_{[i-1]}^{\top} + B_{[i,i-1]}\right)V_{[i-1]} + \exp\left(Q_{[i]}C^{\top}\right)U_{i-2}\\[5pt] 
U_i =&\, U_{i-1} + \Delta_{[i]}^{\top}V_{[i]} 
\end{aligned}\end{equation}

$B$の導入はTransformer-VQが他のKernelized Attentionと差をつける要素となります。パラメータ量を減らし、可変長生成をサポートするため、$B$の非ゼロ部分は「トプリッツ行列」として制約されます。これにより、$B_{i,j}$が$i−j$の関数となる場合、$B$は加算的な相対位置エンコーディング⁵に相当することになります。

勾配逆伝播の問題

$\tilde{K}$の各ベクトルがエンコーダーの表$C$のベクトルであることは、順伝播、つまり前向きの計算フェーズでのみ発生します。逆伝播、すなわち勾配の計算の際には元の$K$が使用されます。異なる位置の$\tilde{K}_j$が同じ$C_k$にマッピングされている場合でも、各位置で勾配は異なります。これは、Straight-Through Estimator（STE）と呼ばれる手法です。STEの存在により、「pick out」は理論上、推論フェーズでのみ有効であり、トレーニングフェーズではその利点を完全に活用できないとされています。

では、トレーニングフェーズにおける線形化の効率を向上させるための解決策はあるのでしょうか？厳密な勾配を追求することは、効率的な線形化手法の発見を難しくします。しかし、VQの勾配は元々近似的なものであり、attentionメカニズムが完全に正確な勾配を必要とするわけではありません。そこで次のようなアプローチが提案されています：

\begin{equation}\begin{aligned} 
O_{[i]} =&\, \exp\left(Q_{[i]}\tilde{K}_{[i]}^{\top} + B_{[i,i]}\right)V_{[i]} + \exp\left(Q_{[i]}\tilde{K}_{[i-1]}^{\top} + B_{[i,i-1]}\right)V_{[i-1]} + \exp\left(Q_{[i]}C^{\top}\right)U_{i-2}\\[5pt] 
U_i =&\, U_{i-1} + \Delta_{[i]}^{\top}V_{[i]} 
\end{aligned}\end{equation}

この式では、最初の2項でSTEを使用し、キーシーケンスから勾配を得ます。一方、$U_{i-1}$の勾配は直接止めます（stop_gradient演算子を使用）。これにより、モデルの線形性を保ちながら、隣接する2ブロックの重要な勾配を維持する、合理的な近似を実現します。これはTransformer-XLの手法に似ており、Transformer-XLでは逐次計算を行いつつ、履歴ウィンドウの勾配を止めています。つまり、履歴ウィンドウは計算には参加しますが、勾配は伝播しません。

勾配逆伝播の問題を解決した後、Autoregressive Cross-Entropy Loss⁶に基づく全体的なトレーニング目標を設定できます。ここには、VQがもたらすエンコーディングテーブルのAuxiliary Loss⁷も追加されます。

トレーニングと実験結果

このプロジェクトのソースコードへのリンクは以下の通りです：
Transformer-VQ GitHub

特筆すべき点は、この研究が従来のMulti-Head Attentionではなく、Gated Attention Unit⁸とSoftmaxを基礎として採用していることです。

言語モデリング（例：ENWIK8やPG-19）や画像生成（例：IMAGENET64）に関する実験が行われました。これらの実験において使用されたエンコーディングテーブルのサイズは512であり、モデルの最大パラメータ数は1.3ビリオンに達します。これは、主要な大規模モデルに比べれば小さいかもしれませんが、研究目的には十分適しています。

実験結果は全体的に良好で、以下にその結果のスクリーンショットがあります：

TRANSFORMER-VQ: LINEAR-TIME TRANSFORMERS VIA VECTOR QUANTIZATION Page 8

この結果からは、Transformer-VQが言語モデリングおよび画像生成タスクにおいて高い効率を発揮していることが明らかです。特に、モデルサイズに対するパフォーマンスのバランスが秀逸であることが注目されます。

Transformer-VQの特徴と展望

Transformer-VQの重要な特徴は、VQを使用することで、Transformerの計算複雑性を線形レベルまで抑えることができる点です。これにより、従来のattentionメカニズムから線形なものへの移行が容易になり、さらに、Attention Biasの導入により、一般的なKernelized Attentionメソッドよりも効果的なモデルの構築が可能になります。

VQによる「クラスタリング」手法は、LinformerやNyströmformerなどの従来のアプローチを超えています。これは、未来のデータの漏洩を防ぎ、因果関係のある言語モデルに自然に適合するためです。

VQは基本的にシーケンスを離散的なIDに変換する作業であり、Tokenizerに似た機能を果たします。この観点から、Transformer-VQは、例えばMegaByteモデルのように、モデル内にTokenizerを統合していますが、VQは従来のTokenizerに近い直感的なものです。特に、ENWIK8の実験ではBytesを入力として使用し、MegaByteモデルよりも優れたパフォーマンスを示しています。

RetNetとの比較では、Transformer-VQは明示的な遠距離減衰がなく、長期的なコンテキストを扱う能力に優れている可能性があります。また、VQ化されたKeyにより、未学習のKeyが生じないため、長さの外挿能力にも優れています。Gated Attention Unitが基本構造で、シングルヘッドですが、逐次計算中のモデルの記憶状態は十分に大きく、記憶容量も理論上十分です。

Finite Scalar Quantization⁹の使用案もありますが、VQはFSQよりも広範囲のエンコード数で効果的であるため、FSQへの切り替えは効果が低下するリスクがあります。加えて、各層のattention KeyがVQ化されることで、VQの近似精度が向上し、FSQはエンコーダーとデコーダーが強力なシナリオに適しています。

VQの利用により、学習済みモデルからの微調整が可能になります。VQは明確な幾何学的意味を持ち、K-Meansと多くの共通点を持っています。学習済みモデルからのサンプルを使用してKeyにK-Meansを適用し、エンコーディングテーブルの初期化に使用する中心ベクトルを得ることができます。ただし、Transformer-VQはRoPEには適さないため、RoPEモデルはReRoPEに変更してからVQを実施する方が良いでしょう。

総じて、Transformer-VQは多くの効率的なTransformer研究の中でも独自の特徴を持ち、多様な可能性を秘めています。

最後に

この記事では、ベクトル量子化（VQ）をKeyに適用することで、Transformerの計算複雑性を線形レベルに抑えるという、独自の特性を持つ「Transformer-VQ」モデルに焦点を当てました。このアプローチは従来の方法とは異なり、ユニークな特徴を持っています。

Transformer-VQのテスト結果は、その有効性を示しており、このモデルは洗練された線形AttentionやRNNモデルの機能を持つとも解釈できます。また、トレーニング可能なTokenizerを有するAttentionモデルとしても理解できるため、言語モデリングや画像生成など、幅広い分野での高いパフォーマンスが期待されます。

特に、Transformer-VQは他の大型モデルと比べても劣らない成果を、より少ないパラメータで達成しています。これは、効率とパフォーマンスを同時に追求する研究において顕著な進展を示しています。VQの採用は、モデル内にトークン化プロセスを統合することで、より直観的で柔軟な表現方法を可能にしています。

Transformer-VQの導入は、効率的なTransformerの分野に新たな視点と潜在力をもたらしており、今後の発展と応用に対する期待が高まっています。

Transformer-VQ

出典：Transformer-VQ: Linear-Time Transformers via Vector Quantization
著者：Lucas D. Lingle
発行年：2023
eprint：arXiv:2309.16354 [cs.LG]
DOI：10.48550/arXiv.2309.16354

https://japan.zdnet.com/glossary/exp/%E8%A8%88%E7%AE%97%E8%A4%87%E9%9B%91%E5%BA%A6/?s=4 ↩
https://ejje.weblio.jp/content/linear+complexity ↩
https://jp.morgenrot.cloud/blog/quantization-of-deep-learning/ ↩
特定の要素や操作を選択するプロセスを表します。例えば、複雑な計算の中で重要な部分を「ピックアウト」することで、効率を高めることができます。 ↩
https://www.morisawa.co.jp/culture/dictionary/1902 ↩
https://en.wikipedia.org/wiki/Cross-entropy ↩
https://magattaca.hatenablog.com/entry/2021/09/19/003636 ↩
https://link.springer.com/chapter/10.1007/978-981-99-2789-0_23 ↩
https://arxiv.org/abs/2309.15505 ↩