TransformerのFFNとは？Attentionの後で各トークンを変換する層
Transformerの図ではAttentionが目立つ。しかし、原論文の各エンコーダー層・デコーダー層には、Attentionの後にFeed-Forward Network（FFN）も入っている。Attentionが別の位置から情報を集めるのに対し、FFNは各位置のベクトルを同じ規則で個別に変換する。
「Attentionだけ理解すればTransformerの計算は終わり」と思って図を読むと、この層を見落としやすい。FFNは次トークンを選ぶ最後のsoftmaxでも、過去の単語を再検索する機構でもない。Attentionで更新された各位置の表現を、次の層へ渡す前に変換する部品だ。
別の位置を混ぜず、同じ重みを各位置に使う
原論文『Attention Is All You Need』の3.3節は、FFNを各位置に「別々かつ同じように」適用すると説明している。同じ層にあるトークン位置には同じ重みを使う一方、層が変われば別の重みを使う。
原論文のFFNは、線形変換、ReLU、もう一度線形変換の順だ。行ベクトル表記では次のようになる。
ReLUは負の成分を0にする活性化関数である。2回の線形変換だけなら1回の線形変換へまとめられるが、間にReLUを挟むことで非線形の変換になる。原論文の基本モデルでは、入力と出力が512次元で、中間だけ2048次元へ広げる。PyTorchの TransformerEncoderLayer にも、この中間幅を指定する dim_feedforward がある。
2つの位置を小さな計算で比べる
次の例は、2次元の入力を2次元の中間表現へ変換する説明用のFFNだ。原論文の512次元・2048次元を再現するものではない。2つの位置に同じ計算を適用する。
tokens = [[2, -1], [0, 2]]
for x0, x1 in tokens: hidden = [max(0, x0 + x1), max(0, x0 - x1)] output = [hidden[0] + hidden[1], hidden[0] - hidden[1]] print(hidden, output)実行結果:
[1, 3] [4, -2][2, 0] [2, 2]第1位置 [2, -1] はReLU後に [1, 3]、最終出力は [4, -2] となる。第2位置 [0, 2] も同じ式を通るが、入力が違うため出力は [2, 2] になる。ここには位置間の参照がない。隣のトークンの情報が反映されるとすれば、それはこのFFNの前にあるAttentionなどを経て、各位置の入力ベクトルへすでに入っているからだ。
AttentionとFFNは役割が違う
Multi-Head Attentionは、ある位置が他の位置をどれだけ参照するかを計算し、Valueを重み付きで集める。FFNは、その結果を位置ごとに変換する。どちらか一方をもう一方の別名と考えると、Transformerブロックの流れを見誤る。
原論文では両サブレイヤーの周囲に残差接続とLayerNormもある。つまり、ブロック全体は「Attention → FFN」だけを直結したものではない。また、近年のモデルがすべて原論文と同じ活性化関数や正規化の順番を採るわけではない。PyTorchの参照実装でも、FFNの活性化関数やLayerNormを先に置くかどうかを設定できる。個別モデルの構造を確認する際には、その実装仕様を見る必要がある。
FFNの要点は、位置間の情報を集める層ではなく、各位置の表現に共通の学習済み変換をかける層だということだ。Attentionで文脈を取り込んだベクトルが、FFNを経て次のブロックへ進む。