TRL が視覚言語モデル向けに Direct Preference Optimization をサポート

TL;DR

Hugging Face は TRL ライブラリが視覚言語モデル (VLM) に対して Direct Preference Optimization (DPO) をサポートするようになったと発表しました。これにより、Idefics‑2、Llava 1.5、PaliGemma などのモデルを、bfloat16 量子化と LoRA アダプターを活用して、メモリ使用量を抑えながら、好みのデータでファインチューニングできるようになります。

VLM における好みベースのファインチューニング

好み最適化は、高コストなラベルごとの監視を、二値比較に置き換えます。各訓練例には、プロンプト、選ばれた回答、および 拒否された回答が含まれます。モデルは選ばれた回答に高い確率を割り当てるよう学習します。このアプローチは、人間の微妙な判断を捉え、言語モデルでは広く採用されていますが、新しい TRL 統合により、マルチモーダルな VLM にも拡張されました。

例のデータセット

ブログでは openbmb/RLAIF‑V‑Dataset を使用しており、83 k 以上の画像-質問ペアと、選ばれた・拒否されたテキスト回答が提供されています。サンプルエントリは以下の通りです:

Question: "How many families?"
Rejected: "The image does not provide any information about families."
Chosen:   "The image shows a Union Organization table setup with 18,000 families."

選ばれた回答が事実上誤っている可能性もありますが、拒否された回答よりも「誤りが少ない」ことが、好み学習の核心的な前提です。

チャット形式の VLM 用フォーマット

データセットは、ユーザーが画像とテキストクエリを提供し、アシスタントが選ばれたまたは拒否されたテキストで応答するチャット形式に再構成する必要があります。Hugging Face の AutoProcessor(例:HuggingFaceM4/idefics2-8b)を使用して、チャットテンプレートを適用し、プロセッサの最大エッジ長に画像をリサイズすることで、メモリ不足(OOM)エラーを回避します。以下のコードスニペットは変換を示しています:

from datasets import features
from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b", do_image_splitting=False)

def format(example):
    prompt = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": example["question"]}]}]
    chosen = [{"role": "assistant", "content": [{"type": "text", "text": example["chosen"]}]}]
    rejected = [{"role": "assistant", "content": [{"type": "text", "text": example["rejected"]}]}]
    prompt   = processor.apply_chat_template(prompt,   tokenize=False)
    chosen   = processor.apply_chat_template(chosen,   tokenize=False)
    rejected = processor.apply_chat_template(rejected, tokenize=False)
    max_size = processor.image_processor.size["longest_edge"]
    example["image"].thumbnail((max_size, max_size))
    return {"images": [example["image"]], "prompt": prompt, "chosen": chosen, "rejected": rejected}

この関数をデータセットにマッピングし、images 列をデコードされた PIL.Image オブジェクトに変換することで、訓練に必要なデータが準備されます。

DPO による VLM の訓練

ブログでは、Idefics‑2‑8b を参照モデルとしてファインチューニングする例を示していますが、同じパイプラインは Llava 1.5 および PaliGemma にも適用可能です。

メモリ予算

フル精度の 8 B パラメータモデルを訓練するには、約 160 GB の VRAM(モデル、参照コピー、勾配、AdamW 状態)が必要です。著者たちはステップバイステップの計算を示しています:

コンポーネント パラメータあたりのバイト数 合計 (GB)
モデル (訓練用) 4 (float32) 32
参照モデル 4 32
勾配 4 32
オプティマイザ状態 (2×) 4 64
合計 160

ほとんどの GPU はこれよりもはるかに小さいため、ブログでは2つの補完的な技術を推奨しています。

bfloat16 への量子化

torch.bfloat16 に切り替えることで、パラメータあたりのストレージを 4 バイトから 2 バイトに半減し、モデルメモリを 32 GB から 16 GB に削減します。この変更はモデルとオプティマイザの両方に適用されます:

model = AutoModelForVision2Seq.from_pretrained(..., torch_dtype=torch.bfloat16)
training_args = DPOConfig(..., bf16=True)

PEFT を通じた LoRA アダプター

低ランク適応(LoRA)は、ベースモデルを固定し、線形層にトレーナブルなランク分解行列を挿入します。peft.LoraConfig(target_modules="all-linear") を使用すると、トレーナブルパラメータは 8 B から約 55 M(全体の約 0.65 %)に減少します。勾配とオプティマイザ状態のメモリは、数百メガバイトにまで低下します。

量子化と LoRA を適用した後の予算再計算では、約 32 GB の合計となり、80 GB GPU に余裕を持って収まります。

バッチサイズと活性化メモリ

活性化は静的予算には含まれていません。著者たちは経験的なアプローチを提案します:希望するバッチサイズ(例:64)から始め、OOM を観察し、バッチサイズを半分にし、gradient_accumulation_steps を2倍にして有効バッチサイズを一定に保ちます。実験では per_device_train_batch_size=2gradient_accumulation_steps=32 に落ち着きました。gradient_checkpointing=True を有効にすることで、活性化メモリをさらに削減できますが、計算コストが増加します。

完全な訓練スクリプト

自己完結型のスクリプト(dpo_idefics2-8b.py)が、モデルの読み込み、データセットのフォーマット、LoRA 設定、および DPOTrainer を統合しています。主な引数には以下が含まれます:

  • bf16=True および gradient_checkpointing=True
  • per_device_train_batch_size=2gradient_accumulation_steps=32
  • dataset_num_proc=32 および dataloader_num_workers=32 を用いた並列前処理
  • DPOTrainer に渡される LoraConfig(target_modules="all-linear")

accelerate launch dpo_idefics2-8b.py でスクリプトを実行すると、1エポック分の DPO ファインチューニングが開始されます。

訓練結果

損失曲線は、2つの DPO 特有の指標において継続的な改善を示しています:

  • 正解率 – モデルが選ばれた回答に高い確率を割り当てるサンプルの割合。
  • 報酬マージン – 選ばれた回答と拒否された回答の報酬(対数確率)の差;増加するマージンは、成功した好み学習を示しています。

両方の指標が訓練期間中に増加しており、DPO が VLM を望ましい応答に効果的に導くことができることを確認しています。

ハルシネーション低減の評価

DPO がハルシネーションを軽減するかどうかを評価するために、ファインチューニングされた Idefics‑2 モデルを AMBER ベンチマーク(VLM 専用のハルシネーションテスト)で評価しました。結果(正解率 / F1)は以下の通りです:

モデル 正解率 F1
GPT‑4o 88.8 91.6
Idefics‑2 + DPO 85.9 89.4
Idefics‑2 (ベースライン) 85.8 89.1
GPT‑4v 83.4 87.4
MiniGemini 82.6 87.6

DPO ファインチューニングモデルはベースラインと同等またはわずかに上回っており、ハルシネーションがわずかに減少していることを示しています。

定性的な例

選ばれた AMBER サンプルは変化を示しています:

画像 質問 ベースライン Idefics‑2 Idefics‑2 + DPO
![ships] Are there two ships? Yes No
![ground] Is the ground uneven? No Yes
![shovel] Is there one shovel? Yes No

これらの例は、トレーニングデータが好みを示す場合、モデルがよりハルシネーションの少ない回答を好むよう学習していることを示しています。

他の VLM への DPO の拡張

TRL の DPO 実装は、既に Llava 1.5 および PaliGemma をサポートしています。ブログでは TRL リポジトリ内の例スクリプト(examples/scripts/dpo_vlm.py)を紹介しています。PaliGemma 用の典型的なコマンドラインは以下の通りです:

accelerate launch examples/scripts/dpo_visual.py \
    --dataset_name HuggingFaceH4/rlaif-v_formatted \
    --model_name_or_path google/paligemma-3b-pt-224 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 32 \
    --dataset_num_proc 32 \
    --output_dir dpo_paligemma_rlaif-v \
    --bf16 \
    --torch_dtype bfloat16 \
    --gradient_checkpointing \
    --use_peft \
    --lora_target_modules=all-linear

同じ量子化+LoRA レシピが適用可能であり、中程度の GPU ハードウェアでも DPO を利用可能になります。

意味

TRL に VLM 用 DPO を統合することで、Hugging Face はマルチモーダルモデルの好み駆動型ファインチューニングの障壁を下げました。開発者は、高価なラベル収集をせずに、人間の判断にモデルを合わせられるようになり、1台の高級 GPU のメモリ制限内に収まる形で実現できます。AMBER でのわずかなハルシネーション改善は、好みデータが過信した誤りを低減する有効なシグナルである可能性を示しており、より安全で信頼性の高い視覚言語アシスタントへの道を開いています。


TL;DR – 新しい TRL DPO サポートにより、二値の好みデータを使って視覚言語モデルをファインチューニングできるようになり、bfloat16 量子化と LoRA アダプターを併用することで、1台の 80 GB GPU 上で 8 B パラメータの VLM を訓練でき、好みの正確性の向上とハルシネーションの低減という測定可能な成果が得られます。

Sources

関連