入力埋め込み（Token Embedding） | Transformer 内部構造ガイド

Level 1概要役割と処理の流れ

何をする機能か

[B, S] の整数ID列を、[B, S, D] の実数ベクトル列へ変換します。

処理そのものは**表引き（テーブルルックアップ）**です。語彙数 V 行 × 次元 D 列の行列（埋め込み行列）を持っておき、 ID が指す行をそのまま取り出します。DB でいえば主キーによる 1 行の SELECT に相当します。

なぜ必要か

Token ID は単なる通し番号で、値の大小に意味がありません。 ID 6 と ID 7 が似た意味とは限らず、そのまま足したり掛けたりしても意味のある計算になりません。

そこで、各トークンを「意味を座標で表したベクトル」に置き換えます。この座標（埋め込み行列の中身）は学習で決まるパラメータです。学習後は、似た文脈で使われるトークンどうしが近い方向のベクトルになります。

全体の中でいつ呼ばれるか

Tokenizer の直後、位置情報の付加の直前です。 Encoder 側・Decoder 側の入り口それぞれで実行されます。

Level 2内部処理Tensor形状・内部ステップ・設計上の意味

前段から受け取るもの

入力トークンが出力した input_ids [B, S]（int64）です。

後段へ渡すもの

hidden_states [B, S, D]（float）です。これ以降、Transformer の内部を流れるデータは 最後まで [B, S, D] のままで、形は変わりません。変わるのは中身（各位置のベクトルが持つ情報）だけです。

処理前

input_ids [B, S] = [1, 8] の整数

処理後

hidden_states [B, S, D] = [1, 8, 512] の実数

例文での処理

共通例文

The keys on the table are lost.

概念上のトークン列（S = 8）: [The, keys, on, the, table, are, lost, .]

The
keys
on
the
table
are
lost
.

位置1の keys と位置3の the は、この時点ではまだ「文中での役割」を持ちません。同じ綴りのトークンは、位置0の The と位置3の the のように大文字小文字が違えば別IDですが、同じIDなら必ず同じベクトルになります。

重要なのは、この段階のベクトルは文脈に依存しないという点です。

位置 3 の the と、別の文に出てくる the は、まったく同じベクトルになります
keys（鍵／キー）が「物理的な鍵」なのか「暗号鍵」なのかは、まだ区別されていません

文脈に応じた区別がつくのは、マルチヘッド自己注意を通ったあと、つまり Encoder出力の段階です。

EncoderとDecoderでの違い

処理内容は同じで、入力する系列が違うだけです。

Encoder: 入力文の input_ids [B, S] → [B, S, D]
Decoder: 右シフトした出力の input_ids [B, T] → [B, T, D]

原典 Transformer では、Encoder の埋め込み・Decoder の埋め込み・出力側の Linear 層の 3 つで重みを共有しています（weight tying）。パラメータ数を減らしつつ、「IDからベクトル」と「ベクトルからID」の対応を一貫させる狙いがあります。

入力と出力

入力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`input_ids`	各位置のトークンID	int64 Tensor	`[B, S]`	B = バッチサイズ（同時に処理する文の数） / S = 入力系列長（Encoder側のトークン数）	[[5, 6, 7, 8, 9, 10, 11, 12]]	埋め込み行列の行を引く

出力

項目名	データの意味	データ型	Tensor形状	記号の意味	具体例	次の渡し先
`hidden_states`	各トークンの意味ベクトル（文脈は未反映）	float32 Tensor	`[B, S, D]`	B = バッチサイズ（同時に処理する文の数） / S = 入力系列長（Encoder側のトークン数） / D = モデル内部のベクトル次元（d_model）	S=8, D=512 なら [1, 8, 512]	位置情報の付加（Positional Encoding）

B: バッチサイズ（同時に処理する文の数）
S: 入力系列長（Encoder側のトークン数）
D: モデル内部のベクトル次元（d_model）

このサイトでは Tensor 形状を PyTorch の `batch_first=True` 相当、つまり `[B, S, D]`（バッチ次元が先頭）で統一しています。実装によっては `[S, B, D]` を採用するものもあります。

学習対象のパラメータは埋め込み行列 embedding.weight [V, D] です。 V = 32,000、D = 512 なら約 1,640 万パラメータになります。

よくある誤解

Level 3数式・Python数式・疑似実装・計算量・実装差

数式

埋め込み行列を $E \in \mathbb{R}^{V \times D}$ 、位置 $i$ のトークンIDを $t_i$ とすると、

x_i = E_{t_i} \cdot \sqrt{D}

原典 Transformer では $\sqrt{D}$ を掛けます。これは、後段で加算する Positional Encoding の値の大きさ（おおむね $[-1, 1]$ ）に対して埋め込み側が小さくなりすぎないようにするためのスケール調整です。

Pythonで表したインターフェース

import math
import torch
from torch import Tensor, nn
 
 
class TokenEmbedding(nn.Module):
    """[B, S] の Token ID を [B, S, D] のベクトルへ変換する。"""
 
    def __init__(self, vocab_size: int, d_model: int) -> None:
        super().__init__()
        self.d_model = d_model
        # 学習対象のパラメータ: [V, D]
        self.embedding = nn.Embedding(vocab_size, d_model)
 
    def forward(self, input_ids: Tensor) -> Tensor:
        """
        Args:
            input_ids: [B, S] int64。語彙表のインデックス。
 
        Returns:
            [B, S, D] float32。各位置の意味ベクトル。
        """
        # [B, S] -> [B, S, D]
        hidden_states = self.embedding(input_ids)
        # 原典 Transformer のスケーリング
        return hidden_states * math.sqrt(self.d_model)
 
 
def embed_tokens(input_ids: Tensor, layer: TokenEmbedding) -> Tensor:
    """[B, S] の Token ID を [B, S, D] のベクトルへ変換する。"""
    return layer(input_ids)
 
 
# 最小使用例
V, D = 32_000, 512
layer = TokenEmbedding(vocab_size=V, d_model=D)
 
input_ids = torch.tensor([[5, 6, 7, 8, 9, 10, 11, 12]])  # [B=1, S=8]
hidden_states = embed_tokens(input_ids, layer)
 
print(input_ids.shape)      # torch.Size([1, 8])       == [B, S]
print(hidden_states.shape)  # torch.Size([1, 8, 512])  == [B, S, D]
 
# 同じ ID は必ず同じベクトルになる（文脈非依存）
print(torch.allclose(hidden_states[0, 0], layer.embedding.weight[5] * math.sqrt(D)))  # True

この例で単純化している点

重み初期化・学習・weight tying は省略し、順伝播の形だけを示しています。
パディング位置の扱い（padding_idx）は省略しています。
sqrt(D) のスケーリングは原典の仕様で、実装によっては行いません。

型と Tensor 形状の見え方を掴むための擬似実装です。そのまま学習・推論に使えるものではありません。

計算量とメモリ

計算量: 表引きのみなので $O(B \cdot S)$ 。Attention や Feed-Forward に比べて無視できます。
パラメータ数: $V \times D$ 。語彙数が大きいモデルでは全体の中でも大きな割合を占めます。
メモリ: 出力は $B \times S \times D$ 個の浮動小数点数です。

派生と実装差

weight tying: 入力側 Embedding と出力側 Linear で $E$ を共有する方式。原典 Transformer および多くの LLM が採用しています。
factorized embedding: $V \times D$ を $V \times D'$ と $D' \times D$ に分解してパラメータを削減する方式（ALBERT など）。
Segment Embedding: BERT のように、文の区別を表す埋め込みを加算するモデルもあります。原典 Transformer にはありません。

入力埋め込みToken Embedding