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

論文: 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は線形
  - 他と同様にlinear attentionは収束
- 画像生成: 自己回帰的にピクセル単位で画像を予測する
  - MNIST:
    - bits/dimは同じだが、300倍以上のスループット
  - CIFAR-10:
    - 学習で3倍、生成時に4000倍の速さ、perplexityも良い
- 自動音声認識: 若干精度が下がった？