雑記 - クッキーの日記

最終更新日： 2022-01-30　お気付きの点がありましたらご指摘いただけますと幸いです。

Stable, Fast and Accurate: Kernelized Attention with Relative Positional Encoding

💜

f:id:cookie-box:20211229152010p:plain:w70

💛

f:id:cookie-box:20211229162343p:plain:w70

［1］のアブストラクトは、「セルフアテンションの計算量を削減する既存研究の多くは『内積をとってからソフトマックスする』方式にしか対応できない」と主張しているようにみえます。そして、それだと「相対位置エンコーディング（RPE）に対応できない」と……これまでに提案されているセルフアテンション計算量削減ってそんなに制約があったんですか？？
まず RPE の手続きをちゃんと確認しよう。2節「事前準備」の 2.2 が RPE だからここを確認してみようか。この論文の表記ではセルフアテンション行列は ${\rm softmax} \bigl( Q W^Q (K W^K)^\top / \sqrt{d} \bigr)$ みたいだね。つまり、 $Q \in {\mathbb R}^{n \times d_{\rm in}}, \;$ $W^Q \in {\mathbb R}^{d_{\rm in} \times d}$ みたいな定義だね（ソフトマックス行列のソフトマックス内の分母にくるのは $Q W^Q$ の列数のルートのはずだから）。まあ今回はエンコーダ層への入力 $x=(x_1, x_2, \cdots, x_n)$ が $x_i \in {\mathbb R}^d$ とされているから $d_{\rm in}$ も $d$ みたいだけど、マルチヘッドだったらこれはイコールじゃなくて $d_{\rm in} = n_{\rm head} \times d$ になると思うんだけどね。ともかく、セルフアテンション後を $z=(z_1, z_2, \cdots, z_n)$ とすると $z_i$ を求めるのに $\mathcal{O}(n)$ かかるから、 $z$ を求めるのは $\mathcal{O}(n^2)$ であるといっている。
ん？を求めるのがだからを求めるのはといえるんですか？ちゃんと考えてみます。……そうですね、セルフアテンションを1層経るとは単語列内の各単語が異世界転生すると捉えることができるのではないでしょうか。つまり、
- $Q W^Q$ の $i$ 行目＝ $i$ 単語目が「異世界ではこんな風に転生したい」というプロフィール
- $K W^K$ の $i$ 行目＝ $i$ 単語目が異世界転生した場合のプロフィール
- $V W^V$ の $i$ 行目＝ $i$ 単語目の異世界での遺伝子
- $Q W^Q (K W^K)^\top$ の $i$ 行目＝ $i$ 単語目からみて各単語がどれくらい理想の転生姿か
- $z_i$ ＝ ${\rm softmax} \bigl( Q W^Q (K W^K)^\top / \sqrt{d} \bigr) V W^V$ の $i$ 行目＝ $i$ 単語目が異世界転生した遺伝子
ということです。このときにかかわってくるのは下図のピンク色のセルですが、これらは全て以下ですね。しかし、全ての単語の異世界転生した遺伝子を求めるときは、の全行が必要になります。これがになりますね。確認できました。
そ、そっか……まあそれはともかく、この論文以前の研究でアテンションにカーネルとしての見方が導入されたみたいだね。式 (2) のように。「内積をとって $\sqrt{d}$ で割ってエクスポネンシャルをとる」の操作をカーネル関数と考えて、この操作でなくても任意の正定値カーネルでいいのでは？ということなのかな。
そうですね、「希望プロフィール」と「プロフィール」の類似度を出すのに「内積をとって $\sqrt{d}$ で割ってエクスポネンシャルをとる」に拘る必要はないと思います。
う、うん。それで、式 (3) は式 (2) をカーネル関数 $\kappa$ ではなく対応する特徴写像 $\phi$ でかいた形式だね。カーネル法では $\phi$ は陽に手に入らない（手に入れる必要もない）と思うけどね。それで、セルフアテンションではむしろこちらの形式にすれば計算量が $\mathcal{O}(n)$ になるといっている。