はじめに

本日は脆弱エンジニアのAdvent Calendar22日目です。

こんにちは、今秋から社会人博士という人生最高級の暇つぶしをしているセキュリティーホールこと、いっせいです。
実はその研究テーマはTransformerがコアになっています。決して、好んでTransformerを選んだ訳ではなく、あまりにもTransformerが深層学習界隈を席巻しているゆえ、使わざるを得ない状況になったのです。

今更Transformerネタの記事なんて散々こすられてはいるが、勉強と備忘を兼ねて特に小生が気になる部分についてまとめていきます。

Transformerの歴史

2017年に登場しました。
「Attention Is All You Need」

ネットワークアーキテクチャ

まずはTransformerとは何ぞや。中身をしっかり理解します。

Vaswani, A. et al.(2017) Attention Is All You Need

左側がエンコーダ、右側がデコーダになります。この機構は紀元前からあるので、飛ばします。
さて、以下からこの図の色の付いている層ひとつひとつを解説します。

Input Embedding

トークン（単語）を、実ベクトルに変換します

埋め込み行列 $\boldsymbol{E}$ というものを使い、入力トークン $\rm token$ に対応するベクトルを引き出します。

\begin{align}
    \boldsymbol{x}_t
    =
    \boldsymbol{E}
    [
        {\rm token}_t
    ]
\end{align}

想像できるように、この行列 $\boldsymbol{E}$ はクソでか行列である。
次元は「語彙数×埋め込み次元」、しかもE2E学習パラメータである。
以下からは、この埋め込み次元を $d_{\rm model}$ と置くことにしよう。

自然言語の文脈で埋め込み、と聞くとword2vecやGloVeのようなものを想像するが、これらは静的Embeddingという別物らしい。
対して、現在の主流は動的Embeddingと言われる、文脈に依存する（更新される）方法らしい。

そして、細かいがこうして埋め込まれたベクトル $\boldsymbol{x}_t$ は、$\sqrt{d_{\rm model}}$ 倍される。

    \boldsymbol{x}_t
    \leftarrow
    \sqrt{
        d_{\rm model}
    }
    \,
    \boldsymbol{x}_t

実装・論文の慣習ではこの埋め込みベクトルは行ベクトルとして扱われる。
つまり、$\boldsymbol{x}_t\in\mathbb{R}^{1\times d_{\rm model}}$。

行ベクトルなんて高校以来扱ってないので、少し違和感がある。

Scaled Dot-Product Attention; SDPA

本題のMulti-Head Attentionについて記述する前に、Multi-Head Attentionに組み込まれている、Scaled Dot-Product Attentionについて説明します。

どのトークンに注目するか、を確率的に推定し、情報を集約します。

説明だけではどういうことかピンとこないが、数式を見れば解像度が格段に上がる。

    \text{Attention}
    \left(
        \boldsymbol{Q},
        \boldsymbol{K},
        \boldsymbol{V}
    \right)
    =
    \text{softmax}
    \left(
        \cfrac{
            \boldsymbol{Q}
            \boldsymbol{K}^\top
        }{
            \sqrt{d_k}
        }
    \right)
    \,
    \boldsymbol{V}

ここで、

$n$：トークン数
$d_k$：$d_{\rm model}$ から線形変換された後の適当な次元数（Multi-Head Attentionで言及）
$d_v$：$d_{\rm model}$ から線形変換された後の適当な次元数（Multi-Head Attentionで言及）
$\boldsymbol{Q}\in\mathbb{R}^{n\times d_k}$：入力クエリ。$n$個のトークンが行方向に束ねられている行列
$\boldsymbol{K}\in\mathbb{R}^{n\times d_k}$：キー。各トークンが持つ、比較用の特徴量。$\boldsymbol{Q}$ と内積（Dot Product）を取ることにより、どのトークンに注目すべきかの情報を得る。学習パラメータ
$\boldsymbol{V}\in\mathbb{R}^{n\times d_v}$：バリュー。各トークンが持つ、出力用の特徴量。内積 $\boldsymbol{QK}^\top$ で得られた重み（softmax）で $\boldsymbol{V}$ が加重和される。学習パラメータ

$\sqrt{d_k}$ で割る（Scaled）ことで、エネルギー正規化を行なっています。

softmax関数については、以下の通り。

\begin{align}
    \text{softmax}
    \left(
        \boldsymbol{x}
    \right)
    _i
    :=
    \cfrac{
        e^{x_i}
    }{
      \sum_j
      e^{x_j}
    }
\end{align}

この関数の出力ベクトルは、各成分が区間 $(0,1)$ に収まり、全ての成分の和が1となるため、「確率」として解釈できる。

さて、本項冒頭では「SPDAとは、どのトークンに注目するか、を確率的に推定し、情報を集約するものです」なんて、シャバい記述をしたが、小生の解釈は以下である。

SDPAとは、sotfmax正規化されたカーネルによる積分作用素である。

まず、以下の様にひとつのトークン $\boldsymbol{q}_t$ に注目して、正規化内積を考える。

\begin{align}
    \boldsymbol{p}_t
    =
    \cfrac{
        \boldsymbol{q}_t
        \boldsymbol{K}^\top
    }{
        \sqrt{d_k}
    }
    \in
    \mathbb{R}^{1\times d_k}
    ,
    \qquad
    t=1,\cdots,n
\end{align}

さらに、ひとつのキー $\boldsymbol{k}_s$ に着目すると、

\begin{align}
    \boldsymbol{p}_{t,s}
    =
    \cfrac{
        \boldsymbol{q}_t
        \boldsymbol{k}^\top_s
    }{
        \sqrt{d_k}
    }
    \in
    \mathbb{R}
    ,
    \qquad
    s=1,\cdots,n
\end{align}

これをスコア関数 $p(t,s)$ とみなすことができる。
そして、SDPAは次のように表現できる。

\begin{align}
    \boldsymbol{y}_t
    =
    \sum_s
    \alpha(t,s)
    \,
    \boldsymbol{v}_s
    ,
    \qquad
    \alpha(t,s)
    :=
    \text{softmax}_t
    \left(
        p(t,s)
    \right)
\end{align}

この形を見ると、連続積分の離散近似が浮かんでくる。
入力が定義される領域 $\Omega\in\mathbb{R}^{d_k}$ を与えれば、

\begin{align}
    y(t)
    =
    \int_\Omega
    \alpha(t,s)
    \,
    v_s
    \,
    d\mu(s)
\end{align}

という、SDPAを積分作用素として表すことができる。
連続版の $\alpha_c(t,s)$ はsoftmax関数を単純に連続にすればよく、

\begin{align}
    \alpha_c(t,s)
    =
    \cfrac{
        e^{p(t,s)}
    }{
        \int_\Omega
        e^{p(t,s^\prime)}
        d\mu(s^\prime)
    }
\end{align}

ここで、測度 $\mu$ は例えば非構造メッシュなど、領域の重み $w_s$ を考えたい場合に用いる。通常の空間積分や一様格子であれば、

d\mu(s)=ds

として良い。
逆に、重み $w_s$ を扱いたい場合のSDPAは、積分の離散近似として考えればよかったので、

\begin{align}
    \int_\Omega
    \alpha_c(t,s)
    \,
    v_s
    \,
    d\mu(s)
    \simeq
    \sum_s
    w_s
    \alpha^\prime(t,s)
    \,
    \boldsymbol{v}_s
\end{align}

右辺の $w_s\alpha^\prime(t,s)$ はsoftmax正規化されている必要があるため、$\alpha^\prime(t,s)$ を以下のように再定義する。

\begin{align}
    \alpha^\prime(t,s)
    &=
    \cfrac{
        e^{p(t,s)}
    }{
      \sum_s
      w_s
      \,
      e^{p(t,s)}
    }
    \\
    &=
    \text{softmax}_t
    \left(
        p(t,s)
        +
        \log{w_s}
    \right)
\end{align}

つまり、スコア関数 $+\log{w_s}$ とするのが自然で筋が良さそうだ。

実際に、Transformerを用いた作用素学習ではこのようにAttentionを学習カーネルの積分変換とみなす流れが強い。

Multi-Head Attention

さて、ようやく本稿の本題である。

Multi-Head Attentionは、同じ入力を複数のheadに分割し、どのトークンがどのトークンに注目すべきかを並列に学習し、再結合する。

早速、定義を記述する。

\begin{gather}
    \text{Multi-Head}
    \left(
        \boldsymbol{Q},
        \boldsymbol{K},
        \boldsymbol{V}
    \right)
    =
    \text{concat}
    \left(
        \text{head}_1,
        \cdots
        \text{head}_h
    \right)
    \boldsymbol{W}_O
    \\
    \text{where}
    \quad
    \text{head}_i
    =
    \text{Attention}
    \left(
        \boldsymbol{QW}_i^Q,
        \boldsymbol{KW}_i^K,
        \boldsymbol{VW}_i^V
    \right)
\end{gather}

ここで、

$h$：ヘッド数。元論文では $h=8$
$d_k=d_{\rm model}/h$：$\boldsymbol{Q},\boldsymbol{K}$ の射影次元
$d_v=d_{\rm model}/h$：$\boldsymbol{V}$ の射影次元
$\boldsymbol{W}_i^W\in\mathbb{R}^{d_{\rm model}\times (hd_k)}$：クエリ空間への線形写像。学習パラメータ
$\boldsymbol{W}_i^K\in\mathbb{R}^{d_{\rm model}\times (hd_k)}$：キー空間への線形写像。学習パラメータ
$\boldsymbol{W}_i^V\in\mathbb{R}^{d_{\rm model}\times (hd_v)}$：バリュー空間への線形写像。学習パラメータ
$\boldsymbol{W}^O\in\mathbb{R}^{(hd_v)\times d_{\rm model}}$：モデル空間への線形写像。学習パラメータ

ややこしいのが、この $\boldsymbol{Q},\boldsymbol{K},\boldsymbol{V}$ はそのAttentionブロックに入力される行列であり、ブロック入力を $\boldsymbol{X}$ とすると、

Self-Attentionでは、$\boldsymbol{X}=\boldsymbol{Q}=\boldsymbol{K}=\boldsymbol{V}$
Decorder Masked Self-Attentionでは、デコーダ前層出力を $\boldsymbol{Y}$ とおけば、$\boldsymbol{Y}=\boldsymbol{Q}=\boldsymbol{K}=\boldsymbol{V}$

正解をマスクしないと、ただのカンニングマシンになるため。

Cross Attentionでは、エンコーダ出力を $\boldsymbol{Z}$ とおけば、

    \boldsymbol{Y}=\boldsymbol{Q}
    ,\quad
    \boldsymbol{Z}=\boldsymbol{K}=\boldsymbol{V}

になる。

実は、Multi-Head AttentionはSDPAと比べてパラメータ数は同じなのだが、並列に複数のAttention分布（softmax）を作れるので、表現力を増すことができる。

良く言われている例えとして、head1→色を見る、head2→エッジを見る、head3→位置関係を見る、など人間的に並行して別々な特徴注目を行なっている、などがある。

Masked Multi-Head Attention

Masked Multi-Head Attentionはデコーダ側の最初のサブレイヤに位置するブロックである。ほとんどMulti-Head Attentionと同じです。

Multi-Head Attentionのスコア関数にマスクを追加し、参照してはならない位置（例えば未来、次の正解トークン）への重みをゼロにする。

数式で書くと次のようになる。

\begin{align}
    \boldsymbol{P}^\prime
    =
    \cfrac{
        \boldsymbol{QK}^\top
    }{
        \sqrt{d}
    }
    +
    \boldsymbol{M}
\end{align}

マスク $\boldsymbol{M}$ の決め方は色々あるが、代表的なものとしてCausal MaskやPadding Maskがある。

Causal Mask

\begin{align}
    \boldsymbol{M}_{ij}
    =
    \begin{cases}
        0\quad &j\leq i\\
        -\inf\quad &j>i
    \end{cases}
\end{align}

次トークンを予測する自己回帰生成モデルでは、位置 $t$ の出力が $t+1$ 以降を見るとカンニングになるので、隠すようにします。

t\j  1 2 3 4
 1   o x x x
 2   o o x x
 3   o o o x
 4   o o o o

このような上三角を禁止するようなマスク $\boldsymbol{M}$ になります。

Padding Mask

\begin{align}
    \boldsymbol{M}_{ij}
    =
    \begin{cases}
        0\quad &\text{有効トークン}\\
        -\inf\quad &\text{無効トークン}
    \end{cases}
\end{align}

Causal maskが未来禁止なのに対し、Padding maskはダミー要素禁止のマスキングを行う。
例として、次の2つのセンテンスを考えてみる。

A：[I, like, cats, <pad>, <pad>]
B：[I, like, cats, a, lot]

このとき、<pad>は情報を持たないので、どのクエリからもPAD位置（Key/Valueの列）を見ないようにします。

Add & Norm

Transformerの各サブレイヤ後に入ってきます。学習の安定性を担っています。

AddとNormを分けて見ていきます。

Add（Residual Connection）

数学的に書くと以下の通り。

\boldsymbol{Y}
=
\boldsymbol{X}
+
\mathcal{F}(
    \boldsymbol{X}
)

$\mathcal{F}(\boldsymbol{X})$ はサブレイヤからの出力である。
このAdd層は、ResNetに採用されているような残差学習であり、勾配消失を防ぐ効果がある。

物理のサロゲート的な観点では、「ベース状態+相互作用補正」となる。

Layer Normalization

これは、レイヤー内正規化を行なっている。

\text{LN}(
    \boldsymbol{X}
)
=
\gamma
\,
\cfrac{
    \boldsymbol{X}
    -
    \boldsymbol{\mu}
}{
    \sqrt{
        \boldsymbol{\sigma}
        +
        \boldsymbol{\varepsilon}
    }
}
+
\boldsymbol{\beta}

ここで、

$\boldsymbol{\mu}$：特徴次元方向の平均ベクトル
$\boldsymbol{\sigma}^2$：特徴次元方向の分散ベクトル
$\gamma,\beta$：学習可能パラメータ

Post-LNとPre-LN

元祖Transformerでは、

\boldsymbol{Y}
=
\text{LN}
\left(
    \boldsymbol{X}
    +
    \mathcal{F}(\boldsymbol{X})
\right)

と $\text{Norm}\circ\text{Add}\circ\mathcal{F}$ の合成関数（Post-LN）であったが現在の主流は、

\boldsymbol{Y}
=
\boldsymbol{X}
+
\mathcal{F}
\left(
    \text{LN}
    (
        \boldsymbol{X}
    )
\right)

という、$\text{Add}\circ\mathcal{F}\circ\text{Norm}$ の合成関数（Pre-LN）である。

恒等経路を保った方が、勾配が安定するから。

小生も博士研究ではこちらを用いる。（小並）

Feed Forward

冒頭の図の青色で示されたブロックである。

ただの2層FCです。表現力ドーピング

\text{FNN}
\left(
    \boldsymbol{x}
\right)
=
\text{ReLU}
\left(
    \boldsymbol{x}
    \boldsymbol{W}_2
    +
    \boldsymbol{b}_1
\right)
\boldsymbol{W}_2
\boldsymbol{b}_2

$\boldsymbol{x}$は行ベクトルなので、右から行列が掛かります。

Positional-Encoding

最後に、Positioal-Encodingについてまとめます。
これは、ネットワークアーキテクチャの太極図みたいな記号のやつだが、調べてみるとなかなかややこしかったので、大トリとすることにした。

陰陽魚太極図（Wikipedia）

Positional-Encodingは、入力に位置情報を埋め込みます。

Attentionの式を再喝すると、

    \text{Attention}
    \left(
        \boldsymbol{Q},
        \boldsymbol{K},
        \boldsymbol{V}
    \right)
    =
    \text{softmax}
    \left(
        \cfrac{
            \boldsymbol{Q}
            \boldsymbol{K}^\top
        }{
            \sqrt{d_k}
        }
    \right)
    \,
    \boldsymbol{V}

ここで、射影カーネル $\boldsymbol{QK}^\top$ を取っているが、これは系列の絶対的な位置情報を含まない。そのため、TransformerではPositional-Encodingを用いて、位置情報をAttentionの入力に組み込むことを図る。

位置情報を組み込む際に、以下のような条件を満足する必要がある。

明確であること：位置情報と値が相互に変換できること
決定的であること：ルールベース、つまり学習対象外であること
距離が推定できること：モデルが2トークン間の位置の相対距離を推定できること
長さについての外挿性：学習時よりも長いトークン列を推論できること

元祖Transformerは、以下に示す Sinusoidal Positional Encorder を用いておりこれは上の条件をしっかり満たす。

Sinusoidal Positional Encorder

まず、トークン埋め込み $\boldsymbol{x}_t$ と同じ次元数の位置ベクトル $\boldsymbol{p}_t$ を用意します。（※スミマセン、2回目の $\boldsymbol{p}$です...。
これを加算することで、位置情報を組み込みます。

\boldsymbol{x}_t
\leftarrow
\boldsymbol{x}_t
+
\boldsymbol{p}_t

一見、加算すると位置ベクトルの情報が落ちてしまいそうだが、Googleの天才はものすごい。

各トークンについて、各周波数 $k$ ごとに正弦・余弦を交互に並べたベクトル $\boldsymbol{p}_t$ を作成する。

\boldsymbol{p}_t
^\top
=
\begin{bmatrix}
    \\
    \sin(t/\lambda_0) \\
    \cos(t/\lambda_0) \\
    \sin(t/\lambda_1) \\
    \cos(t/\lambda_1) \\
    \vdots \\
    \sin(
        t
        /
        \lambda_{
            \tilde{d}_{\rm model}
        }
    ) \\
    \cos(
        t
        /
        \lambda_{
            \tilde{d}_{\rm model}
        }
    ) \\\\
\end{bmatrix}

ここで、

$\tilde{d}_{\rm model}=d_{\rm model}/2-1$（偶数なので）
$\lambda_j=10000^{2j/d_{\rm model}},\quad j=0,\cdots,\tilde{d}_{\rm model}$

である。

なぜ sin / cos を交互に使うのか？

周波数 $\lambda_j$ に対応する 2 次元ブロック

\begin{bmatrix}
    \sin(t/\lambda_j)\\
    \cos(t/\lambda_j)
\end{bmatrix}

は、位置 $t$ を角度とみなした回転表現になっている。この構造により、異なる位置 $t$ と $s$ に対して、

\sin(t/\lambda_j)\sin(s/\lambda_j)+\cos(t/\lambda_j)\cos(s/\lambda_j)
= \cos\bigl((t-s)/\lambda_j\bigr)

が成り立ち、内積が相対距離 $(t-s)$ の関数として表現できるようになる。

ちゃんと自作して確認しました。
2つの位置の差が各周波数ごとに現れ、距離を一意に推定することができる。

Attention内部での推定

先述した通り、Attentionは射影カーネル $QK^\top$ が中心的な役割を果す。入力に位置ベクトルを加えると、スコア行列は、

(X+P)W^Q {W^K}^\top (X+P)^\top

の形になり、

内容 × 内容
内容 × 位置
位置 × 位置

といった項が自然に現れる。

さて、上記で説明したPositional Encorderは絶対位置エンコーディングと称されるもので、これに対して、相対位置エンコーディング、Rotaryエンコーディングなどがある。それぞれ、

絶対PE：非一様媒質
相対PE：並進不変物理
RoPE：距離依存相互作用

という解釈ができ、問題設定に合わせたPEを設計するのが良さそうである。

例えば、並進運動群 $T(n)$ に対して不変なPEを設計することなども可能
応用すれば、ユークリッド群 $SE(3)$ に不変なPEを考えることができそうだ。

おわりに

今回はTransfomerについて、大まかなイメージを掴むことができ、研究余地なども見えて来ました。スクラッチでの実装も是非取り組んでいきたいです。

では、最後に。
脆弱、脆弱っ！