Self-Attentionにとって文はただの「単語の集合」だった。位置エンコーディングを実装して、順序という情報がどう埋め込まれるかを確かめてみた
約16分で読めます

Self-Attentionにとって文はただの「単語の集合」だった。位置エンコーディングを実装して、順序という情報がどう埋め込まれるかを確かめてみた


前回、Self-Attentionを自分で実装して、有名な「it」が「animal」に注目するという説明図が、学習済みの射影行列があって初めて成り立つ話だと確認した。あのときcausal maskも実装して、「未来のトークンを見てはいけない」という制約は入れた。

ただ、実装しながらずっと引っかかっていたことがある。causal maskは「どこまで見ていいか」を制限するだけで、「このトークンが何番目にあるか」という情報自体は、QKVの計算のどこにも出てこない。 単語ベクトルを並べて QKQK^\top を計算しているだけなら、その「並べる順番」を変えても計算結果は変わらないのではないか。今回はそこを実際に確かめたうえで、Transformerが実際に使っている「位置エンコーディング(Positional Encoding)」を自分で実装し、なぜあの独特なsin/cosの形をしているのかを手を動かして検証してみた。

実験1: Self-Attentionは本当に語順を見ていないのか

前回と同じ、GloVeの本物の単語ベクトルを使った例文で試す。位置エンコーディングを一切加えず、生の単語ベクトルだけでattentionを計算したあと、単語の並び順をランダムにシャッフルして、もう一度同じ計算をする。

def self_attention(X):
    d_k = X.shape[1]
    scores = X @ X.T / np.sqrt(d_k)
    return softmax(scores), softmax(scores) @ X

attn_orig, out_orig = self_attention(X)          # 元の語順
attn_shuf, out_shuf = self_attention(X[perm])     # シャッフルした語順

# シャッフル後のattentionは、元のattentionを同じ並び替えで並べ直したものと一致するはず
diff = np.abs(attn_shuf - attn_orig[np.ix_(perm, perm)]).max()
shuffled order: ['it', 'tired', 'cross', 'because', 'the', 'the', 'animal', 'was', 'did', 'too', 'street', 'not']
max |attn_shuffled - attn_original[perm,perm]| = 5.55e-17
max |output_shuffled - output_original[perm]|   = 8.88e-16

結果は誤差5.55×10⁻¹⁷、実質ゼロ。つまりシャッフルした語順でattentionを計算した結果は、元の語順の結果を単に並べ替えただけのものと完全に一致する。 数字の羅列だけだと実感しづらいので、元の語順で計算したattention行列と、シャッフルして計算してから元の並びに戻した行列を、そのまま並べて比べてみた。

3つのヒートマップを並べた図。左は元の語順で計算したattention行列、中央はシャッフルして計算してから元の語順に並び戻した行列、右はその差分の絶対値。左と中央はピクセル単位で見分けがつかないほど同一で、右の差分行列は1e-17オーダーのごく薄い色しかついていない

左と中央、目で見て違いが分からない。右の差分ヒートマップも、色のスケールが10⁻¹⁷のオーダーで、事実上真っ白(=浮動小数点の丸め誤差でしかない)だ。Self-Attentionという仕組みは、入力を「順序のある列」ではなく「順序のない集合」として扱っている。数式を読んで「置換不変(permutation invariant)」と説明を受けたことはあったが、実際に自分の目で「シャッフルしても数値が寸分違わず一致する」のを見ると、思っていた以上に徹底していた。「the animal did not cross the street」も「street the not did cross animal the」も、Self-Attention単体にとっては同じ入力ということになる。これは明らかにまずい。言葉は語順が変わればほぼ別の文になるのに、仕組みの根っこにそれを区別する手がかりが何もない。

sin/cosの位置エンコーディングを実装する

そこでTransformerの原論文が採用したのが、位置ごとに固定のベクトルを足し込む方法だ。位置 pospos、次元 ii に対して、

PE(pos,2i)=sin(pos100002i/dmodel),PE(pos,2i+1)=cos(pos100002i/dmodel)PE_{(pos,\,2i)} = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right), \quad PE_{(pos,\,2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right)

という値を単語埋め込みにそのまま足す。実装は数行で済む。

def sinusoidal_pe(max_pos, d_model):
    pos = np.arange(max_pos)[:, None]
    i = np.arange(d_model)[None, :]
    angle_rates = 1.0 / np.power(10000, (2 * (i // 2)) / d_model)
    angles = pos * angle_rates
    pe = np.zeros((max_pos, d_model))
    pe[:, 0::2] = np.sin(angles[:, 0::2])
    pe[:, 1::2] = np.cos(angles[:, 1::2])
    return pe

dmodel=64d_{\text{model}}=64、位置0〜99で実際に計算し、位置×次元のヒートマップにしてみた。

sin/cos位置エンコーディングのヒートマップ。横軸が位置0〜99、縦軸が次元0〜63。次元番号が小さいほど縞模様が細かく高速に変化し、次元番号が大きいほど縞がゆったりと変化している

次元番号が小さいほど細かい縞(高周波)、大きいほど緩やかな縞(低周波)になっている。これは2進数の桁と同じ構造だと気づいた。一番下の桁(1の位)は0,1,0,1と高速に切り替わり、上の桁ほどゆっくり切り替わる。位置エンコーディングも同じで、高周波の次元が「今どのあたりか」を細かく刻み、低周波の次元が「大まかにどのブロックか」を表す。異なる周波数を束ねることで、限られた次元数でも広い範囲の位置をほぼ一意に表現できる。

実験2: 位置が近いほど似ているのか、単調に減っていくのか

位置エンコーディングが「近さ」の情報として機能するなら、コサイン類似度は距離が離れるほど単調に下がっていきそうな気がする。全ての位置ペアの組み合わせを一度に見たかったので、平均を取った折れ線グラフではなく、位置×位置のコサイン類似度をそのまま行列にしてヒートマップにしてみた。

def cos_sim(a, b):
    return (a @ b) / (np.linalg.norm(a) * np.linalg.norm(b))

n = 150
Pn = PE[:n] / np.linalg.norm(PE[:n], axis=1, keepdims=True)
sim_matrix = Pn @ Pn.T   # (n, n) 全ての位置ペアのコサイン類似度

位置エンコーディング同士のコサイン類似度行列のヒートマップ。位置0〜149の全ペアについて計算。対角線上は濃い赤(類似度1.0)で、対角線から離れるほど色が薄くなっていくが、対角線に平行な縞模様が何本も規則的に浮き出ており、単純な同心円状の減衰にはなっていない

対角線(同じ位置同士、類似度1.0)を中心に色が薄くなっていくのは予想通りだが、よく見ると対角線と平行な縞模様が何本も規則的に浮き出ている。単純に「対角線から離れるほど均一に薄くなる」のではなく、一定の間隔で類似度が持ち直す帯が繰り返し現れている。 実際、Δpos\Delta pos ごとの平均類似度を数値で見ても、単調減少にはなっていなかった。

delta=  1  avg_sim=0.9662
delta=  5  avg_sim=0.7345
delta= 10  avg_sim=0.6579
delta= 20  avg_sim=0.5900
delta= 40  avg_sim=0.4812
delta= 79  avg_sim=0.4664

これは複数の周波数のsin/cosを重ね合わせているせいで、ある距離では複数の高周波成分がたまたま似た位相に戻ってきて、局所的に類似度が持ち直す瞬間があるからだ。「近いほど単調に似ている」という単純な仮定は、少なくともこの生の位置エンコーディング単体を見る限りでは正確ではなかった。学習済みの重みを通せばこの縞模様は吸収されるはずだが、位置情報そのものの形が思っていたより素朴ではないというのは、行列を丸ごと可視化して初めて気づいた点だった。

実験3: 「相対位置は線形変換で表現できる」は本当か

位置エンコーディングにsin/cosが採用されている一番の理由としてよく紹介されるのが、「PEpos+kPE_{pos+k}PEposPE_{pos} の線形変換として表現できるので、モデルが相対位置に注目する仕組みを学習しやすい」という主張だ。三角関数の加法定理を使えば理屈の上では成り立つのはわかるが、実際どれくらい正確に成り立つのか、自分で確かめたことはなかった。

やることはシンプルで、ある距離 kk について、いくつかの位置 pospos から (PEpos, PEpos+k)(PE_{pos},\ PE_{pos+k}) のペアを集め、最小二乗法で行列 MkM_k を求める。そのうえで、学習には使っていない別の位置でどれだけ正確に PEpos+k=MkPEposPE_{pos+k} = M_k \, PE_{pos} が成り立つかを検証した。

M_k, *_ = np.linalg.lstsq(PE[train_pos], PE[train_pos + k], rcond=None)
Y_pred = PE[test_pos] @ M_k
residual = np.abs(Y_pred - PE[test_pos + k]).max()   # held-outな位置での誤差
k=  1  residual=2.02e-12
k=  2  residual=8.38e-14
k=  5  residual=4.30e-14
k= 10  residual=3.27e-12
k= 20  residual=7.83e-14
k= 30  residual=5.37e-14
k= 50  residual=2.24e-13

kk を1から50までどう変えても、誤差は10⁻¹²〜10⁻¹⁴のオーダー、つまり浮動小数点演算の誤差レベルでしかない。実質ゼロ、完全に一致していると言っていい。k=7k=7 の場合について、学習に使っていない位置を1つずつスキャンしながら、実際の PEpos+7PE_{pos+7} (青の実線)と M7PEposM_7 \cdot PE_{pos} による予測(赤の点線)がどれだけ重なるかをアニメーションにした。

k=7の相対位置を表す線形変換M_kについて、実際のPE(pos+7)(青の実線)と、M_kによる予測(赤の点線)を、学習に使っていない位置を左から右へスキャンしながら重ねていくアニメーション。4つの次元それぞれで、赤の点線が青の実線にぴったり重なったまま最後までずれない

学習に使った60個の位置とは別の場所を左から順にスキャンしていっても、赤の点線は最後まで青い実線からずれない。4つの次元どれを見ても同じだった。「相対位置は固定の線形変換で表現できる」という、論文の脚注レベルの主張が、実際に数値でここまで綺麗に成り立つとは思っていなかった。MkM_kpospos に依存せず、kk だけで決まる行列になっている、というのがポイントで、これがあるから、attentionの射影行列がこの線形変換をうまく学習できれば「7個前のトークンだけを見る」といった相対位置ベースの注目の仕方を、位置ごとに別々のルールを覚えることなく獲得できる。

実験4: 位置情報を足すだけで「it」は「animal」を見るようになるのか

最後に、前回の「it」の実験に位置エンコーディングを足して、話がつながるか確かめた。GloVeの単語ベクトルに、対応する位置の PEPE をそのまま加算するだけ(射影なし)で計算し直す。

X_pos = X + PE[:len(sentence), :d_model]
attn_with_pe = self_attention(X_pos)

単語ごとの前後変化を、「itからの相対位置」順に並べ替えた棒グラフにしてみた。

位置エンコーディングなし(青)とあり(赤)で、'it'からの各単語へのattention重みを比較した棒グラフ。横軸はitからの相対位置順(it自身、+1のbecause、-1のwas...)に並んでいる。it自身とbecause・tooなど近い単語の赤い棒が青より高くなっている一方、animalの棒だけは青と赤がほぼ完全に重なって高さが変わっていない

赤(位置エンコーディングあり)の棒は、it自身や近くの「because」「too」で青(なし)より明らかに高くなっている。一方、一番右寄りに位置する「animal」だけは、青と赤の棒がほぼ完全に重なっている。数値で見ても、

'it' attention WITHOUT positional encoding:  animal 0.0310
'it' attention WITH positional encoding added: animal 0.0310   (変化 -0.0000)

位置情報を足した途端、「it」の注目は自分自身(0.18→0.34)と、すぐ近くにある「because」「too」に強く偏った。位置エンコーディングは、素の状態でも「近くのトークンほど似ている」というバイアスを持っているので、これは実験2の結果と整合する。ただし肝心の**「animal」への注目は小数点以下まで一切変わらなかった。** 文中でだいぶ離れた位置にある「animal」は、位置エンコーディングを足しただけでは何も恩恵を受けない。位置情報は「相対的な近さ」を持ち込むだけで、意味的に重要な遠くの単語と結びつける力はない。前回の結論(「itがanimalに注目する」のはQKVの射影行列を学習した結果であって、仕組みそのものの性質ではない)は、位置情報を足しても変わらなかった。位置エンコーディングは語順という情報を「使える形にして渡す」だけで、それをどう使うかは、結局学習に委ねられている。

手を動かして意外だったこと

一番印象に残ったのは、Self-Attentionが語順を「知らない」のではなく、語順という概念自体が計算の外側にあるという点だった。QKVの数式そのものは、入力が集合であろうと列であろうと気にしない。順序という情報は、位置エンコーディングという形で人間側が明示的に「足してあげて」初めて存在するようになる。

これは、コミュニケーションの多くの場面にも近い気がする。同じ事実や言葉でも、それがどんな順番・タイミングで出てきたかによって受け取られ方は大きく変わる。ただ、その「順番」や「文脈」は、話す側が意識的に示さない限り、聞く側には自動的には伝わらない。中身(単語の意味、事実)と、それがいつ・どんな順で出てきたか(位置)は、本質的に別の情報で、後者は誰かが明示的に運んでこない限りどこにも存在しない。もう一つ面白かったのは、相対位置が線形変換として厳密に取り出せるという結果だ。「絶対的に何番目か」よりも「何個前・何個後か」という相対的な関係の方が、数学的にはずっと扱いやすい形で埋め込まれている。人間関係や記憶も、絶対的な日付そのものより、「あの出来事の少し後」「あれよりずっと前」という相対的な位置関係で覚えていることの方が多い気がして、妙に納得感があった。

まとめ

  • Self-Attentionは入力を「順序のある列」ではなく「順序のない集合」として扱う。単語をシャッフルして計算しても、attentionの重みは元の結果を並べ替えただけのものと誤差5.55×10⁻¹⁷で完全一致する
  • sin/cos位置エンコーディングは、次元番号が小さいほど高周波、大きいほど低周波になっており、2進数の桁のような構造で位置を表現している
  • 位置エンコーディング同士のコサイン類似度は、距離が離れるほど全体としては下がるが単調ではなく、複数の周波数が重なり合うことで山と谷を繰り返す波打った形になる
  • PEpos+kPE_{pos+k}PEposPE_{pos} の固定された線形変換 MkM_k で表現できる」という主張は、held-outの位置で検証しても誤差10⁻¹²〜10⁻¹⁴(浮動小数点誤差レベル)で成立しており、kk が1でも50でも変わらなかった
  • GloVe埋め込みに位置エンコーディングを射影なしで足しただけでは、「it」の注目は自分自身と近くの単語に偏るだけで、遠くにある意味的に重要な「animal」への注目は変化しなかった。位置情報の活用も、結局は学習された射影行列に委ねられている