Transformer の概要

シェイプ

$B$：バッチ数
$C$：コンテキスト長
$E$：埋め込みベクトル次元数
$D$：クエリ、キー次元数

概要

LLM の流れは C 個のトークン ID 列を入力すると、入力と同じ C 個の確率分布が得られます。 C の各行の確率分布は、その位置の次に来るトークンの予想を表します。

Embed：トークン ID 列は、 Embed （埋め込み）によって埋め込みベクトルに変換されます
Attention と FFN （ Feed-Forward Network ）による変換が繰り返し実行されます
- Attention と FFN を合わせて Transformer ブロックと呼びます
Linear （線形変換）：各トークンの埋め込みベクトルを語彙サイズと同じ次元数に変換します（出力をスコアと呼びます）
Softmax 関数：スコアから次に出現するトークンの確率分布を生成します

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

Embed

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

重み行列

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

順伝播

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

学習との関係

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

補足：バッチ処理の場合

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

Attention と FFN

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

詳細は Attention を参照。

Linear と Softmax

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

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

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 . の次にくるトークンの確率分布を表します