学も衒うし、奇も衒う

自己注意機構はローパスフィルタである．

メモ書き程度の考察．全く別のことを考えていた際に下の考察を思いついたが，[Wang+ ICLR 2022]を見つけてしまったので供養を込めて投稿する．

自己注意機構は本質的にローパスフィルタである．この事実を2通りの数式で眺める．

(1) 確率行列としての固有値構造 [Wang+ ICLR 2022]

(2) 熱核近似によるラプラシアン拡散

なお，先に論文紹介．

[Wang+ ICLR 2022] Theorem 1

openreview.net

[Choi+ NeurIPS 2024]
深層での過平滑化（oversmoothing）現象を指摘している．

[Wang+ ICLR 2022]の結果を引用する．

$A=\text{softmax}(P),\ P\in \mathbb{R}^{n\times n}$

ここで、確率行列であることから，ペロン＝フロベニウスの定理より，最大固有値1の直流モードが存在し，

$\lim_{t \to \infty} \frac{\| \text{HC}(A^t z)\|}{\| \text{DC}(A^t z) \|}=0$

が成り立つことを示した．

ここで，離散フーリエ変換 $F$ で

$\hat{z}=Fz,\ DC(z)=(F^{-1}e_1)\hat{z}_1,\ HC(z)=z-DC(z)$ .

DCは直流成分，HCは高周波成分を表す．

ここで別の見方を示す．

ある層での入力系列を行列 $X\in \mathbb{R}^{n\times d}$ とし，重み行列でクエリ・キー・バリューを作ると，

$Q=XW_Q,\ K=XW_K,\ V=XW_V$

であり，自己注意の出力 $Y\in \mathbb{R}^{n\times d}$ はsoftmax関数を用いて

$Y=\text{softmax} \left( \frac{QK^\top}{\sqrt{d}} \right)V$

と書ける．ここで，簡略のために

$A=\text{softmax} \left(\frac{QK^\top}{\sqrt{d}}\right)$

と置く．

トークン間のグラム行列を，グラフの隣接行列ように捉え，

$S=\frac{QK^\top}{\sqrt{d}}$

とする．ここで， $\tilde{S}=\frac{1}{2}(S+S^{\top})$ のように対称化して，対称行列として扱うことにする．そして，その行列指数関数を取り，固有分解しよう．行列指数関数の性質から次のように書ける．

$\exp(\tilde{S})=U \exp(\Lambda) U^\top$

注意， $\exp(\Lambda)=\text{diag} (e^{\lambda_1}, e^{\lambda_2},\cdots, e^{\lambda_n})$ である．

ここで非正規化ラプラシアンをとる．当然，正規化してもよい．

$L=D-\tilde{S},\ D_{ii}=\sum_j (\tilde{S})_ij$

すると，

$\exp(-\tau L)=U\exp(-\tau \Lambda)U^\top$

となる．高周波成分（大きな $\lambda_i$ ）は $e^{-\tau \lambda_i}$ は急激に小さくなり，高周波成分は抑制される．

注意機構では距離が大きいところほど減衰させたい．熱核による近似として，行ごとの正規化を忘れずに

$A \approx \text{diag}(\exp(-\tau L) \mathbf{1})^{-1} \exp(-\tau L)$

とすれば，高周波ほど小さくなるので結果的に低周波のみを残すローパスフィルタとして振舞う．

もう少し直感的な理解を考える．この行列指数関数をテイラー展開して1次近似すると，

$\exp(-\tau L)= I-\tau L+O(\tau^2)\approx I-\tau L$

であり，この層の出力は

$Y=AV \approx (I-\tau L)V$

となり，入力をそのまま通す項と，少しだけ隣接トークンと混ぜ合わせる項の足し合わせとして解釈できる．

実際のブロック単位では，もとの行列をそのまま足し合わせる残差接続が入る．

$\text{LN}(X+Y)=\text{LN}(X+AV) \approx \text{LN}(X+(I-\tau L )XW_V)$

である．ここで層正規化などを無視して，伝達関数だけを見れば，

$(I+I-\tau L)W_V=(2I-\tau L)W_V$

とみることができるので，バイパス＋ローパスとみることができる．

残差接続と層正規化を，さらに陽なオイラー離散化と見做して，第 $l$ 層を連続時間 $t$ に対応付けると

$\frac{\partial X(t)}{\partial t}=-\tau L(t)X(t)W_V+O(h^2)$

という時間依存な係数が付いた拡散方程式を得ることができる．ただし，層幅 $h$ であり， $h\to 0$ で厳密にPDEとして解釈できる．これもいろいろ考えられそう．