기술 요약: 불확실성 가이드를 위한 샘플 적응형 잠재 보상 디퓨전 사후 학습 (SURE)

1. 문제 정의

인간의 선호도에 맞게 시각적 디퓨전 모델을 정렬(alignment)하는 사후 학습(post-training)은 일반적으로 디코딩된 RGB 출력물을 평가하는 데 의존하며, 이는 피드백 루프 내에서 반복적이고 비용이 많이 드는 VAE 디코딩을 필요로 합니다. 노이즈가 섞인 잠재 상태(latent)를 직접 점수화하여 디코딩 비용을 방지하고 중간 단계의 감독을 가능하게 하는 "디퓨전 네이티브" 잠재 보상 모델(LRM)이 등장했음에도 불구하고, 기존 방식들은 두 가지 결정적인 한계점을 가집니다:

1. 스칼라 출력의 불확실성 결여: 현재의 잠재 보상 모델은 스칼라 점수만을 출력합니다. 즉, 각 예측의 불확실성을 추정하지 못합니다. 결과적으로 생성기(generator)는 신뢰할 수 있는 피드백과 신뢰할 수 없는 피드백을 구분할 수 없습니다.
2. 정적 불확실성 가정: DiNa-LRM과 같은 이전 연구들은 노이즈 수준(타임스텝)에 따라서만 불확실성을 할당합니다. 이는 동일한 타임스텝의 모든 샘플이 동일한 신뢰도를 공유한다고 가정합니다. 그러나 본 논문은 보상의 난이도가 콘텐츠에 따라 달라질 수 있다고 주장합니다. 예를 들어, 명확한 피사체가 있는 "노이즈가 많은" 샘플이 모호한 콘텐츠를 가진 "깨끗한" 샘플보다 더 예측 가능할 수 있습니다.
3. 최적화 위험: 불확실성 추정이 없으면 생성기가 신뢰할 수 없는 로컬 피드백을 바탕으로 최적화될 수 있으며, 이는 "보상 해킹(reward hacking)"(모델이 품질을 개선하는 대신 보상 함수를 악용하는 현상)과 불안정한 학습 역학으로 이어집니다.

2. 방법론

저자들은 SURE(Sample-Adaptive Latent Reward for Uncertainty-Guided Diffusion Post-Training)를 제안합니다. 이는 SURE-LRM(보상 모델)과 SURE-REFL(사후 학습 알고리즘)이라는 두 가지 구성 요소로 이루어진 통합 프레임워크입니다.

2.1 SURE-LRM: 샘플 적응형 잠재 보상 모델

SURE-LRM은 단순히 보상 점수를 예측하는 것이 아니라, 프롬프트와 타임스텝에 따라 각 노이즈 잠재 상태에 대한 전체 가우시안 효용 분포(Gaussian utility distribution)를 예측하도록 설계되었습니다.

• 아키텍처: VAE 디코딩 없이 노이즈 잠재 상태, 프롬프트, 타임스텝을 처리하기 위해 대상 생성기의 디퓨전 트랜스포머(DiT) 백본을 재사용합니다(이미지 및 비디오 모두 해당).
• 예측 대상: 노이즈 잠재 상태 $z_\sigma$에 대해 모델은 다음을 예측합니다:
  • 평균 보상 $\mu_\theta$.
  • 샘플 적응형 잔차 분산 파라미터 $a_\theta$.
• 분산 모델링: 총 예측 분산 $v_\eta(\sigma)$는 노이즈 의존적 사전 분산 $v_0(\sigma)$(노이즈 증가에 따른 평균적인 정보 손실을 포착)와 학습된 샘플별 잔차 분산 $v_{\eta, res} = \text{softplus}(a_\theta) + \epsilon_v$의 합입니다.
• 학습 목적 함수: 모델은 표준 쌍별 선호도($z^+ \succ z^-$)를 사용하여 헤테로스케다스틱(heteroscedastic) 선호도 목적 함수로 학습됩니다:
  • 우도 손실 ($L_{pref}$): Thurstone의 무작위 효용 모델(random-utility model)에 기반하며, 선호되는 샘플이 더 높은 효용을 갖도록 최적화합니다. 그래디언트는 모델이 순위가 틀리거나 모호할 때 분산(불확실성)을 높이고, 순위가 확신 있고 정확할 때 분산을 낮추도록 유도합니다.
  • 정규화:
    • 분산 정규화 ($L_{var}$): 모델이 단순히 불확실성을 팽창시켜 모든 순위 오류를 설명하려는 현상(분산 팽창)을 방지합니다.
    • 마진 정규화 ($L_{margin}$): 선호되는 점수와 거부된 점수 사이의 최소한의 격차를 보장합니다.
    • 앵커 손실 ($L_{anchor}$): 전체적인 점수 척도를 안정화합니다.

2.2 SURE-REFL: 불확실성 가이드 밀집 보상 피드백 학습

SURE-LRM이 학습되면 그 파라미터는 고정됩니다. SURE-REFL은 이 모델을 활용하여 잠재 공간에서 생성기의 사후 학습을 가이드합니다.

• 로컬 피드백 메커니즘: 전체 디노이징 그래프를 역전파하는 대신, SURE-REFL은 디노이징 궤적 상의 선택된 전이(transition, $j$)에서 고정된 보상 모델을 쿼리합니다.
• 불확실성 가이드 가중치 부여:
  1. 모델은 전이 $j$에서의 상태에 대한 평균 보상 $\mu$와 잔차 분산을 예측합니다.
  2. 총 분산은 신뢰도 가중치 $\rho = V^{-1/2}$(역 표준 편차)로 변환됩니다.
  3. 이 가중치는 detachment(stop-gradient) 처리됩니다. 즉, 생성기는 가중치를 통해 학습하지만, 이 단계에서 보상 모델의 불확실성 추정치를 업데이트하지는 않습니다.
  4. 공정한 비교를 위해 동일한 전이에서 평가된 샘플들 간에 가중치를 정규화합니다.
• 최적화: 생성기는 예측된 분산이 낮은(신뢰도가 높은) 샘플이 그래디언트에 더 크게 기여하도록 하는 가중 보상 목적 함수를 사용하여 업데이트됩니다. 이 밀집 피드백은 각 전이에 로컬하게 적용되어, 전체 그래프 유지나 픽셀 공간 디코딩의 필요성을 제거합니다.

3. 주요 기여

1. 샘플 적응형 불확실성 공식화: 추가적인 인간의 신뢰도 주석 없이도 표준 쌍별 선호도로부터 유도되는 학습 가능한 분포로서 잠재 공간의 보상 불확실성을 공식화했습니다.
2. SURE-LRM: 동일한 타임스텝의 샘플이라도 콘텐츠에 따라 서로 다른 신뢰도 추정치를 가질 수 있음을 입증하며, 샘플 의존적 불확실성을 학습하는 잠재 보상 모델을 제시했습니다.
3. SURE-REFL: 디코딩된 픽셀 공간의 VAE나 전체 디노이징 그래프 유지가 필요 없이, 분리된(detached) 불확실성을 사용하여 밀집 로컬 피드백에 가중치를 부여하는 이미지 및 비디오를 위한 통합 사후 학습 방법을 제공합니다.
4. 이론적 통찰: 저자들은 분산 그래디언트를 도출하여 선호도 손실이 어떻게 불확실성 학습을 위한 내재적 신호를 제공하는지 보여주었습니다.

4. 실험 결과

4.1 보상 모델 정확도 (SURE-LRM)

• 이미지 벤치마크: 세 가지 백본(SD3.5, Z-Image-Turbo, FLUX.1-dev)에 대해 네 가지 벤치마크(ImageReward, HPD v2, HPD v3, GenAI-Bench)에서 SOTA 수준의 쌍별 선호도 정확도를 달성하였으며, DiNa-LRM 및 VLM 기반 픽셀 보상을 포함한 강력한 베이스라인들을 능가했습니다.
• 비디오 벤치마크: VisionRewardDB 및 T2V Ranking에서 최고 정확도를 기록하며, 디코딩된 비디오 평가기 및 비디오 특화 보상 모델들을 앞질렀습니다.
• 불확실성 진단: 실험을 통해 학습된 잔차 분산이 예측 오류와 상관관계가 있음을 확인했습니다. 순위가 잘못된 샘플은 평균적으로 훨씬 높은 잔차 분산(24.26% 더 높음)을 보였습니다. 낮은 예측 분산을 가진 샘플로 필터링했을 때 쌍별 정확도가 79.40%에서 90.10%로 향상되었습니다.

4.2 사후 학습 성능 (SURE-REFL)

• 이미지 생성: SURE-REFL은 SD3.5, Z-Image-Turbo, FLUX.1-dev에서 HPSv2.1, ImageReward, PickScore 지표에 대해 최고의 성능을 보였으며, Flow-GRPO, ViPO, DiNa-LRM과 같은 베이스라인보다 우수한 성능을 나타냈습니다.
• 비디오 생성: VBench에서 SURE-REFL은 Wan2.1-14B 비디오 생성기에 대해 Quality, Semantic, Total 지표에서 가장 높은 점수를 기록하며 DanceGRPO, ViPO, Video-DPO, HY-PRFL을 능가했습니다.
• 안정성 및 보상 해킹: 최적화 역학을 보면, 베이스라인(예: DiNa-LRM-REFL)은 독립적인 품질 지표(HPSv3)가 하락하는 동안 보상 점수는 상승하는 보상 해킹 현상을 보인 반면, SURE-REFL은 최적화된 보상과 독립적 품질 지표 간의 정렬을 더 오랫동안 유지하며 더 높은 최종 점수에 도상했습니다.

4.3 절제 연구 (Ablation Studies)

보상 모델의 손실 구성 요소에 대한 절제 연구 결과, 모든 항이 필수적임이 확인되었습니다. 분산 정규화($L_{var}$)를 제거했을 때 정확도가 가장 크게 떨어졌으며, 이는 분산 팽창을 제어해야 할 필요성을 입증합니다.

5. 의의 및 주장

본 논문은 SURE가 잠재 보상 모델링의 "신뢰도 격차"를 해결함으로써 시각적 디퓨전 모델을 정렬하기 위한 효율적이고 통합된 프레임워크를 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

• 효율성: VAE 디코딩 없이 잠재 공간에서 완전히 작동함으로써 피드백 루프 중의 계산 오버헤드를 줄입니다.
• 강건성: 샘플 적응형 불확실성을 명시적으로 모델링하고 활용함으로써, 노이즈가 많거나 모호한 피드백에 과적합되는 것을 방지하여 보상 해킹을 줄이고 학습 안정성을 높입니다.
• 범용성: 프레임워크는 모달리티에 구애받지 않으며, 동일한 불확실성 가중치 로직을 이미지 및 비디오 생성 작업 모두에 성공적으로 적용합니다.

저자들은 표준 쌍별 선호도로부터 샘플 의존적 불확실성을 학습하는 것이 디퓨전 사후 학습에서 밀집한 보상 피드백의 신뢰도를 개선하는 매우 효과적인 전략임을 결론짓습니다.