MiniMind — 小さな言語モデルを最初から最後まで自分で学習する
DPO は報酬モデルなしに「こちらが良い」を教える
一言でいうと
SFTで答え方を学んだモデルに、どちらの答えがよりよいかを教える段階です。DPO(Direct Preference Optimization)は、同じ質問に対するよい答え(chosen)と悪い答え(rejected)のペアだけで、報酬モデルも強化学習もなく、分類損失1つでこれを行います。MiniMindのtrain_dpo.pyは、その損失を10行で実装しています。このモジュールでは、事実は同じで口調だけが違うペア(입니다と이다。韓国語の語尾で、順に「です」と「だ」にあたります)でDPOを行い、口調が変わるかと、内容は守られるかを一緒に測ります。
なぜ必要なのか
RLHFは、人の選好で報酬モデルを先に学習し、その報酬を最大化しつつ、元のモデルから離れすぎないように、強化学習(PPO)でモデルを押します。モデルを4つ(ポリシー・基準・報酬・価値)立ち上げる必要があり、学習が不安定です。DPOの論文は、この目標の最適解を閉じた形で解いて、同じ問題を、ポリシーと基準の2つと、単純な分類損失で解けるようにしました。MiniMindのREADMEは、DPOが静的な選好データを何度も回すオフポリシーの方式なので安定していますが、自分で探索しないため、「問題を当てる能力」よりも、選好・安全のようなアライメントに合うと書いています。このモジュールの実験が、まさにその境界を見せてくれます。
どう動くのか
文1つの対数確率は、答えトークンごとのlog pを足したものです。MiniMindは、SFTと同じマスクで、答えトークンだけを残します。
log_probs = torch.gather(F.log_softmax(logits, dim=2), 2, labels.unsqueeze(2)).squeeze(-1)
seq = (log_probs * mask).sum(dim=1) # 문장마다 답 토큰의 합
chosen, rejected = seq[:B // 2], seq[B // 2:] # 배치 앞 절반이 chosen, 뒤 절반이 rejected
logits = (π_chosen − π_rejected) − (ref_chosen − ref_rejected)
loss = −F.logsigmoid(β * logits).mean()
- 基準モデル: SFTモデルをもう1つ読み込んで、凍結します(
eval()、requires_grad_(False))。ポリシーが基準よりchosenを相対的にもっと好むようになるほど、損失が減ります。ポリシーが基準と同じなら、かっこの中が0なので、損失は−log σ(0) = ln 2 ≈ 0.693です。 - β: 基準からどれだけ離れてよいかを決めます。論文は、β·log(π/π_ref)を暗黙の報酬と解釈します。MiniMindのデフォルトは0.15です。
- 学習率: MiniMindのデフォルトは4e-8で、コードのコメントが「5e-8以下を推奨、忘却を避けるため」と書いています。SFT(1e-5)よりも数百分の1です。選好学習は、モデルを少しだけ押す必要があります。
落とし穴が1つあります。DPO損失は、chosenとrejectedの差だけを見ます。そのため、学習率が大きいと、rejectedを大きく引き下げながら、chosenまで一緒に引き下げることがよくあります。差は開いて損失は減りますが、モデルが実際に出す答えは、どちらでもないものになります。このコースのモデルでlrを1e-4に上げると、chosenの平均対数確率が−3.5から−16.5に崩れ、丁寧な語尾(韓国語の「です」にあたる語尾)で答える割合は、むしろ下がりました。そのため、損失だけを見ずに、生成結果と、ほかの能力が守られたかを一緒に見る必要があります。
現場での姿
社内チャットボットがタメ口を混ぜる、禁止された話題に答える、というフィードバックがたまったら、選好ペアを作ってDPOを行います。このときよく起こるのが、「口調は直ったが、正確さが下がった」ということです(よくalignment taxと呼ばれます)。そのため、DPOの前後で、選好の指標(選好精度・望む口調の割合)と、元の能力の指標(事実の正答率)を、両方測るのが標準の手順です。選好データが事実を教えられないことも、覚えておく必要があります。モデルが知らない事実は、chosenをいくら見せても、新しく生まれません。選好ペアを作ること自体もコストです。人が2つの答えを読んで選ぶのに時間がかかり、選ぶ人ごとに基準が違って、データにノイズが混ざります。そのため、「何を好むのか」を先に1文で決めておいて(例: 敬語、根拠のない断定の禁止)、その基準1つだけを違えたペアから始めるのが、結果を解釈しやすくなります。このラボのペアが、口調1か所だけが違う理由です。
MiniMindのオリジナルとこのコースの違い
MiniMindの選好データは、DPO-En-Zh-20kから取った実際の人の選好ペアで、1組の2つの答えは、内容が違います。このコースのペアは、内容が同じで、口調1か所だけを違えて作りました。そうすれば、「選好学習が何を移し、何を守るのか」を、1つの変数で測れます。学習率も違います。MiniMindは、長いデータを回す64Mモデルに4e-8を使い、このコースは、100学習ステップだけ回す1Mモデルに1e-5を使います。数字は違っても、原則は同じです。SFTの学習率よりずっと小さくすること、そして、指標がよくなる間、元の能力も一緒に測ることです。
次のラボですること
SFTの基準モデルで選好ペア1組の対数確率を測り、DPO損失を論文の式どおりに実装して、ポリシーが基準と同じときにln 2が出るかを確認します。DPOの前の選好精度と丁寧語の割合を測り、100学習ステップのDPOのあとに、同じ指標がどう変わるか、口調を無視した正答率が守られたかを測ります。