Transformer(GPT / ViT / MAE)

Transformer の中心にあるのは自己注意です。系列の各要素が、他の要素をどれだけ参照するかを自分で決め、その重み付き和で文脈表現を作ります。

このノートでは、自己注意の最小式から始めて、因果マスク、位置情報、ViT のパッチ列、MAE のマスク再構成までを一つの視点でつなぎます。

Transformer は「順に読む」代わりに「どこを見るかを毎回選ぶ」

RNN が前から順に情報を渡していくのに対し、Transformer は各トークンが他のトークンをどれだけ参照するかを毎回計算します。この視点に立つと、GPT も ViT も MAE も、入力は違っても同じ骨格で読めます。

まず小さな attention 行列を作って、「各トークンがどこを見たか」を可視化します。そこへ未来を隠すマスクを入れると GPT 風になり、画像パッチをトークンとして並べると ViT 風になります。後半の話題が散らばって見えないよう、最初から同じ骨格で読みます。

ここでは巨大モデルを追うのではなく、attention の見方を身につける

重要なのはパラメータ数ではなく、attention weight が何を表しているかです。各行がどこを見たか、mask を入れると何が遮られるか、位置情報がないと何が崩れるかを順に確かめます。

GPT, ViT, MAE を別物として覚えない

自己回帰、画像分類、自己教師あり再構成という目的は違っても、入力をトークン列へし、attention をかけるという骨格は共通です。この notebook では差分より先に共通骨格を押さえます。

mask の使い方が、モデルの役割を決める

未来を隠せば GPT になり、一部パッチを隠して復元させれば MAE になります。Transformer 自体より、どこを見せてどこを隠すかの設計がタスクを規定していることを意識して読み進めてください。

後半は vision の例ですが、注目すべきは入力の違いよりトークン化の設計です

画像を patch に切ると、Transformer は文章以外にもそのまま使えます。ここでは NLP 専用のモデルと見るのではなく、可変な参照機構として読み直します。

まずは Q, K, V から attention を作る

最初に、入力埋め込みから Q, K, V を作り、重み行列がどうできるかを最小コードで見ます。ここが Transformer の核です。

import math
import numpy as np
import matplotlib.pyplot as plt

try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    TORCH_AVAILABLE = True
except ModuleNotFoundError:
    torch = None
    nn = None
    optim = None
    TORCH_AVAILABLE = False

自己注意では、入力埋め込み X から Q, K, V を作り、次で重みを計算します。

Attention(Q, K, V) = softmax((QK^T) / sqrt(d_k)) V

ここで X の形状を (T, d_model)、射影行列 W_Q, W_K, W_V(d_model, d_k) とすると、
Q, K, V の形状は (T, d_k) になります。
要素で見ると score_{i,j} = (q_i ・ k_j) / sqrt(d_k) で、行ごとに softmax して V を重み付き和します。

Q は「何を探しているか」、K は「何を持っているか」、V は「実際に受け渡す情報」と見ると理解しやすくなります。

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


# 4トークン, 埋め込み次元3の玩具例
X = np.array([
    [1.0, 0.2, 0.0],   # token 0
    [0.9, 0.1, 0.1],   # token 1
    [0.1, 0.8, 0.2],   # token 2
    [0.0, 0.7, 1.0],   # token 3
], dtype=np.float64)

W_Q = np.array([[0.8, 0.0, 0.1], [0.2, 0.9, 0.1], [0.1, 0.2, 0.7]])
W_K = np.array([[0.7, 0.1, 0.1], [0.1, 0.8, 0.2], [0.2, 0.1, 0.6]])
W_V = np.array([[1.0, 0.2, 0.0], [0.1, 0.9, 0.1], [0.0, 0.2, 0.8]])

Q = X @ W_Q
K = X @ W_K
V = X @ W_V

scores = (Q @ K.T) / math.sqrt(Q.shape[-1])
weights = softmax_rowwise(scores)
context = weights @ V

print('attention weights:')
print(np.round(weights, 4))
print('\ncontext vectors:')
print(np.round(context, 4))

GPT では未来を見ないようにする

次は因果マスクです。自己回帰モデルでは、未来トークンを見えてしまうと次トークン予測の意味が壊れるので、ここを強制的に隠します。

plt.figure(figsize=(5.0, 4.2))
plt.imshow(weights, cmap='Blues')
plt.colorbar(label='attention weight')
plt.xticks(range(len(X)), [f'k{j}' for j in range(len(X))])
plt.yticks(range(len(X)), [f'q{i}' for i in range(len(X))])
plt.title('Self-Attention Weight Matrix')
plt.tight_layout()
plt.show()

GPTのような自己回帰モデルでは、未来トークンを見ないように因果マスク(causal mask)を入れます。
j > i(未来位置)をマスクし、スコアを -∞ 相当に落として softmax 後の重みをほぼ0にします。

下のコードでは、同じ入力で「マスクなし」と「マスクあり」を比較します。

def causal_masked_attention(Q, K, V):
    T, d = Q.shape
    scores = (Q @ K.T) / math.sqrt(d)
    mask = np.triu(np.ones((T, T), dtype=bool), k=1)  # future positions
    scores_masked = scores.copy()
    scores_masked[mask] = -np.inf
    w = softmax_rowwise(scores_masked)
    return w @ V, w


context_full = softmax_rowwise((Q @ K.T) / math.sqrt(Q.shape[-1])) @ V
context_causal, w_causal = causal_masked_attention(Q, K, V)

print('full attention (token 1):', np.round(context_full[1], 4))
print('causal attention (token 1):', np.round(context_causal[1], 4))
print('\ncausal weight matrix:')
print(np.round(w_causal, 4))

順序情報がないと何が困るか

注意機構だけでは、単語の順番やパッチの並びを区別できません。そこで位置情報を足して、並びそのものを表現へ埋め込みます。

def sinusoidal_positional_encoding(seq_len, d_model):
    pe = np.zeros((seq_len, d_model), dtype=np.float64)
    pos = np.arange(seq_len)[:, None]
    i = np.arange(d_model)[None, :]
    angle_rates = 1.0 / np.power(10000, (2 * (i // 2)) / d_model)
    angles = pos * angle_rates

    pe[:, 0::2] = np.sin(angles[:, 0::2])
    pe[:, 1::2] = np.cos(angles[:, 1::2])
    return pe


pe = sinusoidal_positional_encoding(seq_len=24, d_model=16)
print('positional encoding shape:', pe.shape)

plt.figure(figsize=(7.2, 3.6))
plt.imshow(pe.T, aspect='auto', cmap='coolwarm')
plt.colorbar(label='value')
plt.xlabel('position')
plt.ylabel('channel')
plt.title('Sinusoidal Positional Encoding')
plt.tight_layout()
plt.show()

画像もトークン列にしてしまう

ViT は画像をパッチ列に分けて、文章のトークン列と同じように扱います。ここでは「画像でもトークン列を作れば Transformer が使える」という発想を押さえてください。

def patchify(image, patch_size=2):
    h, w = image.shape
    assert h % patch_size == 0 and w % patch_size == 0
    patches = []
    for y in range(0, h, patch_size):
        for x in range(0, w, patch_size):
            p = image[y:y+patch_size, x:x+patch_size].reshape(-1)
            patches.append(p)
    return np.stack(patches, axis=0)


img = np.arange(64, dtype=np.float64).reshape(8, 8) / 63.0
patches = patchify(img, patch_size=2)  # (16, 4)

W_patch = np.random.default_rng(0).normal(0, 0.4, size=(4, 6))
patch_tokens = patches @ W_patch  # (16, 6)

# 実際のViTは学習可能な位置埋め込みを加える
pos_embed = np.random.default_rng(1).normal(0, 0.1, size=patch_tokens.shape)
patch_tokens = patch_tokens + pos_embed

# CLSトークンも通常は学習可能ベクトル(ここでは最小例として0初期化)
cls_token = np.zeros((1, 6), dtype=np.float64)
vit_tokens = np.concatenate([cls_token, patch_tokens], axis=0)  # (17, 6)

print('image shape      :', img.shape)
print('patches shape    :', patches.shape)
print('patch token shape:', patch_tokens.shape)
print('ViT token shape  :', vit_tokens.shape, '(CLS + patches)')

plt.figure(figsize=(3.8, 3.8))
plt.imshow(img, cmap='gray')
plt.title('Toy image (8x8)')
plt.axis('off')
plt.show()

MAE は隠したパッチを復元する

MAE では、入力の一部を隠し、その欠けた部分を当てるように学習します。生成モデルというより、表現学習のための再構成タスクとして読むと整理しやすくなります。

def make_random_mask(n_tokens, mask_ratio=0.75, seed=0):
    rng = np.random.default_rng(seed)
    n_mask = int(n_tokens * mask_ratio)
    perm = rng.permutation(n_tokens)
    mask_idx = perm[:n_mask]
    keep_idx = np.sort(perm[n_mask:])
    return keep_idx, np.sort(mask_idx)


n_patches = patches.shape[0]
keep_idx, mask_idx = make_random_mask(n_patches, mask_ratio=0.75, seed=1)
kept_tokens = patch_tokens[keep_idx]

print('all patches :', n_patches)
print('kept patches:', len(keep_idx), 'indices=', keep_idx)
print('masked      :', len(mask_idx), 'indices=', mask_idx)
print('encoder input token shape (without CLS):', kept_tokens.shape)

# デコーダ入力側で元の長さに戻す(masked位置にはmask tokenを置く)
mask_token = np.zeros((len(mask_idx), patch_tokens.shape[1]))
decoder_input = np.zeros_like(patch_tokens)
decoder_input[keep_idx] = kept_tokens
decoder_input[mask_idx] = mask_token

# 最小デモ: 線形デコーダでパッチを復元し、masked部分のみ誤差を計算
W_rec = np.random.default_rng(2).normal(0, 0.3, size=(patch_tokens.shape[1], patches.shape[1]))
recon_patches = decoder_input @ W_rec
masked_recon_mse = np.mean((recon_patches[mask_idx] - patches[mask_idx]) ** 2)
print('masked reconstruction MSE (toy):', round(float(masked_recon_mse), 6))

masked_view = img.copy()
patch_size = 2
for idx in mask_idx:
    gy = idx // (img.shape[1] // patch_size)
    gx = idx % (img.shape[1] // patch_size)
    y0, x0 = gy * patch_size, gx * patch_size
    masked_view[y0:y0+patch_size, x0:x0+patch_size] = 0.0

fig, axes = plt.subplots(1, 2, figsize=(7.2, 3.4))
axes[0].imshow(img, cmap='gray', vmin=0, vmax=1)
axes[0].set_title('original')
axes[0].axis('off')
axes[1].imshow(masked_view, cmap='gray', vmin=0, vmax=1)
axes[1].set_title('MAE masking (75%)')
axes[1].axis('off')
plt.tight_layout()
plt.show()

VAE の再構成と何が違うか

VAE も再構成を使いますが、潜在分布を学ぶ生成モデルです。MAE は可視パッチから隠したパッチを復元することで良い表現を作る、という点が違います。

# VAEの再パラメータ化トリックと損失項の最小計算
mu = np.array([0.2, -0.4, 0.1], dtype=np.float64)
logvar = np.array([-0.2, 0.3, -0.5], dtype=np.float64)

rng = np.random.default_rng(42)
eps = rng.normal(0.0, 1.0, size=mu.shape)
std = np.exp(0.5 * logvar)
z = mu + std * eps

# KL項
kl = -0.5 * np.sum(1 + logvar - mu**2 - np.exp(logvar))

# 再構成項(ここではMSEの玩具例)
x_true = np.array([0.9, 0.2, 0.7], dtype=np.float64)
x_recon = np.array([0.8, 0.3, 0.6], dtype=np.float64)
recon_loss = np.mean((x_true - x_recon) ** 2)

beta = 1.0
vae_loss = recon_loss + beta * kl

print('mu    =', np.round(mu, 4))
print('logvar=', np.round(logvar, 4))
print('z sample =', np.round(z, 4))
print('recon loss =', round(float(recon_loss), 6))
print('KL(q(z|x) || p(z)) =', round(float(kl), 6))
print('VAE loss = recon + beta*KL =', round(float(vae_loss), 6))

最後に GPT 型の最小学習へつなぐ

最後は小さな decoder-only Transformer で次トークン予測を行います。ここまで見た attention, causal mask, token 列の考え方が、そのまま言語モデルの学習へ入っていくことを確認します。

if TORCH_AVAILABLE:
    torch.manual_seed(0)

    vocab_size = 20
    seq_len = 14
    d_model = 48

    def sample_batch(batch_size=64):
        # Fibonacci-like modulo sequence: x_t = (x_{t-1} + x_{t-2}) mod vocab
        seq = torch.zeros(batch_size, seq_len + 1, dtype=torch.long)
        seq[:, 0] = torch.randint(0, vocab_size, (batch_size,))
        seq[:, 1] = torch.randint(0, vocab_size, (batch_size,))
        for t in range(2, seq_len + 1):
            seq[:, t] = (seq[:, t - 1] + seq[:, t - 2]) % vocab_size
        return seq[:, :-1], seq[:, 1:]

    class TinyGPT(nn.Module):
        def __init__(self, vocab_size, d_model=48, nhead=4, num_layers=2, seq_len=14):
            super().__init__()
            self.seq_len = seq_len
            self.token_emb = nn.Embedding(vocab_size, d_model)
            self.pos_emb = nn.Parameter(torch.zeros(1, seq_len, d_model))
            layer = nn.TransformerEncoderLayer(
                d_model=d_model,
                nhead=nhead,
                dim_feedforward=4 * d_model,
                dropout=0.0,
                batch_first=True,
                activation='gelu',
            )
            self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
            self.norm = nn.LayerNorm(d_model)
            self.head = nn.Linear(d_model, vocab_size)

        def forward(self, x):
            bsz, t = x.shape
            h = self.token_emb(x) + self.pos_emb[:, :t, :]
            mask = torch.triu(torch.ones(t, t, device=x.device), diagonal=1).bool()
            h = self.encoder(h, mask=mask)
            h = self.norm(h)
            return self.head(h)

    model = TinyGPT(vocab_size=vocab_size, d_model=d_model, nhead=4, num_layers=2, seq_len=seq_len)
    optimizer = optim.AdamW(model.parameters(), lr=3e-3)
    criterion = nn.CrossEntropyLoss()

    for step in range(180):
        x, y = sample_batch(batch_size=64)
        logits = model(x)

        # t=0 の予測は初期2トークン中1つ目だけでは不確定なので除外
        loss = criterion(logits[:, 1:, :].reshape(-1, vocab_size), y[:, 1:].reshape(-1))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if step % 45 == 0:
            print(f'step={step:>3d}, loss={loss.item():.4f}')

    model.eval()
    seed = torch.tensor([[3, 7]], dtype=torch.long)
    generated = seed.clone()
    for _ in range(10):
        cur = generated[:, -seq_len:]
        logits = model(cur)
        next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
        generated = torch.cat([generated, next_token], dim=1)

    print('generated tokens:', generated.squeeze(0).tolist())
else:
    print('PyTorch未導入のためGPTミニ実験セルはスキップしました。')

Transformer を使うときに本当に中心になるのは、「どの入力をトークン列として作るか」と「どこを見せてどこを隠すか」です。文章でも画像でも、その設計がモデルの役割を決めます。