RでTransformerのAttention機構:スケーリング済み内積スコア

\(QK^{\top}\) が意味すること

\(Q\)(Query) と \(K\)(Key) は、それぞれ複数のベクトルを行として持つ行列です。

Transformerのattention機構では、入力系列の各トークン(単語)ごとに、同じ埋め込みベクトルから2種類の異なる役割を持つベクトルを、別々の学習可能な重み行列(\(W_Q\), \(W_K\))を使って作り出します。

なお、埋め込みベクトルとは、各単語(トークン)を、連続値の数値の並び(ベクトル)に変換したものであり、BERT-baseの場合、768次元の実数値ベクトルになります。

\[
Q = X W_Q, \qquad K = X W_K
\]

ここで、\(X\)は入力トークンの埋め込みを並べた行列です。

• Query(Q):
  • そのトークンが「何を知りたいか」を表すベクトル。「自分の意味を明確にするために、他のどんな情報が欲しいか」という問いかけに相当します
• Key(K):
  • そのトークンが「何を提供できるか」を表すベクトル。他のトークンからの問いかけに対して、「自分はこういう情報を持っている」と示す見出しやラベルに相当します

同じ埋め込みベクトル \(X\) から出発しているにもかかわらず、\(W_Q\)と\(W_K\)という別々の重み行列をかけることで、\(Q\) と \(K\) は異なる空間に射影され、それぞれ異なる役割(問いかける側/答える側)を担うようになります。

この \(W_Q\)、\(W_K\) 自体は、モデル全体の学習を通じて獲得されたパラメータです。

\(QK^{\top}\) を計算するということは、\(Q\) の各行ベクトルと、\(K\) の各行ベクトルの内積を、全ての組み合わせについて一斉に計算するということです。

内積は「2つのベクトルがどれだけ同じ方向を向いているか」を表す量です。

ベクトルの向きが近いほど内積は大きくなり、直交していれば0、逆方向なら負になります。

つまり \(QK^{\top}\) は、\(Q\) のどのベクトルと \(K\) のどのベクトルが”似ている(近い方向を向いている)かを表す、類似度の一覧表になります。

なぜ \(\sqrt{d_k}\) で割るのか

「内積の値は、ベクトルの次元数 \(d_k\) が大きくなるほど、大きくばらつくようになる」という統計的な性質にあります。

\(Q\) と \(K\) の各成分が、平均0・分散1程度の(独立な)乱数だとすると、内積

\[
q \cdot k = \sum_{i=1}^{d_k} q_i k_i
\]

は \(d_k\) 個の項の合計になるため、その分散は\(d_k\)に比例して大きくなることが期待値の計算から分かります(項が独立ならVarの加法性が使える)。

つまり \(d_k\)(embeddingの次元数)が大きいモデルほど、\(QK^{\top}\) の値そのものが極端に大きく(または小さく)なりやすいのです。

その何が問題かというと、この後 softmax に通すためです。

softmax は入力の値が極端に大きいと、出力がほぼ0か1に「飽和」してしまい(差が大きいほど勝者に極端に確率を寄せる性質があります)、学習時の勾配がほとんどゼロになってしまいます(勾配消失)。

\(\sqrt{d_k}\) で割ることで、内積の分散を \(d_k\) の値に関係なくほぼ1程度に保ち、この問題を防いでいます。

以下のコードで、この「次元数と分散の関係」と「スケーリングの効果」を実際に確認します。

Part 1: \(QK^T\) は「内積による類似度の一覧表」

簡単な2次元・3次元ベクトルで、\(QK^{\top}\) が「向きの近さ(内積)の一覧表」であることを確認します。

library(ggplot2)

seed <- 20260714
set.seed(seed)

# 2次元ベクトルで直感的に確認する
q1 <- c(1, 0) # 右向き
k1 <- c(1, 0) # 右向き (q1と同じ方向)
k2 <- c(0, 1) # 上向き (q1と直交)
k3 <- c(-1, 0) # 左向き (q1と逆方向)

cat("=== 内積は「向きの近さ」を表す ===\n")
cat(sprintf("q1・k1 (同じ方向)  = %.2f\n", sum(q1 * k1)))
cat(sprintf("q1・k2 (直交)      = %.2f\n", sum(q1 * k2)))
cat(sprintf("q1・k3 (逆方向)    = %.2f\n", sum(q1 * k3)))

# 少し現実的な例: 複数のQuery行、複数のKey行を持つ行列同士の
# QK^T を計算する
Q <- rbind(
  q_a = c(1.0, 0.2, 0.0),
  q_b = c(0.0, 1.0, 0.3)
)
K <- rbind(
  k_x = c(0.9, 0.1, 0.0), # q_a に似ている
  k_y = c(0.1, 0.9, 0.2), # q_b に似ている
  k_z = c(-1.0, -1.0, 0.0) # どちらとも似ていない
)

QKt <- Q %*% t(K)
cat("\n=== QK^T (Query x Key の類似度一覧表) ===\n")
print(round(QKt, 3))

=== 内積は「向きの近さ」を表す ===
q1・k1 (同じ方向)  = 1.00
q1・k2 (直交)      = 0.00
q1・k3 (逆方向)    = -1.00

=== QK^T (Query x Key の類似度一覧表) ===
     k_x  k_y  k_z
q_a 0.92 0.28 -1.2
q_b 0.10 0.96 -1.0

\(q_a\) は \(k_x\) と、\(q_b\) は \(k_y\) と、内積(類似度)が高くなっており、\(QK^T\) の結果は「誰と誰が似ているか」を一斉に数値化していることになります。

Part 2: 次元数 \(d_k\) が大きいほど、内積の分散が大きくなる

モンテカルロ・シミュレーションで、次元数 \(d_k\)(4〜1024)を変化させ、内積 \(q \cdot k\) の分散が \(d_k\) にほぼ比例して増加することを理論値 \(d_k\) と比較しながら確認します。

\(Q\), \(K\) の各成分が平均0・分散1の乱数だとすると、内積 \(q\cdot k\) の分散は理論上 \(d_k\) に比例して大きくなります (\(Var(q\cdot k) = d_k\), 各項が独立であれば分散が単純に加算されるため)。

d_k_values <- c(4, 16, 64, 256, 1024)
n_trials <- 5000

simulate_dot_product_variance <- function(d_k, n_trials) {
  # 平均0, 分散1の乱数ベクトルを大量に作り、内積を計算する
  dot_products <- replicate(n_trials, {
    q <- rnorm(d_k, mean = 0, sd = 1)
    k <- rnorm(d_k, mean = 0, sd = 1)
    sum(q * k)
  })
  var(dot_products)
}

variance_results <- sapply(d_k_values, simulate_dot_product_variance, n_trials = n_trials)

variance_df <- data.frame(
  d_k = d_k_values,
  empirical_variance = variance_results,
  theoretical_variance = d_k_values # 理論値: Var(q・k) = d_k
)

cat("=== 次元数 d_k と 内積 q・k の分散の関係 ===\n")
print(variance_df)

=== 次元数 d_k と 内積 q・k の分散の関係 ===
   d_k empirical_variance theoretical_variance
1    4           3.947771                    4
2   16          15.549082                   16
3   64          62.586327                   64
4  256         254.483085                  256
5 1024        1052.891052                 1024

empirical_variance が \(d_k\) にほぼ比例して増えており、理論値 (\(d_k\) そのもの) にほぼ一致しています。

つまり \(d_k\) が大きいモデルほど、\(QK^T\) の値は大きくばらつくことになります(標準偏差は \(\sqrt{d_k}\) 程度になる)。

Part 3: \(\sqrt{d_k}\) で割ると、分散が \(d_k\) に依らず一定になる

\(\sqrt{d_k}\) で割った場合、分散が \(d_k\) の値に関係なくほぼ1で一定に保たれることを確認します。

simulate_scaled_variance <- function(d_k, n_trials) {
  scaled_dot_products <- replicate(n_trials, {
    q <- rnorm(d_k, mean = 0, sd = 1)
    k <- rnorm(d_k, mean = 0, sd = 1)
    sum(q * k) / sqrt(d_k) # ここで sqrt(d_k) で割る
  })
  var(scaled_dot_products)
}

scaled_variance_results <- sapply(d_k_values, simulate_scaled_variance, n_trials = n_trials)

cat("=== sqrt(d_k) でスケーリングした後の分散 ===\n")
scaled_df <- data.frame(d_k = d_k_values, scaled_variance = scaled_variance_results)
print(scaled_df)

=== sqrt(d_k) でスケーリングした後の分散 ===
   d_k scaled_variance
1    4       0.9668607
2   16       0.9856271
3   64       0.9650066
4  256       1.0077857
5 1024       0.9983623

\(d_k\) がいくら変わっても、分散はほぼ1程度で一定になっており、次元数に関係なく、内積のスケールが一定に保たれています。

Part 4: スケーリングしないと softmax が飽和し、勾配が消える

\(d_k\) が大きいときの \(QK^T\) をそのまま softmax に通しますと、値のばらつきが大きすぎて、softmax の出力がほぼ0か1に「飽和」してしまいます。

これは学習時の勾配がほぼゼロになる(=学習が進みにくくなる) ことを意味します。

\(d_k=1024\) を例に、スケーリングの有無で softmax 出力がどう変わるか、また \(p\times(1-p)\)(勾配の大きさの目安)を使って勾配消失の様子を具体的な数値で確認します。

softmax <- function(x) exp(x - max(x)) / sum(exp(x - max(x)))

d_k <- 1024
q <- rnorm(d_k)
k1 <- rnorm(d_k)
k2 <- rnorm(d_k)
k3 <- rnorm(d_k)

raw_scores <- c(sum(q * k1), sum(q * k2), sum(q * k3))
scaled_scores <- raw_scores / sqrt(d_k)

cat("=== d_k = 1024 のときの、スケーリングの有無による違い ===\n")
cat("生のスコア (スケーリングなし):", round(raw_scores, 2), "\n")
cat("  -> softmax:", round(softmax(raw_scores), 4), "\n\n")
cat("スケーリング後のスコア:", round(scaled_scores, 2), "\n")
cat("  -> softmax:", round(softmax(scaled_scores), 4), "\n")

# softmaxの勾配 (softmax'(x) = softmax(x) * (1 - softmax(x)) の
# 対角成分に相当する量) がどれだけ小さくなるかも確認する
gradient_magnitude <- function(probs) probs * (1 - probs)

cat("\n=== softmax出力における勾配の大きさの目安 (p * (1-p)) ===\n")
cat("スケーリングなしの場合:", round(gradient_magnitude(softmax(raw_scores)), 6), "\n")
cat("スケーリングありの場合:", round(gradient_magnitude(softmax(scaled_scores)), 6), "\n")

=== d_k = 1024 のときの、スケーリングの有無による違い ===
生のスコア (スケーリングなし): -20.88 16.94 -28.92 
  -> softmax: 0 1 0 

スケーリング後のスコア: -0.65 0.53 -0.9 
  -> softmax: 0.1985 0.6472 0.1544 

=== softmax出力における勾配の大きさの目安 (p * (1-p)) ===
スケーリングなしの場合: 0 0 0 
スケーリングありの場合: 0.159078 0.228345 0.130545 

スケーリングなし

softmax: 0  1  0
勾配:    0  0  0

3つとも完全に0か1に張り付き、勾配が全てゼロになっています。

これは学習が完全に停止してしまう状態そのもので、勾配消失を最も分かりやすい形で示しています。

スケーリングあり

softmax: 0.1985  0.6472  0.1544
勾配:    0.159    0.228    0.131

k2が最も高い確率(0.6472)を持ちつつも、他の2つも0.15〜0.20程度の確率を保っており、極端な一強状態にはなっていません。

勾配も0.13〜0.23の範囲で、実質的な大きさを保っています。

また、生スコアの範囲は約-29〜17(差は最大で約46)でしたが、スケーリング後は約-0.9〜0.53(差は最大で約1.4)まで縮小しています。

\(\sqrt{1024}=32\) で割っていますので、\(46/32 \approx 1.4\) と、計算通りの縮小率になっています。

Part 5: ggplot2 による可視化

# --- 図1: 次元数 d_k と 内積の標準偏差の関係 ---
sd_df <- data.frame(
  d_k = rep(d_k_values, 2),
  sd_value = c(sqrt(variance_results), sqrt(scaled_variance_results)),
  type = rep(c(
    "スケーリングなし (標準偏差 ≈ sqrt(d_k))",
    "スケーリングあり (標準偏差 ≈ 1 で一定)"
  ), each = length(d_k_values))
)

p1 <- ggplot(sd_df, aes(x = d_k, y = sd_value, color = type)) +
  geom_line(linewidth = 1) +
  geom_point(size = 2.5) +
  scale_x_log10() +
  labs(
    title = "次元数 d_k と QK^T の標準偏差の関係",
    subtitle = "sqrt(d_k) でスケーリングすると、標準偏差が d_k に依らず一定になる",
    x = "d_k (対数スケール)", y = "標準偏差", color = NULL
  ) +
  theme_minimal(base_size = 12) +
  theme(legend.position = "top")

print(p1)

# --- 図2: d_k=1024 における、スケーリングの有無でのsoftmax出力比較 ---
compare_df <- data.frame(
  key = factor(rep(c("k1", "k2", "k3"), 2)),
  probability = c(softmax(raw_scores), softmax(scaled_scores)),
  type = rep(c("スケーリングなし", "スケーリングあり"), each = 3)
)

p2 <- ggplot(compare_df, aes(x = key, y = probability, fill = type)) +
  geom_col(position = "dodge") +
  geom_text(aes(label = round(probability, 3)),
    position = position_dodge(width = 0.9), vjust = -0.5, size = 3.3
  ) +
  labs(
    title = "スケーリングの有無による softmax 出力の違い (d_k = 1024)",
    subtitle = "スケーリングなしだと、わずかな内積の差でも極端に一強になりやすい",
    x = "Key", y = "Attention Weight", fill = NULL
  ) +
  theme_minimal(base_size = 12) +
  theme(legend.position = "top") +
  coord_cartesian(ylim = c(0, 1))

print(p2)

以上です。