【論文要約・解釈】Key-Value Transformer：Attention機構のQueryとKeyは一緒でもいいらしい

ソリングです．今回はAli Borjiさんの論文「Key-Value Transformer」を要約し，自分なりに解釈します．原文を知りたい方は，以下のリンクから読んでみてください．以下，特に断りのない場合，画像や表などは以下の論文を出典とします．

なお，この論文はVision Transformerを理解していることを前提としています．必要に応じてこの記事内でも解説はしたいと思っていますが，もし理解できていない場合は，先にこちらを理解することをお勧めします．

3行要約

・Query=Keyにして，KeyとValueだけを使うTransformerだよ
・KeyとValueを一緒にしても，そこまで精度は落ちなかったよ
・QKVが必要な場面はよくわからないよ

論文のAbstract

まずいきなりですが，アブストラクトを示します．

Transformers have emerged as the prevailing standard solution for various AI tasks, including computer vision and natural language processing. The widely adopted Query, Key, and Value formulation (QKV) has played a significant role in this. Nevertheless, no research has examined the essentiality of these three components for transformer performance. Therefore, we conducted an evaluation of the key-value formulation (KV), which generates symmetric attention maps, along with an asymmetric version that incorporates a 2D positional encoding into the attention matrix. Remarkably, this transformer requires fewer parameters and computation than the original one. Through experiments encompassing three task types -- synthetics (such as reversing or sorting a list), vision (mnist or cifar classification), and NLP (character generation and translation) -- we discovered that the KV transformer occasionally outperforms the QKV transformer. However, it also exhibits instances of underperformance compared to QKV, making it challenging to draw a definitive conclusion. Nonetheless, we consider the reported results to be encouraging and anticipate that they may pave the way for more efficient transformers in the future.
(Adopted from: Key-Value Transformer, https://arxiv.org/abs/2305.19129, Ali Borji, 2023)

和訳は以下の通りです．(DeepLをもとに，必要に応じて加筆修正)

Transforerは、コンピュータビジョンや自然言語処理を含む様々なAIタスクの標準的な解決方法として広く普及している。これには、広く採用されているQuery, Key, Valueの定式化（QKV）が重要な役割を果たしている。しかしながら、Transformerの性能にこれら3つの構成要素が不可欠であることを検証した研究はない。そこで我々は、対称的な注意マップを生成するKey-Value定式化（KV）と、注意行列に2次元位置エンコーディングを組み込んだ非対称バージョンの評価を行った。驚くべきことに、このTransformerは元のTransformerよりも少ないパラメータと計算量で済む。我々は、3つのタスク、すなわち合成（リストの反転や並べ替えなど）、視覚（mnistやcifarの分類）、NLP（文字の生成と翻訳）の実験を通して、KV-TransformerがQKV-Transformerを上回ることがあることを発見した。しかし、QKVに比べて性能が劣る場合もあり、決定的な結論を出すのは困難である。しかしながら、報告された結果は勇気づけられるものであり、将来、より効率的なTransformerへの道が開かれることが期待される．

初めに

既存のTransformerやその派生のようなモデルは，多くの場合Query, Key, Valueを使用しています．しかしながら，Query, Key, ValueがTransformer，特にAttention機構に必要なものかという点について論じている研究は多くありません．

そこで，BorjiさんはQuery, Key, Valueが本当に必要かどうかという観点を考えるために，Key, Valueのみを使用したTransformerを考えました．

Key-Value Transformer

上で述べたように，通常のTransformerのSelf-Attention機構には，Query, Key, Valueが存在しています．入力のトークンの数を$n$個，トークンの次元数を$d$として，入力を$x \in \mathbb{R}^{n\times d}$とすると，通常のSelf-Attentionは以下の式で書けます．

$$ \mathrm{SA}(x)=\mathrm{softmax}\left(\dfrac{xW_q (xW_k)^{\top}}{\sqrt{d}}\right) (xW_v) $$

ここで，$W_q, W_k, W_v \in \mathbb{R}^{d\times d}$は入力からそれぞれQuery, Key, Valueを得るための行列を表します．Borjiさんは，$W_q=W_k$とおいた以下のKey-Value Self-Attentionを考えました．通常のTransformerに比べ，Key-Value Transformerはパラメータが少なくなっています．

$$ \mathrm{KV}(x)=\mathrm{softmax}\left(\dfrac{xW_k (xW_k)^{\top}}{\sqrt{d}}\right) (xW_v) $$

式の形からわかるように，この設定だとAttention Weightが対称行列になります．そこで，softmaxの内部に学習可能な位置エンコーディング$B \in \mathbb{R}^{n\times n \times m}$を置き，非対称性を導入しました．実装上は$\dfrac{xW_k (xW_k)^{\top}}{\sqrt{d}}$をブロードキャストして$B$を足したのち，それを線形変換してsoftmaxの引数とします．$m=1$の場合は以下の式で書けます：

$$ \mathrm{KV}_{pos}(x)=\mathrm{softmax}\left(\dfrac{xW_k (xW_k)^{\top}+B}{\sqrt{d}}\right) (xW_v) $$

なお，タスクによってはCross Attention（QueryのソースとがKey, Valueのソースが異なるもの）が必要になることがあるため，Cross Attentionは通常のAttentionを行いますが，Self AttentionはKVに置き換えます．

実験

Ali Borjiさんは3つの実験「Synthetic Task」「Vision Task」「NLP Task」を行いました．それぞれの実験についての結果および考察を以下にまとめます．

Synthetic Task

0から9までの数字のリストが与えられて，それに所定の操作（何もせずに出力，逆順にする，9からそれぞれの要素を引く，など）を行うタスクです．reverseは逆順，sortは小さい順に並べる，swapは前半と後半を入れ替える，subは9からそれぞれの要素を引く，copyは配列をそのまま複製するタスクです．

数字はone-hotベクトルとしてTransformerに与えられ，それを埋め込む次元は(32, 64, 256)の中から，Layer数は2または4，Head数も2または4，入力の長さは(16, 64, 128)としています．また，$m=10$としています．以下の表は，それぞれの設定について実験を行った平均と標準偏差です．

みてわかる通り，KV+Pos Transformerの性能が一番良くなっています．

Attention Mapを可視化すると以下のようになります．上から順にreverse, sort, swap, sub, copyで，左から順にQKV, KV, KV+posを表しています．論文内では，Key-Value Transformerがy=xあたりに対称性を持つことおよび，QKV Transformerのreverseは対角成分に注意重みをかなり大きめに持つ一方で，対角成分に近い部分にも重みを少し持っているのは，厳密なアテンションではなく大雑把なアテンションが必要だからなのではないかと論じられています．

以下，これに対する僕個人の考察としては，

Reverse：[3, 5, 2, 4]を[4, 2, 5, 3]のように並び替えるタスクなので，対角成分の重みが大きくなるのは自然だと思います．KV TransformerとKV+posでは対角成分およびy=x上にも大きめの重みが乗っている理由はよくわかりません．

Sort：これを解釈するのはかなり難しいですが，みた感じQKVでは序盤に来て欲しい小さい数ほどアテンションが下に来て，大きい数ほど上のほうに行っているように見えます．KVとKV+posについては全く解釈できませんでした．

Swap：これは前半と後半を入れ替えるタスクなので，例えば[1, 4, 2, 3]は[2, 3, 1, 4]になります．QKVとKVを見ると，1に対応するAttentionは2が大きくなっています．これは自然だと思います．

Sub：これに関しては全く解釈できませんでした．

Copy：これも解釈が困難でした．

となっています．

Vision Task

様々なデータセットを用いて実験を行いました．Anomalyとは，CIFAR100のあるクラスから持ってきた9枚の画像と，それとは別のクラスから持ってきた画像1枚の計10枚のうち，別のクラスから持ってきた画像を検出するタスクです．その結果，画像分類タスクではKV, QKV, KV+posにほとんど差のないことがわかります．

AnomalyタスクではKV+posがの度が良くなっていると報告されていますが，この理由については論文内では言及さておらず，謎です．

Language Task

Language Taskとして，「英単語で書かれた数字から次の数字をアルファベットで出力する(character generation task)」「シェイクスピアの文章から次の単語を推測する(number generation task)」「ドイツ語と英語を翻訳する(translation task)」というTaskを設定した結果，以下の結果を得たと報告されています．

みてわかるように，ほぼ全てのタスクでqkvの性能が良い一方，numeber generation taskではkv+posもほぼ同じ精度を出しています．また，他のタスクでも，ほぼqkvに迫る性能を出しているということができます．このことから，kv transformerもqkv transformerに近い性能を出せることがわかります．

結論

TransformerのQueryとKeyを同じにしてもほぼ性能の差が出ませんでした．また，対称性が重要な場面も存在していることがわかりました．その一方で，どのような場合に対称性が重要なのかはわかっていない点も多いのですが，QKVが必要という固定観念を払拭できたのは新規性と言えそうです．

個人的な考え

・Query=Key=Valueでもうまくいきそうな気がします．

・言語のTaskではqkvの精度がよく，画像ではkvの方が精度が良いことも多いのは不思議です．個人的には，音声などでどのような結果が得られるかも気になります．

・（自分は画像の研究をしているため）今回画像に使われたデータセットはかなり小さなデータセットでした．そのため，大きなデータセットになるとどのような傾向の差が出るのか気になります．

・IdentityFormerという，TransformerのAttention Matrixを単位行列に置き換えてもあまり精度が落ちないというモデルがあります．このことから，Transformer内部のFFNが重要であるような気もします．
(参考：https://qiita.com/omiita/items/dfb021be491974a4149a)

最後に

ここまで読んでいただきありがとうございました．Transformerのquery, key, valueが必ずしも全て必要ではないというのは新しい知見だと思います．それではまた次回の記事でお会いしましょう！