RLHFとは？人間のフィードバックによる強化学習でAIを改善する仕組み

Hugging Faceのtrlライブラリを使って、RLHF（人間のフィードバックによる強化学習）の3段階パイプライン（SFT → Reward Model → PPO）の概念コードを示します。

from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig, RewardTrainer, RewardConfig
from datasets import load_dataset
import torch

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# =============================
# Stage 1: SFT（教師あり微調整）
# =============================
# 高品質な指示-応答ペアでモデルを微調整
sft_dataset = load_dataset("tatsu-lab/alpaca", split="train[:1000]")

sft_model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.float16, device_map="auto"
)

sft_config = SFTConfig(
    output_dir="./sft_output",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=2e-5,
    logging_steps=10,
    max_seq_length=512,
)

sft_trainer = SFTTrainer(
    model=sft_model,
    args=sft_config,
    train_dataset=sft_dataset,
    tokenizer=tokenizer,
)
# sft_trainer.train()  # 実行時はコメント解除

# =============================
# Stage 2: Reward Model（報酬モデル）
# =============================
# 人間が「良い回答」と「悪い回答」を比較評価したデータで学習
from transformers import AutoModelForSequenceClassification

reward_model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=1,               # スカラー報酬を出力
    torch_dtype=torch.float16,
    device_map="auto",
)

# 比較データの形式: chosen（好ましい応答）vs rejected（好ましくない応答）
# reward_dataset = load_dataset("Anthropic/hh-rlhf", split="train[:1000]")

reward_config = RewardConfig(
    output_dir="./reward_output",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    learning_rate=1e-5,
)

# =============================
# Stage 3: PPO（強化学習による最適化）
# =============================
from trl import PPOTrainer, PPOConfig

ppo_config = PPOConfig(
    model_name=model_name,
    learning_rate=1.41e-5,
    batch_size=16,
    mini_batch_size=4,
    ppo_epochs=4,               # 1バッチあたりのPPO更新回数
    kl_penalty="kl",            # KLダイバージェンス制約
    init_kl_coef=0.2,           # KLペナルティの初期係数
)

# PPOの学習ループの概念:
# for batch in dataloader:
#     query = batch["query"]
#     response = model.generate(query)       # SFTモデルで応答生成
#     reward = reward_model(query + response) # 報酬モデルでスコア算出
#     ppo_trainer.step(query, response, reward) # PPOで方策を更新

print("RLHF Pipeline構成完了")
print("Stage 1: SFT -> 指示に従う基本能力を獲得")
print("Stage 2: Reward Model -> 人間の好みを数値化")
print("Stage 3: PPO -> 報酬を最大化する方策に最適化")

RLHFの3段階はそれぞれ異なる役割を担います。SFTで「指示に従う能力」を身につけ、Reward Modelで「人間の好みを数値化」し、PPOで「報酬が高い応答を生成するように方策を最適化」します。PPOではKLダイバージェンスのペナルティが重要で、これがないとモデルが報酬ハッキング（報酬モデルを騙す不自然な応答）に陥ります。最近ではDPO（Direct Preference Optimization）がReward Model + PPOの2段階を1段階に統合する手法として注目されています。