LLM は入力として C 個のトークン ID 列を受け取り、入力と同じ C 個の（語彙数の次元数を持つ）確率分布を出力します。
C の各行の確率分布は、その位置の次にくるトークンの予想を表します。

説明は以下のシェイプを仮定します。

  • $C$：コンテキスト長
  • $E$：埋め込みベクトル次元数
  • $D$：クエリ、キー次元数
  1. Embedding レイヤー：トークン ID 列を埋め込みベクトルに変換します
    • Embedding レイヤーの重み行列は、語彙数 $V$ 、埋め込みベクトルの次元数 $E$ の行列 $(V, E)$です
  2. Attention と FFN （ Feed-Forward Network ）による変換が繰り返し実行されます
    • Attention と FFN を合わせて Transformer ブロックと呼びます
  3. Linear （線形変換）： 各トークンの埋め込みベクトルを語彙サイズと同じ次元数に変換します（出力をスコアと呼びます）
  4. Softmax 関数： スコアから次に出現するトークンの確率分布を生成します

Transformer ブロック（ Attention と FFN ）を通過するデータは、一貫して同じ形状を維持します。
例では $(C, E)$ になります。
この一貫性により Transforer ブロックによる複数回の処理が可能となります。

Embedding レイヤー

Embedding レイヤー（埋め込み層）は、トークン ID（入力）に対応する重み行列に格納された分散表現ベクトルを出力します。

重み行列

  • 重み行列（ $W$ ）のシェイプ： $V \times E$（語彙数 $V$, 埋め込みベクトルの次元数 $E$）
  • 各行がトークンの分散表現ベクトルに対応

順伝播

  • 入力：トークン ID（整数）、例えば id = 42
  • 出力： W[42]、つまり 42 行目のベクトル（分散表現ベクトルの次元数 $E$ ）
  • 実装上は（入力トークンの） one-hot ベクトルとの行列積と数学的に等価ですが、実際は単なるテーブルルックアップ（行の抜き出し）として実装されます

学習との関係

この重み行列 $W$ は学習によって更新されます。
結果として、意味的・文法的に近いトークンが空間上で近くなるよう分散表現が形成されます。

補足：バッチ処理の場合

入力が (バッチサイズ $B$, コンテキスト長 $C$) なら出力は (バッチサイズ, コンテキスト長, $E$ ) のテンソルになります。

Attention と FFN （Transformer ブロック）

Attention と FFN を合わせて Transformer ブロックと呼びます。 Transformer ブロックを通過するテンソルの形状は一貫しています（$(C, E)$）。
そのため Transformer ブロックを複数回適用することができます。

詳細は Attention を参照。

Linear と Softmax

形状が $(C, E)$ のテンソルは、線形変換によって $(C, V)$ の形状に変換されます。
$V$ はトークンの語彙サイズを表します。

Softmax 関数によって C 個の確率分布が得られます。 各位置の確率分布は、その位置の次にくるトークンの予想を表します。

NOTE
末尾の単語の次にくるトークンは、学習データ的には多くの場合文末トークン（EOS）やパディングトークンになります。
推論時においては、次の新しい文章の始まりの単語を予測する挙動になります。

以下の解釈は Attention にマスクが設定（つまり未来を考慮しない）されている場合に限ります。

C = 5 ( This is a pen .) 
V = 1000
         0     1             999  
This [0.001, 0.021, ......, 0.002]
is   [0.020, 0.001, ......, 0.060]
a    [0.020, 0.001, ......, 0.060]
pen  [0.100, 0.001, ......, 0.001]
.    [0.034, 0.021, ......, 0.071]

This の確率分布 [0.001, 0.021, ......, 0.002] は This の次にくるトークンの確率分布を表します
is の確率分布 [0.020, 0.001, ......, 0.060] は This is の次にくるトークンの確率を表します
a の確率分布 [0.020, 0.001, ......, 0.060] は This is a の次にくるトークンの確率分布を表す
pen の確率分布 [0.100, 0.001, ......, 0.001] は This is a pen の次にくるトークンの確率分布を表します
. の確率分布 [0.034, 0.021, ......, 0.071] は This is a pen . の次にくるトークンの確率分布を表します

※ . の次にくるトークンは、学習データ的には多くの場合文末トークン（EOS）やパディングトークンになります。推論時においては、次の新しい文章の始まりの単語を予測する挙動になります。