第7章 スケーリング:なぜルートd_kで割るのか

こんにちは。ゆうせいです。

第6章では、QueryとKeyの内積によって、注意スコアが計算されることを確認しました。自己注意の数式を見ると、この内積の結果を、そのまま次の処理に渡してはいません。必ず、ある値で割り算をしています。

Attention(Q, K, V) = softmax\left(\frac{QK^T}{\sqrt{d_k}}\right)V

分母にある \sqrt{d_k} が、この章の主題です。d_kは、1ヘッドあたりのベクトルの次元数を表します。GPT-2(small)では64ですから、その平方根である8で割ることになります。多くの解説では、この処理が「値が大きくなりすぎるのを防ぐため」と一行で片付けられています。しかし、なぜ値が大きくなると困るのか、なぜ平方根なのかまで説明されることは、あまりありません。第7章では、この二つの疑問に答えます。

スケーリングとは何をしているのか

処理そのものは、単純です。

第6章で求めたスコア行列のすべての値を、\sqrt{d_k} という定数で割ります。

GPT-2(small)の場合、1ヘッドあたりの次元数d_kは64です。

\sqrt{64} = 8

したがって、すべてのスコアを8で割ります。スコアが40であれば5に、スコアがマイナス24であればマイナス3になります。

すべての値を同じ数で割るだけですから、値どうしの大小関係は変わりません。40のほうが24より大きいという関係は、5のほうが3より大きいという関係として、そのまま保たれます。

大小関係が変わらないのであれば、この処理には意味がないように思えます。しかし、次の章で扱うSoftmax関数を通すと、この定数倍の違いが、決定的な差を生むのです。

割らないと何が起きるのか:Softmaxの飽和

Softmax関数の性質を先に確認する

第9章で詳しく扱いますが、ここで必要な範囲だけ、Softmax関数の性質を先に確認します。

Softmax関数は、複数の数値を、合計が1になる確率に変換する関数です。

softmax(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}

重要なのは、指数関数 e^x が使われている点です。指数関数は、入力が大きくなると、爆発的に大きな値を返します。

e^1 \approx 2.72

e^5 \approx 148.4

e^{10} \approx 22026

e^{20} \approx 485165195

入力が1から10に増えると、値は約8000倍になります。10から20に増えると、さらに約22000倍になります。

この性質が、スケーリングの必要性を生みます。

数値例で確認する

三つのトークンに対するスコアが、次の値だったとします。

スケーリングなしの場合:[10, 8, 6]

Softmaxを計算します。

e^{10} \approx 22026

e^{8} \approx 2981

e^{6} \approx 403

合計は約25410です。それぞれの確率は、次のようになります。

1番目:\frac{22026}{25410} \approx 0.867

2番目:\frac{2981}{25410} \approx 0.117

3番目:\frac{403}{25410} \approx 0.016

1番目のトークンに、約87パーセントの注目が集中しています。

次に、同じスコアを8で割った場合を計算します。

スケーリングありの場合:[1.25, 1.0, 0.75]

e^{1.25} \approx 3.49

e^{1.0} \approx 2.72

e^{0.75} \approx 2.12

合計は約8.33です。それぞれの確率は、次のようになります。

1番目:\frac{3.49}{8.33} \approx 0.419

2番目:\frac{2.72}{8.33} \approx 0.327

3番目:\frac{2.12}{8.33} \approx 0.254

こちらでは、注目が複数のトークンに分散しています。

分布が尖ることの何が問題なのか

スケーリングをしないと、注意が一つのトークンにほぼ集中してしまいます。これには、二つの問題があります。

一つ目の問題は、複数の情報を統合できなくなることです。

自己注意の目的は、周囲の複数のトークンから、必要な情報を集めることです。ほぼ一つのトークンだけを見る状態になると、複数の手がかりを組み合わせた判断ができなくなります。

二つ目の問題は、学習が進まなくなることです。こちらのほうが深刻です。

Softmaxの出力が0または1に極端に近づくと、その勾配がほぼゼロになります。第9章で数式を示しますが、Softmaxの微分には、出力の値そのものが含まれる形になっており、出力が0や1に張り付くと、微分値が消えてしまうのです。

勾配がゼロになれば、誤差逆伝播法によるパラメータの更新が行われません。学習が事実上停止します。

この現象を、Softmaxの飽和と呼びます。スケーリングは、この飽和を避けるための処理なのです。

なぜ平方根なのか:分散の観点からの説明

ここからが、この章の核心です。なぜ、次元数d_kそのものではなく、その平方根で割るのでしょうか。

内積の値は次元数とともに大きくなる

内積は、d_k個の掛け算の結果を、すべて足し合わせた値です。

\mathbf{q} \cdot \mathbf{k} = q_1 k_1 + q_2 k_2 + \cdots + q_{d_k} k_{d_k}

足し合わせる項の数が増えれば、結果として得られる値のばらつきも大きくなります。次元数が64のときと、次元数が1024のときでは、内積の値のばらつきが大きく異なるのです。

統計的な性質から導く

ここで、簡単な仮定を置きます。QueryとKeyの各要素が、平均0、分散1の独立な確率変数だとします。

このとき、一つの掛け算 q_i k_i について、次のことが言えます。

平均は0です。独立な二つの平均0の変数の積ですから、期待値は0になります。

分散は1です。独立な二つの変数の積の分散は、それぞれの分散の積になるため、1 \times 1 = 1 です。

次に、これらをd_k個足し合わせます。独立な確率変数の和では、分散が足し算されるという性質があります。

Var(q_1k_1 + q_2k_2 + \cdots + q_{d_k}k_{d_k}) = 1 + 1 + \cdots + 1 = d_k

つまり、内積の分散は、次元数d_kに等しくなります。

分散がd_kであるということは、標準偏差は \sqrt{d_k} です。標準偏差は、値がどれくらいばらつくかの目安を表します。

標準偏差で割れば、ばらつきが一定になる

ここで、内積を \sqrt{d_k} で割ってみます。

確率変数を定数cで割ると、分散は c^2 で割られます。したがって、\sqrt{d_k} で割ると、分散は次のようになります。

\frac{d_k}{(\sqrt{d_k})^2} = \frac{d_k}{d_k} = 1

分散が1に戻ります。

つまり、\sqrt{d_k} で割るという操作は、次元数がいくつであっても、スコアのばらつきを一定の大きさに揃えるための正規化なのです。

これが、d_kそのものではなく、平方根で割る理由です。d_kで割ってしまうと、分散が \frac{1}{d_k} になり、今度は値が小さくなりすぎます。すべてのスコアが0に近づくと、Softmaxの出力が均一に近づき、どのトークンにも同じように注目することになってしまいます。これでは、注意機構が機能しません。

大きすぎず、小さすぎず、ちょうどよいばらつきに揃える。そのための値が、標準偏差である \sqrt{d_k} なのです。

次元数による違いを具体的に確認する

異なる次元数で、スケーリングの有無がどう影響するかを整理します。

次元数 d_k内積の標準偏差の目安スケーリング係数スケーリング後の標準偏差
64(GPT-2 small)881
128約11.3約11.31
25616161
102432321

次元数が16倍(64から1024)になると、内積の標準偏差は4倍になります。スケーリングをしなければ、大きなモデルほどSoftmaxが飽和しやすくなる、という関係です。

モデルを大規模化する際に、スケーリングがなければ、次元数を増やすほど学習が不安定になっていたことになります。この一つの割り算が、大規模モデルの学習を可能にした、地味ながら重要な工夫だと言えます。

Transformer Explainerで観察できること

Transformer Explainerでは、スコア行列に対してスケーリングが適用される段階が、表示されています。

次の点を確認してみてください。

一つ目は、スケーリング前後で、値の大小関係が変わっていないという点です。最も大きい値は、割り算の後も最も大きいままです。

二つ目は、値の範囲が縮まっているという点です。極端に大きな値や小さな値が、より中央に寄った範囲に収まります。

三つ目は、この段階でも、まだ値がマイナスになり得るという点です。スケーリングは、値の範囲を狭めるだけであり、符号を変えることはありません。確率への変換は、まだ先の処理です。

初学者がつまずきやすい点

つまずき1:d_kが768だと思ってしまう

\sqrt{d_k} のd_kは、埋め込み次元の768ではありません。1ヘッドあたりの次元数です。

GPT-2(small)では、768次元を12ヘッドに分割するため、1ヘッドあたりは64次元です。

\frac{768}{12} = 64

したがって、割る値は \sqrt{768} \approx 27.7 ではなく、\sqrt{64} = 8 です。

この点を間違えると、実装時に値が合わなくなります。マルチヘッドの仕組みは第10章で扱いますが、スケーリングの時点では、すでにヘッドごとに分割された後の次元数を使う、と押さえてください。

つまずき2:大小関係が変わらないのに意味があるのか、という疑問

「同じ数で割るだけなら、順位は変わらないのだから、意味がないのではないか」という疑問は、非常に自然です。

答えは、Softmaxが非線形な関数だからです。

線形な処理であれば、定数倍は後で調整できます。しかし、指数関数を含むSoftmaxでは、入力の絶対的な大きさが、出力の分布の形を決めます。入力を8で割ることは、出力の確率分布を「なだらかにする」ことに対応します。

この関係は、第14章で扱う温度パラメータとまったく同じ構造です。温度も、Softmaxに入れる前の値を定数で割ることで、確率分布の尖り具合を調整します。スケーリングは、いわば固定された温度である、と理解しておくと、第14章の内容がつながりやすくなります。

つまずき3:分散の計算で使った仮定を絶対視する

「QueryとKeyの各要素が、平均0、分散1の独立な確率変数である」という仮定は、あくまで理論的な説明のための単純化です。

実際のモデルでは、学習が進むにつれて、QueryとKeyの分布はこの仮定から外れていきます。要素どうしが独立であるとも限りません。

したがって、\sqrt{d_k} で割れば分散が厳密に1になる、というわけではありません。あくまで、おおよその目安として妥当な値である、という理解が正確です。

数式による導出は、なぜその値が選ばれたかを説明するものであり、実際の値を保証するものではない、という点を押さえておいてください。

演習

演習1:スケーリング係数を計算する

次の条件のモデルについて、スケーリングに使う値を計算してください。

埋め込み次元が4096、ヘッド数が32のモデル。

演習2:Softmaxの飽和を確認する

スケーリングの有無で、Softmaxの出力がどう変わるかを、実際に確認してください。

import torch
import torch.nn.functional as F
import math

# スコアの例(次元数64を想定した大きさ)
scores = torch.tensor([12.0, 9.0, 6.0, 3.0])

print("=== スケーリングなし ===")
probs_raw = F.softmax(scores, dim=0)
print(f"スコア: {scores.tolist()}")
print(f"確率: {[round(p, 4) for p in probs_raw.tolist()]}")
print(f"最大値: {probs_raw.max().item():.4f}")

print("\n=== スケーリングあり(d_k=64) ===")
d_k = 64
scaled = scores / math.sqrt(d_k)
probs_scaled = F.softmax(scaled, dim=0)
print(f"スケーリング後のスコア: {[round(s, 4) for s in scaled.tolist()]}")
print(f"確率: {[round(p, 4) for p in probs_scaled.tolist()]}")
print(f"最大値: {probs_scaled.max().item():.4f}")

print("\n=== 次元数を変えた比較 ===")
for d in [64, 256, 1024]:
    scaled = scores / math.sqrt(d)
    probs = F.softmax(scaled, dim=0)
    print(f"d_k={d:5d}: 最大確率 {probs.max().item():.4f}, スケーリング係数 {math.sqrt(d):.2f}")

演習3:内積の分散を実験で確かめる

理論どおりに、内積の分散が次元数に一致するかを確認してください。

import torch

torch.manual_seed(0)

for d_k in [16, 64, 256, 1024]:
    # 平均0、分散1のランダムなベクトルを大量に生成
    q = torch.randn(10000, d_k)
    k = torch.randn(10000, d_k)
    
    # 対応するペアの内積を計算
    dots = (q * k).sum(dim=1)
    
    print(f"d_k={d_k:5d}: 内積の分散 {dots.var().item():8.2f}, 標準偏差 {dots.std().item():6.2f}, sqrt(d_k)={d_k**0.5:6.2f}")

演習の解答例と解説

演習1の解答は、次の通りです。

1ヘッドあたりの次元数は、次のように計算されます。

\frac{4096}{32} = 128

スケーリングに使う値は、その平方根です。

\sqrt{128} \approx 11.31

埋め込み次元の4096ではなく、ヘッドで分割した後の128を使う点が、この演習の要点です。

演習2では、スケーリングなしの場合に、最大確率が0.95を超えるほど極端に偏ることが確認できます。一方、スケーリングありでは、確率がより均等に分散します。

また、次元数を変えた比較では、d_kが大きいほどスケーリング係数も大きくなり、確率分布がなだらかになることが分かります。

演習3では、内積の分散が、おおむねd_kの値に一致することが確認できます。d_kが64であれば分散は約64、標準偏差は約8になるはずです。この実験結果が、本文で説明した理論的な導出を裏付けています。

講師向けの補足

この章は、数式の導出が中心になるため、受講者の数学的な背景によって、進め方を変える必要があります。

数学に苦手意識がある受講者が多い場合は、分散の議論を省略し、「次元数が大きいほど内積が大きくなるので、その分だけ割り戻している」という説明にとどめても構いません。そのうえで、演習3のプログラムを実際に動かし、数値で確認してもらうほうが、理解が定着します。

E資格対策として扱う場合は、分散の計算を丁寧に追ってください。試験では、「なぜ \sqrt{d_k} で割るのか」という問いが、選択式で出題されることがあります。「値が大きくなりすぎるのを防ぐため」という選択肢と、「分散を一定に保つため」という選択肢が並んだ場合、後者がより正確な説明です。

理解度を確認する問い

第7章の内容を、次の問いで確認してみてください。

スケーリングをしないと、Softmaxの出力にどのような問題が起きるでしょうか。また、それがなぜ学習の妨げになるのでしょうか。

d_kそのものではなく、その平方根で割る理由を、分散の観点から説明してみてください。

d_kで割ってしまった場合には、どのような問題が起きるでしょうか。

まとめ

スケーリングとは、QueryとKeyの内積によって得られたスコアを、\sqrt{d_k} という定数で割る処理です。d_kは、1ヘッドあたりの次元数であり、GPT-2(small)では64、したがって割る値は8になります。

この処理が必要な理由は、スコアの値が大きいまま Softmax関数に入力されると、指数関数の性質によって、注意が一つのトークンにほぼ集中してしまうためです。出力が0や1に張り付くと、勾配がほぼゼロになり、学習が進まなくなります。

平方根で割る理由は、統計的な性質にあります。QueryとKeyの各要素が平均0、分散1であると仮定すると、内積の分散は次元数d_kに等しくなります。その標準偏差である \sqrt{d_k} で割ることにより、次元数がいくつであっても、スコアのばらつきを一定に保てます。

d_kそのもので割ると、今度は値が小さくなりすぎ、すべてのトークンに均等に注目する状態になってしまいます。大きすぎず小さすぎない、ちょうどよい値が、標準偏差なのです。

次の第8章では、因果マスクを扱います。GPT-2のような文章生成モデルでは、あるトークンを処理する際に、それより後ろにあるトークンを見てはいけないという制約があります。この制約を、どのように実装しているのかを解説します。

第8章に進む前に、Transformer Explainerでスケーリング後のスコア行列を確認し、値の範囲がどの程度に収まっているかを見ておいてください。次章では、この行列の一部が、意図的にマイナス無限大に置き換えられる様子を観察することになります。

セイ・コンサルティング・グループでは新人エンジニア研修のアシスタント講師を募集しています。

投稿者プロフィール

山崎講師
山崎講師代表取締役
セイ・コンサルティング・グループ株式会社代表取締役。
岐阜県出身。
海外放浪の末、2000年創業、2004年会社設立。
IT企業向け人材育成研修歴業界歴20年以上。
すべての無駄を省いた費用対効果の高い「筋肉質」な研修を提供します!
この記事に間違い等ありましたらぜひお知らせください。

学生時代は趣味と実益を兼ねてリゾートバイトにいそしむ。長野県白馬村に始まり、志賀高原でのスキーインストラクター、沖縄石垣島、北海道トマム。高じてオーストラリアのゴールドコーストでツアーガイドなど。現在は野菜作りにはまっている。