Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

読んだときのメモ。間違いがあるかもしれません。図表はすべて論文より引用しています。

論文: https://proceedings.mlr.press/v119/katharopoulos20a.html
参考にした記事: https://qiita.com/Yosemat1/items/802a41588e8bffaca992
参考にした記事2: https://scrapbox.io/yuwd/Linear_Attention:_Transformers_are_RNNs



- **Title**: Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
- **Authors**: Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François Fleuret
- **Venue**: ICML 2020
- 概要: Self-attentionをカーネル特徴関数の線形ドット積として表現することで計算量を$O(N)$に削減
- 問題設定 (Transformer)
	- $x\in\mathbb{R}^{N\times F}$: $F$次元の特徴ベクトル$N$個
	- Transformer: $T:\mathbb{R}^{N\times F}\to\mathbb{R}^{N\times F}$
	- $L$層のTransformerについて、各層$T_1(\cdot),\dots,T_L(\cdot)$は$$T_l=f_l\left(A_l\left(x\right)+x\right)$$
	- $f_l\left(\cdot\right)$: 各特徴を他の特徴から独立に変換 (通常は2層のFNN)
	- $A_l\left(\cdot\right)$: Self-attention (系列を通して作用する) $W_Q\in\mathbb{R}^{F\times D}, W_K\in\mathbb{R}^{F\times D}, W_V\in\mathbb{R}^{F\times M}$$$\begin{align*}
	Q&=xW_Q,\\\\
	K&=xW_K,\\\\
	V&=xW_V,\\\\
	A_l(x)&=V^\prime=\text{softmax}\left(\frac{QK^\top}{\sqrt{D}}\right)V.\\\\
	\end{align*}$$
- カーネルに基づいたSelf-attentionの定式化 (行列の結合法則を使用)$$V^\prime=\frac{\sum_{j=1}^N \phi(Q_i)^\top \phi(K_j) V_j}{\sum_{j=1}^N \phi(Q_i)^\top \phi(K_j)}=\frac{\phi(Q_i)^\top \sum_{j=1}^N \phi(K_j) V_j^\top}{\phi(Q_i)^\top \sum_{j=1}^N \phi(K_j)}$$
	- 元のSoftmax Attentionは$$V^\prime=\frac{\sum_{j=1}^N \text{sim}\left(Q_i, K_j\right) V_j}{\sum_{j=1}^N \text{sim}\left(Q_i, K_j\right)}$$の類似度関数を$$\text{sim}\left(q, k\right)=\exp\left(\frac{q^\top k}{\sqrt{D}}\right)$$とする（ただし非負）
		- 時間・空間ともに$O(N^2)$
	- $\phi(x)\in\mathbb{R}^C$は特徴表現
		- $\phi(K_j)$はクエリごとに再利用できるので計算量は時間・空間ともに$O(N)$
		- ただし、指数カーネルに対応する$\phi$は無限次元なので、厳密なattentionの線形化は実行できない
    - 代わりに多項式カーネルとしたときは指数カーネル・RBFカーネルと同等に機能する (tsai2019transformer) ことから、提案法では$\phi$を$$\phi(x)=\text{elu}(x)+1$$とする
        - $\text{elu}$(指数線形ユニット)は$x$が負でも勾配が0にならない$$\text{elu}(x) =
\begin{cases} 
x, & \text{if } x > 0 \\\\
\alpha \cdot (\exp(x) - 1), & \text{if } x \leq 0
\end{cases}$$
- 因果マスキング
	- $i$番目は$j\leq i$となる$j$番目からしか影響を受けないから$$V^\prime=\frac{\sum_{j=1}^i \text{sim}\left(Q_i, K_j\right) V_j}{\sum_{j=1}^i \text{sim}\left(Q_i, K_j\right)}=\frac{\phi(Q_i)^\top \sum_{j=1}^N \phi(K_j) V_j^\top}{\phi(Q_i)^\top \sum_{j=1}^N \phi(K_j)}=\frac{\phi(Q_i)^\top S_i}{\phi(Q_i)^\top Z_i}$$となり、$S_i, Z_i$はそれぞれ$S_{i-1},Z_{i-1}$から定数時間で計算できるので、これの計算量は系列長に対して線形
		- ただし$$\begin{align*}
	S_i&=\sum_{j=1}^N \phi(K_j) V_j^\top\\\\
	Z_i&=\sum_{j=1}^N \phi(K_j)
\end{align*}$$
	- 勾配計算に必要なメモリは定数オーダー
	- 学習・推論時は、$\phi(K_j)V_j^\top$行列を内部状態として保存してRNNのように時間ステップごとに更新すればよい（つまり$O(N)$）
- TransformerとRNNの関係
	- Transformerの層は、前の式から、Attentionのメモリ$s$と正規化のメモリ$z$を持つRNNとして定式化できる (添え字を再帰のタイムスタンプとする)$$\begin{align*}
	s_0&=0,\\\\
	z_0&=0,\\\\
	s_i&=s_{i-1}+\phi\left(x_iW_K\right)\left(x_iW_V\right)^\top,\\\\
	z_i&=z_{i-1}+\phi\left(x_iW_K\right),\\\\
	y_i&=f_l\left(\frac{\phi\left(x_iW_Q\right)^\top s_i}{\phi\left(x_iW_Q\right)^\top z_i}+x_i\right).
\end{align*}$$
		- 特徴関数$\phi$に制約がない
- 実験
	- 標準のTransformer, Reformerと比較
	- 合成データ: kitaev2019reformerの配列複製タスク（一連の記号をコピーする（最大長128、10種類の記号））
		- Softmax attentionは時間・空間で2乗に比例, Reformerとlinear attentionは線形
![fig1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/227494/3e05242b-d3d6-5dd2-fbd1-5d4bbedae2b7.png)
		- 他と同様にlinear attentionは収束![fig2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/227494/8d1465b0-6c2e-9887-420c-ec3a9d9c0bdc.png)
	- 画像生成: 自己回帰的にピクセル単位で画像を予測する
		- MNIST:![fig3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/227494/6c437bd3-675a-b6b8-0e71-d4306b095f61.png)
			- bits/dimは同じだが、300倍以上のスループット![tab1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/227494/705341ab-8eed-5eb0-a4d5-4e44bed7d76f.png)
		- CIFAR-10:![fig4.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/227494/f3d5ee75-eed3-5413-6423-aed1478f2a28.png)
			- 学習で3倍、生成時に4000倍の速さ、perplexityも良い![tab2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/227494/53bd68a5-7737-e6dc-05ac-b6216defebe7.png)
	- 自動音声認識: 若干精度が下がった？![tab3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/227494/138b458f-e73d-058c-7d9f-21828beca7ff.png)