TRL, Vision-Language Models를 위한 Direct Preference Optimization 지원 추가

TL;DR

Hugging Face는 TRL 라이브러리가 이제 Vision‑Language Models(VLMs)를 위한 Direct Preference Optimization(DPO)을 지원한다고 발표했습니다. 이를 통해 개발자들은 bfloat16 양자화와 LoRA 어댑터를 사용하여 메모리 요구 사항을 관리 가능한 수준으로 유지하면서 Idefics‑2, Llava 1.5, PaliGemma와 같은 모델을 선호도 데이터로 미세 조정할 수 있습니다.

VLMs를 위한 선호도 기반 미세 조정

선호도 최적화(Preference optimization)는 비용이 많이 드는 레이블 단위 감독을 이진 비교로 대체합니다. 각 학습 예제에는 프롬프트, 선택된(chosen) 답변, 그리고 거부된(rejected) 답변이 포함됩니다. 모델은 선택된 응답에 더 높은 확률을 할당하는 법을 배웁니다. 이 접근 방식은 미묘한 인간의 판단을 포착하며 언어 모델에 널리 채택되어 왔습니다. 새로운 TRL 통합은 이를 멀티모달 VLM으로 확장합니다.

예시 데이터셋

이 블로그는 openbmb/RLAIF‑V‑Dataset을 사용하며, 이는 선택된 답변과 거부된 텍스트 답변이 포함된 83k개 이상의 이미지-질문 쌍을 제공합니다. 샘플 항목은 다음과 같습니다:

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에도 적용됩니다.

메모리 예산

전정밀도(full-precision) 8B 파라미터 모델을 학습하려면 약 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바이트로 절반으로 줄어들어 모델 메모리가 32GB에서 16GB로 감소합니다. 이 변경 사항은 모델과 옵티마이저 모두에 적용됩니다:

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

PEFT를 통한 LoRA 어댑터

Low-Rank Adaptation (LoRA)는 기본 모델을 동결하고 선형 레이어에 학습 가능한 랭크 분해 행렬을 주입합니다. peft.LoraConfig(target_modules="all-linear")를 사용하면 학습 가능한 파라미터가 8B에서 약 55M(전체의 ≈0.65%)로 줄어듭니다. 이제 그래디언트와 옵티마이저 상태를 위한 메모리는 수백 메가바이트로 떨어집니다.

양자화와 LoRA를 적용한 후 예산을 다시 계산하면 총 ≈32 GB가 되어 80 GB GPU에 여유롭게 들어갑니다.

배치 크기 및 활성화 메모리

활성화(Activations)는 정적 예산에 포함되지 않습니다. 저자들은 경험적인 접근 방식을 제안합니다: 원하는 배치 크기(예: 64)로 시작하여 OOM이 발생하면 배치 크기를 절반으로 줄이고, 유효 배치 크기를 일정하게 유지하기 위해 gradient_accumulation_steps를 두 배로 늘립니다. 그들의 실험에서는 per_device_train_batch_size=2와 gradient_accumulation_steps=32로 결정했습니다. gradient_checkpointing=True를 활성화하면 추가 연산 비용을 대가로 활성화 메모리를 더욱 줄일 수 있습니다.

전체 학습 스크립트

독립 실행 가능한 스크립트(dpo_idefics2-8b.py)는 모델 로딩, 데이터셋 포맷팅, LoRA 구성 및 DPOTrainer를 결합합니다. 주요 인수는 다음과 같습니다:

  • bf16=True 및 gradient_checkpointing=True
  • per_device_train_batch_size=2, gradient_accumulation_steps=32
  • dataset_num_proc=32 및 dataloader_num_workers=32를 이용한 병렬 전처리
  • DPOTrainer에 전달되는 LoraConfig(target_modules="all-linear")

accelerate launch dpo_idefics2-8b.py로 스크립트를 실행하면 단일 에포크 DPO 미세 조정이 시작됩니다.

학습 결과

손실 곡선은 두 가지 DPO 전용 지표에서 꾸준한 개선을 보여줍니다:

  • Accuracy (정확도) – 모델이 선택된 답변에 더 높은 확률을 할당하는 샘플의 비율.
  • Reward margin (보상 마진) – 선택된 답변과 거부된 답변의 보상(로그 확률) 차이; 마진이 커지는 것은 성공적인 선호도 학습을 나타냅니다.

두 지표 모두 학습 과정 동안 증가하며, 이는 DPO가 VLMs를 선호하는 응답으로 효과적으로 유도할 수 있음을 확인시켜 줍니다.

환각 감소에 대한 평가

DPO가 환각을 완화하는지 평가하기 위해, 미세 조정된 Idefics-2 모델을 AMBER 벤치마크(VLM 전용 환각 테스트)에서 평가했습니다. 결과(정확도 / F1)는 다음과 같습니다:

모델 Accuracy 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 샘플은 변화를 보여줍니다:

이미지 질문 Baseline 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.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에 VLMs를 위한 DPO를 통합함으로써, Hugging Face는 멀티모달 모델의 선호도 기반 미세 조정에 대한 장벽을 낮춥니다. 개발자들은 이제 값비싼 레이블 수집 없이도 단일 고성능 GPU의 메모리 제한 내에서 VLMs를 인간의 판단에 맞게 정렬할 수 있습니다. AMBER에서의 완만한 환각 개선은 선호도 데이터가 과도하게 확신하는 오류를 줄이는 데 효과적인 신호가 될 수 있음을 시사하며, 더 안전하고 신뢰할 수 있는 시각-언어 어시스턴트로 나아가는 길을 열어줍니다.

TL;DR – 새로운 TRL DPO 지원을 통해 이진 선호도 데이터를 사용하여 시각-언어 모델을 미세 조정할 수 있으며, bfloat16 양자화와 LoRA 어댑터를 사용하면 단일 80GB GPU에서 8B 파라미터 VLM을 학습시켜 선호도 정확도 향상 및 환각 감소라는 측정 가능한 이득을 얻을 수 있습니다.

Sources