[論文メモ] Hydra Attention: Efficient Attention with Many Heads

CADL2022

効率的なmulti-head attentionの提案

transformerのattentionはトークンの数の2乗オーダーの計算コストを必要とする。
そのためVision Transformer(ViT)などで高解像度の画像を扱うとトークン数が膨大になり、計算のほとんどをattention matrixの生成と適用に費やすことになる。
これをなんとかしたいというお気持ち。

提案手法

一般的なsoftmax self-attentionは以下の式(1)

トークン数を $T$ 、特徴量の次元数を $D$ とすると計算量は $O(T^2D)$ 。
softmaxを $Q$ と $K$ の類似度を図る関数 $sim$ として一般化したのが以下の式

$A(Q, K, V) = sim(Q, K)V$

非線形関数 $\phi(\cdot)$ で $sim$ を分解する。

$A(Q, K, V; \phi) = (\phi(Q) \phi(K))V = \phi(Q)(\phi(K))V$

$\phi(K)^T V$ を先に計算することで、計算量は $O(TD^2)$ になる。

これが1つのヘッドに相当。計算量はTについて線形になったが一般に $D \geq 768$ なのでまだ高コスト。

基本self-attentionはmulti-headで扱われる(MSA)。
ヘッド数を $H$ は大体6~16で、 $Q, K$ の特徴量を $D/H$ に分割して行う。

$A(Q_h, K_h, V_h) = softmax(\frac{Q_h K^{T}_{h}}{\sqrt{D}})V_h \quad \forall_h \in \{1,...,H\}$

$Q_h,K_h,V_h \in \mathbb{R}^{T \times \frac{D}{H}}$

コレがmulti-head linear attention(MLA)
multi-headはもとのattentionと計算量は変わらないが、先程のように非線形関数で分解することで計算量を $O(HT(D/H)^2) = O(TD^2/H)$ に抑えられる。

$A(Q_h, K_h, V_h; \phi) = \phi(Q_h)(\phi(K_h)^T V_h) \quad \forall_h \in \{1,...,H\}$

$O(TD^2/H)$ なのでヘッド数を増やすと高速化できるが精度とのトレードオフで、実際いくつくらいまで増やしていいのか。

調査のためImageNet-1kをDeiT-Bで学習した結果が以下の図2。横軸が $H$ 。

MSAは $H \gt 96$ で、MLAは $H \lt 3$ でメモリ不足。

MLAは $H=768$ でもある程度精度を保っているが、これは $H = D$ でただのスカラ特徴。

類似度関数としてsoftmaxを使わなければ $H$ をスケールアップできそう(ここではcosine similarityを採用)。

そこで $H = D$ とした hydra trick を導入する。

$A(Q_h, K_h, V_h; \phi) = \phi(Q_h)(\phi(K_h)^T V_h) \quad \forall_h \in \{1,...,D\}$
なお、 $Q_h,K_h,V_h \in \mathbb{R}^{T \times 1}$

$\odot$ をアダマール積として
$Hydra(Q, K, V; \phi) = \phi(Q) \odot \Sigma^{T}_{t=1} \phi(K)^t \odot V^t$

$\phi$ は $Q, K$ 全体に適用することに注意( $Q_h,K_h$ は列ベクトルなので)。

HydraはMSAとは全く異なる動作で、すべてのトークンを集約したグローバルな特徴ベクトル $\Sigma^{T}_{t=1} \phi(K)^t \odot V^t$ に対して $\phi(Q)$ でゲーティングしている。

計算量は $O(TD(D/H)) = O(TD)$ 。

その他の $O(TD)$ の手法でAttention-Free TransformerやPloyNLなどがあるが、Hydra Attentionはこれらの一般化と捉えることができる(論文参照)。

実験・結果

アーキテクチャは基本的にDeiT-B、データセットはImageNet-1k。
$sim(\cdot, \cdot)$ としてcosine similarityを採用( $\phi$ はL2 normになる)。

cosine similarity以外について調査した結果が表1。

cosine similarityが最もよく、MSAのそもそもの性質を変化させてるのが原因と考えられる。
MSAは重みの和が1になるようになっているがそれらがそもそも望ましい性質では無いのかもしれない。

Hydra Attentionでの置き換え位置の調査。
すべてを置き換えるのではなく一部を置き換えた方がいいのではというお気持ち(よくあるグローバルを扱うAttention系は後半の層を置き換えると良い的なのが多い)。
Hydra Attentionはグローバルな情報を扱うためその可能性は高い。
実験結果が図4。