Tech カテゴリのイメージ
Tech

Transformerの残差接続とLayerNormとは？Attentionの前後で何を足し、何を整えるのか

Transformerのブロックは、AttentionやFFNの出力をそのまま次へ送るだけではない。原論文では、それぞれのサブレイヤーへ入ったベクトルを出力に足し、LayerNormを適用する。残差接続が元の入力を通す経路を作り、LayerNormが各位置の特徴量を正規化する。

図にある「Add & Norm」を一つの謎の演算として覚えると、モデルによってLayerNormの位置が違う理由を読み解けない。足し算と正規化を分けて見ると、Transformerブロックのデータの流れが明確になる。

残差接続は、入力を変換結果に足す

サブレイヤーの入力を xx、AttentionやFFNによる変換を F(x)F(x) とすると、残差接続の基本形は x+F(x)x+F(x) だ。元の xx を捨てず、変換結果を加える。足し算をするため、両者の次元は合っていなければならない。原論文のTransformerは、この経路の表現を512次元にそろえていた。

たとえば入力が [1, 2, 3]、サブレイヤーの出力が [0.5, -0.5, 1] なら、足し合わせた結果は [1.5, 1.5, 4] だ。残差接続が「元の情報を必ずそのまま最終出力に残す」とまでは言えない。その後に正規化や別の層が続く。しかし、サブレイヤーを迂回する加算の経路があることが重要だ。

LayerNormは、同じ位置の特徴量から平均と分散を求める

LayerNormは、一般的なTransformerのベクトル表現では、各トークン位置にある特徴量を対象に平均と分散を求める。別の文章や別のトークン位置の値をまとめて平均する操作ではない。PyTorchの LayerNorm も、指定した末尾の次元に沿って統計量を計算すると説明している。

単純化した形では、要素 xix_i の正規化は次の通りだ。

x^i=xi−μσ2+ε\hat{x}_i=\frac{x_i-\mu}{\sqrt{\sigma^2+\varepsilon}}

μ\mu は対象となる特徴量の平均、σ2\sigma^2 は分散、ε\varepsilon は分母を安定させる小さな正の値だ。実際のLayerNormには、正規化後に学習可能なスケールとバイアスを適用する設定もある。したがって、最終出力の平均が常に0、分散が常に1になると断定しない。

Addの後にNormをかける小さな例

次の標準Pythonコードは、1位置・3次元の入力へ残差加算を行い、学習可能なスケールとバイアスを省略して正規化する。原論文のPost-LNという順序の一部分を取り出した説明用の値であり、学習済みモデルの出力ではない。

import math


x = [1.0, 2.0, 3.0]
update = [0.5, -0.5, 1.0]
added = [a + b for a, b in zip(x, update)]
mean = sum(added) / len(added)
variance = sum((value - mean) ** 2 for value in added) / len(added)
normalized = [
    (value - mean) / math.sqrt(variance + 1e-5)
    for value in added
]


print(added)
print([round(value, 4) for value in normalized])

実行結果:

[1.5, 1.5, 4.0]
[-0.7071, -0.7071, 1.4142]

足し算の後、同じベクトル内の3成分から平均と分散を求めている。別のトークン位置の値は使っていない。1e-5 はゼロに近い分散でも分母を安定させるための値で、この例では結果の丸め値にほとんど影響しない。

原論文はPost-LNだが、順番は固定ではない

『Attention Is All You Need』の3.1節は、各サブレイヤーの周囲に残差接続を置き、その後にLayerNormを適用する。式で書けば LayerNorm⁡(x+F(x))\operatorname{LayerNorm}(x+F(x)) で、Post-LNと呼ばれる配置だ。

一方、正規化をサブレイヤーの前へ置くPre-LNもある。この場合の基本形は x+F(LayerNorm⁡(x))x+F(\operatorname{LayerNorm}(x)) となり、計算順序は同じではない。PyTorchの TransformerEncoderLayer は norm_first で正規化を先にするかどうかを切り替えられる。どちらの配置かを確かめずに、図中の「Norm」を常に同じ場所だと扱わないほうがよい。

残差接続とLayerNormは、Attentionが何を見るかを決める機構ではない。元の表現にサブレイヤーの変換を加える経路と、ベクトルの数値スケールを整える操作だ。Attention、FFN、Add、Normを順に追うことで、1つのTransformerブロックの計算がつながる。

参照資料

生成AI・機械学習講座の一覧へ

関連記事