[論文メモ] SELF-ATTENTION DOES NOT NEED O(n^2) MEMORY

self-attentionの計算にメモリ $O(n^2)$ は必要ない

self-attentionはクエリ $q \in \mathbb{R}^d$ 、長さ $n$ のキーとバリューをそれぞれ $k_1,...k_n$ 、 $v_1,...v_n \in \mathbb{R}^d$ として次の式で表せる（ただしクエリが1つのとき）。
f:id:Ninhydrin:20211216092016p:plain

普通に実装すると $s_i$ の計算・保存ために $O(n)$ の計算量とメモリが必要。そしてself-attentionは $O(n^2)$ 必要。
これを改善し、

attentionの計算について。
まずsoftmax(式の $s'_i$ の部分)で $\Sigma_{j} e^{s_j}$ をattentionの最後に移動する。
f:id:Ninhydrin:20211216092728p:plain

これは定数メモリで計算出来る。
attentionの除算部分のために $v^* \in \mathbb{R}^d$ と $s^* \in \mathbb{R}$ を用意する。
キーとバリューのペア $k_i, v_i$ を取り出したら、 $s_i$ を計算し $v^*$ と $s^*$ を蓄積していく( $v^* \leftarrow v^* + v_ie^{s_i}$ 、 $s^* \leftarrow s^* + e^{s_i}$ )。
そして最後に $\frac{v^*}{s^*}$ で割ればよい。