みんなのためのAI 第1回 — 1600万パラメータの言語モデルを15分でゼロから学習させる

• はじめに — なぜ小さなモデルを自分で作るのか
• 実験環境
• データセット — TinyStories
• モデル構造 — デコーダ専用トランスフォーマー
  • 因果マスクがなければ学習は成立しない
  • 重み共有 — パラメータの40%を節約する
  • Pre-norm 配置
• 学習ループ
• 学習ログ
• 生成結果
• perplexity 8 が意味すること
• まとめ
• 🧠 理解度チェッククイズ
• 参考資料

はじめに — なぜ小さなモデルを自分で作るのか

巨大言語モデルの使い方を扱う記事は数多くあります。しかし「中で何が起きているのか」は API の背後に隠れています。このシリーズは逆方向に進みます。 自分で学習させてみる のです。

目標は性能ではありません。条件はこう設定しました。

• パラメータ 1億個未満 (実際にはずっと少ないです)
• GPU 1枚で 20分以内 の学習
• データは Hugging Face に公開されているものだけ
• 学習ログと成果物を すべて残す

小さく作ると全体の構造が一目で分かります。アテンションがなぜ必要か、損失関数が何を測るのか、学習が失敗するときにどんな兆候が現れるのかを、数分単位で観察できます。

今回はその第一歩、 テキスト生成モデル です。

実験環境

実際に使用した機材とバージョンです。

データセット — TinyStories

roneneldan/TinyStories は、3〜4歳の子どもが理解できる語彙だけで書かれた短い英語の童話集です。このデータセットが小さなモデルの実験に理想的な理由があります。

一般的なウェブテキストは語彙が膨大なため、小さなモデルは文法を身につける前に単語を覚えることで容量を使い切ってしまいます。TinyStories は語彙を意図的に制限しているので、 小さなモデルでも「文を作る方法」そのものに容量を使える ようになります。

トークナイザーは新しく作らず、GPT-2 のものをそのまま借りました。ハブに公開されたモデルの一部だけを使うのも有用な方法です。

from transformers import AutoTokenizer
from datasets import load_dataset

tok = AutoTokenizer.from_pretrained("openai-community/gpt2")
ds = load_dataset("roneneldan/TinyStories", split="train[:40000]")

ids = []
for ex in ds:
    ids.extend(tok(ex["text"]).input_ids + [tok.eos_token_id])
    if len(ids) >= 4_000_000:
        break
data = torch.tensor(ids[:4_000_000], dtype=torch.long)

トークンを一つの長いストリームにつなげました。学習時にはこのストリームから任意の位置を切り取って使います。文書の境界をまたぐサンプルが生じますが、 eos_token_id が境界を示しているので、モデルは「ここで物語が終わる」ことを学びます。

モデル構造 — デコーダ専用トランスフォーマー

全体の構造はこれだけです。4層、埋め込み次元 256、ヘッド 8個です。

CTX, DIM, LAYERS, HEADS = 128, 256, 4, 8

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(DIM), nn.LayerNorm(DIM)
        self.attn = nn.MultiheadAttention(DIM, HEADS, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(DIM, 4 * DIM), nn.GELU(), nn.Linear(4 * DIM, DIM)
        )

    def forward(self, x, mask):
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))


class TinyGPT(nn.Module):
    def __init__(self, vocab):
        super().__init__()
        self.emb = nn.Embedding(vocab, DIM)
        self.pos = nn.Embedding(CTX, DIM)
        self.blocks = nn.ModuleList(Block() for _ in range(LAYERS))
        self.ln = nn.LayerNorm(DIM)
        mask = torch.triu(torch.full((CTX, CTX), float("-inf")), 1)
        self.register_buffer("mask", mask)

    def forward(self, x):
        T = x.shape[1]
        h = self.emb(x) + self.pos(torch.arange(T, device=x.device))
        for b in self.blocks:
            h = b(h, self.mask[:T, :T])
        return self.ln(h) @ self.emb.weight.T  # weight tying

ここで確認すべき設計判断が3つあります。

因果マスクがなければ学習は成立しない

torch.triu(..., 1) で作った上三角行列は、対角線より上をすべて負の無限大で埋めます。ソフトマックスを通すとこれらの位置の重みが 0 になります。つまり 各位置が自分より後ろのトークンを見られないようにする 仕組みです。

これがないとどうなるでしょうか。モデルは次のトークンを当てる問題を解くのに、正解である次のトークンをすでに入力として見ていることになります。損失は 0 に近く落ちますが、生成はまったくできません。 「損失が速く落ちすぎる」のは、たいてい正解が漏れているという信号 です。

重み共有 — パラメータの40%を節約する

最後の行の @ self.emb.weight.T が重み共有 (weight tying) です。出力層を別に持たず、入力埋め込み行列を転置して再利用します。

GPT-2 の語彙は 50,257 個です。出力層を別に持つと 50257 × 256 = 約1287万 個のパラメータが追加で必要になります。全体が 1606万個なので、 共有しなければ 2893万個 になっていたはずです。さらに「トークンをベクトルに変える仕事」と「ベクトルをトークンに戻す仕事」は互いに逆方向の同じ作業なので、同じ行列を使うのが自然です。

Pre-norm 配置

self.ln1(x) をアテンションの 前に 適用し、残差接続は正規化されていない x に加えます。この順序 (pre-norm) が post-norm より深いモデルで学習が安定すると知られています。4層では差は大きくありませんが、慣例に従いました。

学習ループ

model = TinyGPT(len(tok)).to("cuda")
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)

step = 0
while not run.over_budget():          # 15分の予算
    ix = torch.randint(0, len(data) - CTX - 1, (64,))
    x = torch.stack([data[i:i + CTX] for i in ix]).to("cuda")
    y = torch.stack([data[i + 1:i + CTX + 1] for i in ix]).to("cuda")
    with torch.autocast("cuda", torch.bfloat16):
        loss = F.cross_entropy(model(x).flatten(0, 1), y.flatten())
    opt.zero_grad(); loss.backward(); opt.step()
    step += 1

x と y が 1つずれているのがすべてです。入力が "Once upon a" なら正解は "upon a time" です。すべての位置で同時に次のトークンを当てるので、長さ 128 のサンプル 1つが 128個の学習信号を作ります。

torch.autocast で bfloat16 を使います。3090 ではメモリと速度の両方で得があり、bf16 は float16 と違って指数範囲が float32 と同じなので、損失スケーリングなしでも安定します。

学習ログ

実際に記録された値です。

[    0.6s] step      0 loss=168.4468 ppl=...
[   38.2s] step   1000 loss=5.2213  ppl=185.13
[  114.7s] step   3000 loss=3.6841  ppl=39.81
[  305.9s] step   6700 loss=3.0350  ppl=20.80
[  662.6s] step  16100 loss=2.2882  ppl=9.86
[  901.2s] step  22358 loss=2.0948  ppl=8.12

• 合計 22,358 ステップ、バッチ 64、コンテキスト 128 で約1億8千万トークンを処理しました
• 損失 168 から始まるのは初期化直後の値です。一様分布なら ln(50257) ≈ 10.8 になるはずですが、LayerNorm と重み共有の相互作用で最初のステップに大きな値が出て、すぐ正常範囲に下がります
• 最終 perplexity 8.12 は 次の単語の候補を平均8個程度に絞った という意味です

生成結果

温度 0.8 でサンプリングした実際の出力です。手を加えていません。

PROMPT: Once upon a time
Once upon a time, there was a little boy named Timmy. Timmy liked to play
with his toy cars instead. One day, Timmy wanted to play with the red cars,
but he couldn't find it.

As Timmy was leaving a fight on the slide, he saw a big tree. Timmy thought
it looked like fun, so he started to tremble. He mixed the leaves and

PROMPT: Lily went to the
Lily went to the park with her mom. They saw a big tree, a hole and a hole.
They were curious and wanted to see what was inside.

They heard a hole in the hole. It was a zipper, but it was in the hole.

この結果を正直に読んでみましょう。

うまくいったこと — 文法がほぼ完璧です。時制の一致、冠詞、代名詞の指示が合っています。登場人物の名前を段落を通して一貫して保ちます。引用符を開けば閉じます。段落の区切りも自然です。

うまくいかなかったこと — 意味が崩れます。"a big tree, a hole and a hole" のように同じ単語を繰り返し、"leaving a fight on the slide" は文法的ですが意味が通りません。"he started to tremble" は前後とつながりません。

この対比が重要です。 文法は局所的なパターンなので小さなモデルでも学べますが、一貫性は長い依存関係なので容量とコンテキストが必要です。 コンテキストが 128 トークンなので、モデルが記憶できる範囲そのものが短いのです。モデルを大きくすると最初に良くなるのが、まさにこの部分です。

perplexity 8 が意味すること

perplexity は exp(cross_entropy) です。直感的には 「モデルが次のトークンを何個の候補に絞ったか」 です。

• 学習前 (一様分布): 50,257
• 1,000ステップ後: 185
• 最終: 8.12

50,257個の中から選んでいたものを 8個に絞りました。ただしこの数値を他のモデルと直接比較してはいけません。perplexity はトークナイザーとデータ分布に強く依存するからです。TinyStories のように語彙が制限されたデータでは低い値が出やすくなります。 同じデータ・同じトークナイザーの中で比較するときだけ意味がある指標 です。

まとめ

1600万パラメータ、15分、電気代にすれば100円足らずです。それで文法に合った英語の童話を書くモデルができました。

学んだことを3行にまとめるとこうなります。因果マスクは正解の漏れを防ぐ仕組みであり、損失が異常に速く落ちたらこれをまず疑うべきです。重み共有は語彙が大きいモデルでパラメータの半分近くを節約します。そして文法と意味は難易度の異なる問題であり、小さなモデルは前者を先に学びます。

次回は 一つのモデルがテキストと画像を同時に扱う方法 を見ます。ピクセルをトークンとして扱えば、同じトランスフォーマーが画像生成と画像説明の両方をこなせます。

🧠 理解度チェッククイズ

参考資料

• TinyStories データセット
• TinyStories 論文 — How Small Can Language Models Be and Still Speak Coherent English?
• Attention Is All You Need
• Using the Output Embedding to Improve Language Models — 重み共有