TRL を用いた Direct Preference Optimization (DPO) による Llama 2 のファインチューニング

Hugging Face は TRL ライブラリに Direct Preference Optimization (DPO) を統合し、開発者が Llama 2 などの Large Language Models (LLMs) を簡素化されたバイナリクロスエントロピー損失を使用してアラインできるようにしました。この方法は、Reinforcement Learning from Human Feedback (RLHF) に通常必要とされる補助報酬モデルと複雑な強化学習 (RL) の機構の必要性をなくします。

DPO vs. PPO: アラインメントパイプラインの簡素化

Direct Preference Optimization (DPO) は、報酬モデリングと RL 最適化のステップをバイパスすることでアラインメントプロセスを簡素化します。Proximal Policy Optimization (PPO) を使用するような従来の RLHF パイプラインでは、報酬モデルが人間の好みを推定するために訓練され、凍結された参照モデルが KL ペナルティとともに使用され、ポリシーモデルが過度に逸脱して意味のない出力を生成するのを防ぎます。

DPO はこれを、報酬関数から最適な RL ポリシーへの解析的マッピングに置き換えます。これにより、報酬モデルと参照モデルに対する RL 損失が参照モデルに対する損失に直接変換されます。その結果、言語モデルは好みのデータに対する直接的な尤度目的を使用して最適化され、別個の報酬モデルや手間のかかる RL ベースの最適化の必要性がなくなります。

TRL ライブラリを用いた DPO の実装

TRL ライブラリは、DPO ワークフローを実装するためのヘルパーを提供します。従来の RLHF パイプラインは、教師ありファインチューニング (SFT)、好みデータのアノテーション、報酬モデルのトレーニング、および RL 最適化という 4 つのステップを含みますが、DPO は後者の 2 つを 1 つのステップにまとめます。

データ要件

TRL で DPOTrainer を使用するには、好みのデータを次の 3 つの特定のキーを含む辞書形式で提供する必要があります。

  • prompt: 推論時にモデルに提供されるコンテキストプロンプト。
  • chosen: 優先される生成レスポンス。
  • rejected: 優先されない生成レスポンス。

トレーナーの設定

DPOTrainer は、ベースモデル（SFT パイプラインから）、参照モデル（通常は SFT トレーニング済みベースモデルのコピー）、および温度ハイパーパラメータ beta （通常は 0.1 から 0.5 の間）を必要とします。beta パラメータは参照モデルの影響を制御し、beta の値が小さいほど参照モデルは無視されやすくなります。

ケーススタディ: Llama 2 7B のアラインメント

Hugging Face は、質問に対するランク付けされた回答を含む stack-exchange 好みデータセットで Llama 2 7B パラメータモデルをファインチューニングすることにより、DPO の適用を実証しました。

ステップ 1: 教師ありファインチューニング (SFT)

モデルはまず、bitsandbytes ライブラリを介した SFTTrainer と QLoRA テクニックを使用して SFT を行います。これにより、ベースモデルが 4 ビット量子化でロードされ、ターゲットモジュール (q_proj, v_proj) に LoRA レイヤーが追加されます。

ステップ 2: DPO トレーニング

SFT の後、その結果得られたモデルは DPO トレーニングにおいてベースモデルおよび参照モデルの両方として使用されます。AutoPeftModelForCausalLM を使用して、モデルは 4 ビット構成でロードされ、QLoRA メソッドでトレーニングされます。DPOTrainer は評価データセットを使用して進捗を評価し、暗黙の報酬メトリックを報告します。

DPO パフォーマンスのモニタリング

トレーニングと評価中、DPOTrainer はアラインメントの進捗を測定するために 4 つの主要な報酬メトリックを追跡します。

  • rewards/chosen: beta でスケーリングされた、選択されたレスポンスに対するポリシーモデルと参照モデルの対数確率の平均差。
  • rewards/rejected: beta でスケーリングされた、拒否されたレスポンスに対するポリシーモデルと参照モデルの対数確率の平均差。
  • rewards/accuracies: 選択された報酬が対応する拒否された報酬よりも高い割合のパーセンテージ。
  • rewards/margins: 選択された報酬と拒否された報酬の平均差。

成功したトレーニングは、マージンが増加し、 accuracies が 1.0 に近づくことで示され、これはモデルが優先されるレスポンスに一貫してより高い報酬を割り当てていることを意味します。

Sources