セルフアテンション

アテンションをConvの代わりに使おうというアイデア。Transformerで提唱された。

このページではセルフアテンション自体と、それを使った1 Transformerブロックの、Convなどの他のブロックとの比較を行う。

演習 3.4 QとかKを書き下してアテンションの式になるのを確認

原論文から解き明かす生成AIの演習3.4。Q, K, Vでの記述が便利なので、この定式化の計算を確かめておく。

Fully-connectedなパーセプトロンとのパラメータ数の比較

セルフアテンションとConvやRNNの比較は論文や原論文から解き明かす生成AIにあるが、 この辺は当時の状況からの比較であって、今から新しくこの辺を学ぶ人にとっては不要に難しい比較に思う。

素人の視点としては、Fully-connectedなパーセプトロンとの違いを見てみるのが教育的だろう。

Fully-connectedなパーセプトロンのパラメータ数

一つのoutputにつき、トークン数、トークンの次元を512とする。

一つのトークンあたり、512次元のアウトプットを出すには512x512となる。 一つのアウトプットにつき512トークンを足し合わせるのだから、それが512個必要になる。

これが512個あるのだから、全体では512の4乗で、だいたい64G個（64Billion）となる。

Transformerブロックのパラメータ数

MultiHeadとLayerNormとFFNで出来ている。LayerNormは大したこと無いので無視しよう。また多くの演算はW+Bの形式になるがWが512x512の時Bは512のオーダーなので、Wだけ見ていけば丼勘定としては十分。

という事で、64Giと2Mi、くらいの差（32*1024倍）がある。

なお、パラメータ数を直感的に感じるには、どの行列は同じものを掛けているかに注目すると良い。

計算量の比較

論文のセクション4や原論文から解き明かす生成AIに話題があるが、少し自分でも計算してみる。

内積とWとの積

まず大前提として、内積と、Wとの積の計算量から見る。

内積はO(d)となる。 Wとの積は出力の次元によるが、入力と同じ次元を出力するのが基本とすると$O(d^2)$

マルチヘッドではd/8を8個出すので、d/8のケースも考えておくと、計算量は1/8だが、オーダーとしては結局d/8の項はdとなる。512程度でビッグO記法はどうなんだ、という話でもある。

Q, K, VにWを掛ける計算は全て同じような演算なので、KWだけ考えておくと上記のようになる。

マルチヘッドアテンションのQ, K, V

マルチヘッドのQ, K, VにWを掛ける計算は先にも述べた通りビッグO記法では定数倍は影響が無いが、厳密に考えても1/8したものが8個あるのでやはり通常のアテンションと同じくらいになる。

K, VとWとの積の計算量は、上の図でも書いたように、KとVは行列との積がn個あると解釈して、$O(nd^2)$となる。

QとWの積は、アテンションの種類によってQがベクトル1つかn個かの違いがある。

• クロスアテンション(Qはベクトル1つ）: $O(d^2)$
• セルフアテンション(Qはベクトルn個）: $O(nd^2)$

これらがアテンションの入力を作るのに必要な計算量で、この後にこれらの結果を使ってアテンションを計算する計算量が掛かる。

アテンションの計算量

FFNの計算量

セルフアテンション全体

以上をまとめると以下のようになる。

Convの計算量

比較のため、1次元Convolutionの計算量も見ておく。カーネルサイズは便宜上5としておく（k=5）。なお、ConvS2Sはk=3なのでもっと小さい。

テンソルは慣れないと分かりにくいので、dxdの行列を5個持っている、として計算すると良い。 一つのdxd行列との積の計算量が$O(d^2)$で、それがk個ある。

なお、convolutionはinputの個数だけこの演算を行うので、最後にn倍する事に注意。

r個にrestrictしたケースの計算量等（演習3.6）

原論文から解き明かす生成AIの演習3.6に、アテンションを近隣のrに絞った場合の計算をせよ、というのがあるのでやってみる。

最大パス長は、nの間の相互作用を得られるまでに何層つなげるか、という事だと思う。 微妙な話ではあるが、一層でr個分の相互作用があると思うと、n/r 個のレイヤーを積み重ねれば一応つながるという話だろう。 softmaxがrの端をちゃんと考慮に入れるかどうかは自明では無いので正直n/r個のレイヤーを重ねてもn離れた相互作用を計算出来るかは微妙だが。