TRL 為視覺語言模型新增 Direct Preference Optimization 支援

TL;DR

Hugging Face 宣布 TRL 函式庫現在支援視覺語言模型 (VLMs) 的 Direct Preference Optimization (DPO),讓開發者能夠使用偏好數據對 Idefics-2、Llava 1.5 和 PaliGemma 等模型進行微調,並透過 bfloat16 量化和 LoRA 适配器保持可控的記憶體需求。

針對 VLMs 的基於偏好的微調

偏好優化取代了成本高昂的逐標籤監督,改用二元比較:每個訓練範例包含一個提示詞 (prompt)、一個被選中 (chosen) 的答案和一個被拒絕 (rejected) 的答案。模型學習為被選中的回答分配更高的機率。這種方法捕捉了細微的人類判斷,並已廣泛應用於語言模型;新的 TRL 整合將其擴展到了多模態 VLMs。

範例數據集

本部落格使用了 openbmb/RLAIF-V-Dataset,它提供了 8.3 萬多行包含被選中與被拒絕文本答案的圖像-問題對。範例條目如下:

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."

被選中的答案可能在事實上仍有錯誤,但它比被拒絕的答案錯誤程度較低,這是偏好學習的核心前提。

聊天風格 VLMs 的格式化

數據集必須重新整理成聊天格式,其中使用者提供圖像和文本查詢,而助手則以被選中或被拒絕的文本進行回覆。使用 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)
Model (train) 4 (float32) 32
Reference model 4 32
Gradients 4 32
Optimizer states (2×) 4 64
Total 160

由於大多數 GPU 的容量小得多,部落格建議使用兩種互補的技術。

量化至 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% 的總量)。梯度和優化器狀態的記憶體現在降至幾百 MB。

在量化和 LoRA 之後重新計算預算,總計約為 ≈32 GB,可以輕鬆放入 80 GB 的 GPU 中。

Batch size 與激活記憶體

激活值 (Activations) 未計入靜態預算中。作者建議一種經驗方法:從理想的 batch size(例如 64)開始,觀察是否發生 OOM,若發生則將 batch size 減半,並加倍 gradient_accumulation_steps 以保持有效 batch size 不變。在他們的實驗中,他們決定使用 per_device_train_batch_size=2 搭配 gradient_accumulation_steps=32。啟用 gradient_checkpointing=True 可以進一步減少激活記憶體,但代價是增加計算量。

完整的訓練腳本

一個自包含的腳本 (dpo_idefics2-8b.py) 將模型載入、數據集格式化、LoRA 配置和 DPOTrainer 整合在一起。關鍵參數包括:

  • bf16=Truegradient_checkpointing=True
  • per_device_train_batch_size=2, gradient_accumulation_steps=32
  • 使用 dataset_num_proc=32dataloader_num_workers=32 進行並行預處理
  • 傳遞給 DPOTrainerLoraConfig(target_modules="all-linear")

使用 accelerate launch dpo_idefics2-8b.py 執行腳本可啟動單輪 (single-epoch) DPO 微調。

訓練成果

損失曲線顯示兩個 DPO 特定指標穩定提升:

  • 準確度 (Accuracy) – 模型將較高機率分配給被選中答案的樣本比例。
  • 獎勵邊際 (Reward margin) – 被選中答案與被拒絕答案之間的獎勵(對數機率)之差;邊際增加表示偏好學習成功。

這兩個指標在訓練過程中都在增加,證實了 DPO 可以有效地引導 VLMs 朝向偏好的回覆發展。

幻覺減少的評估

為了評估 DPO 是否減輕了幻覺,微調後的 Idefics-2 模型在 AMBER 基準測試(專門針對 VLM 的幻覺測試)上進行了評估。結果 (準確度 / F1) 為:

模型 準確度 F1
GPT-4o 88.8 91.6
Idefics-2 + DPO 85.9 89.4
Idefics-2 (baseline) 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

這些範例顯示,當訓練數據發出偏好信號時,模型學會了偏好錯誤率較低的答案。

將 DPO 擴展到其他 VLMs

TRL 的 DPO 實作已經支援 Llava 1.5PaliGemma。部落格指出了 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。

影響

透過將 DPO 整合進 TRL 以用於 VLMs,Hugging Face 降低了多模態模型進行偏好驅動微調的門檻。開發者現在可以在不進行昂貴標籤收集的情況下,將 VLMs 與人類判斷對齊,同時保持在單個高端 GPU 的記憶體限制內。AMBER 上微小的幻覺改善表明,偏好數據可以成為減少過度自信錯誤的有效信號,為開發更安全、更可靠的視覺語言助手開闢了道路。


TL;DR – 新的 TRL DPO 支援讓您可以使用二元偏好數據來微調視覺語言模型,透過 bfloat16 量化加上 LoRA 适配器,您可以在單個 80 GB GPU 上訓練 8 B 參數的 VLM,並在偏好準確度和減少幻覺方面取得顯著進步。

Sources