AI・機械学習
上級

報酬モデル訓練パイプライン(ホウシュウモデルクンレンパイプライン)

LLMの選好データ収集からモデル訓練・評価・デプロイまでの一連のワークフロー。SFT済みモデルをバックボーンとし、ペアワイズ比較データで報酬ヘッドを訓練する標準的なMLパイプラインである。

0 回閲覧
0 いいね

報酬モデル訓練パイプラインとは

報酬モデル訓練パイプラインは、LLMの出力品質を評価する報酬モデル(Reward Model)を構築するための体系的なワークフローである。データ収集・前処理・モデル訓練・評価・デプロイの各段階を経て、RLHF や Best-of-N サンプリングで実用される RM を完成させる。

2026年現在、HuggingFace TRL・OpenRLHF・DeepSpeed-Chat・Axolotl など複数のフレームワークが RM 訓練パイプラインを提供しており、A100/H100 クラスの GPU があれば 8B〜70B パラメータの RM を数時間〜数日で訓練可能である。

パイプラインの全体構成

報酬モデル訓練パイプラインは以下の5段階で構成される。

Stage 1: SFT(Supervised Fine-Tuning)

RM のバックボーンとなるモデルを instruction-following データで事前微調整する。

  • データ: OpenAssistant、ShareGPT、UltraChat(約140万会話)等
  • 訓練量: 1〜3エポック、学習率 2e-5、コサインスケジューラ
  • 目的: 応答品質の基準線を確立し、RM が「良い応答」と「悪い応答」を識別できる基盤を作る

Stage 2: 選好データ収集

SFT モデルから複数の応答を生成し、人間または AI が比較ランク付けを行う。

  • サンプリング: temperature 0.7〜1.0 で同一プロンプトから 4〜8 個の応答を生成
  • 比較形式: ペアワイズ(2択)が最も一般的。K個のランキングからは C(K,2) 個のペアを生成可能
  • 品質管理: Inter-Annotator Agreement(IAA)を測定し、κ > 0.6 を基準とする

Stage 3: RM 訓練

ペアワイズ比較データを用いて報酬ヘッドを訓練する。

パラメータ推奨値備考
学習率1e-5 〜 5e-6SFT より低く設定
バッチサイズ64〜128(ペア数)GPU メモリに依存
エポック数1〜2過学習防止のため少なめ
重み減衰0.01正則化
最大系列長2048〜4096長い応答も評価するため
損失関数Bradley-Terryペアワイズランキング損失

Stage 4: 評価

RewardBench・MT-Bench・人間評価の3軸で RM の品質を検証する。

Stage 5: デプロイ

訓練済み RM を推論サーバーにデプロイし、RL 訓練ループまたは Best-of-N リランキングに組み込む。vLLM や TensorRT-LLM でバッチ推論を最適化することで、RM 推論のスループットを 10 倍以上向上できる。

フレームワーク比較

フレームワーク言語RM 訓練RL 統合マルチGPUライセンス
TRL (HuggingFace)PythonRewardTrainerPPOTrainer/DPOTrainerDeepSpeed/FSDPApache 2.0
OpenRLHFPython専用 RM TrainerPPO/DPO/KTORay + vLLMApache 2.0
DeepSpeed-ChatPythonStep 2: RMStep 3: RLHFDeepSpeed ZeRO-3Apache 2.0
AxolotlPythonYAML設定DPO/ORPOFSDPApache 2.0
LLaMA-FactoryPythonGUI/CLIPPO/DPO/KTO/ORPODeepSpeedApache 2.0

実装例(TRL)

from trl import RewardTrainer, RewardConfig
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct", num_labels=1
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")

config = RewardConfig(
    output_dir="./reward_model",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=16,
    learning_rate=1e-5,
    num_train_epochs=1,
    bf16=True,
    max_length=2048,
)

trainer = RewardTrainer(
    model=model,
    tokenizer=tokenizer,
    args=config,
    train_dataset=preference_dataset,
)
trainer.train()

コスト見積もり

構成GPU8B RM 訓練時間70B RM 訓練時間概算コスト(クラウド)
シングル GPUA100 80GB約4時間不可(VRAM不足)$8〜$12
4x GPUA100 80GB×4約1.5時間約12時間$12〜$100
8x GPUH100 80GB×8約30分約4時間$30〜$120

よくある質問(FAQ)

Q1: RM のバックボーンはポリシーモデルと同じサイズにすべきですか? A: 必ずしも同じサイズである必要はない。一般的にポリシーモデルの1/4〜1/2サイズで十分な精度が得られる。ただし、Nemotron-4 340B RM のようにポリシーモデルと同等サイズの RM が最高精度を達成した例もある。

Q2: 訓練データはどのくらいの量が必要ですか? A: 最低10,000ペア、推奨50,000〜100,000ペアである。UltraFeedback(約64,000件)で訓練した 8B RM が RewardBench で90%以上のスコアを達成している実績がある。

Q3: RM 訓練で過学習を防ぐにはどうすればよいですか? A: エポック数を1〜2に制限し、学習率を SFT より低く(1e-5以下)設定する。また、訓練データの多様性確保(複数ドメイン・複数アノテーター)と、バリデーションセットでの早期停止が有効である。

まとめ

  • RM 訓練パイプラインは SFT → 選好データ収集 → RM 訓練 → 評価 → デプロイの5段階
  • TRL・OpenRLHF 等のフレームワークで数行のコードから実装可能
  • 8B RM は A100 1枚で約4時間、概算$10以下で訓練可能
  • 過学習防止にはエポック数制限・低学習率・データ多様性確保が重要