← ブログに戻る

チャットモデルは損失マスクで生まれる — nanochatのSFTを「採点表」として読む

事前学習を終えたばかりのGPTに「こんにちは」と話しかけても、相槌が返ってくる保証はどこにもありません。手に入るのは大量のWebテキストの「続き」を予測する文章生成器で、会話という概念をまだ持っていないからです。

ではチャットモデルはどこで生まれるのか。答えは損失マスクという仕組みにあります。名前は仰々しいですが、実体はトークンごとに0か1を立てるだけの配列です。この0と1の列が、文章生成器をチャット相手に変えます。

この記事では、その現場をAndrej Karpathy公開の nanochat のコードで読みます。nanochatはtokenizer訓練からチャットCLIまで、LLMの製造工程一式を実装本体およそ8,159行に収めたリポジトリです。訓練費用の話(2019年の$43,000が2026年に$48になった経緯)は以前書いたので、今回は工程の中身、それも一番面白い「チャットが生まれる瞬間」に絞ります。読むのはコミット 92d63d4 時点のコードです。

用語を3つだけ

  • 事前学習: 大量のWebテキストで「次のトークン(単語の断片)を当てるクイズ」をひたすら解かせる工程。これを終えたモデルは文章の続きがうまくなりますが、それだけです
  • SFT (Supervised Fine-Tuning): 事前学習済みのモデルに会話の模範解答を見せて追加学習する工程。nanochatでは scripts/chat_sft.py が担当します
  • 損失: モデルの予測と正解のズレを測った点数。「損失を計算する」は「答え合わせをして採点する」とほぼ同じ意味です

主役の損失マスクは、要するに 「どのトークンを採点するか」の指定表 です。

事前学習は、会話用トークンを1つも使わない

nanochat/tokenizer.py の特殊トークン定義から始めます。

SPECIAL_TOKENS = [
    # every document begins with the Beginning of Sequence (BOS) token that delimits documents
    "<|bos|>",
    # tokens below are only used during finetuning to render Conversations into token ids
    "<|user_start|>", # user messages
    "<|user_end|>",
    "<|assistant_start|>", # assistant messages
    "<|assistant_end|>",
    "<|python_start|>", # assistant invokes python REPL tool
    "<|python_end|>",
    "<|output_start|>", # python REPL outputs back to assistant
    "<|output_end|>",
]

コメントに「only used during finetuning」とある通り、<|bos|>(文書の区切り印)を除く8個は事前学習の文章に一度も登場しません。語彙表に席だけ用意されていて、誰も座っていない予約席です。

事前学習を終えた時点のモデルにとって、<|user_start|> はほぼ意味を持たない記号です。会話の区切りも発言者の交代も白紙。この予約席に意味を書き込む工程がSFTです。

会話は「トークン列」と「採点表」のセットに変換される

会話データをトークン列へ変換するのは nanochat/tokenizer.pyrender_conversation です。この関数はトークン列と同じ長さの mask を一緒に返します。1が「採点する」、0が「採点しない」。ツール呼び出し部分の実物がこれです。

elif part["type"] == "python":
    # python tool call => add the tokens inside <|python_start|> and <|python_end|>
    add_tokens(python_start, 1)
    add_tokens(value_ids, 1)
    add_tokens(python_end, 1)
elif part["type"] == "python_output":
    # python output => add the tokens inside <|output_start|> and <|output_end|>
    # none of these tokens are supervised because the tokens come from Python at test time
    add_tokens(output_start, 0)
    add_tokens(value_ids, 0)
    add_tokens(output_end, 0)

add_tokens の第2引数がmaskです。render_conversation の分岐を全部追うと、採点方針はこう整理できます。

会話の部分mask意味
`<bos>`
ユーザー発言(枠トークン含む)0採点しない
アシスタント発言の本体1ここを採点する
`<assistant_end>`
ツール呼び出しの式1採点する
ツール実行結果0採点しない

教科書は全ページ読ませるけれど、テストに出すのはアシスタントの発言だけ。そういう採点方針です。

mask=0のトークンは、採点欄から消える

この採点表を学習につなぐのが chat_sft.py です。言語モデルの「次のトークン当てクイズ」では正解は常に「1つ右のトークン」なので、maskも1つ右にずらして正解ラベル(targets)に重ねます。

# Apply the loss mask from render_conversation (mask=1 for assistant completions,
# mask=0 for user prompts, BOS, special tokens, tool outputs). mask[1:] aligns
# with targets (shifted by 1). Unmasked positions get -1 (ignore_index).
mask_tensor = torch.tensor(mask_rows, dtype=torch.int8)
mask_targets = mask_tensor[:, 1:].to(device=device)
targets[mask_targets == 0] = -1

最後の1行がすべてです。mask=0の位置は正解ラベルを-1に書き換える。nanochatの損失計算(nanochat/gpt.py)は F.cross_entropy(..., ignore_index=-1) で、-1の位置を採点から除外します。答案用紙に「この問題は採点対象外」とスタンプを押すイメージです。

SFTで採点されるのはアシスタントの発言本体、<|assistant_end|>、ツール呼び出しの式だけ。ユーザーの発言もツールの実行結果も、モデルは文脈として読んではいますが、点数には入りません。

事前学習が「すべてのトークンを採点する」工程だったのに対し、SFTは「どのトークンを採点するかを選ぶ」工程です。チャットモデルを生むのは新しいアーキテクチャでも追加の魔法でもない、採点範囲の指定です。期末テストの範囲表くらいの地味さですが、これが効きます。

電卓の答えは、暗記させない

このmask設計の面白さが一番よく出ているのが、GSM8K(算数の文章題データセット)の扱いです。

tasks/gsm8k.py は解答に埋め込まれた電卓注釈 <<式=結果>> を正規表現で切り出し、式を python タイプ、結果を python_output タイプとしてアシスタントメッセージに埋め込みます。式はmask=1、結果はmask=0。

つまり「17×3を計算したくなったら電卓を呼ぶ」という振る舞いは採点する。でも「51」という答えそのものは採点しない。本番の会話では結果はPythonが返してくれるので、モデルが覚える必要がないからです。もし結果までmask=1にすると、モデルは計算せずに答えを暗記する方向に育ちます。九九を理解せずに答えの表を丸暗記して乗り切る、あの勉強法です。人間で失敗済みの学習戦略を、わざわざモデルに実装することはありません。

SFTが教えるのは「振る舞い」で、知識ではない

runs/speedrun.sh のコメントが端的です。

# SFT (teach the model conversation special tokens, tool use, multiple choice)

chat_sft.py の学習データは3本柱。SmolTalk(訓練460K行)で「相手が話し終えたら自分が話す」というターン交代を、MMLU(100K行×3エポック)で多肢選択の答え方を、GSM8K(8K行×4エポック)でツール使用を覚えます。MMLUの模範解答は「A」のような1文字で、教えているのは実質マークシートの塗り方です。

どれも模範解答を真似る学び方です。知識そのものは事前学習で仕込まれたものがすべてで、SFTはその知識を会話として取り出す型を教えているに過ぎません。実際、READMEのspeedrunモデルとの会話例では空が青い理由にそれらしく答えつつ、README本文は「空はなぜ緑色かとも聞いてみろ」とけしかけています。会話の型は完璧、中身は別問題。誰の周りにもいるタイプです。

まとめ

  • 事前学習は会話用の特殊トークンを一切使わない。会話はSFTで上書きされる
  • render_conversation が会話を「トークン列」と「採点表(mask)」のセットに変換する
  • 採点されるのはアシスタント発言本体・<|assistant_end|>・ツール呼び出しの式だけ
  • maskは1つ右にずらしてtargetsに重ね、0の位置は ignore_index=-1 で採点から消える
  • 電卓の式は採点するが答えは採点しない。答えの暗記は人間で失敗済みだから

損失マスクは、配列に0と1を立てるだけの仕組みです。それでも「どのトークンを採点するか」という選択そのものがチャットモデルを形づくっている、というのがSFTのコードを読んで一番面白かったところです。

スライドでおさらい

nanochatの全体工程(費用崩壊・—depthダイヤル・損失マスク・CPUでの一周)を12枚のスライドにまとめています。

この記事の続き

この記事は、書籍「nanochatで理解するLLM製造工程」の第4章を抜粋・再構成したものです。tokenizer訓練から事前学習、強化学習、評価、推論エンジン、MacBookでの全工程一周まで、全8章+トリビア付録をKindle(¥250)で読めます。

nanochatで理解するLLM製造工程 関連書籍 nanochatで理解するLLM製造工程 nanochat 解説 · LLM 作り方 · Karpathy · $48でGPT-2級を訓練する全工程を読む (Zennで全章無料) 書籍ページを見る →