Transformer攻略(2) Q・K・VからAttention scoreを作る
はじめに
こちらは第二回目です。
前回、AIに仕事を取られないエンジニアになるぞ、という気持ちでTransformerを触り始めました。
前回は、文章をtokenに分けて、IDにして、Embeddingして、位置情報を足すところまでやりました。
全体の流れでいうと、ここまでです。
文章
↓
トークン化
↓
トークンIDに変換
↓
単語Embeddingに変換
↓
位置Embeddingを足す
↓
Query・Key・Valueを作る
↓
Self-Attention
↓
Multi-Head Attention
↓
Feed Forward
↓
Transformer Blockを通す
↓
次のトークンを予測する
↓
小さな文章生成モデルになる
今回はその続きで、Attentionに入ります。
正直、ここから急に難しくなりました。
Q, K, V とか、q @ k.transpose(-2, -1) とか、最初に見たときは「何をやっているんだこれは」という感じでした。
なので今回は、完璧な理論理解というより、まずはshapeと処理の流れを追って、Attentionで何が起きているのかを掴むことを目標にしました。
前回の振り返り
前回やったことをかなりざっくり書くと、以下です。
文章はそのままだとモデルが扱えないので、まず数値にします。
例えば今回のコードでは、以下の文を使いました。
I like small models
これをtoken IDにして、Embeddingして、位置情報を足すと、最終的にこういうshapeになります。
positioned: (1, 4, 8)
意味はこうです。
1 = batch数
4 = token数
8 = 1 tokenあたりのベクトル次元
今回のAttentionは、この positioned を入力として使います。
今回やること
今回から、Scaled Dot-Product Attentionを実装します。
この記事では、その前半として、Q・K・Vを用意し、Attention scoreを作るところまで進めます。
名前は強そうですが、まずはこう考えました。
各tokenが、他のtokenをどれくらい見るかを決める処理
流れはこんな感じです。
一言でいうと、
QとKで「どれを見るか」を決める
決まった割合でVを混ぜる
です。
この記事では、図のうちscoresを作るところまで、ひとつずつ見ていきます。
Q, K, Vって何
最初に分からなかったのが Q, K, V です。
名前はそれぞれ以下です。
| 記号 | 名前 | ざっくりした役割 |
|---|---|---|
Q |
Query | 自分が何を探しているか |
K |
Key | 自分がどう見つけられたいか |
V |
Value | 見つけられたときに渡す中身 |
検索で考えると少し分かりやすかったです。
Q = 検索欄に入れる言葉
K = 検索対象についたタグ
V = 実際に読む本文
まず Q と K を比べて、どの情報を見るかを決めます。
そのあと、決まった割合で V の中身を受け取ります。
今回の実装では、理解を優先してかなり単純化しています。
q = positioned
k = positioned
v = positioned
本来はLinear層を通して Q, K, V を別々に作るようですが、今回はAttentionの計算を追いたいので、同じTensorに別名をつけています。
ここは最初、「同じものを代入しているだけなら意味あるのか？」と思いました。
ただ、後続の計算では役割が違います。まずは名前と役割に慣れることを優先し、今回やりたいことからもずれるため、単純化します。
scoreを作る
次に、Q と K を比べてscoreを作ります。
コードはこれです。
d_k = q.size(-1)
scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)
さっぱりわかりませんでした。
なので、まずは意味よりshapeから見ます。
q: (1, 4, 8)
k: (1, 4, 8)
k.transpose: (1, 8, 4)
q @ k.T: (1, 4, 4)
4 x 4 になるのは、4個のtokenが、それぞれ4個のtoken全部と比べられるからです。
key
I like small models
query I ? ? ? ?
query like ? ? ? ?
query small ? ? ? ?
query models ? ? ? ?
この ? がscoreです。
つまりscoreは、
各token同士の「関係ありそう度」の表
みたいなものだと理解しました。
図にするとこうです。
なぜ sqrt(d_k) で割るのか
ここも最初はよく分かりませんでした。
q @ k.T だけでもscoreは作れます。
では、なぜわざわざ sqrt(d_k) で割るのか。
自分の理解では、scoreが大きくなりすぎるのを防ぐためです。
内積は、ベクトルの次元数が増えるほど、足し合わせる数も増えます。
そうするとscoreが大きくなりやすいです。
scoreが大きすぎると、次に出てくるsoftmaxがかなり極端になります。
[0.9999, 0.0000, 0.0000, 0.0001]
みたいになると、ほぼ1つのtokenしか見なくなります。
それを少しならすために、sqrt(d_k) で割っています。
ここはまだ完全に腹落ちしているというより、
softmaxが極端になりすぎるのを防ぐ調整
くらいの理解です。
実装
今回のコードです。
前半は前回と同じで、token ID、Embedding、位置情報を作っています。
後半でQ・K・Vを用意し、Attention scoreを計算します。
import math
import torch
from torch import nn
vocab = {
"<pad>": 0,
"<unk>": 1,
"I": 2,
"like": 3,
"small": 4,
"models": 5,
}
sentence = "I like small models"
tokens = sentence.split()
token_ids = [vocab[token] for token in tokens]
input_ids = torch.tensor([token_ids])
embedding_dim = 8
embedding = nn.Embedding(
num_embeddings=len(vocab),
embedding_dim=embedding_dim,
)
embedded = embedding(input_ids)
max_len = 20
position = torch.arange(max_len, dtype=torch.float32).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, embedding_dim, 2).float()
* (-math.log(10000.0) / embedding_dim)
)
pe = torch.zeros(max_len, embedding_dim)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
positioned = embedded + pe[:, :embedded.size(1)]
q = positioned
k = positioned
v = positioned
d_k = q.size(-1)
scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)
print("tokens:", tokens)
print("positioned.shape:", positioned.shape)
print("q.shape:", q.shape)
print("k.transpose.shape:", k.transpose(-2, -1).shape)
print("scores.shape:", scores.shape)
実行すると、shapeはこうなりました。
tokens: ['I', 'like', 'small', 'models']
positioned.shape: torch.Size([1, 4, 8])
q.shape: torch.Size([1, 4, 8])
k.transpose.shape: torch.Size([1, 8, 4])
scores.shape: torch.Size([1, 4, 4])
今回ハマったところ
今回ハマったところを、自分用にもう一回短くまとめます。
scores のshapeがなぜ (1, 4, 4) なのか
最初、positioned が (1, 4, 8) なので、scoreも (1, 4, 8) っぽくなるのかなと思っていました。
でもscoreはtoken同士の関係表です。
4 token が 4 token を見る
ので、4 x 4 になります。
sqrt(d_k) で割る理由
ここはまだ完全には腹落ちしていません。
ただ、今の理解では、softmaxが極端になりすぎるのを防ぐための調整です。
こういう「まだ怪しいけど、まずはこの理解で進む」ポイントが、Transformerには結構ありそうです。
全体まとめ
今回の流れをまとめると、以下です。
shapeだけで追うと、こうです。
positioned: (1, 4, 8)
q: (1, 4, 8)
k: (1, 4, 8)
k.T: (1, 8, 4)
scores: (1, 4, 4)
今回の理解を一言でまとめるなら、
QとKを使って、各token同士の「関係ありそう度」を表すscoreを作る
です。
scoreはまだ生の点数です。
次回はsoftmaxで「各tokenをどれくらい見るか」という割合に変換し、その割合でVを混ぜます。