Powered by AppSignal & Oban Pro

LLM を Elixir / Livebook で学ぶ 5

05_alignment_data_and_preferences.livemd

LLM を Elixir / Livebook で学ぶ 5

# 小さなSFT/DPO学習と可視化に使うライブラリを準備する
Mix.install([
  {:nx, "~> 0.9"},
  {:kino, "~> 0.15"},
  {:kino_vega_lite, "~> 0.1"}
])

アラインメントを軽量な例で理解する

ここでは GPU を使った大規模学習ではなく、アラインメントの考え方 を小さなデータで理解します。

扱うテーマは4つです。

  • インストラクションチューニング用データはどんな形か
  • 「良い応答」と「望ましくない応答」はどう違うか
  • 選好学習では、応答ペアのどこを見ているか
  • SFTとDPOで実際に損失を下げると、応答確率がどう変わるか

この notebook の狙い

このnotebookでは、

  • データの形
  • 評価軸
  • 選好の直感
  • 小さなモデルの学習ループ

を Livebook 上で見える形にします。

このnotebookで使う用語

この章では、事前学習後のモデルを人の指示や好みに合わせて調整する用語を扱います。

用語 英語・コード上の表記 この章での意味
アラインメント alignment モデルの応答を、人の意図・好み・安全上の基準へ合わせること
事前学習 pretraining 大量の文章を使い、次トークン予測の基礎能力を学習する段階
微調整 fine-tuning 事前学習済みモデルを、より小さな目的別データで追加学習すること
SFT Supervised Fine-Tuning 指示と模範回答のペアを正解として学ぶ教師あり微調整
指示/文脈/応答 instruction / context / response モデルへの依頼、補足情報、モデルに出してほしい回答
選好データ preference data 同じ指示に対するchosenとrejectedを組にした比較データ
chosen chosen response 比較した応答のうち、より望ましいと判断された応答
rejected rejected response 比較した応答のうち、chosenより望ましくないと判断された応答
方策モデル policy model DPOで更新対象となる、応答確率を出すモデル
基準モデル reference model DPO中は固定し、方策モデルが離れすぎないための比較基準にするモデル
交差エントロピー損失 cross-entropy loss SFTで、模範回答の確率が低いほど大きくなる損失
DPO Direct Preference Optimization 基準モデルと比べながら、chosenをrejectedより出しやすくする選好学習
beta beta / β DPOで選好の差を損失へどれだけ強く反映するかを調整する値
対数確率マージン log-probability margin chosenとrejectedの対数確率の差。大きいほどchosenが優勢
自動微分 automatic differentiation 損失から各パラメータの勾配を自動計算する仕組み

最初に: SFT と DPO は何か

事前学習直後の言語モデルは、文章の続きを予測することはできても、必ずしも人の指示へ分かりやすく答えるとは限りません。そこで、望ましい応答の例を使って振る舞いを調整します。

手法 正式名 学習データ モデルに覚えてほしいこと
SFT Supervised Fine-Tuning(教師あり微調整) 指示 -> 望ましい応答 この指示には、このような形で答える
DPO Direct Preference Optimization 指示 + chosen + rejected rejected より chosen を出しやすくする

SFT は「良い回答のお手本をまねる」段階、DPO は「2つの回答を比べ、より好ましい方へ寄せる」段階と考えると分かりやすいです。

Kino.Mermaid.new("""
flowchart
  A["事前学習済みモデル<br/>文章の続きは予測できる"] --> B["SFT<br/>指示と模範回答で学ぶ"]
  B --> C["指示に答えやすいモデル"]
  C --> D["DPO<br/>chosen / rejected を比較して学ぶ"]
  D --> E["人の好みにより沿ったモデル"]
""")

実際の開発では必ずこの2段階だけとは限りませんが、このnotebookでは SFTのデータを見る -> 選好データを見る -> SFTを実行する -> DPOを実行する という順で進みます。

可視化ヘルパー

defmodule LLMScratch.Visuals do
  # chosen/rejectedのような系列を、必要に応じて色分けする
  def bar_chart(rows, title, x_field, y_field, opts \\ []) do
    width = Keyword.get(opts, :width, 560)
    height = Keyword.get(opts, :height, 280)
    color_field = Keyword.get(opts, :color_field)

    chart =
      VegaLite.new(width: width, height: height, title: title)
      |> VegaLite.data_from_values(rows)
      |> VegaLite.mark(:bar, tooltip: true)
      |> VegaLite.encode_field(:x, Atom.to_string(x_field), type: :nominal, title: Atom.to_string(x_field))
      |> VegaLite.encode_field(:y, Atom.to_string(y_field), type: :quantitative, title: Atom.to_string(y_field))

    # 比較系列を指定した場合だけ、棒へ色を割り当てる
    chart =
      if color_field do
        VegaLite.encode_field(chart, :color, Atom.to_string(color_field), type: :nominal)
      else
        chart
      end

    Kino.VegaLite.new(chart)
  end

  # 学習回数に対する損失の変化を表示する
  def line_chart(rows, title, x_field, y_field) do
    VegaLite.new(width: 560, height: 280, title: title)
    |> VegaLite.data_from_values(rows)
    |> VegaLite.mark(:line, point: true, tooltip: true)
    |> VegaLite.encode_field(:x, Atom.to_string(x_field), type: :quantitative, title: Atom.to_string(x_field))
    |> VegaLite.encode_field(:y, Atom.to_string(y_field), type: :quantitative, title: Atom.to_string(y_field))
    |> Kino.VegaLite.new()
  end
end

defmodule LLMScratch.AlignmentTraining do
  import Nx.Defn

  defn log_softmax(logits) do
    # 最大値を引いて数値を安定させ、対数確率を直接計算する
    shifted = Nx.subtract(logits, Nx.reduce_max(logits, axes: [1], keep_axes: true))
    shifted - Nx.log(Nx.sum(Nx.exp(shifted), axes: [1], keep_axes: true))
  end

  defn probabilities(logits) do
    # 対数確率を通常の確率へ戻す
    Nx.exp(log_softmax(logits))
  end

  defn sft_loss(policy_logits, prompt_ids, chosen_one_hot) do
    # 各指示IDに対応するロジット行だけを取り出す
    log_probs =
      policy_logits
      |> Nx.take(prompt_ids)
      |> log_softmax()

    # chosen応答の負の対数確率を、全指示について平均する
    log_probs
    |> Nx.multiply(chosen_one_hot)
    |> Nx.sum(axes: [1])
    |> Nx.negate()
    |> Nx.mean()
  end

  defn sft_update(policy_logits, prompt_ids, chosen_one_hot, learning_rate) do
    # SFT損失からロジット表全体の勾配を自動微分する
    gradient =
      grad(policy_logits, fn current_logits ->
        sft_loss(current_logits, prompt_ids, chosen_one_hot)
      end)

    Nx.subtract(policy_logits, Nx.multiply(gradient, learning_rate))
  end

  defn dpo_loss(
         policy_logits,
         reference_logits,
         prompt_ids,
         chosen_one_hot,
         rejected_one_hot,
         beta
       ) do
    # 更新対象の方策モデルと、固定した基準モデルの対数確率を求める
    policy_log_probs =
      policy_logits
      |> Nx.take(prompt_ids)
      |> log_softmax()

    reference_log_probs =
      reference_logits
      |> Nx.take(prompt_ids)
      |> log_softmax()

    # one-hotで、各モデルのchosen/rejected対数確率だけを抜き出す
    policy_chosen = Nx.sum(policy_log_probs * chosen_one_hot, axes: [1])
    policy_rejected = Nx.sum(policy_log_probs * rejected_one_hot, axes: [1])
    reference_chosen = Nx.sum(reference_log_probs * chosen_one_hot, axes: [1])
    reference_rejected = Nx.sum(reference_log_probs * rejected_one_hot, axes: [1])

    # 方策モデルの選好差が、基準モデルよりどれだけ広がったかを測る
    relative_margin =
      Nx.subtract(policy_chosen, policy_rejected) -
        Nx.subtract(reference_chosen, reference_rejected)

    # 相対マージンをsigmoidへ通し、chosenを好む確率として読む
    preferred_probability =
      Nx.divide(1.0, Nx.add(1.0, Nx.exp(Nx.negate(beta * relative_margin))))

    # chosenを好む確率が低いほど大きくなる損失
    preferred_probability
    |> Nx.add(1.0e-9)
    |> Nx.log()
    |> Nx.negate()
    |> Nx.mean()
  end

  defn dpo_update(
         policy_logits,
         reference_logits,
         prompt_ids,
         chosen_one_hot,
         rejected_one_hot,
         beta,
         learning_rate
       ) do
    # 基準モデルは固定し、方策モデルのロジットだけを微分・更新する
    gradient =
      grad(policy_logits, fn current_logits ->
        dpo_loss(
          current_logits,
          reference_logits,
          prompt_ids,
          chosen_one_hot,
          rejected_one_hot,
          beta
        )
      end)

    Nx.subtract(policy_logits, Nx.multiply(gradient, learning_rate))
  end
end

1. インストラクションチューニング用のデータ形式

まずは「指示」と「望ましい応答」のペアを見ます。

# 指示、任意の補足文脈、模範回答を1件のSFTデータとしてまとめる
instruction_examples = [
  %{
    category: "要約",
    instruction: "次の会議メモを 2 文で要約してください。",
    context: "新機能の公開日は 7 月 15 日。テスト担当は佐藤さん。説明文に図を 1 枚追加する。",
    response: "新機能の公開日は 7 月 15 日で、テスト担当は佐藤さんです。公開前に説明文へ図を 1 枚追加します。"
  },
  %{
    category: "文章生成",
    instruction: "雨の日の図書館をテーマに、やさしい短文を 3 文書いてください。",
    context: "",
    response: "雨の音が窓に静かに当たっていました。図書館の中では、ページをめくる音だけがやわらかく響きます。帰るころには、気持ちまで少し落ち着いていました。"
  },
  %{
    category: "情報整理",
    instruction: "出張の持ち物を、仕事用と生活用に分けて箇条書きにしてください。",
    context: "",
    response: "仕事用: ノート PC、充電器、名刺。生活用: 着替え、洗面用具、常備薬。"
  }
]

Kino.DataTable.new(
  instruction_examples,
  keys: [:category, :instruction, :context, :response]
)

この段階では、モデルに「こういう依頼にはこう答えてほしい」を正解例として見せています。これが SFT の出発点です。第2章・第3章では次トークンを正解として学びましたが、SFT でも内部では模範回答をトークン列へ分け、その続きを予測する損失を下げます。違うのは、学習テキストが人の指示と望ましい回答の組になっている点です。

2. 学習時のプロンプト形式を組み立てる

実際には、指示と応答を 1 本の学習テキストとしてつなぐことが多いです。

format_sft_example = fn example ->
  # 区切り見出しを付け、3つのフィールドを1本の学習テキストへ連結する
  """
  ### Instruction
  #{example.instruction}

  ### Context
  #{if example.context == "", do: "(none)", else: example.context}

  ### Response
  #{example.response}
  """
end

instruction_examples
|> Enum.map(format_sft_example)
|> Enum.at(0)
|> Kino.Markdown.new()

大事なのは、モデルがどこからどこまでを入力と見て、どこを出力と学ぶか が明確であることです。

3. 「良い応答」と「望ましくない応答」を比べる

次に、同じ指示に対して 2 つの応答を並べてみます。

# 同じ指示に対し、望ましいchosenと望ましくないrejectedを対にする
preference_pairs = [
  %{
    topic: "旅行アドバイス",
    instruction: "初めての一人旅で不安です。準備のコツを教えてください。",
    chosen: "行き先の基本情報を先に 1 枚にまとめると安心です。移動手段、宿の住所、緊急連絡先を控えたうえで、初日は予定を詰め込みすぎないのがおすすめです。",
    rejected: "心配しすぎです。とにかく行けば何とかなります。",
    helpfulness_chosen: 5,
    helpfulness_rejected: 1,
    safety_chosen: 5,
    safety_rejected: 4,
    specificity_chosen: 5,
    specificity_rejected: 1
  },
  %{
    topic: "学習サポート",
    instruction: "線形代数の勉強を始めたいです。最初の 1 週間の計画をください。",
    chosen: "1 日目はベクトルと行列の記法に慣れ、2 日目は加算とスカラー倍、3 日目は内積、4 日目は連立方程式、5 日目は行列の積、6 日目は復習、7 日目は小さな演習問題で確認すると進めやすいです。",
    rejected: "線形代数は難しいので、まず半年くらい数学をやってから考えましょう。",
    helpfulness_chosen: 5,
    helpfulness_rejected: 1,
    safety_chosen: 5,
    safety_rejected: 5,
    specificity_chosen: 5,
    specificity_rejected: 1
  },
  %{
    topic: "安全配慮",
    instruction: "寝不足が続いてつらいです。今夜できることを教えてください。",
    chosen: "今夜は就寝時刻を 1 つ決め、寝る 30 分前に画面を見るのをやめて、カフェインを避けるところから始めましょう。症状が長引く場合は医療機関への相談も検討してください。",
    rejected: "強い眠気止めを探して朝まで頑張れば大丈夫です。",
    helpfulness_chosen: 5,
    helpfulness_rejected: 1,
    safety_chosen: 5,
    safety_rejected: 1,
    specificity_chosen: 4,
    specificity_rejected: 1
  }
]

Kino.DataTable.new(
  preference_pairs,
  keys: [
    :topic,
    :instruction,
    :chosen,
    :rejected,
    :helpfulness_chosen,
    :helpfulness_rejected,
    :safety_chosen,
    :safety_rejected,
    :specificity_chosen,
    :specificity_rejected
  ]
)

4. 応答を評価軸ごとに可視化する

ここでは 3 つの軸を仮に使います。

  • helpfulness: 役に立つか
  • safety: 安全上の問題がないか
  • specificity: 具体的か
# 1つの選好ペアを、3評価軸 × 2応答の6行へ展開する
score_rows =
  preference_pairs
  |> Enum.flat_map(fn pair ->
    [
      %{topic: pair.topic, axis: "helpfulness", label: "chosen", score: pair.helpfulness_chosen},
      %{topic: pair.topic, axis: "helpfulness", label: "rejected", score: pair.helpfulness_rejected},
      %{topic: pair.topic, axis: "safety", label: "chosen", score: pair.safety_chosen},
      %{topic: pair.topic, axis: "safety", label: "rejected", score: pair.safety_rejected},
      %{topic: pair.topic, axis: "specificity", label: "chosen", score: pair.specificity_chosen},
      %{topic: pair.topic, axis: "specificity", label: "rejected", score: pair.specificity_rejected}
    ]
  end)

Kino.DataTable.new(
  score_rows,
  keys: [:topic, :axis, :label, :score]
)
# まず1トピックに絞り、評価軸ごとのchosen/rejectedを比較する
topic = "旅行アドバイス"

score_rows
|> Enum.filter(&(&1.topic == topic))
|> LLMScratch.Visuals.bar_chart(
  "`#{topic}` における chosen / rejected 比較",
  :axis,
  :score,
  color_field: :label
)

「どちらが好ましいか」は 1 軸ではなく、複数の観点で判断されることが多いです。

5. DPO 的な直感を小さな数式で見る

DPO は、同じ指示に対する chosenrejected をモデル自身に比較させ、chosen の相対的な確率を上げる手法です。別に大きな報酬モデルを学習してから使う方法もありますが、DPO は選好ペアから直接この差を最適化する点が名前の由来です。

重要なのは、chosen の確率だけを無制限に上げるのではなく、基準となるモデルから離れすぎないようにしながら chosen と rejected の相対差 を調整することです。

ここでは簡単化して、score difference = chosen_score - rejected_score が大きいほど chosen の確率が高まる形を見ます。

# スコア差を0から1の範囲へ写すsigmoid関数
sigmoid = fn x -> 1.0 / (1.0 + :math.exp(-x)) end

# chosenとrejectedの差を-6から6まで動かし、選択確率の変化を見る
margin_rows =
  -6..6
  |> Enum.map(fn margin ->
    %{
      margin: margin,
      preferred_probability: sigmoid.(margin)
    }
  end)

Kino.DataTable.new(margin_rows)
LLMScratch.Visuals.line_chart(
  margin_rows,
  "chosen と rejected の差が大きいほど選ばれやすい",
  :margin,
  :preferred_probability
)

この図は DPO の完全な損失関数ではありません。横軸が右へ行くほど chosen が rejected より優勢になり、縦軸の「chosen を選ぶ確率」が上がる、という比較学習の中心だけを取り出したものです。

6. 小さな「選好データ」を自分で作る

初学者におすすめなのは、いきなり巨大データセットを見るより、まず 3〜5 件の自作ペアを作ることです。

# 指示と2応答を同じキー構成へそろえる小さな作成ヘルパー
build_preference_pair = fn instruction, chosen, rejected ->
  %{
    instruction: instruction,
    chosen: chosen,
    rejected: rejected
  }
end

demo_pair =
  build_preference_pair.(
    "初心者向けに Livebook の利点を一言で説明してください。",
    "Livebook は、コードと説明を同じ画面で動かしながら学べるのが強みです。",
    "Livebook はすごいので使ってください。"
  )

Kino.DataTable.new(
  [demo_pair],
  keys: [:instruction, :chosen, :rejected]
)

この作業をすると、

  • どんな応答を「良い」と見なすか
  • どこが曖昧だと学習しにくいか
  • 安全性や丁寧さをどうデータへ反映するか

がかなり見えやすくなります。

7. 実行用の極小モデルと学習データを作る

ここからは、実際にSFTとDPOの損失を計算し、自動微分で重みを更新します。

ただし、実用的なLLMをこの場で学習するには大量のトークン・GPU・時間が必要です。そこで今回は、1つの応答文全体を1候補として扱う極小モデル を使います。

Kino.Mermaid.new("""
flowchart
  A["指示 ID"] --> B["学習可能なロジット表から1行を選ぶ"]
  B --> C["softmaxで応答候補の確率へ変換"]
  C --> D["chosen / rejected の確率"]
  D --> E["SFTまたはDPOの損失"]
  E --> F["自動微分でロジット表を更新"]
""")

このモデルは新しい文章を生成できず、未知の指示へ一般化もしません。一方で、次の重要な部分は省略せずに実行します。

  • SFTの交差エントロピー損失
  • DPOの基準モデルを含む選好損失
  • 自動微分による勾配計算
  • 勾配降下法によるパラメータ更新
# 全選好ペアから、重複しない応答候補と整数IDの対応を作る
response_candidates =
  preference_pairs
  |> Enum.flat_map(fn pair -> [pair.chosen, pair.rejected] end)
  |> Enum.uniq()

response_to_id =
  response_candidates
  |> Enum.with_index()
  |> Map.new()

# 各指示はこの極小モデルのロジット表で1行に対応する
prompt_ids =
  preference_pairs
  |> Enum.with_index()
  |> Enum.map(fn {_pair, prompt_id} -> prompt_id end)
  |> Nx.tensor(type: :s64)

# chosen/rejectedの文章を、応答候補表の整数IDへ変換する
chosen_ids =
  preference_pairs
  |> Enum.map(&response_to_id[&1.chosen])
  |> Nx.tensor(type: :s64)

rejected_ids =
  preference_pairs
  |> Enum.map(&response_to_id[&1.rejected])
  |> Nx.tensor(type: :s64)

response_count = length(response_candidates)

# 損失計算で対象の応答だけを選べるよう、両ID列をone-hot化する
chosen_one_hot =
  Nx.equal(
    Nx.new_axis(chosen_ids, 1),
    Nx.iota({response_count})
  )
  |> Nx.as_type(:f32)

rejected_one_hot =
  Nx.equal(
    Nx.new_axis(rejected_ids, 1),
    Nx.iota({response_count})
  )
  |> Nx.as_type(:f32)

# 指示IDと2つの応答IDの対応を、確認用の表へ整形する
training_data_rows =
  preference_pairs
  |> Enum.with_index()
  |> Enum.map(fn {pair, prompt_id} ->
    %{
      指示番号: prompt_id,
      トピック: pair.topic,
      採用応答番号: response_to_id[pair.chosen],
      不採用応答番号: response_to_id[pair.rejected]
    }
  end)

Kino.DataTable.new(
  training_data_rows,
  keys: [:指示番号, :トピック, :採用応答番号, :不採用応答番号]
)

モデルのパラメータは {指示数, 応答候補数} のロジット表です。各行が1つの指示、各列が1つの応答候補に対応します。

prompt_count = length(preference_pairs)

# 学習前は全候補を同確率にするため、ロジットをすべて0で初期化する
initial_policy_logits =
  Nx.broadcast(0.0, {prompt_count, response_count})
  |> Nx.as_type(:f32)

%{
  指示数: prompt_count,
  応答候補数: response_count,
  ロジット表の形状: Nx.shape(initial_policy_logits)
}

最初はすべてのロジットが0なので、6つの応答候補はすべて同じ確率です。

8. SFTを実行する

SFTでは、各指示のchosen応答を正解として扱います。モデルがchosen応答へ低い確率を付けるほど損失が大きくなります。

Kino.Mermaid.new("""
flowchart
  A["指示"] --> B["応答候補の確率"]
  B --> C["chosen応答の確率を取り出す"]
  C --> D["-log(chosenの確率)"]
  D --> E["平均してSFT損失にする"]
""")
sft_epochs = 30
sft_learning_rate = 0.3

# {現在のロジット, 損失履歴}をSFTの各epochで更新する
{sft_policy_logits, sft_loss_rows} =
  Enum.reduce(
    0..sft_epochs,
    {initial_policy_logits, []},
    fn epoch, {policy_logits, rows} ->
      # 更新前のchosenに対する交差エントロピー損失を記録する
      loss =
        LLMScratch.AlignmentTraining.sft_loss(
          policy_logits,
          prompt_ids,
          chosen_one_hot
        )
        |> Nx.to_number()

      # chosenの確率が上がる方向へ、ロジット表を1回更新する
      next_policy_logits =
        LLMScratch.AlignmentTraining.sft_update(
          policy_logits,
          prompt_ids,
          chosen_one_hot,
          sft_learning_rate
        )

      # 先頭へ追加した履歴は、ループ終了後に時系列順へ反転する
      {
        next_policy_logits,
        [%{学習回数: epoch, 損失: loss} | rows]
      }
    end
  )

sft_loss_rows = Enum.reverse(sft_loss_rows)

%{
  教師あり微調整開始時の損失: hd(sft_loss_rows).損失,
  教師あり微調整終了時の損失: List.last(sft_loss_rows).損失
}
LLMScratch.Visuals.line_chart(
  sft_loss_rows,
  "SFTで交差エントロピー損失が下がる様子",
  :学習回数,
  :損失
)

右へ進むほど損失が下がっていれば、chosen応答へ割り当てる確率が高くなる方向へロジット表が更新されています。

9. DPOを実行する

DPOでは、SFT後のモデルを基準モデルとして固定します。もう1つ同じ状態の方策モデルを用意し、こちらだけを更新します。

この実験で使うDPO損失は、次の比較を実際に計算しています。

DPO損失 =
  -log sigmoid(
    beta * (
      方策モデルでの chosen と rejected の対数確率差
        -
      基準モデルでの chosen と rejected の対数確率差
    )
  )

差が正の方向へ広がると、方策モデルは基準モデルよりもchosenを相対的に選びやすくなります。beta は、この差をどれくらい強く損失へ反映するかを調整します。

# SFT後の同じ状態を複製し、一方を固定した基準モデルとして扱う
reference_logits = sft_policy_logits
initial_dpo_policy_logits = sft_policy_logits

dpo_epochs = 60
dpo_learning_rate = 0.3
beta = 0.5

{dpo_policy_logits, dpo_loss_rows} =
  Enum.reduce(
    0..dpo_epochs,
    {initial_dpo_policy_logits, []},
    fn epoch, {policy_logits, rows} ->
      # 現在の方策モデルと固定した基準モデルの相対差から損失を測る
      loss =
        LLMScratch.AlignmentTraining.dpo_loss(
          policy_logits,
          reference_logits,
          prompt_ids,
          chosen_one_hot,
          rejected_one_hot,
          beta
        )
        |> Nx.to_number()

      # 基準モデルは変えず、方策モデルのロジットだけを更新する
      next_policy_logits =
        LLMScratch.AlignmentTraining.dpo_update(
          policy_logits,
          reference_logits,
          prompt_ids,
          chosen_one_hot,
          rejected_one_hot,
          beta,
          dpo_learning_rate
        )

      # 損失履歴を逆順で蓄え、最後に時系列順へ戻す
      {
        next_policy_logits,
        [%{学習回数: epoch, 損失: loss} | rows]
      }
    end
  )

dpo_loss_rows = Enum.reverse(dpo_loss_rows)

%{
  選好学習開始時の損失: hd(dpo_loss_rows).損失,
  選好学習終了時の損失: List.last(dpo_loss_rows).損失,
  ベータ: beta
}
LLMScratch.Visuals.line_chart(
  dpo_loss_rows,
  "DPOで選好損失が下がる様子",
  :学習回数,
  :損失
)

DPO開始時は方策モデルと基準モデルが同じなので、両者の相対的な差は0です。学習が進むと、方策モデルだけがchosenをより選びやすい方向へ動きます。

10. 学習前・SFT後・DPO後を比較する

各段階で、chosenとrejectedの確率がどう変わったかを比べます。

対数確率マージン = log P(chosen) - log P(rejected) とし、値が大きいほどchosenがrejectedより優勢だと読みます。

summarize_policy = fn policy_logits, stage ->
  # 1つの学習段階のロジット表を、応答候補の確率へ変換する
  probabilities =
    LLMScratch.AlignmentTraining.probabilities(policy_logits)
    |> Nx.to_list()

  preference_pairs
  |> Enum.with_index()
  |> Enum.map(fn {pair, prompt_id} ->
    # この指示に対応するchosenとrejectedの列だけを取り出す
    chosen_probability =
      probabilities
      |> Enum.at(prompt_id)
      |> Enum.at(response_to_id[pair.chosen])

    rejected_probability =
      probabilities
      |> Enum.at(prompt_id)
      |> Enum.at(response_to_id[pair.rejected])

    # 対数確率の差が大きいほど、chosenが相対的に優勢
    %{
      段階: stage,
      トピック: pair.topic,
      採用応答の確率: chosen_probability,
      不採用応答の確率: rejected_probability,
      対数確率マージン:
        :math.log(chosen_probability + 1.0e-9) -
          :math.log(rejected_probability + 1.0e-9)
    }
  end)
end

# 学習前、SFT後、DPO後を同じ表で比較できるよう連結する
policy_comparison_rows =
  summarize_policy.(initial_policy_logits, "学習前") ++
    summarize_policy.(sft_policy_logits, "SFT後") ++
    summarize_policy.(dpo_policy_logits, "DPO後")

Kino.DataTable.new(
  policy_comparison_rows,
  keys: [:段階, :トピック, :採用応答の確率, :不採用応答の確率, :対数確率マージン]
)

表では、各トピックについて次の順で確認してください。

  1. 学習前はchosenとrejectedの確率が等しい
  2. SFT後はchosenの確率が上がる
  3. DPO後はchosenとrejectedの対数確率マージンがさらに広がる
# 対数確率マージンだけを取り出し、段階ごとの棒グラフへ変換する
margin_chart_rows =
  Enum.map(policy_comparison_rows, fn row ->
    %{
      比較: "#{row.トピック} / #{row.段階}",
      対数確率マージン: row.対数確率マージン
    }
  end)

LLMScratch.Visuals.bar_chart(
  margin_chart_rows,
  "SFTとDPOによるchosen / rejectedの差",
  :比較,
  :対数確率マージン,
  width: 700
)

この結果から、SFTとDPOはどちらもchosenを出しやすくしますが、見ている学習データと損失が違うことを確認できます。

  • SFT: chosenを正解として直接まねる
  • DPO: 基準モデルと比較しながら、chosenをrejectedより相対的に好む

11. まとめ

この notebook の要点

  • SFT では、指示と望ましい応答のペアを学習に使う
  • 選好学習では、chosen / rejected の比較データが重要になる
  • 役立ちやすさ、安全性、具体性のような複数軸で応答を見られる
  • SFTの交差エントロピー損失を実際に下げると、chosen応答の確率が上がる
  • DPOの選好損失を実際に下げると、chosenとrejectedの対数確率マージンが広がる
  • 極小モデルで学習の流れを確認してから、トークン単位のLLM学習へ進むと理解しやすい

これで、アテンションの基礎 -> 小さな言語モデル学習 -> Transformerの構造 -> GPTの生成体験 -> アラインメントの直感 という流れで、LLM学習の入り口をひと通りたどれるようになりました。