メインコンテンツまでスキップ

Chapter 2: Transformer

前章では、ベクトル・内積・行列積・softmax といった「LLM を支える数学」を積み上げました。この章では、いよいよその数学が組み上がってできる主役、Transformer(トランスフォーマー) に迫ります。

この章は 前半(概念)→ 後半(実装) の2部構成です。まず「そもそも Transformer って何なの?」「なぜ必要だったのか」をつかみ、心臓部の Attention(注意機構) の仕組みを前章の数学とつなげて理解します。そのうえで後半では、その Attention を 実際に数値とコードで計算し、手を動かしながら腑に落とします。

この章の読み方

数式や用語が出てきても身構えなくて大丈夫です。新しい記号や式は 使う直前に意味を渡します。特に Attention は「前章で学んだ内積・softmax をくり返しているだけ」だと分かると、一気に見通しが良くなります。後半の実装も、前半で理解した3ステップをそのままコードにするだけです。


1. ざっくり一言で言うと

Transformer をひとことで言うと、こうです。

文章の中で「どの単語が、どの単語に注目すべきか」を一気に計算する仕組み(ニューラルネットのアーキテクチャ)

2017 年の論文 "Attention Is All You Need" で提案され、いまの GPT・Claude をはじめとする大規模言語モデルが、ほぼ例外なくこの Transformer をベースにしています。LLM を学ぶうえで「ここだけは外せない」という中心的な構造です。

そして大事なのは、Transformer は決して魔法の箱ではない、ということです。中身は前章で学んだ数学の組み合わせでできています。

Transformer の部品中身(前章で学んだ数学)
単語を数で表すベクトル(単語埋め込み)
単語どうしの関連度を測る内積・コサイン類似度
全単語ペアの関連度を一括計算行列積 QKTQK^T
関連度を「注目度(確率)」に変えるsoftmax
表現を別空間へ変換する線形変換(重み行列)
「何番目の単語か」を表す三角関数(位置エンコーディング)

つまり Transformer の理解とは、「前章の数学が、どういう順番で、何のために組み合わさるのか」を理解することにほかなりません。

アーキテクチャって何?

「アーキテクチャ(architecture)」は、ここでは ニューラルネットワークの設計図・構造の型のことです。「どんな計算ブロックを、どんな順番でつなぐか」という骨組みを指します。Transformer は、その骨組みの一種の名前だと思ってください。


2. なんで生まれたの?(背景がアツい🔥)

新しい技術は、たいてい「前のやり方の不満」から生まれます。Transformer も同じです。ここを押さえると、Attention という仕組みが「なぜその形なのか」が腑に落ちます。

2.1 むかしのやり方:1単語ずつ順番に読む(RNN)

Transformer 以前、文章を扱う主役は RNN(再帰型ニューラルネット) やその改良版の LSTM でした。これらは、人間が文章を読むように 単語を1つずつ、順番に処理していきます。

前の単語を読んだ「記憶」を次の単語へ受け渡しながら、左から右へ進んでいくイメージです。一見、自然な方法に見えます。しかし、ここには大きな問題が2つありました。

問題1:遅い(並列化できない)

順番に処理するということは、2単語目は1単語目が終わるまで計算できないということです。100 単語あれば 100 回の処理を直列に待たなければなりません。

現代の GPU は「大量の計算を同時に(並列に)こなす」のが得意なのに、RNN はその強みを活かせません。順番待ちが発生してしまうからです。

問題2:遠い単語を忘れる(長距離依存が苦手)

記憶を1単語ずつバケツリレーで受け渡していくと、遠く離れた単語の情報は薄れていきます

昨日、駅前の本屋で友達とばったり会って立ち話をして…(中略)…それはとても楽しかった」

「それ」が「昨日の出来事」を指すと分かるには、ずっと前の情報を保持し続ける必要があります。ところが RNN ではリレーの途中で情報がぼやけてしまい、長い文ほど関係をとらえそこねます。

2.2 Transformer の発想:順番に読むのをやめる

そこで Transformer は、発想を逆転させました。

1単語ずつ順番に読むのをやめて、文章全体を一度に見渡し、すべての単語ペアの関連を一気に計算する。

この転換が、上の2つの問題を同時に解決します。

  • 並列化できる:単語を順番待ちさせないので、全単語の計算を GPU でいっせいに走らせられる → 速い
  • 遠い単語と直接つながる:バケツリレーをやめ、どの単語ともダイレクトに関連度を計算する → どんなに離れていても1ステップでつながる

下の図は、RNN(順番に渡す)と Transformer(全単語を直接つなぐ)の情報の流れ方の違いです。

昨日本屋それ順番にリレー:遠い「昨日」は薄れる
① RNN:左から右へ1単語ずつ。並列化できず、遠い単語の情報は薄れる
昨日本屋それ全ペアを直接計算:「それ」と「昨日」も1ステップ
② Transformer:全単語どうしを直接つなぐ。並列化でき、遠い単語ともすぐつながる
論文タイトル "Attention Is All You Need" の意味

「必要なのは Attention だけだ」という、やや挑発的なタイトルです。それまで主流だった RNN の「順番に読む」仕組みを 完全に取り払い、後述する Attention だけで文章を処理してみせた——という宣言になっています。実際、それがうまくいったことが、その後の LLM 時代の幕開けになりました。

「全ペアを一気に」のコストには注意

全単語ペアの関連を計算するということは、単語数が nn なら n×nn \times n 個の組み合わせを計算する、ということです。文章が長くなると計算量が急増する(おおよそ n2n^2 に比例する)という弱点もあり、これは後の章で扱う「効率化」の大きなテーマになります。まずは「全ペアを一気に見るからこそ速くて賢い」という利点を押さえておけば十分です。


3. 一番の主役:Attention(注意機構)

ここからが Transformer の心臓部、Attention(アテンション/注意機構) です。名前のとおり「どの単語に注意を向けるか」を計算する仕組みで、Transformer の賢さはほぼここから来ています。

3.1 直感:その単語は、どの単語を見るべきか

次の文を読んでください。

「その動物は疲れていたので、それは道を渡らなかった」

「それ」は何を指しているでしょうか。もちろん「動物」ですね。人間は無意識に分かりますが、機械にこの判断をさせるのが Attention です。

Attention のアイデアはこうです。

各単語が、文中の全単語に対して「私はあなたにどれくらい注目すべき?」という関連度を計算し、関連度の高い単語の情報をより多く受け取る。

「それ」という単語を処理するとき、「動物」への関連度を高く、「疲れて」や「道」への関連度をそれなりに、関係ない単語への関連度を低く——というふうに 注目度の配分を決め、その配分にしたがって情報を集めてくる、というわけです。

3.2 Query・Key・Value:検索にたとえる

Attention では、各単語から 3つのベクトルを作ります。これがいちばんの肝です。図書館での「検索」にたとえると分かりやすいです。

名前役割検索のたとえ
Query(クエリ) qq「私はこういう情報を探している」という問い合わせ検索キーワード
Key(キー) kk「私はこういう特徴を持つ単語だよ」という見出し本の背表紙ラベル
Value(バリュー) vv実際に持っている中身の情報本の中身

この3つはどこから出てくるのか

「Query・Key・Value という3つのベクトルを作る」と言いましたが、ではこれらはどこから湧いてくるのでしょうか。答えは、各単語のベクトルに、それぞれ専用の変換をかけて作るです。

各単語はまず、前章で見たような数字の並び——埋め込みベクトル(embedding) xx——で表されています。この同じ xx に、3種類の重み行列 WQ,WK,WVW_Q, W_K, W_V をかけることで、3つのベクトルを取り出します。

埋め込みベクトル(embedding)とは

コンピュータは「猫」「犬」といった文字そのものを計算できません。そこで、各単語を 意味を表す数字の並び(ベクトル) に変換します。これが 埋め込みベクトル(embedding) です。

「猫」(0.21, 0.84, 0.05, , 0.63)\text{「猫」} \rightarrow (0.21,\ -0.84,\ 0.05,\ \dots,\ 0.63)

ベクトルの長さ(数字の個数)は、たとえば 512 や 768 のように決められた 次元で、1本が単語1個ぶんの「意味の座標」にあたります。

うれしいのは、意味が近い単語ほどベクトルも近い向きになるように学習される点です。「猫」と「犬」はどちらも動物なので近く、「猫」と「車」は遠くなります。だからこそ、Chapter 1 で学んだ「内積が大きい=似た方向=関連が強い」を使って、単語どうしの関連度を計算できるのです。Attention が動くのは、単語がこの埋め込みベクトルになっているおかげ、というわけです。

「単語」というバラバラな記号を、数字が連続する空間の中に "埋め込む(置く)" ——意味の似た単語が近所に集まる巨大な「意味の地図」に配置するイメージから、embedding と呼ばれます。

q=xWQ,k=xWK,v=xWVq = x W_Q, \quad k = x W_K, \quad v = x W_V
記号中身
xxその単語の埋め込みベクトル(Attention への入力)
WQ,WK,WVW_Q, W_K, W_VQuery・Key・Value 用の重み行列(変換の道具)

つまり、同じ1つの単語から、3つの違う「見方」で別々のベクトルを取り出しているわけです。これは Chapter 1 で学んだ「行列をかけることはベクトルを変換すること」がそのまま効いている箇所です。

そして大事なのは、WQ,WK,WVW_Q, W_K, W_V学習で決まるという点です。最初はでたらめな数字ですが、「どんな問い合わせ(Query)を投げ、どんな見出し(Key)を掲げれば、関連語をうまく見つけられるか」を、学習を通してモデル自身が獲得していきます。Query をどう作るか自体が、モデルの賢さの一部になっているのです。

Attention 1回の流れ

ここまでが準備です。実際の流れはこうです。ある単語の Query を、全単語の Key と照らし合わせて「相性(関連度)」を測り、相性の良い単語の Value をたくさん受け取る——これが Attention の1回の計算です。

そして「Query と Key の相性をどう測るか」こそ、前章で学んだ 内積です。

前章とのつながり:相性 = 内積

Chapter 1 で「内積が大きいほど、2つのベクトルは似た方向を向いている(=関連が強い)」と学びました。Attention はまさにこれを使います。Query と Key の内積を取れば、それがそのまま「その単語ペアの関連度スコア」になるのです。Attention は、内積による関連度の計算を全単語ペアにわたってくり返しているだけ、とも言えます。

3.3 3ステップで見る Attention

Query・Key・Value がそろえば、Attention は3ステップで計算できます。

  1. 関連度を測る(内積):注目元の単語の Query qq と、各単語の Key kk の内積を取る。値が大きいほど関連が強い。
  2. 注目度に変える(softmax):内積で出たスコアを softmax にかけ、すべて正・合計1の「注目度の配分(重み)」にする。
  3. 情報を集める(加重和):各単語の Value vv を、注目度を重みにして足し合わせる。注目した単語の情報が濃く混ざる。
Queryそれ各単語の Key動物疲れ内積9.12.00.5softmax注目度0.910.070.02出力ΣValue の加重和≈ 動物「それ」は「動物」へ最も強く注目し、その情報を濃く受け取る
Attention の3ステップ:内積で関連度 → softmax で注目度(合計1)→ Value の加重和

この図では、「それ」の Query が「動物」の Key と最も相性が良く(内積 9.1)、softmax を通すと注目度が 0.91 に集中します。結果として出力は「動物」の Value がほぼそのまま——つまり「それ = 動物」という対応を、計算だけで取り出せたことになります。

前章とのつながり:注目度 = softmax

ステップ2で使う softmax も Chapter 1 で学んだ関数です。内積で出たスコアはバラバラな大きさの数ですが、softmax を通すと すべて正・合計1 の確率分布になります。これがそのまま「どの単語にどれだけ注目するか」の配分になる、というのが美しいところです。

3.4 まとめて1つの式に:Scaled Dot-Product Attention

ここまでの3ステップを、全単語ぶんまとめて1つの式で書いたものが、Transformer の中核をなす式です。単語ごとの q,k,vq, k, v を縦に積んで行列 Q,K,VQ, K, V にまとめると、

Attention(Q,K,V)=softmax ⁣(QKdk)V\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^{\top}}{\sqrt{d_k}}\right) V

と書けます。記号が多いですが、やっていることは3ステップそのままです。

式の部分対応するステップ中身
QKQ K^{\top}① 関連度を測る全 Query と全 Key の 内積を一括計算(前章の行列積!)
    dk\dfrac{\;\cdot\;}{\sqrt{d_k}}(調整)値が大きくなりすぎないよう、Key の次元 dkd_k の平方根で割る
softmax()\text{softmax}(\cdots)② 注目度に変えるスコアを合計1の注目度へ
()V(\cdots)\,V③ 情報を集める注目度を重みに Value を加重和

注目してほしいのは QKQ K^{\top} の部分です。これは前章の「行列積は内積の集まり」がそのまま効いている箇所で、全単語ペアの内積(関連度)を1回の行列積でまとめて計算しています。RNN のような順番待ちなしに、全ペアを一気に求められる——これが「速い」の正体です。

なぜ dk\sqrt{d_k} で割るのか(スケーリング)

Query と Key の次元 dkd_k が大きいほど、内積はたくさんの項を足すので値が大きくなりがちです。値が大きすぎると softmax の出力が一箇所に極端に偏り、学習がうまく進まなくなります。そこで dk\sqrt{d_k} で割って大きさをならし、安定させています。この「割る」操作があるので Scaled(スケール済み)Dot-Product Attention と呼ばれます。詳しい理由は後の章でも触れます。

3.5 Self-Attention と Multi-Head Attention

最後に、よく出てくる2つの言葉だけ押さえておきましょう。詳しくは後の章で扱いますが、名前の意味が分かるだけで理解がぐっと楽になります。

  • Self-Attention(自己注意) … Query・Key・Value を すべて同じ文章(同じ単語列)から作る Attention のこと。文中の単語どうしが互いに注目し合い、「それ = 動物」のような 文の内部の関係をとらえます。Transformer の主役はこの Self-Attention です。
  • Multi-Head Attention(マルチヘッドアテンション) … Attention を 複数セット並列に行う仕組み。1つの「ヘッド(head)」だけだと1種類の見方しかできませんが、ヘッドを複数用意することで「文法的なつながりを見るヘッド」「意味的な近さを見るヘッド」のように 別々の観点から同時に注目でき、最後にそれらを統合します。
入力ヘッド1:文法を見るヘッド2:意味を見るヘッド3:…統合
Multi-Head Attention:複数のヘッドが別々の観点で注目し、最後に統合する
前半(概念)のまとめ
  • Transformer は「全単語ペアの注目度を一気に計算する」アーキテクチャで、いまの LLM の土台。
  • RNN の「遅い・遠い単語を忘れる」を、「順番に読むのをやめる」ことで解決した。
  • 心臓部の Attention は、内積で関連度 → softmax で注目度 → Value の加重和 という、前章の数学そのもの。
  • まとめると Attention(Q,K,V)=softmax ⁣(QKdk)V\text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)VQKQK^{\top} は全ペア内積の一括計算。

ここまでで、Attention が「内積 → softmax → 加重和」の3ステップでできていることが分かりました。ここからは後半(実装)パートです。その3ステップを、実際の数値で手を動かし、最後は Python(NumPy)で実装して確かめます。ゴールは1つ。アテンション機構の出力 oo を自分の手で計算できるようになることです。

ここからは実践パート

仕組みは前半でつかんだので、後半は「結局どう計算するのか/コードでどう書くのか」に集中します。手元で動かしたいときは Google Colab が手軽です(pip install numpy だけで動きます)。

4. 注目度から出力 oo までを計算する

4.1 注目度 aia_i を「計算する形」に書き下す

前半の 3 節では、クエリ qq とキー kik_i の内積を softmax に通したものが注目度になる、と見ました。計算に移る前に、これを1つの式として書いておきます。キーが全部で nn 個あるとして、スコアは内積 ai=qkia'_i = q\,k_i^{\top}、それを softmax に通した各成分を aia_i と書くと、

ai=exp(qki)jexp(qkj)a_i = \frac{\exp(q k_i^{\top})}{\sum_{j} \exp(q k_j^{\top})}

これが クエリ qq の、キー kik_i に対する注目度です。各 aia_iスカラ(ただの1つの数) で、0ai10 \le a_i \le 1、かつ a1+a2++an=1a_1 + a_2 + \dots + a_n = 1 を満たします。「どの位置をどれくらい重視するか」を表す割合だと思ってください。

なぜ kik_i を転置(\top)するの?

スコアの式に出てくる qkiq\,k_i^{\top} の「\top(転置)」は、横ベクトル同士を掛けて「1つの数(関連度スコア)」を取り出すために必要です。

このテキストでは qqkik_i横に数を並べた行ベクトルとして扱います。次元を d=3d=3 とすると、どちらも 1×31 \times 3 の行列です。

q=[q1 q2 q3],ki=[ki1 ki2 ki3]q = [\,q_1\ q_2\ q_3\,], \qquad k_i = [\,k_{i1}\ k_{i2}\ k_{i3}\,]

行列のかけ算には「左の列数と右の行数が一致しないと掛けられない」というルールがあります。qq1×31 \times 3、列数 3)と kik_i1×31 \times 3、行数 1)はそのままでは 313 \ne 1 で掛けられません。そこで kik_i を転置して 3×13 \times 1縦ベクトルにします。

ki=[ki1ki2ki3]k_i^{\top} = \begin{bmatrix} k_{i1} \\ k_{i2} \\ k_{i3} \end{bmatrix}

すると qkiq\,k_i^{\top}(1×3)×(3×1)=1×1(1 \times 3) \times (3 \times 1) = 1 \times 1、つまりスカラになり、これがちょうど内積(=関連度)になります。

qki=q1ki1+q2ki2+q3ki3q\,k_i^{\top} = q_1 k_{i1} + q_2 k_{i2} + q_3 k_{i3}

上で「各 aia_i はスカラ」と書けるのは、この \top のおかげです。前半まとめ(3.4 節)の QKQK^{\top} も同じ理屈で、こちらは「全クエリ × 全キー」の内積を行列で一括計算した版です。

4.2 バリューベクトル vv を用意する

出力を求めるには、前半 3.2 で出てきた バリューベクトル vv(実際に受け取る中身の情報)を、いよいよ式に登場させます。本来は Key とは別物ですが、計算を追いやすくするため、この章の範囲では、

バリューベクトル v=キーベクトル k\text{バリューベクトル } v = \text{キーベクトル } k

とします。つまり vi=kiv_i = k_i。いったんは「キーと同じものをもう一度使うだけ」と思っておけば十分です。

なぜ Value と Key を別々に用意するの?

理由は、「探すための見出し」と「実際に渡す中身」では役割がまったく違うからです。Attention の中で各単語の情報は、2つの別の場面で使われます。

  • Key kik_i:クエリと内積を取って関連度を測るための見出し(「これは関係ある?」)
  • Value viv_i:関連度が決まったあと、実際に受け取る中身(「で、結局どの情報をもらう?」)

前半 3.2 の図書館のたとえに乗せると、はっきりします。

  • Query = 探したいキーワード(「猫の飼い方」)
  • Key = 本の背表紙・タイトル(探すときにスキャンする部分)
  • Value = 本の中身そのもの(実際に読む情報)

「タイトル」と「中身」が別物であるように、kik_iviv_i も役割が違うので、一般には vikiv_i \ne k_i です。実際の Transformer では、同じ入力 xix_i から別々の重み行列で作ります。

ki=xiWK,vi=xiWVk_i = x_i W_K, \qquad v_i = x_i W_V

こうして分けておくと、モデルは「ある特徴でマッチさせて、でも渡すのは別の情報」という使い分けを学習でき、表現力が上がります。

ただし計算を最初に追う段階では、両者を同じにしておくと式が見通しやすくなります。そこでこの章では vi=kiv_i = k_i で進め、両者を分ける一般形は後の章で扱います。

4.3 出力 oo は「Value の加重和」

アテンション機構の出力 oo は、各バリューベクトル viv_i を注目度 aia_i で重みづけして、すべて足し合わせたものです。

o=a1v1+a2v2++anvno = a_1 v_1 + a_2 v_2 + \dots + a_n v_n

oovv と同じ次元数のベクトルになります(スカラ × ベクトルを足しているだけなので、次元は変わりません)。

やっていることを一言で言うと、これだけです。

ii 番目のバリューベクトルを、ii 番目の注目度倍して、全部足す。

  • 注目度が大きいもの = クエリと強く関連しているもの。
  • 注目度が大きい Value は oo に強く反映され、注目度がゼロに近い Value はほとんど無視される。

つまり oo とは、

クエリ qq が入力されたときの、qq に関連のある内容を強く反映した「新しいベクトル表現」

なのです。前半 3.3 の「それ → 動物」の例で言えば、出力 oo は「動物」の情報を濃く受け取ったベクトルになります。

注目度 × Value0.665 · v₁0.090 · v₂0.245 · v₃+出力 o≈ v₁ に近いベクトル
注目度が高い v₁ の成分が濃く混ざり、出力 o は v₁ に引き寄せられる(数値は次節 4.4 の計算例)
結局 oo ってなに? なぜこれがゴールなの?

oooutput(出力)の oo。アテンション機構が「内積 → softmax → 加重和」の3ステップの最後に吐き出す答えのベクトルで、この計算全体のゴールです。中身を3つの視点で整理すると:

  • 作り方:各 Value viv_i を注目度 aia_i で重みづけして全部足した「加重和」。注目度が大きい Value ほど濃く混ざり、ゼロに近い Value はほぼ無視される(要は関連度で重みづけした平均)。
  • 意味:クエリ qq関連する情報を濃く取り込んだ「新しいベクトル表現」。「それ → 動物」の例なら、oo は「動物」の情報を強く受け取ったベクトルになる。
  • 大きさaia_i はスカラ、viv_i はベクトルなので、oovv と同じ次元のベクトル。入力と同じサイズで出てくるから、そのまま次の層へ渡していける。

ひとことで言えば、oo は「qq の視点で文脈を読み直して作り直した、その単語の新しい意味ベクトル」です。この oo を積み重ねていくのが Transformer の本体になります。

4.4 ダミーの数値で計算してみる

言葉だけだとピンとこないので、小さな例で最後まで計算してみましょう。3 次元のベクトルで、キーは 3 個(n=3n=3)とします。vi=kiv_i = k_i でしたね。

q=(1, 0, 1),k1=(1, 0, 1),k2=(0, 1, 0),k3=(1, 1, 0)q = (1,\ 0,\ 1), \quad k_1 = (1,\ 0,\ 1), \quad k_2 = (0,\ 1,\ 0), \quad k_3 = (1,\ 1,\ 0)

ステップ① 内積でスコア ai=qkia'_i = q k_i^{\top} を出す

a1=11+00+11=2,a2=10+01+10=0,a3=11+01+10=1a'_1 = 1{\cdot}1 + 0{\cdot}0 + 1{\cdot}1 = 2, \quad a'_2 = 1{\cdot}0 + 0{\cdot}1 + 1{\cdot}0 = 0, \quad a'_3 = 1{\cdot}1 + 0{\cdot}1 + 1{\cdot}0 = 1

qq と向きがそっくりな k1k_1 がいちばん高いスコア(2)になりました。

ステップ② softmax で注目度 aia_i にする

exp(2)7.389, exp(0)=1, exp(1)2.718\exp(2) \approx 7.389,\ \exp(0) = 1,\ \exp(1) \approx 2.718、合計は 11.107\approx 11.107 なので、

a1=7.38911.1070.665,a2=1.00011.1070.090,a3=2.71811.1070.245a_1 = \frac{7.389}{11.107} \approx 0.665, \quad a_2 = \frac{1.000}{11.107} \approx 0.090, \quad a_3 = \frac{2.718}{11.107} \approx 0.245

合計はちゃんと 1(0.665+0.090+0.245=1.0000.665 + 0.090 + 0.245 = 1.000)になっています。

ステップ③ Value の加重和で出力 oo を作る

o=0.665(1,0,1)+0.090(0,1,0)+0.245(1,1,0)=(0.665+0+0.245,0+0.090+0.245,0.665+0+0)(0.910, 0.335, 0.665)\begin{aligned} o &= 0.665\,(1,0,1) + 0.090\,(0,1,0) + 0.245\,(1,1,0) \\ &= (0.665 + 0 + 0.245,\quad 0 + 0.090 + 0.245,\quad 0.665 + 0 + 0) \\ &\approx (0.910,\ 0.335,\ 0.665) \end{aligned}

できました! 出力 o(0.910, 0.335, 0.665)o \approx (0.910,\ 0.335,\ 0.665) は、最も注目された k1=(1,0,1)k_1 = (1,0,1) にかなり近いベクトルになっています。「qq に関連のある内容を強く反映した新しい表現」という説明が、数値でも確かめられました。

5. 行列で一気に書く

nn 個のキー・バリューを縦に積んで行列 K,VK, V にまとめると、ステップ①〜③の足し算をすべて 行列積1本ずつにまとめられます。

o=softmax(qK)Vo = \text{softmax}(q K^{\top})\, V

足し算の記号 \sum が消えているのがポイントです。qKqK^{\top} が全キーとの内積(スコア a1,,ana'_1, \dots, a'_n)を一括で計算し、最後の VV との積が加重和(iaivi\sum_i a_i v_i)をまとめて行ってくれます。

式の部分対応するステップ中身
qKq K^{\top}① スコア全キーとの内積をまとめて計算
softmax()\text{softmax}(\cdots)② 注目度スコアを合計1の注目度へ
()V(\cdots)\,V③ 出力注目度を重みに Value を加重和
前半の式とのちがい:dk\sqrt{d_k}

前半 3.4 で出てきた式 softmax ⁣(QKdk)V\text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V には dk\sqrt{d_k} で割る項がありました。ここではまず 割らない素朴な形で計算の流れをつかみ、7 節でこのスケーリングを足します。

これが「速い」の正体

ステップ①〜③をループで1個ずつ回す代わりに、行列積で一気に計算できる——これが前半で触れた「全単語ペアを並列に処理できる」という Transformer の強みそのものです。次の実装でも、ループ版と行列版の両方を書いて、結果が一致することを確かめます。

6. NumPy で実装する

それでは、5 節までの計算をそのままコードにします。Colab や手元の Python で動かしてみてください。

コード 2.1:入力ベクトルを用意する

import numpy as np

# クエリ(1個)とキー(3個)。本章では Value = Key とする。
q = np.array([1.0, 0.0, 1.0])
K = np.array([
[1.0, 0.0, 1.0], # k1
[0.0, 1.0, 0.0], # k2
[1.0, 1.0, 0.0], # k3
])
V = K # バリューベクトル = キーベクトル

コード 2.2:スコア a=qKa' = qK^{\top} を計算する

scores = q @ K.T # 内積をまとめて計算 → [2.0, 0.0, 1.0]
print(scores)

@ は行列積(ベクトルどうしなら内積)の演算子です。手計算の a1,a2,a3=2,0,1a'_1, a'_2, a'_3 = 2, 0, 1 と一致します。

コード 2.3:softmax で注目度にする

def softmax(x):
e = np.exp(x - np.max(x)) # オーバーフロー防止に最大値を引く
return e / e.sum()

a = softmax(scores) # → [0.665, 0.090, 0.245]
print(a, a.sum()) # 合計は 1.0
なぜ最大値を引くの?

数学的には softmax(x)\text{softmax}(x)softmax(xc)\text{softmax}(x - c) は完全に同じ値になります(分子・分母に同じ exp(c)\exp(-c) が掛かって打ち消し合うため)。一方コンピュータでは exp\exp が大きな入力で簡単にオーバーフローするので、c=max(x)c = \max(x) を引いて指数を 00 以下に抑えるのが定番のテクニックです。結果は変わらず、計算だけ安定します。

コード 2.4:出力 oo を加重和で求める

o = a @ V # softmax(qKᵀ)V を1行で
print(o) # → [0.910, 0.335, 0.665]

手計算で出した o(0.910, 0.335, 0.665)o \approx (0.910,\ 0.335,\ 0.665) とぴたり一致します🎉

ループ版とまとめて答え合わせ(クリックで開く)

行列版が正しいことを、定義どおりのループ版と比べて確認します。

# ループ版:o = Σ aᵢ vᵢ
o_loop = np.zeros_like(V[0])
for i in range(len(V)):
o_loop += a[i] * V[i]

print(np.allclose(o_loop, a @ V)) # True

np.allcloseTrue を返せば、ループ版と行列版の出力が(浮動小数点の誤差の範囲で)一致しているということです。

7. (発展)スケール化内積アテンション

最後に、前半 3.4 で出てきた dk\sqrt{d_k} で割るスケーリングを、ちゃんと理由づけしてから実装に足します。やること自体は「スコアをキーの次元数 dkd_k の平方根で割るだけ」です。

o=softmax ⁣(qKdk)Vo = \text{softmax}\!\left(\frac{q K^{\top}}{\sqrt{d_k}}\right) V

これが スケール化内積アテンション(Scaled Dot-Product Attention) で、実用上はこちらが標準形です。dkd_k はキー(とクエリ)ベクトルの次元数で、たとえば dk=64d_k = 64512512 のような値です。

なぜわざわざ割り算を1つ足すのか——順番に見ていきましょう。

7.1 なぜ割るのか:softmax の「偏りすぎ」を防ぐ

思い出してほしいのが softmax の性質です。softmax は入力のを指数関数で強調します。入力の値が全体的に大きくなると、いちばん大きい成分だけが極端に勝ち、出力がほぼ「1, 0, 0, …」に張り付いてしまいます。

これが起きると困ることが2つあります。

  • 注目がガチガチに偏る:本当は「動物に強め、疲れに少し」と配分したいのに、「動物にほぼ全部」になってしまい、他の単語の情報を拾えない。
  • 学習が進まなくなる:出力が 1 や 0 に張り付いた softmax は、入力が少し変わっても出力がほとんど動きません。これは「勾配がほぼ 0(傾きが消える)」状態で、モデルが「どっちに修正すればいいか」を受け取れず、学習が止まってしまいます。

下の図の左(①)は、スコア(注目元と注目先の相性)を横軸に、softmax が返す注目度を縦軸にとったものです。中央のスコアがほどよい範囲ではカーブに傾きがあり、スコアが少し動けば注目度もちゃんと動きます(=「どっちに修正すべきか」が伝わる=学習できる)。ところがスコアが大きい右側はカーブが平らに寝てしまい、スコアが動いても注目度はほぼ 1 のまま変わりません。この「平らな領域」が、注目が偏りきって勾配がほぼ 0 になる飽和ゾーンです。

右(②)は、この2つの動作点が実際に生む注目度の配分(棒グラフ)です。飽和側(赤)は1か所に全振りして他をいっさい拾えず、ほどよい側(青)は強弱をつけつつ他の単語の情報も残しています。

00.51スコア大:平ら勾配ほぼ0ほどよい:傾きあり学習できる÷√dₖ で引き戻すスコア(Q と K の相性)大きい →注目度
① 注目度のカーブ:スコアが大きいほど平らに飽和し、傾き(勾配)が消える。dk\sqrt{d_k} で割るのは、動作点を平らな右側から傾きのある中央へ引き戻す操作にあたる
スコア大(飽和)a₁a₂a₃1か所に全振りほどよいa₁a₂a₃強弱つき・他も残る
② 実際の注目度の配分:飽和(赤)は1か所に全振りして他を拾えない。ほどよい(青)は強弱をつけつつ他の情報も残る(具体的な数値は 7.3 で確認)

つまりスコアが大きくなりすぎると softmax が飽和して使いものにならない。だから スコアの大きさをほどよく抑えたい——これがスケーリングの目的です。後の 7.2 で見るように、その「ほどよい大きさ」に引き戻す割り算がちょうど dk\sqrt{d_k} になります。

7.2 なぜ「平方根」なのか:スコアの“ばらつき”をそろえる

ここがいちばんの肝です。「大きさを抑えたい」だけなら適当な定数で割ってもよさそうなのに、なぜちょうど dk\sqrt{d_k} なのでしょうか。理由は、スコアのばらつき(標準偏差)がちょうど dk\sqrt{d_k} に比例して大きくなるからです。

ざっくり次のように考えます。クエリ qq とキー kk の各成分が、平均 00・ばらつき(分散)11 くらいでバラバラな値だとします。スコアは内積なので、dkd_k 個の項の足し算です。

qk=q1k1+q2k2++qdkkdkq\,k^{\top} = q_1 k_1 + q_2 k_2 + \dots + q_{d_k} k_{d_k}

内積の計算じたいは、対応する成分どうしを掛けて、全部足すだけです。11 より大きい値が混じっても手順は変わりません。たとえば 33 次元(dk=3d_k=3)で

q=(2, 1, 3),k=(1, 2, 2)q = (2,\ 1,\ 3), \qquad k = (1,\ 2,\ 2)

なら、

qk=21+12+32=2+2+6=10q\,k^{\top} = 2\cdot 1 + 1\cdot 2 + 3\cdot 2 = 2 + 2 + 6 = 10

のように、位置ごとの積 21, 12, 322{\cdot}1,\ 1{\cdot}2,\ 3{\cdot}2 を足し合わせた 1つの数(スコア) になります。

さて、ここからは「次元 dkd_k が大きいほどスコアが大きくなる」様子を見やすくするため、各成分を +1+11-1 に単純化します(平均をちょうど 00、分散をちょうど 11 にそろえ、上の仮定にぴったり合わせるためです)。すると内積の各項 qikiq_i k_i+1+11-1 のどちらかになり、スコアは「±1\pm 1dkd_k 個足し合わせたもの」になります。

まず dk=4d_k = 4 で計算してみます。

位置 ii1234
クエリ qiq_i+1+1+1+11-1+1+1
キー kik_i+1+11-11-1+1+1
qikiq_i k_i+1+11-1+1+1+1+1

スコアは各項の合計なので、

qk=(+1)+(1)+(+1)+(+1)=2q\,k^{\top} = (+1) + (-1) + (+1) + (+1) = 2

目安どおり 4=2\sqrt{4}=2 くらいの大きさになりました。次に次元を dk=16d_k = 16 に増やすと、今度は ±1\pm 116 個足すことになります。たとえば +1+1 が 10 個・1-1 が 6 個そろえば、スコア =106=4= 10 - 6 = 4。やはり目安の 16=4\sqrt{16}=4 くらいです。

次元が 4164 \to 16 と 4 倍になると、スコアの大きさの目安は 242 \to 4 と 2 倍(=4\sqrt{4} 倍)に増えました。 次元が増えるほど足し合わさる項が増えて、スコアが大きく振れるわけです。

この背景にあるのが「独立な数をたくさん足すと、ばらつきは足した個数だけ大きくなる」という統計の性質です。各項のばらつきが 11 なら、dkd_k 個足したスコアの分散は dkd_k、その平方根である 標準偏差(ばらつきの目安)は dk\sqrt{d_k} になります。

コラム:分散と標準偏差ってなに?

どちらも「データが平均のまわりにどれくらい散らばっているか」を表す数です。同じ平均でも、ぎゅっと固まっているか、広く散っているかを1つの数で言い表したいときに使います。

たとえば 3 人のテストの点を考えます。平均はどれも 50 点ですが、散らばり方が違います。

  • {50, 50, 50}\{50,\ 50,\ 50\} … 全員ぴったり同じ。散らばり ゼロ
  • {40, 50, 60}\{40,\ 50,\ 60\} … 少し散らばっている。
  • {0, 50, 100}\{0,\ 50,\ 100\} … 大きく散らばっている。

これを数値にする手順はこうです。

  1. 各データの「平均からのズレ」を求める。
  2. ズレを 2乗 して、その平均をとる。これが 分散
  3. 分散の 平方根(ルート)をとる。これが 標準偏差

{0,50,100}\{0,50,100\} なら、ズレは 50, 0, +50-50,\ 0,\ +50。2乗して平均すると分散は 2500+0+250031667\frac{2500+0+2500}{3} \approx 1667、その平方根の標準偏差は 166741\sqrt{1667} \approx 41 です。「だいたい平均 ±41 点くらいに散らばっている」と読めます。

3つを数直線に並べると、ひと目で違いが分かります。平均(緑の線)はどれも 50 で同じなのに、散らばりを表す青い帯の幅(=標準偏差)が大きく変わります。

平均 50{50, 50, 50}×3(重なり)σ = 0{40, 50, 60}σ ≈ 8.2{0, 50, 100}σ ≈ 41050100
平均(緑)はどれも 50 で同じでも、散らばりを表す青い帯の幅=標準偏差 σ\sigma は大きく違う。分散はこの幅をさらに2乗した量にあたる
分散=(ズレ)2,標準偏差=分散\text{分散} = \overline{(\text{ズレ})^2}, \qquad \text{標準偏差} = \sqrt{\text{分散}}

では、2つは結局それぞれ何を表しているのでしょうか。同じ「散らばり具合」を表す仲間ですが、役割を分けるとこうです。

  • 分散 … 散らばりの大きさそのもの。ただしズレを 2乗 しているので、単位(スケール)が元と変わってしまう(点なら「点²」)。足し算と相性がよく、数式や理論で扱いやすいのが長所。
  • 標準偏差 … その分散を 平方根で元のスケールに戻した もの。「だいたい平均 ±これくらい散らばる」と、人間が実感しやすい目安になる。

ひとことで言えば、分散は計算向き・標準偏差は読み取り向き——同じ散らばりを別の単位で見ているだけ、と思えば十分です。なぜわざわざ分散を経由するかというと、次に出てくる「足すと散らばりがたまる」という性質が、標準偏差ではなく分散でこそきれいに成り立つからです。

このコラムで覚えてほしいのは1点だけ。独立なデータを足し合わせると、分散はそのまま足し算でたまっていく(これを分散の加法性といいます)。だから「11 の項を dkd_k 個足すと分散は dkd_k、標準偏差はその平方根の dk\sqrt{d_k}」という、本文の話につながるのです。

次元 dkd_kスコアの典型的な大きさ(標準偏差 dk\sqrt{d_k}
4422
646488
51251222.6\approx 22.6

つまり次元が大きいモデルほど、何もしなければスコアが自動的に大きくなり、softmax が飽和しやすくなります。そこでスコアを dk\sqrt{d_k} で割ると、ばらつきが 11 前後に戻り、dkd_k がいくつでもスコアの大きさが一定に保たれます。これがちょうど「平方根」を使う理由です。

コラム:なぜ「割る」とばらつきが 1 に戻るの?

カギは、データ全体を同じ数で割ると、散らばりの幅(標準偏差)も同じ数だけ縮むという性質です。数直線に散らばった点を、まるごと 18\frac{1}{8} に縮小コピーすると、点の間隔(=散らばりの幅)も 18\frac{1}{8} になる——それと同じイメージです。

平均(中心)割る前(dₖ = 64)-80+8幅 ≈ √dₖ = 8÷√dₖ(=÷8)で縮小割った後-1+1幅 = 1
dk\sqrt{d_k} で割ると、散らばりの幅が dk\sqrt{d_k} 倍ぶんだけ縮み、ちょうど 1 に揃う(縮小コピーのイメージ)

式で書くと(σ\sigma =標準偏差)こうなります。

σ ⁣(スコアdk)=σ(スコア)dk=dkdk=1\sigma\!\left(\frac{\text{スコア}}{\sqrt{d_k}}\right) = \frac{\sigma(\text{スコア})}{\sqrt{d_k}} = \frac{\sqrt{d_k}}{\sqrt{d_k}} = 1

割る前のばらつきが dk\sqrt{d_k}、それを dk\sqrt{d_k} で割るので、ちょうど 11 になります。

式だけだとイメージしにくいので、小さな数で実際に計算してみましょう。dk=4d_k = 4 のモデルで、あるクエリのスコアが2つ、2-2+2+2 だったとします(このとき dk=4=2\sqrt{d_k} = \sqrt{4} = 2)。スコアを dk=2\sqrt{d_k}=2 で割る前と後で、標準偏差を計算して比べます。

割る前のスコアdk=2\sqrt{d_k}=2 で割った後
スコア2, +2-2,\ +21, +1-1,\ +1
平均0000
平均からのズレ2, +2-2,\ +21, +1-1,\ +1
分散(ズレ²の平均)(2)2+222=4\dfrac{(-2)^2+2^2}{2}=4(1)2+122=1\dfrac{(-1)^2+1^2}{2}=1
標準偏差 σ\sigma4=2\sqrt{4}=21=1\sqrt{1}=1

割る前の標準偏差は 22(=dk\sqrt{d_k})、22 で割った後は 11。たしかに スコアを 22 で割ったら、標準偏差も 22 で割られて、ちょうど 11 になりました。式の dkdk=1\frac{\sqrt{d_k}}{\sqrt{d_k}}=1 は、これを文字に置きかえただけのものです。

次元が大きくても、まったく同じことが起きます。

  • dk=64d_k = 64:ばらつき 8864=8\sqrt{64}=8 で割る → 11
  • dk=512d_k = 512:ばらつき 22.6\approx 22.651222.6\sqrt{512}\approx 22.6 で割る → 1\approx 1

どの次元でも割ったあとのばらつきは 11 に揃います。だから次元 dkd_k を変えても、softmax に入るスコアの振れ幅はいつも同じ手頃なサイズ(およそ ±1\pm 1)に保たれ、7.1 で見た飽和を避けられるのです。

“分散が次元の数だけ足し合わさる”ってどういうこと?

サイコロ1個の出目はばらつきますが、100 個振って合計すると、合計値はもっと大きな幅でばらつきます(だいたい「個数の平方根」倍に広がる)。内積も同じで、独立な項 qikiq_i k_idkd_k 個足すほど、合計であるスコアの振れ幅が大きくなります。分散(ばらつきの2乗)が dkd_k 倍、標準偏差が dk\sqrt{d_k} 倍、というわけです。

なお、この「各成分が平均0・分散1で独立」という前提は厳密には成り立ちませんが、スケーリングの動機を理解するにはこの近似で十分です。元論文 "Attention Is All You Need" でも、この分散の議論を根拠に dk\sqrt{d_k} で割ることが提案されています。

7.3 スケーリングしないとどうなるか(数値で体感)

4.4 節で使ったスコア a=(2, 0, 1)a' = (2,\ 0,\ 1) は、実はすでに「ほどよい大きさ」になっていました。もし dk=64d_k = 64 のモデルでスケーリングをサボった場合、生のスコアは標準偏差 64=8\sqrt{64}=8 くらいまで膨らみ、たとえば (16, 0, 8)(16,\ 0,\ 8) のような値になります。同じ「向きの相性」なのに、大きさだけが 88 倍になったイメージです。

この2つを softmax に通すと、結果はまるで違います。

softmax(16, 0, 8)(1.000, 0.000, 0.000),softmax(2, 0, 1)(0.665, 0.090, 0.245)\text{softmax}(16,\ 0,\ 8) \approx (1.000,\ 0.000,\ 0.000), \qquad \text{softmax}(2,\ 0,\ 1) \approx (0.665,\ 0.090,\ 0.245)

スケーリングしない左は 1か所に全振り(飽和)してしまい、他の単語の情報をいっさい拾えません。dk=8\sqrt{d_k}=8 で割った右(= (2,0,1)(2,0,1))は、強弱をつけつつ他もちゃんと残した、ほどよい配分になっています。

1.0000.0000.000a₁a₂a₃スコア (16, 0, 8):飽和
① スケーリングなし:1か所に全振りし、他の情報を拾えない
0.6650.0900.245a₁a₂a₃÷√64 → (2, 0, 1):ほどよい
dk\sqrt{d_k} でスケーリング:強弱をつけつつ他の情報も残る

同じ「向きの相性」から出発しても、スケーリングの有無で注目度がここまで変わります。だから実用の Transformer は必ず dk\sqrt{d_k} で割る、というわけです。

7.4 実装

理由が分かったところで、コードにします。6 節にわり算を1つ足すだけです。Query が複数(行列 QQ)になっても、同じ式がそのまま動きます。

def softmax_rows(x):
e = np.exp(x - np.max(x, axis=-1, keepdims=True))
return e / e.sum(axis=-1, keepdims=True)

def scaled_dot_product_attention(Q, K, V):
d_k = K.shape[-1]
scores = Q @ K.T / np.sqrt(d_k) # スケール化したスコア
weights = softmax_rows(scores) # 行ごとに softmax
return weights @ V

# 複数クエリ(2個)でも同じ関数でOK
Q = np.array([[1.0, 0.0, 1.0],
[0.0, 1.0, 1.0]])
print(scaled_dot_product_attention(Q, K, V))

softmax_rows は「行ごとに合計1にする softmax」です。クエリが複数あるときは、各クエリ(各行)について独立に注目度を出す必要があるため、こうして行方向に正規化します。

8. (発展)マルチヘッドアテンション

前半 3.5 で Multi-Head Attention(マルチヘッドアテンション) の名前だけ出てきました。ここまでで Attention 1回ぶん——つまり「1つのヘッド(head)」の計算が、数値でもコードでもできるようになりました。実際の Transformer は、このヘッドを 複数並べて同時に走らせます。その仕組みを、これまでと同じ「直感 → 仕組み → 式 → 数値 → コード」の流れで見ていきましょう。

8.1 なぜ1つのヘッドでは足りないのか

思い出してほしいのが softmax の性質です。softmax は注目度を「合計1」に配分するので、1つのヘッドが強く注目できる相手は、せいぜい1〜数か所に限られます。言いかえると、1つのヘッドは「1種類の関係」をとらえるのが得意なのです。

ところが、言葉の関係は1種類ではありません。前半 3.1 の例文をもう一度見てみます。

「その動物は疲れていたので、それは道を渡らなかった」

「それ」という単語を理解するには、本当は複数の観点が同時に必要です。

  • 意味のつながり:「それ」が指すのは「動物」(共参照)。
  • 文法のつながり:「それ」は「渡らなかった」の主語で、「道」とも関係する。

これを1つのヘッドの、たった1つの注目配分に押し込めるのは無理があります。そこでマルチヘッドアテンションは、発想を変えます。

ヘッドを複数用意し、それぞれに別々の観点を担当させ、最後に全部の見方を合体させる。

下の図は、同じ文に対して2つのヘッドが 別々の単語に注目している様子です。ヘッド1は意味の観点から「動物」に、ヘッド2は文法の観点から「道」に強く注目しています。1つのヘッドだけでは片方しか拾えなかった関係を、ヘッドを分けることで 同時にとらえられるのです。

0.700.200.10動物疲れヘッド1:意味 →「動物」に注目
① ヘッド1(意味の観点):「それ」が指す「動物」に強く注目
0.100.250.65動物疲れヘッド2:文法 →「道」に注目
② ヘッド2(文法の観点):述語とつながる「道」に強く注目
1つのヘッド = これまで学んだ Attention そのもの

身構える必要はありません。1つのヘッドの中身は、7 節までに完成させたスケール化内積アテンションそのままです。マルチヘッドは「その計算をいくつも並べて、結果を合体させる」だけ。新しく増えるのは、後述する 連結(Concat)出力射影 WOW_O の2手順だけです。

8.2 仕組み:分ける → 別々に注目 → つなげる → 混ぜる

マルチヘッドアテンションは、次の4ステップでできています。ヘッドの数を hh とします。

  1. 分ける(射影):各ヘッドは専用の重み行列 WQ(i),WK(i),WV(i)W_Q^{(i)}, W_K^{(i)}, W_V^{(i)} を持ち、入力を 小さな次元 dkd_k の Query・Key・Value に変換する。ヘッドごとに違う変換なので、ヘッドごとに違う「見方」が生まれる。
  2. 別々に注目(Attention):各ヘッドで独立に、7 節のスケール化内積アテンションを計算する。出力はヘッドごとに1本のベクトル headi\text{head}_i
  3. つなげる(連結 / Concat):全ヘッドの出力 head1,,headh\text{head}_1, \dots, \text{head}_h横に1本につなげる。これで次元が元の大きさ(dmodeld_{\text{model}})に戻る。
  4. 混ぜる(出力射影):連結したベクトルに重み行列 WOW_O をかけ、ヘッドをまたいだ情報を混ぜ合わせて最終出力にする。
入力 x射影で分けるヘッド1:Attentionヘッド2:Attentionヘッド3:…連結W_O出力つなげて元の次元へヘッドを混ぜる各ヘッドは別々の観点で並列に注目し、最後に連結+W_O で1つに統合する
マルチヘッドアテンションの流れ:分ける → ヘッドごとに注目 → 連結 → WOW_O で混ぜる

8.3 式で書く

4ステップを式にすると、こうなります。

MultiHead(Q,K,V)=Concat(head1, , headh)WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1,\ \dots,\ \text{head}_h)\, W_O

そして、各ヘッド headi\text{head}_i の中身は、7 節のスケール化内積アテンションそのものです。

headi=Attention(QWQ(i), KWK(i), VWV(i))\text{head}_i = \text{Attention}(Q W_Q^{(i)},\ K W_K^{(i)},\ V W_V^{(i)})
記号中身
hhヘッドの数(たとえば 8)
dmodeld_{\text{model}}入力・出力ベクトルの次元(たとえば 512)
WQ(i),WK(i),WV(i)W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}ヘッド ii 専用の射影行列。入力を次元 dkd_k の小さな空間へ変換する
dkd_k各ヘッドの Query・Key(と Value)の次元。ふつう dk=dmodel/hd_k = d_{\text{model}} / h
headi\text{head}_iヘッド ii の Attention 出力(次元 dkd_k
WOW_O連結結果を混ぜて元の次元 dmodeld_{\text{model}} に戻す出力射影行列

ここで効いているのが dk=dmodel/hd_k = d_{\text{model}} / h という割り当てです。たとえば dmodel=512d_{\text{model}} = 512h=8h = 8 なら、各ヘッドは dk=64d_k = 64 次元という 小さな空間で計算します。

ヘッドを増やしても計算量はほぼ変わらない(うれしい設計)

「ヘッドを8個も並べたら8倍重いのでは?」と思うかもしれませんが、そうはなりません。各ヘッドの次元を dk=dmodel/hd_k = d_{\text{model}} / h小さくしているからです。

  • 1つの大きなヘッド(次元 dmodeld_{\text{model}})で計算する場合と、
  • hh 個の小さなヘッド(各 dk=dmodel/hd_k = d_{\text{model}}/h)に分けて計算する場合

で、足し合わせた計算量はだいたい同じになります。つまりマルチヘッドは、同じコストのまま「1つの広い視野」を「複数の専門的な視野」に分け直しているわけです。コストを増やさずに表現力だけ上げられる、よくできた設計です。

なぜ最後に WOW_O が必要なの?

連結しただけの状態は、各ヘッドの出力をただ横に並べただけで、ヘッドどうしの情報がまだ混ざっていません。「ヘッド1は動物に、ヘッド2は道に注目した」という別々の結果が、隣り合って置いてあるだけです。

そこで WOW_O をかけることで、ヘッドをまたいで情報を混ぜ合わせ、「複数の観点を統合した1つの表現」に仕上げます。WOW_O も学習で決まるので、どのヘッドの情報をどう組み合わせると役に立つかを、モデル自身が獲得していきます。

8.4 数値で見る:2つのヘッドをつないでみる

小さな例で、ステップ②〜④(注目 → 連結 → 混ぜる)を追ってみましょう。ヘッドは2つ(h=2h=2)、各ヘッドの Value は2次元(dk=2d_k=2)とします。8.1 の図のとおり、2つのヘッドは別々の単語に注目しているとします。

ヘッド1(意味のヘッド)は「動物」に強く注目し、注目度 a(1)=(0.7, 0.2, 0.1)a^{(1)} = (0.7,\ 0.2,\ 0.1)。3つの Value を v1(1)=(1,0), v2(1)=(0,1), v3(1)=(1,1)v^{(1)}_1=(1,0),\ v^{(1)}_2=(0,1),\ v^{(1)}_3=(1,1) とすると、

head1=0.7(1,0)+0.2(0,1)+0.1(1,1)=(0.8, 0.3)\text{head}_1 = 0.7\,(1,0) + 0.2\,(0,1) + 0.1\,(1,1) = (0.8,\ 0.3)

ヘッド2(文法のヘッド)は「道」に強く注目し、注目度 a(2)=(0.1, 0.2, 0.7)a^{(2)} = (0.1,\ 0.2,\ 0.7)。Value は別の射影なので v1(2)=(0,1), v2(2)=(1,0), v3(2)=(1,1)v^{(2)}_1=(0,1),\ v^{(2)}_2=(1,0),\ v^{(2)}_3=(1,1) だとすると、

head2=0.1(0,1)+0.2(1,0)+0.7(1,1)=(0.9, 0.8)\text{head}_2 = 0.1\,(0,1) + 0.2\,(1,0) + 0.7\,(1,1) = (0.9,\ 0.8)

ステップ③ 連結(Concat):2つのヘッドの出力を、横に1本につなげます。

Concat(head1,head2)=(0.8, 0.3ヘッド1, 0.9, 0.8ヘッド2)=(0.8, 0.3, 0.9, 0.8)\text{Concat}(\text{head}_1, \text{head}_2) = (\underbrace{0.8,\ 0.3}_{\text{ヘッド1}},\ \underbrace{0.9,\ 0.8}_{\text{ヘッド2}}) = (0.8,\ 0.3,\ 0.9,\ 0.8)

2次元 × 2ヘッド = 4次元になり、これが dmodel=4d_{\text{model}} = 4 にあたります。意味の観点(ヘッド1)と文法の観点(ヘッド2)が、1本のベクトルの中に共存しているのがポイントです。

ヘッド10.80.3ヘッド20.90.8連結(4次元 = d_model)W_Oヘッドをまたいで混ぜる最終出力d_model 次元2つの観点を1本にまとめ、W_O で混ぜて次の層へ渡せる形にする
連結で各ヘッドの出力を1本にし、WOW_O で混ぜて最終出力(dmodeld_{\text{model}} 次元)にする

ステップ④ 出力射影:最後に連結ベクトルに WOW_O をかけ、ヘッドをまたいで情報を混ぜます。WOW_O は学習で決まる行列なので、ここでは具体的な数値計算は次のコードに譲りますが、やっていることは「4次元ベクトルに行列をかけて4次元ベクトルにする」だけ——前章で学んだ線形変換そのものです。

8.5 NumPy で実装する

7.4 で作った scaled_dot_product_attention を、そのまま部品として使い回せます。各ヘッドでそれを呼び、出力を連結し、最後に WOW_O をかけるだけです。

# 入力:3トークン、各 d_model = 4 次元
X = np.array([
[1.0, 0.0, 1.0, 0.0],
[0.0, 1.0, 0.0, 1.0],
[1.0, 1.0, 0.0, 0.0],
])

d_model, h = 4, 2
d_k = d_model // h # 各ヘッドの次元 = 2

# 各ヘッド専用の射影行列。本来は学習で決まるが、ここでは例として固定の乱数を使う
rng = np.random.default_rng(0)
W_Q = rng.normal(size=(h, d_model, d_k)) # ヘッドごとの Query 用
W_K = rng.normal(size=(h, d_model, d_k)) # ヘッドごとの Key 用
W_V = rng.normal(size=(h, d_model, d_k)) # ヘッドごとの Value 用
W_O = rng.normal(size=(h * d_k, d_model)) # 連結結果を混ぜる出力射影

# ① 各ヘッドで別々に Attention を計算(中身は 7.4 の関数そのまま)
heads = []
for i in range(h):
Q = X @ W_Q[i] # (3, d_k) … ヘッド i 用に射影
K = X @ W_K[i]
V = X @ W_V[i]
heads.append(scaled_dot_product_attention(Q, K, V)) # (3, d_k)

# ② 連結 → ③ W_O で混ぜる
concat = np.concatenate(heads, axis=-1) # (3, h*d_k) = (3, 4)
output = concat @ W_O # (3, d_model) = (3, 4)

print(output.shape) # (3, 4)

ポイントは2つです。1つめは、各ヘッドの中身は 7.4 の scaled_dot_product_attention をそのまま呼んでいるだけだということ。マルチヘッドで新しく増えたのは、ヘッドごとの射影(X @ W_Q[i] など)と、連結(np.concatenate)、出力射影(@ W_O)の3点だけです。

2つめは、出力の形が入力と同じ (3, 4) になっていること。入力 dmodeld_{\text{model}} 次元 → 出力 dmodeld_{\text{model}} 次元、と形が保たれるので、この出力をそのまま次の層の入力にできます。この「同じ形のブロックを積み重ねる」のが Transformer の本体で、いまの LLM はこのブロックを何十段も重ねてできています。

前半 3.5 とのつながり

3.5 で「複数のヘッドが別々の観点で注目し、最後に統合する」と一言で説明したものが、ここで式とコードになりました。実際の Transformer では、Query・Key・Value をすべて同じ文章から作る Self-Attention(3.5)を、この マルチヘッドにした「Multi-Head Self-Attention」が基本ブロックとして使われています。

9. (発展)PyTorch で実装する

ここまでは仕組みを理解するために NumPy で書いてきました。最後に、実際の LLM 開発で標準的に使われる PyTorch でも、同じ3つ——内積アテンション → スケール化内積アテンション → マルチヘッドアテンション——を書いてみましょう。コードはどれも NumPy 版とほとんど同じ数行です。

なぜ NumPy ではなく PyTorch なのか

NumPy でも計算自体はできますが、実際にモデルを 学習させるには PyTorch(や JAX・TensorFlow)が要ります。理由は2つです。

  • 自動微分(autograd):学習では「重み WQ,WK,W_Q, W_K, \dots をどちらに動かせば良くなるか(勾配)」を計算する必要があります。PyTorch はこの微分を自動でやってくれます(NumPy にはこの機能がありません)。
  • GPU で速い.to("cuda") でテンソルを GPU に載せるだけで、巨大な行列積が一気に並列計算されます。

書き味は NumPy とそっくりです。np.arraytorch.tensor に、@(行列積)はそのまま。最初は「NumPy の置きかえ」くらいの気持ちで読んで大丈夫です。手元で動かすなら Google Colab なら PyTorch がはじめから入っています(ローカルなら pip install torch)。

9.1 内積アテンション(スケーリングなし)

まずは 6 節の素朴な内積アテンション o=softmax(qK)Vo = \text{softmax}(qK^{\top})V を PyTorch で書きます。数値も 4.4 節と同じものを使うので、答えが一致するか確かめられます。

import torch
import torch.nn.functional as F

# 4.4 節と同じ数値。np.array の代わりに torch.tensor を使うだけ
q = torch.tensor([1.0, 0.0, 1.0])
K = torch.tensor([
[1.0, 0.0, 1.0], # k1
[0.0, 1.0, 0.0], # k2
[1.0, 1.0, 0.0], # k3
])
V = K # 本章では Value = Key

scores = q @ K.T # ① 内積でスコア → tensor([2., 0., 1.])
weights = F.softmax(scores, dim=-1) # ② softmax で注目度(合計1)
o = weights @ V # ③ Value の加重和
print(o) # tensor([0.9100, 0.3350, 0.6650])

NumPy 版との違いは、たった3点です。

やることNumPyPyTorch
データの入れ物np.array(...)torch.tensor(...)
softmax自前で def softmax を書いた組み込みの F.softmax を呼ぶだけ
軸の指定axis=-1dim=-1

ポイントは F.softmax(scores, dim=-1)dim=-1 です。「いちばん最後の軸(次元)に沿って合計1にする」という指定で、NumPy の axis=-1 と同じ意味です。スコアが1本のベクトルなら最後の軸はその要素方向なので、[2, 0, 1] 全体が合計1の注目度 [0.665, 0.090, 0.245] に変換され、出力 o は手計算の (0.910, 0.335, 0.665)(0.910,\ 0.335,\ 0.665) とぴたり一致します🎉

torch.tensor ってなに?

PyTorch でのデータの入れ物が テンソル(tensor) です。中身は NumPy 配列とほぼ同じ「多次元の数の並び」ですが、上で触れた 自動微分GPU 実行に対応している点が違います。q @ K.T(行列積)、.shape(形の確認)、.T(転置)など、NumPy で使った操作の多くがそのまま使えます。

別の書き方:nn.Module クラス+バッチ対応

上のコードは「1文・1クエリ」を最短で計算する書き方でした。実際の本やライブラリでは、次のように nn.Module を継承したクラスとして書き、複数の文をまとめて(バッチで)処理できる形にするのが一般的です。9.3 のマルチヘッド実装ともそのままつながる書き方なので、ここで紹介しておきます。

import torch
from torch import Tensor, nn


class DotProductAttention(nn.Module):
def __init__(self):
super().__init__()

def forward(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor:
"""内積アテンションの計算を行う.

Args:
query (Tensor): クエリ. shape は (batch_size, query_len, d_model).
key (Tensor): キー. shape は (batch_size, key_len, d_model).
value (Tensor): バリュー. shape は (batch_size, value_len, d_model).
"""
# 1. query と key から (batch_size, query_len, key_len) のスコアを計算
score = torch.bmm(query, key.transpose(1, 2))
# 2. 重みの和が 1 になるように softmax を計算
weight = torch.softmax(score, dim=-1)
# 3. value の重み付き和を計算
output = torch.bmm(weight, value)
return output

使うときは、まずインスタンスを作って、そこに Query・Key・Value を渡します。

attn = DotProductAttention()
# バッチ2文・各3トークン・d_model=4 のダミー入力
query = torch.randn(2, 3, 4)
key = torch.randn(2, 3, 4)
value = torch.randn(2, 3, 4)
output = attn(query, key, value) # forward が呼ばれる
print(output.shape) # torch.Size([2, 3, 4])

中身は 9.1 の3ステップ(内積 → softmax → 加重和)とまったく同じです。新しく出てくるのは「バッチ」と「torch.bmm」の2つだけなので、そこを押さえれば読めます。

「バッチ」次元 = 何文かをまとめて処理する

9.1 の最初のコードは入力が1文ぶん(2次元)でしたが、こちらは形が (batch_size, 系列長, d_model)3次元になっています。先頭の batch_sizeバッチ——「いちどに処理する文の本数」です。

学習では、1文ずつ計算するより 何十・何百文をまとめて計算したほうが GPU を活かせて圧倒的に速くなります。そのため実用のコードは、ほぼ必ずこのバッチ次元を先頭に付けた形で書かれます。「文が1本増えても、各文に対して同じ計算をするだけ」なので、考え方は 9.1 と変わりません。

torch.bmm(バッチ行列積)とは

torch.bmmbatch matrix multiply(バッチ行列積) の略で、「バッチの中の各文について、行列積を独立にまとめて行う」関数です。形でいうと、

(batch, a, b)  ×  (batch, b, c)    (batch, a, c)(\text{batch},\ a,\ b) \;\times\; (\text{batch},\ b,\ c) \;\rightarrow\; (\text{batch},\ a,\ c)

のように、先頭のバッチ次元はそのまま、後ろの2軸だけで行列積をします。コードの1行目では

  • query(batch, query_len, d_model)
  • key.transpose(1, 2)(batch, d_model, key_len)(後述)

bmm して、(batch, query_len, key_len)スコア行列——全クエリ × 全キーの内積(QKQK^{\top})——を文ごとに一括計算しています。3行目も同様に、注目度 (batch, query_len, key_len) と value (batch, key_len, d_model)bmm して、加重和の出力 (batch, query_len, d_model) を得ます。

9.1 で使った @torch.matmul)でも同じ計算はできます。bmm は「入力がちょうど3次元(バッチ+行列)であることを前提にした専用版」だと思ってください。

key.transpose(1, 2) は、9.2 で出てくる transpose(-2, -1) と同じく 最後の2軸の入れ替えです。key(batch_size, key_len, d_model) の3次元なので、軸 1(key_len)と軸 2(d_model)を入れ替えて (batch_size, d_model, key_len) にし、QKQK^{\top} の「KK^{\top}」を作っています。softmax(score, dim=-1)dim=-1 は最後の軸(key_len)方向の正規化で、9.1 と同じく「各クエリについて、全キーへの注目度を合計1にする」という意味です。

なぜ 1 文版と 2 通り見せるのか

最初の q @ K.T 版は 3ステップの本質を最短で見るため、こちらのクラス版は 実際のコードの形(バッチ・nn.Module に慣れるためのものです。中身は同じなので、「短い版で仕組みを理解 → クラス版で実物に橋渡し」という順で読むと、9.3 のマルチヘッド実装(nn.Module +バッチ+ヘッド軸)がすっと入ってきます。

9.2 スケール化内積アテンション

次に、7 節で足した dk\sqrt{d_k} のスケーリングを入れます。やることは「スコアを dk\sqrt{d_k} で割る」だけです。

import math

def scaled_dot_product_attention(Q, K, V):
d_k = K.shape[-1] # Key の次元
scores = Q @ K.transpose(-2, -1) / math.sqrt(d_k) # スケール化したスコア
weights = F.softmax(scores, dim=-1) # 行ごとに softmax
return weights @ V

# クエリが複数(2個)でも同じ関数でOK
Q = torch.tensor([[1.0, 0.0, 1.0],
[0.0, 1.0, 1.0]])
print(scaled_dot_product_attention(Q, K, V))

7.4 の NumPy 版とほぼ同じですが、転置を K.T ではなく K.transpose(-2, -1) で書いている点に注目してください。

なぜ .T ではなく .transpose(-2, -1) なのか

.T は2次元(行列)の転置には使えますが、後の 9.3 で出てくるような 3次元以上のテンソル(バッチやヘッドの軸が付いたもの)ではうまくいきません。transpose(-2, -1) は「最後の2つの軸だけを入れ替える」という指定で、前にバッチやヘッドの軸がいくつ付いていても、Key の「トークン × 次元」の部分だけを正しく転置できます。最初からこの書き方に慣れておくと、マルチヘッドでもそのまま通用します。

別の書き方:nn.Module クラス版(スケール化)

9.1 の DotProductAttention と同じ要領で、スケール化版も nn.Module クラス+バッチ対応で書けます。9.1 のクラスに スケーリングの1行を足すだけです。

class ScaledDotProductAttention(nn.Module):
def __init__(self):
super().__init__()

def forward(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor:
"""スケール化内積アテンションの計算を行う.

Args:
query (Tensor): クエリ. shape は (batch_size, query_len, d_model).
key (Tensor): キー. shape は (batch_size, key_len, d_model).
value (Tensor): バリュー. shape は (batch_size, value_len, d_model).
"""
# query の次元(= キーの次元)でスケーリング
d_k = query.size(-1)
score = torch.bmm(query, key.transpose(1, 2)) / (d_k ** 0.5)
# 注目度の計算
weight = torch.softmax(score, dim=-1)
# アテンション出力(Value の加重和)の計算
output = torch.bmm(weight, value)
return output

9.1 のクラスと見比べると、増えたのは d_k = query.size(-1) と、スコアを / (d_k ** 0.5) で割る部分だけ。これがまさに 7 節の QKdk\dfrac{QK^{\top}}{\sqrt{d_k}} です。バッチや torch.bmm の意味は 9.1 のコラムのとおりで、変わりません。

query.size(-1)d_k ** 0.5 という書き方

同じことを表す書き方が PyTorch には複数あります。どれを使っても結果は同じなので、本やコードによって見た目が違っても戸惑わなくて大丈夫です。

  • 次元数の取り方query.size(-1) は「最後の軸の大きさ(= dkd_k)」を返します。NumPy ふうの query.shape[-1] と同じ意味です。
  • 平方根d_k ** 0.5 は「dkd_k の 0.5 乗 = dk\sqrt{d_k}」。math.sqrt(d_k)torch.sqrt(...) と同じです。

実は PyTorch 2.0 以降には、これと同じ計算をしてくれる 組み込み関数が用意されています。実用ではこちらを使うのが定番です。

# 自前実装と同じ結果。内部でスケーリング・softmax・加重和をまとめて行う
o = F.scaled_dot_product_attention(Q, K, V)
print(o)
組み込みの F.scaled_dot_product_attention が速い理由

この組み込み関数は、スコア計算・dk\sqrt{d_k} のスケーリング・softmax・Value の加重和を 1つの最適化された処理(FlashAttention などの高速実装)にまとめて実行します。さらに、後の章で出てくる「未来の単語を見ないようにする マスク」や、過学習を防ぐ dropout もオプションで指定できます。仕組みを理解した今なら、この関数が中で何をしているかが分かるはずです——まさに 9.1〜9.2 でやったことそのものです。

9.3 マルチヘッドアテンション

最後に、8 節のマルチヘッドアテンションを PyTorch で書きます。ここでは、学習可能な層をまとめる nn.Module という仕組みを使い、再利用できる部品(クラス)として作ります。

import torch.nn as nn

class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super().__init__()
assert d_model % num_heads == 0 # 割り切れることが前提
self.num_heads = num_heads
self.d_k = d_model // num_heads # 各ヘッドの次元

# Q/K/V/出力の射影。それぞれ「重み行列を持った線形変換」
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model) # 出力射影 W_O

def forward(self, x):
B, T, _ = x.shape # (バッチ, トークン数, d_model)

# ① 射影して、(B, num_heads, T, d_k) の形にヘッド分割
q = self.W_q(x).view(B, T, self.num_heads, self.d_k).transpose(1, 2)
k = self.W_k(x).view(B, T, self.num_heads, self.d_k).transpose(1, 2)
v = self.W_v(x).view(B, T, self.num_heads, self.d_k).transpose(1, 2)

# ② 全ヘッドで一括して Scaled Dot-Product Attention
out = F.scaled_dot_product_attention(q, k, v) # (B, num_heads, T, d_k)

# ③ ヘッドを連結して (B, T, d_model) に戻す
out = out.transpose(1, 2).contiguous().view(B, T, -1)

# ④ 出力射影で、ヘッドをまたいで混ぜる
return self.W_o(out)

# d_model=4・ヘッド2個。入力は バッチ1・3トークン・4次元
mha = MultiHeadAttention(d_model=4, num_heads=2)
x = torch.randn(1, 3, 4)
print(mha(x).shape) # torch.Size([1, 3, 4])

NumPy 版(8.5)と見た目はだいぶ違いますが、やっていることは 8.2 の4ステップそのままです。対応を表で並べます。

8.2 のステップコードの該当部分何をしているか
① 分ける(射影)self.W_q(x).view(...).transpose(1, 2)射影してから、ヘッドごとの軸に分割
② 別々に注目F.scaled_dot_product_attention(q, k, v)全ヘッドぶんの Attention を一括計算
③ つなげる(連結).transpose(1, 2).contiguous().view(B, T, -1)ヘッドの軸を畳んで1本に戻す
④ 混ぜる(出力射影)self.W_o(out)WOW_O をかけてヘッドをまたいで混ぜる

NumPy 版と比べたときの、新しくて分かりにくいところを3つ補足します。

nn.Linear は「重み行列をかける」の PyTorch 版

nn.Linear(d_model, d_model) は、入力に重み行列をかけて変換する層です(前章の線形変換そのもの)。NumPy 版では W_Q[i] をヘッドごとに用意しましたが、PyTorch では nn.Linear 1枚で 全ヘッドぶんの射影をまとめて行い、あとから view でヘッドに切り分けます(こうすると行列積が1回で済んで速い)。nn.Linear が持つ重みは、学習で自動的に更新されます。

viewtranspose でヘッドに分ける

射影直後の形は (B, T, d_model) です。これを view(B, T, num_heads, d_k) で「d_modelnum_heads × d_k に割り直し」、さらに transpose(1, 2) でヘッドの軸を前に出して (B, num_heads, T, d_k) にします。こうすると、F.scaled_dot_product_attentionヘッドごとに独立して(最後の2軸 T × d_k で)Attention を計算してくれます。-1 を使った view(B, T, -1) は「残りの次元はよしなに埋めて」の意味で、ここでは num_heads × d_k = d_model に戻ります。

連結の前に .contiguous() が必要なわけ

transpose は、データを実際に並べ替えるのではなく「見え方(軸の順番)だけ」を変えます。そのため直後に view で形を変えようとすると、メモリの並びが連続していなくてエラーになることがあります。.contiguous() は「いまの見え方どおりにメモリを並べ直す」操作で、これを挟むと view が安全に通ります。お決まりの1手と覚えておけば十分です。

ここでも、出力の形は入力と同じ (1, 3, 4)(バッチ・トークン数・dmodeld_{\text{model}})です。8.5 の NumPy 版と同じく「同じ形のブロック」になっているので、何段も積み重ねられます。

別の書き方:ヘッドをクラスに分けて nn.ModuleList で束ねる

上の MultiHeadAttention は、速さを優先して 全ヘッドを1枚の nn.Linear に詰め込み、viewtranspose でヘッドに切り分ける書き方でした。一方、本やライブラリでは、「1ヘッド=1つのクラス」を作り、それを複数束ねるという、より素直な書き方もよく使われます。8.5 の NumPy 版(for ループで各ヘッドを計算 → np.concatenate)に、いちばん近いのはこちらです。

まず、1つぶんのヘッドAttentionHead クラスにします。中身は「Q・K・V を小さな次元へ射影 → 9.2 の ScaledDotProductAttention に渡す」だけです。

class AttentionHead(nn.Module):
def __init__(self, d_k: int, d_v: int, d_model: int) -> None:
"""MultiHeadAttention のヘッド.

Args:
d_k (int): クエリ・キーの次元数
d_v (int): バリューの次元数
d_model (int): モデルの埋め込み次元数
"""
super().__init__()
# Q・K・V を、このヘッド専用の部分空間へ射影する線形層
self.linear_q = nn.Linear(d_model, d_k)
self.linear_k = nn.Linear(d_model, d_k)
self.linear_v = nn.Linear(d_model, d_v)
self.attention = ScaledDotProductAttention() # 9.2 のクラスを再利用

def forward(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor:
"""単一ヘッドのアテンションを計算する."""
query = self.linear_q(query) # (batch, query_len, d_k)
key = self.linear_k(key) # (batch, key_len, d_k)
value = self.linear_v(value) # (batch, value_len, d_v)
output = self.attention(query, key, value)
return output

次に、この AttentionHeadn_heads 個ならべて束ね、出力を連結して linear_o で混ぜれば、マルチヘッドの完成です。

class MultiHeadAttention(nn.Module):
def __init__(self, n_heads: int, d_k: int, d_v: int, d_model: int) -> None:
super().__init__()
# ヘッドを n_heads 個ぶんリストにして保持
self.heads = nn.ModuleList(
[AttentionHead(d_k, d_v, d_model) for _ in range(n_heads)]
)
# 連結した出力を混ぜる線形層(W_O)
self.linear_o = nn.Linear(n_heads * d_v, d_model)

def forward(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor:
# ① 各ヘッドで別々にアテンションを計算(8.5 の for ループそのもの)
head_out = [head(query, key, value) for head in self.heads]
# ② ヘッドを最後の軸で連結(Concat)
head_out = torch.cat(head_out, dim=-1) # (batch, query_len, n_heads*d_v)
# ③ W_O で混ぜて出力(batch, query_len, d_model)
output = self.linear_o(head_out)
return output

8.2 の4ステップとの対応は、こちらのほうが見た目もそのままです。

8.2 のステップコードの該当部分
① 分ける(射影)AttentionHead 内の linear_q/k/v
② 別々に注目[head(...) for head in self.heads](ヘッドごとに計算)
③ つなげる(連結)torch.cat(head_out, dim=-1)
④ 混ぜる(出力射影)self.linear_o(head_out)
nn.ModuleListtorch.cat ってなに?
  • nn.ModuleList:ただの Python リストではなく、PyTorch が「学習対象の部品」として認識してくれるリストです。これで束ねておくと、中の全ヘッドの重み(各 linear_q/k/v)が学習でちゃんと更新されます。普通の [...] に入れると追跡されないので、サブモジュールを並べるときは nn.ModuleList を使うのがお約束です。
  • torch.cat(..., dim=-1):テンソルを 最後の軸方向につなげる関数で、まさに 8.3 の Concat\text{Concat} そのもの。各ヘッドの出力 (batch, query_len, d_v)n_heads 個つなげて (batch, query_len, n_heads*d_v) にします。9.3 冒頭の版では transposeview でやっていた連結を、ここでは素直に cat 1発でやっています。
2つの書き方、どっちが正解?

どちらも結果は同じで、正解は1つではありません。使い分けの目安はこうです。

  • クラス分割版(こちら)AttentionHead という単位ができて 読みやすく、8 節の概念とそのまま対応。学習や実験で構造をいじりたいときに分かりやすい。
  • 1枚にまとめた版(9.3 冒頭):行列積が1回で済むので 速い。大規模モデルの実装はたいていこちら寄り。

まずはクラス分割版で「何をしているか」をつかみ、速度が要るときにまとめた版へ——という理解でOKです。

実用では nn.MultiheadAttention 一発

上のクラスは仕組みを理解するための手書き版です。実際の開発では、PyTorch 標準の nn.MultiheadAttention を使えば同じことが1行で書けます。

mha = nn.MultiheadAttention(embed_dim=4, num_heads=2, batch_first=True)
out, attn_weights = mha(x, x, x) # Self-Attention は Q=K=V=x(同じ入力を3回渡す)
print(out.shape) # torch.Size([1, 3, 4])

batch_first=True は「入力の形を (バッチ, トークン, 次元) で渡す」指定です。x を3回渡しているのは、Query・Key・Value をすべて同じ文章から作る Self-Attention(3.5)だから。返り値は出力と、おまけの 注目度(attention weights) で、後者を覗くと「どのヘッドがどの単語に注目したか」を可視化できます。

10. この章のまとめ

前半(概念)

  • Transformer は「全単語ペアの注目度を一気に計算する」アーキテクチャで、いまの LLM の土台。RNN の「遅い・遠い単語を忘れる」を、「順番に読むのをやめる」ことで解決した。
  • 心臓部の Attention は、内積で関連度 → softmax で注目度 → Value の加重和 という、前章の数学そのもの。

後半(実装)

  • 注目度は ai=exp(qki)jexp(qkj)a_i = \dfrac{\exp(q k_i^{\top})}{\sum_j \exp(q k_j^{\top})}。内積を softmax に通したスカラで、合計1。
  • 出力は Value の加重和 o=a1v1++anvno = a_1 v_1 + \dots + a_n v_n。注目した Value が濃く混ざった「新しい表現」。
  • 行列でまとめると o=softmax(qK)Vo = \text{softmax}(qK^{\top})V。NumPy では softmax(q @ K.T) @ V の数行で実装でき、手計算と一致する。
  • 実用形は dk\sqrt{d_k} で割る スケール化内積アテンション softmax ⁣(QKdk)V\text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}\right)V

発展(マルチヘッドアテンション)

  • 1つのヘッドは「1種類の関係」しか拾えないので、ヘッドを hh 個並べて別々の観点を担当させる。各ヘッドは次元 dk=dmodel/hd_k = d_{\text{model}}/h の小さな空間で、7 節の Attention をそのまま計算する。
  • MultiHead(Q,K,V)=Concat(head1,,headh)WO\text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1,\dots,\text{head}_h)\,W_O。連結で1本にまとめ、WOW_O でヘッドをまたいで混ぜる。出力は入力と同じ dmodeld_{\text{model}} 次元なので、ブロックとして積み重ねられる。

発展(PyTorch 実装)

  • NumPy で書いた3つ(内積/スケール化/マルチヘッド)は、PyTorch でもほぼ同じ数行で書ける。np.arraytorch.tensor、softmax は F.softmax(..., dim=-1)
  • 実用では自動微分・GPU・高速化に対応した組み込み(F.scaled_dot_product_attentionnn.MultiheadAttention)を使う。仕組みを理解した今なら、その中身が 9.1〜9.3 でやったことだと分かる。

この章では、Transformer の心臓部 Attention を、概念から実装まで一気通貫でつかみました。ただしここまでは「単語はすでにベクトルになっている」ことを前提にしてきました。

次章では、その前提を埋める入り口——テキストをベクトルに変えるトークン埋め込みや、位置エンコーディング・フィードフォワード層・スキップ接続・レイヤー正規化といった Attention 以外の部品を学び、最後に全部を組み合わせて 1つの Transformer ブロックを完成させます。そして、テキストをトークン(数)に変える仕組み(トークナイザー)やモデルの学習については、その先の章で扱っていく予定です。