[논문 리뷰] DiffusionDPO : Diffusion Model Alignment Using Direct Preference Optimization

[2311.12908] Diffusion Model Alignment Using Direct Preference Optimization

Diffusion Model Alignment Using Direct Preference Optimization

 

1. Intro

 

문제점

 

LLM : 2 stage trianing

 

1. Pretraining (large web scale data로)
2. Alignment (SFT + RLHF → pref 반영)

Text-to-Image Diffusion Models : Alignment 단계 없음

• web scale text-img 쌍으로 1 stage training

 

제안 : Diffusion DPO

• Diffusion 모델을 pref에 맞게 안정적으로 alignment 하는 방법 필요
• human pref data로 implicit하게 reward 학습

 

2. Related Work

 

LLM에서의 정렬

 

RLHF

1. Reward Function 학습
   • 모델이 특정 prompt에 대해 답변을 하면 사람이 그 중 더 좋은 것을 선택
   • 이 comparison data로 reward model 학습
2. Policy model 정렬을 위해 RL 사용
   • Reward Model이 준 점수를 최대화하는 방향으로 policy model(=생성 모델)을 RL

but

• 비용이 크고
• hyperparam 조정 필요
• reward hacking 문제

 

대안

1. Reward model based sampling + SFT
   • Base model(정렬 전인 사전 학습된 모델)이 여러 답변 생성하면
   • Reward model로 점수를 매기고
   • 높은 점수 답변만 골라 STF
     • but 여전히 reward model 의존 (reward model 학습하고 그걸 기반으로 또 RL하는 것은 비효율적)
2. Direct Feedback based Fine-Tuning
   • human pref data를 모델 학습에 직접 사용하자! (reward model 필요 없게)
     1. Ranking Loss 
        • A 답변을 B 답변보다 선호했다면 모델이 A의 log-likelihood를 B보다 더 높게 하도록 학습
     2. DPO
        • Ranking Loss와 비슷, but 사전학습된 ref model 분포에서 너무 멀어지지 않도록 제약 추가
        • * Diffusion-DPO는 이 DPO를 diffusion에 일반화한 것!

 

Diffusion에서의 정렬 시도

1. Aesthetics based ft
   • Aesthetics classifier가 보기 좋은 img로 평가한 것 선택해서 ft
   • but text fidelity 같은 다른 기준은 반영 불가
2. 고품질 데이터셋 기반
   • pretrained model을 소규모 고품질 사진 + 수작업 캡션 data로 ft
   • but dataset 제한적
3. Data recaptioning
   • 웹에서 긁어온 기존 img data caption 정밀하게 다시 작성
   • text fidelity 개선
   • but pref 반영 불가
4. Human pref scoring
   • human pref data로 학습한 preference scoring model 사용
   • but 생성 공간 전체의 안정적 정렬은 불가능
5. Inference time alignment
   • DOODL [51]
   • 한 번 생성된 이미지를 inference 단계에서 iterative하게 개선
   • but 추가 연산 → inference 시간 증가
6. Reward based training
   • DRAFT [7], AlignProp [31]
   • 학습 중 reward gradient 사용
   • but 안정성 부족
7. RL based
   • DPOK [6], DDPO [11]
   • RL로 reward 최대화   
   • but open vocabulary 일때 성능 -

비교

2. additional inference time / 3. ref model에서 벗어나지 않게

 

3. Background

 

3.2. Direct Preference Optimization

 

Reward Modeling (reward model 학습!)

• 목표 : pθ​(x0​∣c)가 human pref에 맞게 학습되도록
• comparison data 형태는:

• BT model : human pref를 확률적으로 모델링
• +의 reward가 -보다 높을수록 선호 확률 커짐

BT

• 학습 목표는 :

 

RLHF

 

• 첫번째 항: reward model이 +로 평가한 sample 자주 만들도록
• 두 번째 항: ref 분포에서 벗어나지 않도록 KL penalty
• β: 두 항의 균형을 조절하는 hyperparam

 

DPO Objective

• reward model 학습 필요x, policy vs ref의 likelihood 비율만 필요

• +의 상대적 확률이 -의 상대적 확률보다 크도록
• 단, ref 분포 대비 얼마나 개선되는지도 고려

 

정리

• RLHF
  1. Reward model 학습 (human pref data로)
  2. RL(PPO)로 policy 최적화
• DPO
  • Reward model 학습 생략
  • RL 단계 생략
  • human pref data 직접 likelihood 학습

 

4. DPO for Diffusion Models

 

문제점

• pθ​(x0​∣c) is not tractable
• 왜? : 모든 경로를 marginalize하는 것은 불가능
• 즉, 모든 가능한 노이즈 경로를 고려해야 하는데 계산 불가

 

해결 : ELBO

 

• Latent variable x_1:T(noise step) 도입
• 보상 r(c,x0)를 경로 기반 기대값으로 정의:

 

• KL 항도 직접 계산 어렵기 때문에, joint 분포 KL upper bound 사용:

• 결국, 경로 전체를 대상으로 DPO optimization eq를 다시 세움

 

최종: