SFT は回答だけを教える — 損失マスキングがその境界を引く

一言でいうと

事前学習したモデルは、コーパスの続きを書くだけで、質問には答えません。SFT(教師あり微調整)は、質問と答えの対話をチャット形式の1行にして、同じ次トークン予測で学習しますが、答え(assistant)のトークンでだけ損失を計算します。MiniMindのgenerate_labelsが、その境界を引きます。このモジュールでは、マスキングしたモデルと外したモデルを同じ学習ステップ数で学習して、質問トークンの損失と答えトークンの損失がどう分かれるかを、数字で見ます。

なぜ必要なのか

モデルに「質問を受けたら答える」ことを教えるには、2つのものが必要です。1つは、どこまでが質問で、どこからが答えなのかをモデルが見分けられる目印で、もう1つは、何を学ぶのかの選択です。対話全体で学習すると、モデルは質問の書き方まで学びます。ユーザーが言いそうなことを真似るのに、学習信号の相当な部分を使い、答えが終わったあと、自分で次の質問を作り出すこともあります。私たちが望むのは、質問を条件として見て答えを出すモデルです。

どう動くのか

チャット形式では、MiniMindは、Qwen系と同じChatMLの形を使います。

<|im_start|>user
가람 마을의 특산물은 뭐야?<|im_end|>
<|im_start|>assistant
가람 마을의 특산물은 인삼입니다.<|im_end|>

MiniMindの実際のテンプレートは、assistantの番ごとに空の<think>\n\n</think>\n\nを入れ、学習のとき80%の確率で消します(思考モードとそろえるための仕組みです)。20%の確率で、systemプロンプトを先頭に付けることもあります。このコースは思考モードを扱わないので、両方を外して、上の形だけを使います(mmkit.chat_text)。

ラベルのマスキングでは、generate_labelsは、トークン列から<|im_start|>assistant\nのトークンの断片を探し、その後ろから<|im_end|>\nの断片が終わるまでだけ、ラベルに元のトークンを入れます。残りはすべて-100です。モデルのcross_entropy(..., ignore_index=-100)が、その位置を飛ばします。

labels = [-100] * len(input_ids)
# <|im_start|>assistant\n 을 찾으면 그 뒤부터 <|im_end|>\n 까지 labels[j] = input_ids[j]

終わりの目印<|im_end|>まで正解に入れることが重要です。そうすると、モデルが答えをいつ止めるかを学びます。何度かやりとりする対話では、assistantの番ごとに、この区間ができます。私たちのSFTデータで数えると、実際のトークンのうち損失に入るのは22%ほどです。残りの78%は、条件にすぎません。

ハイパーパラメーターでは、MiniMindのtrain_full_sft.pyは、事前学習の重みから始めて、lr 1e-5(事前学習の5e-4の50分の1)で2エポック回します。すでに学んだ言語能力を大きく揺らさないための選択です。このコースの小さなモデルは1e-3と、ずっと大きくしますが、その代償として、事前学習コーパスの続きを書く能力をかなり忘れます。最後のモジュールで、パープレキシティで測ります。

現場での姿

社内の問い合わせ記録でSFTしたら、モデルが答えのあとに「顧客: …」で始まる偽の質問を書き続けるなら、マスキングが抜けているか、終わりの目印を正解から外しています。逆に、終わりの目印を入れないと、モデルが止まらずに最大長まで書きます。どちらの症状も、ラベルを1行出力して見れば、すぐにわかります。MiniMindのSFTDataset.__getitem__にコメントとして残っているデバッグ出力(入力トークンと次のトークン、ラベルを並べて出力するもの)が、まさにその用途です。

もう1つよくある勘違いは、「SFTで知識を入れる」という期待です。SFTが教えるのは、おもに形式と態度で、質問を受けたらどんな形で答え、どこで止まるのかということです。事実は、ほとんどが事前学習から来ます。データにない事実を、SFT数千件で入れようとすると、モデルはその文を暗記するだけで、別の質問の形で聞くと、取り出せません。このコースの最後のモジュールで、その現象を数字で見ます。

MiniMindのオリジナルとこのコースの違い

MiniMindのSFTデータには、ツール呼び出しの対話(systemにツールの一覧があり、assistantが<tool_call>で呼び出し、toolの役割が結果を返す形)が混ざっていて、pre_processing_chatは、ツールのある対話には手を加えません。テンプレートがツールの結果をuserの番として包むので、その位置もラベルから外れます。モデルはツールの結果を真似るのではなく、それを見て続きを答える方法だけを学びます。このコースは、ツール呼び出しを扱わず、1–2往復の質問と答え、そして挨拶だけを使います。最大長も、MiniMindの推奨値(768)ではなく64です。私たちの対話は、いちばん長いものでも40トークンほどなので、64なら切られるものがなく、パディングも少なくなります。逆に、実際のデータで行うときは、対話の長さの分布を先に測って最大長を決める必要があります。短すぎると、答えの終わりと終わりの目印が切れて、止まる方法を学べません。

次のラボですること

SFTデータの最初の対話をチャット形式に変えて、答えだけを残すラベルを自分で作ります。損失に入るトークンの割合を数え、基準の事前学習の重みから、マスキングしたモデルと外したモデルを同じ条件で学習します。取り分けておいた対話で、2つのモデルの答えの損失と質問の損失を比べ、取り分けておいた質問に対する答えを取り出して見ます。