一文要約: インファレンスとは、モデルの読み込み → プロンプトのエンコード → Autoregressive な生成 → 出力のデコード、です。コードは30行ほどですが、これこそモデルが「話し始める」瞬間です。

完全なコードリポジトリ: github.com/waylandzhang/Transformer-from-scratch

20.1 インファレンスと学習の違い

20.1.1 第16章の復習

20.1.2 インファレンスのコアフロー

1. 学習済みモデルを読み込む
2. プロンプトをトークン ID にエンコードする
3. Autoregressive に生成する(1トークンずつ)
4. トークン ID をテキストに戻す

20.2 モデルの読み込み

20.2.1 チェックポイントを読み込む

# load model
import torch
import tiktoken
from model import Model

# load checkpoint
checkpoint = torch.load('model/model.ckpt')

# restore hyperparameters from checkpoint
h_params = checkpoint['h_params']

# reconstruct model architecture
model = Model(h_params)

# load parameters
model.load_state_dict(checkpoint['model_state_dict'])

# switch to evaluation mode
model.eval()

# move to the correct device
model.to(h_params['device'])

20.2.2 なぜ model.eval() が必要か

model.eval() を呼び出すと、次の2つの効果があります。

1. Dropout を無効化: インファレンスではランダムなドロップは不要です
2. BatchNorm を固定: 学習中に蓄積された統計量を使います

評価モードに切り替えないと、インファレンスのたびに違う出力が出てしまいます。通常それは望ましくありません。

20.3 入力の準備

20.3.1 プロンプトをエンコードする

# encode input
encoding = tiktoken.get_encoding("cl100k_base")

# what do you want the model to continue?
start = "fix(auth): "

# encode to token IDs
start_ids = encoding.encode(start)
print(f"Prompt: {start}")
print(f"Token IDs: {start_ids}")

# convert to Tensor
x = torch.tensor(start_ids, dtype=torch.long, device=h_params['device'])
x = x.unsqueeze(0)  # add batch dimension: [seq_len] -> [1, seq_len]

print(f"Input shape: {x.shape}")

出力例:

Prompt: fix(auth):
Token IDs: [11148, 7, 3997, 1680]
Input shape: torch.Size([1, 4])

モデルは生の文字列を見ることはありません。エンベディングテーブルの行に対応する整数を見ているだけです。unsqueeze(0) はバッチ次元を追加します — モデルは単一プロンプトであっても [batch, seq] 形式を期待します。

20.4 テキストの生成

20.4.1 生成関数を呼び出す

# generate text
with torch.no_grad():  # no gradient computation
    y = model.generate(
        x,
        max_new_tokens=200,   # generate up to 200 new tokens
        temperature=0.5,       # temperature: lower = more deterministic
        top_k=None            # no top-k filtering
    )

# decode
output_text = encoding.decode(y[0].tolist())

print('---------------')
print(output_text)
print('---------------')

20.4.2 生成出力の例

---------------
fix(auth): handle expired refresh tokens before retry
fix(auth): validate JWT signature against rotated key set
fix(auth): clear session cookie on logout in incognito tabs
fix(auth): reject empty Authorization header with 401
---------------

モデルはプルリクエストのタイトル風のテキストを生成するように学習しました — 学習データのスタイルです。これは大事なポイントです: モデルはデータのパターンを学習しているのであって、内容の意味を理解しているわけではありません。

20.5 生成パラメータの詳細

20.5.1 Temperature

y = model.generate(x, temperature=0.5)

Temperature は出力分布の形を変えます。

20.5.2 Top-K サンプリング

y = model.generate(x, top_k=50)

確率上位 K 個のトークンの中からのみサンプリングします。

Original distribution:
  " handle" = 0.30, " validate" = 0.20, " refresh" = 0.15, ...  (100k tokens)

After Top-K=3:
  " handle" = 0.46, " validate" = 0.31, " refresh" = 0.23
  (renormalized over just these 3 tokens)

何が嬉しいか: たまに極端に低確率の奇妙なトークンをサンプリングしてしまい、せっかく一貫した出力を台無しにする、という事態を防げます。

20.5.3 Max New Tokens

y = model.generate(x, max_new_tokens=200)

生成の長さを制御します。

• 短すぎる: 出力が途中で切れる
• 長すぎる: 計算を浪費し、繰り返しが起こりやすくなる

20.6 モデルパラメータの確認

20.6.1 パラメータ数を表示する

# count parameters
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model param size: {total_params:,}")

出力例:

Model param size: 8,234,560

20.6.2 レイヤーごとに確認する

# print each layer's name and shape
for name, param in model.state_dict().items():
    print(f"{name}: {param.shape}")

出力例:

token_embedding_lookup_table.weight: torch.Size([100256, 80])
transformer_blocks.0.ln1.weight: torch.Size([80])
transformer_blocks.0.ln1.bias: torch.Size([80])
transformer_blocks.0.mha.heads.0.Wq.weight: torch.Size([20, 80])
transformer_blocks.0.mha.heads.0.Wk.weight: torch.Size([20, 80])
transformer_blocks.0.mha.heads.0.Wv.weight: torch.Size([20, 80])
...
model_out_linear_layer.weight: torch.Size([100256, 80])
model_out_linear_layer.bias: torch.Size([100256])

これがモデルの完全なパラメータレイアウトです — 各行列に名前と形状があります。これらの名前を第18章のクラス定義と突き合わせて確認できます。

20.7 inference.py の全体像

# -*- coding: utf-8 -*-
"""
Sample from a trained model
"""
import torch
import tiktoken
from model import Model

# load model and hyperparameters
checkpoint = torch.load('model/model.ckpt')
h_params = checkpoint['h_params']
model = Model(h_params)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
model.to(h_params['device'])

# load tokenizer
encoding = tiktoken.get_encoding("cl100k_base")

# input prompt
start = "fix(auth): "
start_ids = encoding.encode(start)
x = torch.tensor(start_ids, dtype=torch.long, device=h_params['device'])[None, ...]

# generate
with torch.no_grad():
    y = model.generate(x, max_new_tokens=200, temperature=0.5, top_k=None)
    print('---------------')
    print(encoding.decode(y[0].tolist()))
    print('---------------')

# print parameter count
total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Model param size: {total_params:,}")

# print parameter shapes
for name in model.state_dict().keys():
    print(name, model.state_dict()[name].shape)

20.8 プロンプトを変えてみる

20.8.1 異なる入力を試す

# try different prompts
prompts = [
    "fix(auth):",
    "feat(api):",
    "refactor(db):",
    "docs(readme):"
]

for prompt in prompts:
    x = torch.tensor(encoding.encode(prompt), dtype=torch.long, device=h_params['device'])[None, ...]
    with torch.no_grad():
        y = model.generate(x, max_new_tokens=50, temperature=0.5)
    print(f"Prompt: {prompt}")
    print(f"Output: {encoding.decode(y[0].tolist())}")
    print("---")

20.8.2 出力の観察

モデルは学習データのスタイルで生成します。

• 学習データが PR タイトルなら、PR タイトル風のテキストを生成する
• 学習データが小説なら、小説風のテキストを生成する
• 学習データがコードなら、コード風のテキストを生成する

モデルはデータの統計的パターンを学習しているのであって、その背後にある「意味」を学習しているわけではありません。 これがモデルの能力の源であると同時に、限界の源でもあります。

20.9 Autoregressive 生成の可視化

20.9.1 ステップごとに見ていく

# visualize generation process
def generate_with_trace(model, x, max_new_tokens=10, temperature=1.0):
    """Generation with step-by-step tracing"""
    encoding = tiktoken.get_encoding("cl100k_base")

    print(f"Initial prompt: {encoding.decode(x[0].tolist())}")
    print("---")

    for i in range(max_new_tokens):
        # forward pass
        with torch.no_grad():
            logits, _ = model(x[:, -model.context_length:])

        # get last-position predictions
        logits = logits[:, -1, :] / temperature
        probs = torch.softmax(logits, dim=-1)

        # get top-5 candidates
        top5_probs, top5_ids = torch.topk(probs[0], 5)
        print(f"Step {i+1} candidates:")
        for prob, idx in zip(top5_probs, top5_ids):
            print(f"  '{encoding.decode([idx.item()])}': {prob.item():.3f}")

        # sample
        idx_next = torch.multinomial(probs, num_samples=1)
        x = torch.cat((x, idx_next), dim=1)

        print(f"  -> selected: '{encoding.decode([idx_next[0].item()])}'")
        print(f"  current sequence: {encoding.decode(x[0].tolist())}")
        print("---")

    return x

20.9.2 トレース出力の例

Initial prompt: fix(auth):
---
Step 1 candidates:
  ' handle': 0.312
  ' validate': 0.198
  ' refresh': 0.087
  ' verify': 0.076
  ' check': 0.065
  -> selected: ' handle'
  current sequence: fix(auth): handle
---
Step 2 candidates:
  ' expired': 0.421
  ' missing': 0.156
  ' invalid': 0.089
  ' empty': 0.067
  ' stale': 0.054
  -> selected: ' expired'
  current sequence: fix(auth): handle expired
---
...

このトレースで Autoregressive のプロセスが具体的にイメージできます。各ステップで蓄積されたシーケンスを見て、確率分布を出力する。サンプリングは確率的なので — 2回実行すれば違うテキストが出ることもあります。

20.10 よくある問題

20.10.1 出力が繰り返しになる

症状: モデルが同じ単語やフレーズを繰り返し続ける。

原因:

• temperature が低すぎる
• 学習データ自体に繰り返しパターンがある
• モデルが過学習している

対策:

• temperature を上げる
• top-k や top-p サンプリングを使う
• 繰り返しペナルティを加える

20.10.2 出力が支離滅裂

症状: 出力が文字化けしているか、意味をなさない。

原因:

• モデルが十分に学習されていない
• プロンプトが学習分布から大きく外れている
• temperature が高すぎる

対策:

• 学習ステップを増やす
• 分布内のプロンプトを使う
• temperature を下げる

20.10.3 生成が遅い

症状: 1トークンの生成に時間がかかる。

原因:

• GPU を使っていない
• KV Cache がない
• モデルが大きすぎる

対策:

• GPU が使えるなら使う
• KV Cache を実装する(第22章)
• 小さいモデルを使う

20.11 章のまとめ

20.11.1 3ステップのインファレンス

1. Load model
   checkpoint = torch.load('model.ckpt')
   model.load_state_dict(checkpoint['model_state_dict'])
   model.eval()

2. Encode prompt
   start_ids = encoding.encode(prompt)
   x = torch.tensor(start_ids)[None, ...]

3. Generate
   with torch.no_grad():
       y = model.generate(x, max_new_tokens=200)
   output = encoding.decode(y[0].tolist())

20.11.2 主要パラメータ

20.11.3 コアの洞察

inference.py は30行のコードに過ぎませんが、これこそ全行程の終着地点 — モデルが「話す」瞬間です。パラメータを読み込み、プロンプトをエンコードし、Autoregressive に生成し、出力をデコードする。私たちの教育用おもちゃモデルから GPT-4 まで、すべての GPT 系モデルがプロンプトに応える方法はこれと同じです。

章末チェックリスト

この章を終えた時点で、次のことができるようになっているはずです。

• 学習済みモデルのチェックポイントを正しく読み込める。
• model.eval() の役割と必要性を説明できる。
• temperature と top-k で生成を制御できる。
• インファレンススクリプトを最初から最後まで実行できる。

第5部のまとめ

これで コード実装 パートを完了しました。

400行未満のコード で、動作する完全な Transformer を実装したことになります。

これらは本番環境の LLM と比べると簡略化されていますが、コアロジックは同じです。このコードを理解できれば、Hugging Face Transformers、LLaMA、GPT-NeoX のソースコードを読んでも、すべての構成要素を見分けられるようになります。

完全なコード

第5部の完全な実装は GitHub にあります。

github.com/waylandzhang/Transformer-from-scratch

含まれているもの:

• model.py — モデル定義一式
• train.py — 学習スクリプト
• inference.py — 推論スクリプト
• step-by-step.ipynb — 注釈付き Jupyter ノートブック

次章でお会いしましょう

私たちのモデルは動きます。しかし遅い。トークンを1つ生成するたびに、シーケンス全体について完全な順伝播を走らせ、毎回すべての Key と Value 行列をゼロから計算し直しているからです。これは無駄です。

第6部では 本番向けの最適化 を扱います。現代の Transformer インファレンスシステムで最も効果の大きい2つの高速化、Flash Attention と KV Cache を見ていきます。KV Cache は、いま見たばかりの冗長な計算をなくします — これまでのトークンすべてについて K と V を再計算する代わりに、それらをメモリに保持しておき、新しいトークン分の寄与だけを計算するのです。この最適化ひとつで、インファレンスは数倍速くなります。

それでは次章で、また一緒に学んでいきましょう。