Beyond Pairwise Preferences: Listwise Reward-Aware Alignment for Diffusion Models

기술적 요약: Diffusion LAIR

문제 제기
직접 선호 최적화 (DPO) 와 확산 모델 전용 변형 (예: Diffusion-DPO, DSPO) 과 같은 텍스트-이미지 확산 모델 정렬을 위한 현재 선호 최적화 방법들은 대부분 쌍별 (pairwise) 감독에 의존합니다. 이러한 방법들은 주어진 프롬프트에 대해 훈련 데이터를 이진 "승자 - 패자" 비교로 축소합니다. 저자들은 이러한 축소가 다음과 같은 두 가지 주요 이유로 제한적이라고 주장합니다:

1. 데이터 구조: 실제 선호 데이터셋 (예: Pick-a-Pic) 은 종종 단일 프롬프트에 대해 여러 후보 이미지를 포함하며, 이는 쌍별 방법이 무시하는 자연스러운 리스트별 (listwise) 구조를 생성합니다.
2. 정보 손실: 현대의 보상 모델은 연속적인 점수를 제공합니다. 이를 단일 이진 레이블로 축소하면 선택되지 않은 후보들의 상대적 품질이 폐기되고, 샘플 간 보상 격차의 크기가 무시됩니다.
   또한, 기존 확산 정렬을 위한 온라인 강화 학습 (RL) 접근 방식은 계산 비용이 많이 들고 병렬화가 어려우며, 온-정책 (on-policy) 특성으로 인해 분포 변화에 민감합니다. 오프라인 방법은 더 효율적이지만, 리스트별 보상 점수 데이터에서 이용 가능한 더 풍부한 신호를 완전히 활용하지는 못했습니다.

방법론: Diffusion LAIR
본 논문은 정렬을 리스트별, 보상 인식 문제로 취급하는 오프라인 선호 최적화 방법인 Diffusion LAIR(Listwise Advantage-weighted Implicit Reward) 을 제안합니다.

• 데이터 표현: 독립적인 쌍 대신, 이 방법은 $N_c$개의 후보 이미지 $\{x^{(i)}_0\}$와 해당 연속 보상 점수 $\{r_i\}$를 가진 프롬프트 $c$를 처리합니다.
• 어드밴티지 가중치: 보상 점수는 중심화된 어드밴티지 가중치($w_i$) 로 변환됩니다. 먼저 온도 $\tau$를 사용하여 소프트맥스 확률 $p_i$를 계산합니다. 그런 다음 가중치를 균일한 기준선에 대해 중심화합니다:
  $$w_i = p_i - \frac{1}{N_c}$$
  이는 $\sum w_i = 0$을 보장하며, 그룹 평균에 비해 높은 보상을 받은 샘플에는 양의 가중치를, 낮은 보상을 받은 샘플에는 음의 가중치를 할당합니다.
• 최적화 목적 함수: 이 방법은 암시적 보상에 어드밴티지 가중 회귀 목적 함수를 최적화합니다. 암시적 보상은 고정된 참조 모델 ($ref$) 대비 현재 모델 ($\theta$) 의 탈노이즈 손실 개선량으로 정의됩니다:
  $$s_\theta(x_0, x_t, t, c, \epsilon) = \omega(\lambda_t) (l_{ref} - l_\theta)$$
  손실 함수는 다음과 같습니다:
  $$L_{LAIR} = \mathbb{E} \left[ -\sum_{i=1}^{N_c} w_i s^{(i)}_\theta + \frac{\lambda}{N_c} \sum_{i=1}^{N_c} (s^{(i)}_\theta)^2 \right]$$
  첫 번째 항은 모델이 높은 보상 샘플에 대해서는 암시적 보상을 높이고 낮은 보상 샘플에 대해서는 낮추도록 장려합니다. 두 번째 항은 참조 모델로부터의 급격한 이탈을 방지하기 위해 암시적 보상의 크기를 정규화하는 2 차 페널티입니다.

이론적 분석
저자들은 목적 함수에 대한 이론적 특성을 제공합니다:

1. 폐형 최적해: 목적 함수는 암시적 보상 공간에서 유일하고 유계인 폐형 최적해를 허용합니다: $s^*_i = \frac{N_c}{2\lambda} w_i$. 이는 정규화 강도 $\lambda$가 선호 업데이트의 크기를 직접 제어함을 명확히 합니다.
2. 대리 KL 경계: 특정 가정 (로그 비율 근사 및 유계 전체 지지 확장) 하에서, 저자들은 유도된 분포와 참조 분포 간의 켈러블 - 라이블 (Kullback-Leibler, KL) 발산에 대한 경계를 유도합니다: $D_{KL} \le \frac{\Delta_c}{\eta}$. 이는 $\lambda$가 선호 기울기의 날카로움을 제어하여 업데이트가 보수적으로 유지됨을 시사합니다.
3. 쌍별 대비 우위: 승자와 패자 간 마진을 극대화하기 위해 암시적 보상을 무한대로 끌어올릴 수 있는 쌍별 목적 함수 (예: Bradley-Terry) 와 달리, LAIR 의 리스트별 공식은 전체 후보 집합에 기반한 유한한 상대적 목표를 할당하여 순위 구조와 보상 간격을 보존합니다.

실험 결과
이 방법은 Pick-a-Pic v2 데이터셋을 사용하여 **Stable Diffusion 1.5 (SD1.5)**와 **Stable Diffusion XL (SDXL)**에서 평가되었습니다.

• 설정: 이미지들은 PickScore 보상 모델을 사용하여 오프라인으로 점수가 매겨졌습니다. 모델들은 텍스트-이미지 생성 (HPD, Parti-prompts), 구성적 생성 (GenEval), 그리고 지시 기반 이미지 편집 (InstructPix2Pix) 에서 파인튜닝 및 평가되었습니다.
• 성능: Diffusion LAIR 은 대부분의 지표에서 강력한 베이스라인 (Diffusion-DPO, DSPO, Diffusion-KTO, MaPO, InPO) 을 능가했습니다.
  • SD1.5에서 Parti-prompt 및 HPD 데이터셋 모두에 대해 PickScore, HPS, Aesthetics, CLIP, ImageReward 에서 최고 점수를 기록했습니다.
  • SDXL에서도 MaPO 및 Diffusion-DPO 를 포함한 베이스라인을 마찬가지로 능가했습니다.
  • 일반화: 주목할 점은 PickScore 로만 훈련되었음에도 불구하고, 모델은 다른 지표들 (예: ImageReward) 에서 상당한 개선을 보여주었는데, 이는 리스트별 목적 함수가 특정 보상 모델을 넘어 일반적인 샘플 품질을 향상시킨다는 것을 시사합니다.
  • 구성 및 편집: 이 방법은 쌍별 베이스라인에 비해 GenEval(구성적 추론) 및 InstructPix2Pix(이미지 편집) 승률에서 일관된 향상을 보였습니다.
• 효율성: 훈련 과정은 온라인 RL 베이스라인보다 훨씬 효율적이었으며, 다른 오프라인 방법과 비슷하거나 더 효율적이었습니다. SD1.5 의 경우 약 24 시간의 GPU 시간이, SDXL 의 경우 140 시간의 A100 GPU 시간이 소요되었습니다.

의의 및 주장
본 논문은 Diffusion LAIR 이 쌍별 축소의 한계를 넘어 확산 모델 정렬을 위한 간단하고 효과적인 방향을 제공한다고 주장합니다. 주요 기여점은 다음과 같습니다:

1. 리스트별 학습: 이는 확산 선호 최적화를 위해 보상 점수가 매겨진 이미지 그룹을 직접 활용하는 첫 번째 방법으로, 상대적 샘플 품질의 전체 구조를 활용합니다.
2. 이론적 명확성: 유계인 폐형 최적해를 제공하여 정규화가 분포 변화와 업데이트 크기를 어떻게 제어하는지에 대한 통찰을 제공합니다.
3. 실증적 우위: 생성, 구성, 편집 작업 전반에 걸쳐 강력한 성능을 입증하여, 리스트별 보상 인식 최적화가 쌍별 방법보다 더 광범위한 개선을 가져옴을 증명합니다.

저자들은 훈련에 사용된 보상 모델의 품질에 대한 의존성과 이론적 KL 경계의 이상화된 가정이라는 한계를 지적합니다. 그러나 인간이 주석한 데이터셋에서 흔히 발생하는 노이즈가 있는 선호 레이블에 대해 이 방법이 견고하다고 주장합니다.