【深層学習】埋め込み(Embedding)や埋め込み層の仕組みと実装を理解する

「単語」や「カテゴリ」は数字ではありません。では、ニューラルネットワークにどうやって渡せばよいでしょうか。

自然言語処理や推薦システムで必ず登場するのが 埋め込み(Embedding) という仕組みです。「りんご」という単語を整数の ID (例: 0) に変換し、そのIDから固定長の密なベクトルを取り出す ── これだけで「単語の意味」がニューラルネットに渡せるようになります。

埋め込みはあらゆる場所で使われています。

• 自然言語処理 (NLP): BERT・GPT・LLaMA など全ての大規模言語モデルは、入力トークンを埋め込みベクトルに変換してから処理します
• 推薦システム: ユーザーIDと商品IDをそれぞれ埋め込んで内積をとることで「好み」を予測します
• テーブルデータ: 都道府県コード・職業コードなどのカテゴリ変数も埋め込みで実数ベクトルに変換できます

本記事では以下の内容を扱います。

• 埋め込みとは何か ── ルックアップテーブルとしての直感的理解
• one-hot ベクトルとの等価性 ── なぜ埋め込みの方が優れているのか
• nn.Embedding の使い方・引数・内部の重み行列
• 学習で「意味」が獲得されるメカニズム
• padding_idx による可変長バッチの処理
• 事前学習済み埋め込みの読み込み
• コサイン類似度で意味の近さを可視化する
• 位置埋め込みとの合算 ── Transformer への接続

埋め込み(Embedding)とは

コンピュータが「単語」を扱うとき、最初にやるべきことは 単語を数値に変換すること です。一番単純な方法は「整数のID」を割り当てることです。たとえば「りんご」→ 0、「バナナ」→ 1、「電車」→ 2 … という辞書を作ります。

しかしただの整数だと、ニューラルネットワークは「2 は 1 より少し大きい」「0 と 1 は近い」といった意味のない関係まで学んでしまいます。「りんご」と「バナナ」が電車より「近い」などという情報は本来ありません。

そこで登場するのが 埋め込み です。整数IDを 密なベクトル (たとえば 64 次元や 128 次元の実数ベクトル) に変換し、そのベクトルをニューラルネットに渡します。

図から読み取れることが2つあります。まず、離散的な整数ID (0, 1, 2, …) がそれぞれ異なる実数ベクトルに変換されています。次に、各ベクトルの値は正負混在の密な実数であり、one-hot とは異なり「ほぼ全ての次元に意味がある」ことが分かります。

これがEmbeddingの本質です。整数IDを「辞書の索引」とみなして、対応するベクトルを取り出す ── まるに本の索引からページを引くような操作です。次に、この「取り出し」が数学的に何を意味するかを見ていきましょう。

one-hot ベクトルとの比較: なぜ埋め込みが優れているのか

埋め込みを理解するうえで、まず one-hot (1-of-K) 表現 と比べることが大切です。

語彙数 $V$ (例: 10,000語) の場合、one-hot ベクトルは $V$ 次元で1か所だけ 1、残りは全て 0 のベクトルです。

$$ \text{one-hot}(\text{りんご}) = [1, 0, 0, \ldots, 0] \quad (V \text{ 次元}) $$

$$ \text{one-hot}(\text{バナナ}) = [0, 1, 0, \ldots, 0] \quad (V \text{ 次元}) $$

この表現には大きな問題が3つあります。

問題1: スパース (メモリ非効率)

10,000語の場合、one-hot ベクトルは 10,000 次元で 9,999 個がゼロです。ゼロだらけの行列を毎回保持するのは非効率です。

問題2: 語の意味的な近さが表現できない

「りんご」と「バナナ」の one-hot ベクトルの内積は 0 です。「自動車」と「電車」の内積も同じく 0 です。どんな単語同士でも等距離なので、意味的な類似度は一切表現されません。

問題3: 次元数が語彙数に比例して増大する

語彙数が 100,000 語になれば one-hot は 100,000 次元になります。一方、埋め込みは次元数をハイパーパラメータ embedding_dim で自由に決めることができます。

左の one-hot 行列を見ると、対角線上にだけ 1 が並ぶスパースな構造が分かります。右の埋め込み行列は正負の実数が密に詰まっており、ずっと小さい次元数で各単語を表現できています。

ではこの「小さい密なベクトル」はどうやって得るのでしょうか。ここで登場するのが 重み行列のルックアップ という考え方です。

one-hot × 重み行列 = ルックアップテーブル (等価性)

埋め込みの仕組みを数学で表すと、実はシンプルです。

語彙数 $V$、埋め込み次元 $D$ の 重み行列 $\bm{W} \in \mathbb{R}^{V \times D}$ を用意します。単語ID $i$ の埋め込みベクトル $\bm{e}_i$ は次のように定義されます。

$$ \bm{e}_i = \bm{W}[i, :] $$

つまり $\bm{W}$ の $i$ 行目をそのまま取り出します。これを「one-hot ベクトルとの行列積」として書くと:

$$ \bm{e}_i = \bm{h}_i^{\top} \bm{W} $$

ここで $\bm{h}_i$ は $i$ 番目だけ 1 の one-hot ベクトルです。展開すると:

$$ \bm{h}_i^{\top} \bm{W} = [0, \ldots, 0, \underbrace{1}_{i\text{番目}}, 0, \ldots, 0] \cdot \bm{W} $$

$i$ 番目だけ 1 なので、結果は $\bm{W}$ の $i$ 行目に他なりません。つまり one-hot との行列積は、行列の特定の行を取り出すだけ という操作と完全に等価です。

この「行の取り出し」を直接行うのがルックアップテーブルです。わざわざ $V \times V$ の one-hot 行列を作って掛け算する必要はなく、インデックスで直接アクセスするだけで済みます。

図の中央を見ると、one-hot の1が立っている行と同じ行 (ID=2の行) が重み行列 $\bm{W}$ からハイライトされています。この1行だけを取り出すことが、one-hot × $\bm{W}$ と数学的に同じであることが視覚的に確認できます。

この等価性を踏まえて、次はPyTorchにおける具体的な実装を見ていきましょう。

nn.Embedding の使い方と引数

PyTorchでは埋め込み層を torch.nn.Embedding で提供しています。

import torch
import torch.nn as nn

# 基本的な作成
# num_embeddings: 語彙数 (ID の最大値 + 1)
# embedding_dim: 埋め込みの次元数
emb = nn.Embedding(num_embeddings=10, embedding_dim=4)

# ID を渡すとベクトルが返る
ids = torch.tensor([0, 3, 7])
out = emb(ids)
print(out.shape)   # torch.Size([3, 4])
print(out)

emb(ids) の戻り値は [バッチサイズ, embedding_dim] の FloatTensor です。内部では emb.weight[ids] という単純なインデックス参照が行われています。

主要な引数を整理します。

コードを実行すると、out.shape が torch.Size([3, 4]) と表示されます。これは 3 個のIDに対して、それぞれ 4 次元のベクトルが対応して返ってきていることを意味します。中身の実数値はランダムに初期化されており、学習を通じて更新されていきます。

重み行列 emb.weight を直接確認する

import torch
import torch.nn as nn

emb = nn.Embedding(num_embeddings=6, embedding_dim=4)
print("weight の形状:", emb.weight.shape)  # torch.Size([6, 4])
print("weight:\n", emb.weight.data)

# ID=2 の埋め込みは weight の 2 行目と一致する
idx = torch.tensor([2])
print("\nemb(2):", emb(idx))
print("weight[2]:", emb.weight[2])

emb.weight は (num_embeddings, embedding_dim) 形状のパラメータです。emb(torch.tensor([2])) と emb.weight[2] は全く同じ値を返します。これがルックアップテーブルとしての等価性をコードで確認したものです。

weight の形状: torch.Size([6, 4])
emb(2): tensor([[-0.1234,  0.5678, -0.3456,  0.9012]], grad_fn=<EmbeddingBackward0>)
weight[2]: tensor([-0.1234,  0.5678, -0.3456,  0.9012], grad_fn=<SelectBackward0>)

(数値は初期化のたびに変わります)

grad_fn=<EmbeddingBackward0> と表示されていることから、この取り出し操作は 自動微分が有効 であり、学習の際には損失から emb.weight まで勾配が流れることが分かります。

ヒートマップを見ると、行がID (単語) に、列が埋め込みの各次元に対応しています。ID=3 の行 (枠で強調) が取り出す対象で、その行のベクトルだけを参照します。学習前の現時点では全IDのベクトルがほぼランダムですが、学習後には意味的に近い単語の行が似た値を持つようになります。

ここまでで「埋め込みとは何か」「内部の重み行列がどう機能するか」を理解しました。次は、この重み行列が学習によってどう変化するかを見ていきます。

学習で「意味」が獲得されるメカニズム

埋め込み層の重み行列 $\bm{W}$ は、他のニューラルネットワークの重みと同じく 誤差逆伝播(backpropagation) で更新されます。

たとえば「単語予測タスク」を考えてみましょう。「私はネコが___」という文の空白に「好き」が来る確率を最大化するように学習します。このとき「ネコ」と「好き」が同じ文脈でよく現れると、両者の埋め込みベクトルは似たような向きに引き寄せられていきます。

より正確に言うと、学習中の勾配更新 $\Delta \bm{W}[i, :]$ は、その単語 $i$ がアクセスされたバッチでのみ行われます。アクセスされなかった行は更新されません。これは埋め込みが「スパースな勾配」を持つことを意味し、大規模語彙でも効率的な学習が可能になっています。

左の「学習前」では8つの単語がランダムに散らばっています。学習後の右では、「王・女王・男・女」(人間カテゴリ)、「犬・猫」(動物カテゴリ)、「車・電車」(乗り物カテゴリ) がそれぞれ近くに集まっています。これは「同じカテゴリの単語は似た文脈に現れる」というデータのパターンを重み行列が学習した結果です。

この学習済みの埋め込みベクトルを用いると、単語同士の意味的な類似度 を数値で測ることができます。次に、その手法であるコサイン類似度を詳しく見ていきましょう。

コサイン類似度で意味の近さを可視化する

2つの埋め込みベクトル $\bm{e}_i$, $\bm{e}_j$ の コサイン類似度 は次のように定義されます。

$$ \text{cos}(\bm{e}_i, \bm{e}_j) = \frac{\bm{e}_i \cdot \bm{e}_j}{\|\bm{e}_i\| \|\bm{e}_j\|} $$

値は $[-1, 1]$ の範囲を取り、1 に近いほど意味が近く、0 に近いほど無関係、-1 に近いほど「逆の意味」を持つことを示します。

ベクトルの大きさ(ノルム)を揃えることで、スケールの違いによる誤差を排除できる点が特徴です。コサイン類似度を使うと「りんごとバナナ」が「りんごと電車」より近いことを数値で確認できます。

以下のコードで埋め込みを学習させてコサイン類似度を計算してみましょう。正の共起ペア(同カテゴリ)と負のペア(異カテゴリ)を使うことで、意味の分離が明確になります。

import torch
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(0)

words = ["王", "女王", "男", "女", "犬", "猫"]

# 正の共起ペア: 同カテゴリの単語は同じ文脈に現れる
pos_pairs = [(0,1),(1,0),(0,2),(2,0),(1,3),(3,1),(2,3),(3,2),(4,5),(5,4)]
# 負のペア: 異カテゴリは共起しない
neg_pairs = [(0,4),(0,5),(1,4),(1,5),(2,4),(2,5),(3,4),(3,5)]

all_pairs = [(c, x,  1.0) for c, x in pos_pairs] + \
            [(c, x, -1.0) for c, x in neg_pairs]

c_ids  = torch.tensor([p[0] for p in all_pairs])
x_ids  = torch.tensor([p[1] for p in all_pairs])
labels = torch.tensor([p[2] for p in all_pairs])

# 4次元の埋め込みで学習
emb = nn.Embedding(num_embeddings=6, embedding_dim=4)
optimizer = optim.Adam(emb.parameters(), lr=0.05)

for step in range(1000):
    cv  = emb(c_ids)
    xv  = emb(x_ids).detach()
    dot = torch.sum(cv * xv, dim=1)
    loss = -torch.mean(labels * dot)        # 正ペアで高く、負ペアで低くする
    optimizer.zero_grad(); loss.backward(); optimizer.step()

# コサイン類似度行列
with torch.no_grad():
    W      = emb.weight
    W_norm = W / (W.norm(dim=1, keepdim=True) + 1e-9)
    sim    = W_norm @ W_norm.T

print("コサイン類似度行列:")
for i, wi in enumerate(words):
    row = [f"{sim[i, j].item():.2f}" for j in range(6)]
    print(f"  {wi}: {row}")

出力:

コサイン類似度行列:
  王:   ['1.00', '1.00', '1.00', '1.00', '-1.00', '-1.00']
  女王: ['1.00', '1.00', '1.00', '1.00', '-1.00', '-1.00']
  男:   ['1.00', '1.00', '1.00', '1.00', '-1.00', '-1.00']
  女:   ['1.00', '1.00', '1.00', '1.00', '-1.00', '-1.00']
  犬:   ['-1.00', '-1.00', '-1.00', '-1.00', '1.00', '1.00']
  猫:   ['-1.00', '-1.00', '-1.00', '-1.00', '1.00', '1.00']

人間カテゴリ (王・女王・男・女) 同士のコサイン類似度が +1.00、動物カテゴリ (犬・猫) との類似度が −1.00 になっています。正のペアと負のペアの両方で学習することで、「同じカテゴリ = 完全に同じ方向」「異カテゴリ = 真逆の方向」という綺麗な構造が埋め込み空間に現れました。

左の類似度行列では、人間グループ内の値が高く(緑)、人間と動物の間が低い(赤)ことが視覚的に明確です。右のベクトル可視化では、同じカテゴリの矢印が近い方向を向いており、コサイン類似度(= ベクトルの向きの近さ)が高くなる理由がよく分かります。

padding_idx ── 可変長バッチを扱う

自然言語処理では、文の長さが各サンプルで異なります。PyTorchのテンソルはすべてのサンプルが同じ長さである必要があるため、短い文の末尾に特別な「パディングトークン」を埋めて長さを揃える のが一般的です。

nn.Embedding の padding_idx 引数にパディング用のIDを指定すると、そのIDの埋め込みベクトルは 常にゼロになり、勾配も流れません。

import torch
import torch.nn as nn

# padding_idx=0 を指定: ID=0 はパディング
emb = nn.Embedding(num_embeddings=10, embedding_dim=4, padding_idx=0)

# バッチ内の 2 文。ID=0 はパディング (PAD)
batch = torch.tensor([
    [3, 5, 7, 0, 0],   # 文1: 3語 + PAD × 2
    [1, 4, 0, 0, 0],   # 文2: 2語 + PAD × 3
])

out = emb(batch)
print("出力形状:", out.shape)              # torch.Size([2, 5, 4])
print("PAD (ID=0) の埋め込み:", emb.weight[0])  # 全て 0
print("PAD 位置の出力:", out[0, 3])            # 全て 0 (パディング部分)

出力形状: torch.Size([2, 5, 4])
PAD (ID=0) の埋め込み: tensor([0., 0., 0., 0.], grad_fn=<SelectBackward0>)
PAD 位置の出力: tensor([0., 0., 0., 0.], grad_fn=<SelectBackward0>)

PAD トークンのベクトルが全てゼロになっています。この出力を後続の Transformer や RNN に渡すとき、アテンションのマスクと組み合わせることで「パディング部分は無視する」という処理が正確に行えます。

図の灰色セルが PAD (ID=0) で、色付きセルが実際のトークンです。文1は3語、文2は2語しかなく、残りを PAD で埋めることで 6 トークンの固定長バッチを作っています。PAD のセルは埋め込みがゼロになるため、後続の演算でも「中身のないゼロ」として扱われます。

embedding_dim の選び方 ── 次元数と表現力のトレードオフ

embedding_dim (埋め込みの次元数) はモデル設計上の重要なハイパーパラメータです。

小さすぎる次元数の場合: ベクトルに詰め込める情報量が少なく、語彙内の微妙な意味の違いを区別できません。極端に言えば 2 次元では「方向」が 2 つしかなく、数万語の違いを表現できません。

大きすぎる次元数の場合: 計算量とメモリが線形に増加します。また語彙が少ない場合はほとんど使われない次元が生まれ、過学習のリスクが高まります。

経験則として、語彙数を $V$ としたとき $D \approx \sqrt[4]{V}$ 程度、あるいはタスクの複雑さに応じて 64〜512 の間 で選ぶことが多いです。

グラフから、表現力 (緑) は次元数に対して対数的に増加し、ある点を超えると飽和していることが分かります。一方、計算コスト (赤) は線形に増加します。点線で示した 64 次元あたりが、多くのタスクで「十分な表現力 / 十分な軽さ」のスイートスポットになっています。

事前学習済み埋め込みの読み込み

大規模なコーパスで学習した Word2Vec や GloVe などの 事前学習済み埋め込み を使うと、少ないデータのタスクでも良い性能を得られることがあります。

import torch
import torch.nn as nn
import numpy as np

# 仮の事前学習済みベクトル (実際は Word2Vec や GloVe のファイルから読み込む)
vocab_size  = 100
embed_dim   = 50
pretrained_vectors = torch.FloatTensor(
    np.random.randn(vocab_size, embed_dim) * 0.01   # 例: GloVe 等を読み込んだもの
)

# nn.Embedding を作成し、重みを事前学習済みベクトルで初期化
emb = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embed_dim)
emb.weight.data.copy_(pretrained_vectors)

# ── fine-tune するか凍結するかを選ぶ ────────────────────────────
# 凍結 (frozen): 埋め込みを変えず下流モデルだけ学習する
emb.weight.requires_grad = False

# 微調整 (fine-tune): 埋め込みも含めて全体を学習する
# emb.weight.requires_grad = True  # デフォルトでは True

print("requires_grad:", emb.weight.requires_grad)
print("weight の一部:", emb.weight[0, :5])

requires_grad = False に設定すると埋め込みは凍結され、学習時に更新されません。語彙が少ない or データが少ないときに有効です。逆に True (デフォルト) のまま学習させれば、事前学習済みを初期値として fine-tune できます。

図のフローを見ると、事前学習済みベクトルの取得から emb.weight.data.copy_() による代入、requires_grad による凍結/微調整の選択、そして下流タスクでの学習という4ステップが確認できます。凍結を選ぶか微調整を選ぶかは、データ量とタスクの性質次第です。

単語埋め込み (Word2Vec) との関係

「Word2Vec の埋め込みと nn.Embedding は何が違うの?」と疑問に思うかもしれません。

本質的には 同じ仕組み です。Word2Vec (Skip-gram や CBOW) も内部に埋め込み行列 $\bm{W}$ を持ち、単語の共起パターンから学習します。学習後の $\bm{W}$ を取り出して PyTorch の nn.Embedding に読み込めば、そのまま使えます。

違いは「何のタスクで学習されたか」です。

つまり Word2Vec はオフラインで学習した「埋め込みの辞書」であり、nn.Embedding はそれを格納して使う「コンテナ」と考えると分かりやすいです。

関連記事: 単語埋め込みの理論 — Word2Vec・GloVe・FastTextの比較

実装: 感情分類で nn.Embedding を使う

ここまでの知識をまとめた実装例として、短いテキストの感情分類モデルを作ります。

import torch
import torch.nn as nn
import torch.optim as optim

# ── データ準備 ───────────────────────────────────────────────────
# 簡易語彙と感情ラベル (0=ネガティブ, 1=ポジティブ)
# 実際は tokenizer で前処理し vocabulary を構築する
vocab = {"<PAD>": 0, "面白い": 1, "楽しい": 2, "退屈": 3,
         "つまらない": 4, "すごい": 5, "悲しい": 6, "映画": 7, "本": 8}
VOCAB_SIZE = len(vocab)
EMBED_DIM  = 8
MAX_LEN    = 4

# (単語IDの列, ラベル)
texts = [
    ([1, 7, 0, 0], 1),   # "面白い 映画 PAD PAD" → ポジティブ
    ([3, 4, 8, 0], 0),   # "退屈 つまらない 本 PAD" → ネガティブ
    ([5, 1, 0, 0], 1),   # "すごい 面白い PAD PAD" → ポジティブ
    ([6, 4, 7, 0], 0),   # "悲しい つまらない 映画 PAD" → ネガティブ
]

X = torch.tensor([t[0] for t in texts])
y = torch.tensor([t[1] for t in texts], dtype=torch.float)

# ── モデル定義 ───────────────────────────────────────────────────
class SentimentModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, max_len):
        super().__init__()
        self.emb   = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.fc    = nn.Linear(embed_dim, 1)

    def forward(self, x):
        # x: (batch, seq_len)
        e = self.emb(x)           # (batch, seq_len, embed_dim)
        mask = (x != 0).float()   # PAD 以外のマスク
        # PAD を除いた平均プーリング
        pooled = (e * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True).clamp(min=1)
        return self.fc(pooled).squeeze(-1)   # (batch,)

model = SentimentModel(VOCAB_SIZE, EMBED_DIM, MAX_LEN)
optimizer = optim.Adam(model.parameters(), lr=0.05)
criterion = nn.BCEWithLogitsLoss()

# ── 学習 ─────────────────────────────────────────────────────────
for epoch in range(100):
    logits = model(X)
    loss   = criterion(logits, y)
    optimizer.zero_grad(); loss.backward(); optimizer.step()

# ── 評価 ─────────────────────────────────────────────────────────
with torch.no_grad():
    probs = torch.sigmoid(model(X))
    preds = (probs > 0.5).long()
    print("予測確率:", probs.numpy().round(3))
    print("予測ラベル:", preds.numpy())
    print("正解ラベル:", y.long().numpy())
    acc = (preds == y.long()).float().mean()
    print(f"正解率: {acc.item():.1%}")

実行結果:

予測確率: [0.921 0.063 0.941 0.058]
予測ラベル: [1 0 1 0]
正解ラベル: [1 0 1 0]
正解率: 100.0%

予測確率は正例 (ポジティブ) で 0.92〜0.94、負例 (ネガティブ) で 0.06 程度となり、4サンプルすべて正しく分類できています。埋め込み層 (self.emb) が感情の正負を区別する方向のベクトルを学習したため、シンプルな線形層 1 枚でも正しく分類できていることが分かります。padding_idx=0 による平均プーリングも、PAD トークンを除外して正しく機能しています。

位置埋め込みとの合算 ── Transformer への接続

Transformer では、トークン埋め込み に加えてトークンの位置情報も埋め込む必要があります。なぜなら、Transformer のアテンション機構は「どのトークンがどの位置にあるか」を意識せず処理するためです。

位置埋め込み (Positional Encoding) には大きく2種類があります。

1. 固定式 (sin/cos) ── BERT のオリジナルで使用

$$ \text{PE}(pos, 2i) = \sin\!\left(\frac{pos}{10000^{2i/d}}\right) $$

$$ \text{PE}(pos, 2i+1) = \cos\!\left(\frac{pos}{10000^{2i/d}}\right) $$

ここで $pos$ は文中の位置 (0始まり)、$i$ は次元のインデックス、$d$ は埋め込み次元です。

2. 学習式 ── GPT-2 や LLaMA で使用。位置ごとにパラメータを学習する

最終的な入力ベクトルは次のようにトークン埋め込みと位置埋め込みの 単純な足し算 です。

$$ \bm{x}_{pos} = \bm{e}_{token\_id} + \bm{p}_{pos} $$

import torch
import torch.nn as nn
import math

class TokenAndPositionalEmbedding(nn.Module):
    def __init__(self, vocab_size, max_len, embed_dim):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, embed_dim, padding_idx=0)

        # 固定式の位置埋め込みを事前計算
        pe = torch.zeros(max_len, embed_dim)
        pos = torch.arange(0, max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, embed_dim, 2).float()
                        * (-math.log(10000.0) / embed_dim))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe)       # 学習しない定数として登録

    def forward(self, x):
        # x: (batch, seq_len) の整数テンソル
        seq_len = x.shape[1]
        tok = self.token_emb(x)                      # (batch, seq, embed_dim)
        return tok + self.pe[:seq_len].unsqueeze(0)  # 位置埋め込みを加算

# 使用例
model = TokenAndPositionalEmbedding(vocab_size=100, max_len=50, embed_dim=16)
ids   = torch.tensor([[1, 3, 5, 7]])    # (1, 4)
out   = model(ids)
print("出力形状:", out.shape)           # torch.Size([1, 4, 16])

出力形状: torch.Size([1, 4, 16])

register_buffer を使うことで、位置埋め込みのテンソルは state_dict に含まれますが parameters() には含まれないため、optimizer に渡したときに更新されません。これが「固定式」と「学習式」の実装上の違いです。

図の左から「Token Embedding」と「Positional Encoding」の2つの行列を足し合わせた結果が「入力表現」として Transformer に渡されます。トークンの意味 (トークン埋め込み) と位置情報 (位置埋め込み) が合算されることで、Transformer はどのトークンがどこにあるかを考慮したアテンション計算ができるようになります。

埋め込みに関するよくある疑問

Q. embedding_dim はどうやって決める?

A. まずはタスクの規模に合わせて 64〜256 を試し、バリデーションスコアを見ながら調整します。あまり小さいと表現力が不足し、大きすぎると学習データが少ない場合は過学習します。事前学習済みを使うときは、その次元数に合わせます (Word2Vec の 300 次元や BERT の 768 次元など)。

Q. nn.Embedding の初期値は?

A. デフォルトでは N(0, 1) (標準正規分布) で初期化されます。nn.init.normal_(emb.weight, mean=0, std=0.1) などで変更できます。

Q. ミニバッチ中の同じ ID に対する勾配はどう扱われる?

A. 同じ ID が複数のサンプルに現れると、それらの勾配が 足し合わされてから 重みに適用されます。scale_grad_by_freq=True にすると、頻出語の勾配が出現頻度で割り引かれます。

Q. 埋め込みベクトルの大きさ (ノルム) は?

A. 学習中に特定の条件がなければノルムは自由に変化します。max_norm 引数を指定すると、取り出したベクトルのノルムが max_norm を超えたときに正規化されます。Transformerではノルムが大きくなりすぎると学習が不安定になるため、スケジュールされた正規化 (LayerNorm など) を後段に置くことが一般的です。

まとめ

本記事では、埋め込み (Embedding) の仕組みを基礎から解説しました。

• 埋め込みとは ── 離散ID を密なベクトルに変換するルックアップテーブル。one-hot × 重み行列と数学的に等価だが、直接インデックス参照で実装される
• nn.Embedding の使い方 ── num_embeddings(語彙数) と embedding_dim(次元数) の2引数が必須。emb.weight が実体の重み行列
• 学習で意味が獲得される ── 誤差逆伝播で重み行列が更新され、意味的に近い語のベクトルが近い向きに揃っていく
• padding_idx ── 可変長バッチのパディングトークンをゼロベクトル固定にする
• 事前学習済み埋め込みの活用 ── emb.weight.data.copy_() で読み込み、requires_grad で凍結/微調整を選択
• コサイン類似度 ── 学習後の埋め込みで意味の近さを数値で評価できる
• 位置埋め込みとの合算 ── Transformer への入力はトークン埋め込み + 位置埋め込みの足し算

埋め込みを理解すれば、BERT・GPT などの大規模言語モデルのアーキテクチャを読む準備が整います。次のステップとして、以下の記事もご参照ください。

単語埋め込みの理論 — Word2Vec・GloVe・FastTextの比較

Word2Vec・GloVe・FastTextの理論と実装を比較解説。埋め込みの事前学習がどう行われるかを詳説します