Mini Transformer Lab ハンズオン

このプロジェクトは、LLM の完成品を作るためではなく、Transformer の主要部品を自作コードで触って理解するための最小実験環境です。

物語形式で概念を追う図解と引用つき解説は LLM Foundations Story にあります。各ハンズオンは、その解説の章と連動しています。

0. セットアップ

devbox run install
devbox run test
devbox run build
devbox run site
devbox run site:build
devbox run site:deploy

以降の任意 uv run ... コマンドは、先に devbox shell に入ってから実行してください。

devbox shell

1. Tokenizer の違いを見る

対応する基礎概念: Text Becomes Tokens

devbox run compare
uv run mini-transformer-compare-tokenizers --file data/natural/tiny.txt
uv run mini-transformer-compare-tokenizers --file data/code/tiny.py.txt

この手順で今動いているコードは、まず文字列を token ID に変える2種類の tokenizer です。

tokenizers.py: char tokenizer and regex tokenizer src/mini_transformer_lab/tokenizers.py L51-L90 Permanent link

@dataclass
class CharTokenizer(_VocabularyMixin):
    """One token per Unicode character."""

    kind: ClassVar[str] = "char"

    @classmethod
    def train(cls, text: str) -> "CharTokenizer":
        return cls(vocab=list(dict.fromkeys(text)))

    def encode(self, text: str) -> list[int]:
        return self._encode_tokens(list(text))

    def decode(self, token_ids: list[int]) -> str:
        return "".join(self.vocab[token_id] for token_id in token_ids)


@dataclass
class RegexTokenizer(_VocabularyMixin):
    """Small regex tokenizer that keeps code identifiers and whitespace chunks intact."""

    kind: ClassVar[str] = "regex"
    pattern: ClassVar[re.Pattern[str]] = re.compile(
        r"\s+|[A-Za-z_][A-Za-z_0-9]*|\d+|[^\sA-Za-z_0-9]",
        flags=re.UNICODE,
    )

    @classmethod
    def train(cls, text: str) -> "RegexTokenizer":
        return cls(vocab=list(dict.fromkeys(cls._tokenize(text))))

    @classmethod
    def _tokenize(cls, text: str) -> list[str]:
        return cls.pattern.findall(text)

    def encode(self, text: str) -> list[int]:
        return self._encode_tokens(self._tokenize(text))

    def decode(self, token_ids: list[int]) -> str:
        return "".join(self.vocab[token_id] for token_id in token_ids)

char は1文字ずつ、regex は識別子・空白・数字・記号をまとまりとして扱います。compare コマンドの差分は、この encode/decode の違いから出ます。

このコマンドで実際に呼ばれる入口は、同じ入力を2つの tokenizer で処理し、語彙数・token 数・先頭 token を表に出すだけの小さな CLI です。

compare_tokenizers.py: compare CLI entrypoint src/mini_transformer_lab/compare_tokenizers.py L17-L27 Permanent link

def main() -> None:
    args = parse_args()
    text = args.text if args.text is not None else Path(args.file).read_text(encoding="utf-8")
    tokenizers = [CharTokenizer.train(text), RegexTokenizer.train(text)]

    print("tokenizer | vocab_size | token_count | first_tokens")
    print("--- | ---: | ---: | ---")
    for tokenizer in tokenizers:
        ids = tokenizer.encode(text)
        tokens = [tokenizer.vocab[token_id].replace("\n", "\\n") for token_id in ids[:16]]
        print(f"{tokenizer.kind} | {tokenizer.vocab_size} | {len(ids)} | {tokens}")

出力の token_count と vocab_size はモデルの性能指標ではなく、同じ文字列がどれだけ違う単位に分解されたかを見るための観察値です。

見るポイント:

char は1文字1トークンなので仕組みが単純
regex は Python の def, return, 空白、記号などをまとまりとして扱う
同じテキストでも token_count と vocab_size が変わる

2. 自然言語を学ぶ

対応する基礎概念: The Only Homework: Predict the Next Token, Dataset Is the Model’s World

devbox run train:natural
uv run mini-transformer-generate --run runs/natural-char --prompt "小さなデータでは"

この手順で今動いているコードは、token ID の列から next-token prediction 用の inputs と targets を作る部分です。

data.py: next-token batch construction src/mini_transformer_lab/data.py L26-L50 Permanent link

def make_lm_batch(
    token_ids: list[int],
    *,
    context_length: int,
    batch_size: int,
    device: str | torch.device = "cpu",
    generator: torch.Generator | None = None,
) -> LanguageModelBatch:
    if context_length < 2:
        raise ValueError("context_length must be at least 2")
    if batch_size < 1:
        raise ValueError("batch_size must be at least 1")
    if len(token_ids) < 2:
        raise ValueError("token_ids must contain at least 2 tokens")

    needed = context_length + 1
    if len(token_ids) < needed:
        repeat_count = (needed // len(token_ids)) + 1
        token_ids = token_ids * repeat_count

    data = torch.tensor(token_ids, dtype=torch.long, device=device)
    max_start = len(data) - needed
    starts = torch.randint(0, max_start + 1, (batch_size,), generator=generator, device=device)
    windows = torch.stack([data[start : start + needed] for start in starts])
    return LanguageModelBatch(inputs=windows[:, :-1], targets=windows[:, 1:])

`inputs` は window の先頭から最後の1つ手前、`targets` は1つ右にずれた列です。モデルは各位置で次の token を当てるように学習します。

学習コマンドの中心は、dataset を読む、tokenizer を作る、モデルを初期化する、loss を計算する、optimizer を進める、という流れです。

train.py: training loop and sample generation src/mini_transformer_lab/train.py L80-L145 Permanent link

    text = read_text(args.dataset)
    tokenizer = build_tokenizer(args.tokenizer, text)
    token_ids = tokenizer.encode(text)
    train_ids, valid_ids = train_valid_split(token_ids)
    if len(valid_ids) < 2:
        valid_ids = train_ids

    config = TransformerConfig(
        vocab_size=tokenizer.vocab_size,
        context_length=args.context_length,
        d_model=args.d_model,
        n_heads=args.n_heads,
        n_layers=args.n_layers,
        d_ff=args.d_ff,
        dropout=args.dropout,
    )
    model = TransformerLanguageModel(config).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
    metrics: list[dict[str, float | int]] = []

    print(json.dumps({
        "dataset": args.dataset,
        "tokenizer": args.tokenizer,
        "tokens": len(token_ids),
        "vocab_size": tokenizer.vocab_size,
        "parameters": count_parameters(model),
        "device": device,
    }, ensure_ascii=False, indent=2))

    for step in range(1, args.steps + 1):
        batch = make_lm_batch(
            train_ids,
            context_length=args.context_length,
            batch_size=args.batch_size,
            device=device,
        )
        _, loss = model(batch.inputs, batch.targets)
        assert loss is not None

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

        if step == 1 or step % args.eval_interval == 0 or step == args.steps:
            valid_loss = estimate_loss(
                model,
                valid_ids,
                context_length=args.context_length,
                batch_size=args.batch_size,
                device=device,
            )
            row = {"step": step, "train_loss": float(loss.item()), "valid_loss": valid_loss}
            metrics.append(row)
            print(json.dumps(row, ensure_ascii=False))

    output_dir = Path(args.output_dir)
    save_experiment(output_dir, model=model, tokenizer=tokenizer, metrics=metrics)

    prompt = args.sample_prompt or text[: min(20, len(text))]
    prompt_ids = tokenizer.encode(prompt)
    generated_ids = model.generate(prompt_ids, max_new_tokens=args.sample_tokens, temperature=0.8)
    sample = tokenizer.decode(generated_ids)
    (output_dir / "sample.txt").write_text(sample, encoding="utf-8")
    print(f"saved: {output_dir}")
    print("--- sample ---")
    print(sample)

`loss.backward()` と `optimizer.step()` が重みを更新する最小単位です。最後に checkpoint と sample.txt を保存するので、generate コマンドが同じ run を読めます。

短いデータなので、自然な文章生成というより「次の文字の癖を覚える」挙動が見えます。

3. コードを学ぶ

対応する基礎概念: Text Becomes Tokens, Dataset Is the Model’s World

devbox run train:code
uv run mini-transformer-generate --run runs/code-regex --prompt "def "

この手順で今動いているコードは、コード向けに識別子や空白をまとまりとして残す regex tokenizer です。

tokenizers.py: regex tokenizer for code-like text src/mini_transformer_lab/tokenizers.py L68-L90 Permanent link

@dataclass
class RegexTokenizer(_VocabularyMixin):
    """Small regex tokenizer that keeps code identifiers and whitespace chunks intact."""

    kind: ClassVar[str] = "regex"
    pattern: ClassVar[re.Pattern[str]] = re.compile(
        r"\s+|[A-Za-z_][A-Za-z_0-9]*|\d+|[^\sA-Za-z_0-9]",
        flags=re.UNICODE,
    )

    @classmethod
    def train(cls, text: str) -> "RegexTokenizer":
        return cls(vocab=list(dict.fromkeys(cls._tokenize(text))))

    @classmethod
    def _tokenize(cls, text: str) -> list[str]:
        return cls.pattern.findall(text)

    def encode(self, text: str) -> list[int]:
        return self._encode_tokens(self._tokenize(text))

    def decode(self, token_ids: list[int]) -> str:
        return "".join(self.vocab[token_id] for token_id in token_ids)

`def` や `return`、改行後の空白、記号を token として残すため、文字単位よりもコードらしい単位で系列を作れます。

見るポイント:

regex tokenizer は def, return, インデントなどの単位を作る
コードデータでは括弧、コロン、改行、インデントが重要なパターンになる
小さいモデルでも、構文らしい断片は比較的早く出る

4. 次元数を変える

対応する基礎概念: Embedding and d_model, Attention and Heads

同じデータで d_model だけ変えます。

uv run mini-transformer-train --dataset data/natural/tiny.txt --tokenizer char --steps 80 --d-model 32 --n-heads 4 --n-layers 2 --context-length 64 --output-dir runs/d32
uv run mini-transformer-train --dataset data/natural/tiny.txt --tokenizer char --steps 80 --d-model 128 --n-heads 4 --n-layers 2 --context-length 64 --output-dir runs/d128

この手順で今動いているコードは、d_model と n_heads から attention head の幅を決める設定です。

model.py: TransformerConfig src/mini_transformer_lab/model.py L11-L28 Permanent link

@dataclass(frozen=True)
class TransformerConfig:
    vocab_size: int
    context_length: int = 128
    d_model: int = 128
    n_heads: int = 4
    n_layers: int = 4
    d_ff: int = 512
    dropout: float = 0.1

    def __post_init__(self) -> None:
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if self.context_length < 2:
            raise ValueError("context_length must be at least 2")

    def to_dict(self) -> dict[str, int | float]:
        return asdict(self)

`d_model` は token 表現の幅です。`d_model % n_heads == 0` が必要なのは、各 head に同じ幅を割り当てるためです。

比較するもの:

parameters
train_loss
valid_loss
runs/*/sample.txt

目安:

d_model が大きいほど1トークンを広いベクトルで表せる
データが小さいと、大きいモデルは汎化より暗記に寄りやすい
d_model は n_heads で割り切れる必要がある

5. 層数を変える

対応する基礎概念: Layers and Feed-Forward Networks

uv run mini-transformer-train --dataset data/code/tiny.py.txt --tokenizer regex --steps 80 --d-model 64 --n-heads 4 --n-layers 1 --context-length 64 --output-dir runs/code-layer1
uv run mini-transformer-train --dataset data/code/tiny.py.txt --tokenizer regex --steps 80 --d-model 64 --n-heads 4 --n-layers 4 --context-length 64 --output-dir runs/code-layer4

この手順で今動いているコードは、attention と feed-forward を1つの block として積み重ねる部分です。

model.py: Transformer blocks src/mini_transformer_lab/model.py L78-L100 Permanent link

class TransformerBlock(nn.Module):
    def __init__(self, config: TransformerConfig) -> None:
        super().__init__()
        self.attention_norm = nn.LayerNorm(config.d_model)
        self.attention = CausalSelfAttention(config)
        self.ffn_norm = nn.LayerNorm(config.d_model)
        self.ffn = FeedForward(config)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.attention(self.attention_norm(x))
        x = x + self.ffn(self.ffn_norm(x))
        return x


class TransformerLanguageModel(nn.Module):
    def __init__(self, config: TransformerConfig) -> None:
        super().__init__()
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.position_embedding = nn.Embedding(config.context_length, config.d_model)
        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
        self.final_norm = nn.LayerNorm(config.d_model)
        self.output = nn.Linear(config.d_model, config.vocab_size)

`n_layers` を変えると、この block の数が変わります。各 block は attention と feed-forward を residual connection で足し戻します。

見るポイント:

層数は「変換を何段重ねるか」
増やすと表現力は上がるが、学習時間も増える
小さいデータでは valid loss が悪化することがある

6. 文脈長を変える

対応する基礎概念: Context Length and the Causal Mask

uv run mini-transformer-train --dataset data/natural/tiny.txt --tokenizer char --steps 80 --context-length 32 --output-dir runs/ctx32
uv run mini-transformer-train --dataset data/natural/tiny.txt --tokenizer char --steps 80 --context-length 128 --output-dir runs/ctx128

この手順で今動いているコードは、未来の token を見ないようにする causal mask です。

model.py: causal self-attention src/mini_transformer_lab/model.py L31-L61 Permanent link

class CausalSelfAttention(nn.Module):
    def __init__(self, config: TransformerConfig) -> None:
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.d_model // config.n_heads
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model)
        self.output = nn.Linear(config.d_model, config.d_model)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        mask = torch.tril(torch.ones(config.context_length, config.context_length))
        self.register_buffer("causal_mask", mask.view(1, 1, config.context_length, config.context_length))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch_size, sequence_length, d_model = x.shape
        qkv = self.qkv(x)
        query, key, value = qkv.split(d_model, dim=2)

        query = query.view(batch_size, sequence_length, self.n_heads, self.head_dim).transpose(1, 2)
        key = key.view(batch_size, sequence_length, self.n_heads, self.head_dim).transpose(1, 2)
        value = value.view(batch_size, sequence_length, self.n_heads, self.head_dim).transpose(1, 2)

        scores = query @ key.transpose(-2, -1)
        scores = scores / math.sqrt(self.head_dim)
        mask = self.causal_mask[:, :, :sequence_length, :sequence_length]
        scores = scores.masked_fill(mask == 0, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        weights = self.attn_dropout(weights)

        attended = weights @ value
        attended = attended.transpose(1, 2).contiguous().view(batch_size, sequence_length, d_model)
        return self.resid_dropout(self.output(attended))

`context_length` は mask と position embedding の最大長にも効きます。長くすると見られる過去は増えますが、attention の計算量も増えます。

見るポイント:

context_length は何トークン前まで見るか
長くすると遠い依存関係を扱える
attention は系列長に対して重くなる

7. データセットを差し替える

対応する基礎概念: Dataset Is the Model’s World, Train Loss and Valid Loss

任意の UTF-8 テキストを置けば使えます。

uv run mini-transformer-train --dataset path/to/your.txt --tokenizer char --steps 200 --output-dir runs/my-text

コードだけを学ばせる場合:

find path/to/code -name "*.py" -maxdepth 4 -print0 | xargs -0 cat > data/code/my-code.txt
uv run mini-transformer-train --dataset data/code/my-code.txt --tokenizer regex --steps 300 --output-dir runs/my-code

注意:

個人情報や秘密鍵を含むデータは入れない
小さいデータでは覚え込みが起きる
データの癖、文体、バグ、偏りはそのまま学習対象になる

8. 生成パラメータを変える

対応する基礎概念: Decoding Changes the Voice, Not the Knowledge

同じ run と prompt で、生成時の選び方だけを変えます。

uv run mini-transformer-generate --run runs/natural-char --prompt "小さな" --temperature 0
uv run mini-transformer-generate --run runs/natural-char --prompt "小さな" --temperature 0.8
uv run mini-transformer-generate --run runs/natural-char --prompt "小さな" --temperature 0.8 --top-k 8

生成時に動くコードは、最後の位置の logits から次 token を選び、その token を列に追加する loop です。

model.py: autoregressive generation src/mini_transformer_lab/model.py L130-L160 Permanent link

    @torch.no_grad()
    def generate(
        self,
        token_ids: list[int],
        *,
        max_new_tokens: int = 80,
        temperature: float = 0.8,
        top_k: int | None = None,
    ) -> list[int]:
        device = next(self.parameters()).device
        generated = torch.tensor([token_ids], dtype=torch.long, device=device)

        for _ in range(max_new_tokens):
            context = generated[:, -self.config.context_length :]
            logits, _ = self(context)
            next_logits = logits[:, -1, :]
            if next_logits.size(-1) > 1:
                next_logits[:, 0] = float("-inf")

            if temperature <= 0:
                next_token = torch.argmax(next_logits, dim=-1, keepdim=True)
            else:
                next_logits = next_logits / temperature
                if top_k is not None:
                    values, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1)))
                    next_logits[next_logits < values[:, [-1]]] = float("-inf")
                probs = F.softmax(next_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            generated = torch.cat([generated, next_token], dim=1)

        return generated[0].tolist()

`temperature=0` は argmax、`temperature>0` は確率分布から sampling します。`top_k` は候補を上位 k 個に絞るだけで、モデルの知識そのものは増えません。

見るポイント:

temperature=0 は安定しやすい
temperature を上げると出力が揺れやすい
top-k は候補を絞る
生成パラメータは、学習済みの能力や知識を増やすものではない

9. Agent 開発へ接続する

対応する基礎概念: Why This Is Not Yet ChatGPT

この lab のモデルは、system prompt や user prompt に従う assistant ではありません。base language model が「続きを予測する」仕組みを理解するための教材です。

agent 開発に進むときは、次を別々に考えます。

model が学習済みの能力
prompt / chat template による入力構造
tool schema による行動範囲
RAG による外部知識
validation による出力検査

この分離ができると、「prompt を強く書けば解決する問題」と「tool 権限、データ、検証で解くべき問題」を切り分けやすくなります。

最小設定の意味

概念と影響の一覧は Concept Impact Checklist にもあります。

パラメータ	意味	まず触る値
`tokenizer`	文字単位か、簡易構文単位か	`char`, `regex`
`context_length`	何トークン前まで見るか	`32`, `64`, `128`
`d_model`	1トークンのベクトル次元	`32`, `64`, `128`
`n_heads`	attention の分割数	`4`
`n_layers`	Transformer block の段数	`1`, `2`, `4`
`d_ff`	feed-forward の中間次元	`4 * d_model` が目安
`steps`	学習更新回数	`80`, `200`, `500`

ファイル対応

src/mini_transformer_lab/model.py: Transformer 本体
src/mini_transformer_lab/tokenizers.py: tokenizer の差分
src/mini_transformer_lab/data.py: next-token prediction 用 batch
src/mini_transformer_lab/train.py: 学習 CLI
src/mini_transformer_lab/generate.py: 生成 CLI
data/natural/tiny.txt: 自然言語の最小データ
data/code/tiny.py.txt: コードの最小データ
docs/LLM_FOUNDATIONS_STORY.md: 図解、概念ごとの影響、外部参考リンク
src/pages/index.astro: Webページ版の入口
src/pages/foundations.astro: docs/LLM_FOUNDATIONS_STORY.md を直接描画する Astro ページ
src/pages/hands-on.astro: docs/HANDS_ON.mdx を直接描画する Astro ページ
src/site/remarkDocLinks.mjs: Markdown 内の文書間リンクを Web ルートへ変換する remark plugin
wrangler.jsonc: Cloudflare Pages の build output 設定
.github/workflows/cloudflare-pages.yml: Cloudflare Pages への GitHub Actions deploy