기술 요약: DiNa-LRM – 확산 기반 잠재 공간 보상 모델링

1. 문제 제기

확산 (diffusion) 및 플로우 매칭 (flow-matching) 모델의 인간 선호도 정렬은 보상 함수에 크게 의존합니다. 현재 최첨단 접근법들은 주로 비전 - 언어 모델 (VLM) 을 보상 백본으로 활용합니다. VLM 은 강력한 판별 능력을 제공하지만, 확산 모델 정렬에 있어 두 가지 치명적인 한계를 지닙니다:

1. 계산 및 메모리 오버헤드: VLM 기반 보상은 최적화 과정에서 반복적으로 쿼리하는 데 비용이 많이 들어 훈련 효율성에 병목 현상을 초래합니다.
2. 도메인 불일치: VLM 은 일반적으로 픽셀 공간에서 작동하는 반면, 잠재 확산 생성기는 잠재 공간에서 훈련되고 최적화됩니다. 이 픽셀 - 잠재 불일치는 정렬, 특히 보상 - 기울기 (reward-gradient) 방법론을 복잡하게 만들며 불필요한 시스템 오버헤드를 야기합니다.

반면, 확산 모델이 판별 백본으로서 가능성을 보여주고 있지만, 이를 보상으로 활용하려는 이전 시도들은 종종 특정 단계별 최적화 알고리즘에 종속되거나 깨끗한 샘플을 채점하기 위한 VLM 과 비교 가능한 범용 프레임워크를 확립하지 못한 채 노이즈가 있는 중간 상태에 집중하는 경향이 있었습니다.

2. 방법론: DiNa-LRM

저자들은 잠재 공간 내의 노이즈가 있는 확산 상태에 직접 선호도 학습을 공식화하는 DiNa-LRM(Diffusion-Native Latent Reward Model, 확산 기반 잠재 보상 모델) 프레임워크를 제안합니다. 이 방법은 세 가지 핵심 구성 요소로 이루어집니다:

A. 확산 기반 선호도 공식화

DiNa-LRM 은 깨끗한 샘플 ($x_0$) 을 직접 평가하는 대신, Thurstone 모델을 노이즈가 있는 확산 상태 ($x_t$) 로 확장합니다.

• 노이즈 보정 불확실성: 저자들은 노이즈 수준이 증가함에 따라 노이즈가 있는 상태에 대한 선호도 판단이 본질적으로 덜 확실해짐을 인식합니다. 그들은 확산 노이즈 수준 $\sigma^2(t)$에 비례하여 조정되는 노이즈 보정 비교 불확실성 $\sigma_u^2(t)$를 도입합니다:
  $$\sigma_u^2(t) = k\sigma^2(t) + \sigma_u^2$$
  이 공식화는 확산 사전 훈련 중 접하게 되는 노이즈 매니폴드와 일관된 보상 목적 함수를 보장합니다. 이는 불확실성 인지 정규화제 역할을 하여, 고노이즈 영역에서 훈련을 불안정하게 만드는 고분산 기울기를 방지합니다.
• 훈련 목적 함수: 모델은 시간 단계 의존적 불확실성에 의해 조절되는 비교 확률을 갖는 선호도 쌍 $(x_t^+, x_t^-)$에 대한 **신뢰도 손실 (fidelity loss)**을 사용하여 최적화됩니다.

B. 잠재 공간 보상 아키텍처

DiNa-LRM 은 경량화된 쿼리 기반 채점 헤드가 탑재된 사전 훈련된 잠재 확산 백본(예: SD3.5-Medium) 위에 구현됩니다.

• 다중 계층 특징 추출: 시각 및 텍스트 특징은 백본의 여러 계층에서 추출됩니다.
• 시간 단계 조건부 적응: 특징은 FiLM(Feature-wise Linear Modulation)을 사용하여 시간 단계 임베딩에 의해 조절되어, 보상 헤드가 노이즈 수준을 명시적으로 인지하도록 합니다.
• 게이트형 Q-Former 채점: 게이트형 Q-Former 는 이러한 다중 계층 및 시간 단계 조건부 특징을 집계하여 스칼라 보상 점수를 생성합니다. 이 아키텍처는 픽셀 공간 디코딩을 피하고 VAE 잠재 공간 전체에서 작동합니다.

C. 노이즈 앙상블을 통한 추론 시간 확장

강건성을 향상시키고 특정 평가 시간 단계에 대한 민감도를 줄이기 위해, DiNa-LRM 은 노이즈 앙상블을 지원합니다.

• 단일 고정 노이즈 수준에서 깨끗한 샘플을 채점하는 대신, 모델은 여러 시간 단계 $\{t_k\}$에서 샘플을 평가합니다.
• 이러한 서로 다른 시간 단계에서 추출된 특징들은 연결되어 동일한 Q-Former 헤드를 통해 처리되어 집계된 점수를 생성합니다. 이는 "확산 기반" 추론 시간 확장 메커니즘으로 작용합니다.

3. 주요 기여

1. 확산 기반 선호도 공식화: 이 논문은 확산 노이즈 수준에 비례하여 조정되는 노이즈 보정 불확실성을 도입함으로써 Thurstone 선호도 모델을 노이즈가 있는 확산 상태로 확장하여, 사전 훈련과 보상 학습 간의 분포적 격차를 해소합니다.
2. 추론 시간 확장: 저자들은 다중 시간 단계 특징을 집계하는 노이즈 앙상블 전략을 제안하여, 단일 점 추정치가 아닌 전역 확산 궤적을 활용하는 추론 시간 확장 메커니즘을 제공합니다.
3. 실증적 검증: 이 연구는 확산 기반 보상이 잠재 공간에서 네이티브로 작동하면서도 강력한 VLM 기반 판정자와 경쟁력 있는 성능을 달성할 수 있음을 보여줍니다.

4. 실험 결과

저자들은 여러 벤치마크 및 정렬 시나리오에서 DiNa-LRM 을 평가했습니다:

• 보상 모델링 정확도:

  • DiNa-LRM 은 ImageReward, HPDv2, HPDv3, GenAI Bench 전반에 걸쳐 기존 확산 기반 보상 베이스라인 (예: LRM-SD1.5, LRM-SDXL) 을 크게 능가합니다.
  • 평균적으로 최첨단 VLM 기반 보상 (예: HPSv3, UnifiedReward) 과 경쟁력 있는 성능을 달성하며, HPDv2 및 HPDv3 와 같은 특정 데이터셋에서는 격차가 눈에 띄게 좁혀집니다.
  • 노이즈 앙상블: 노이즈 앙상블을 적용하면 평균 정확도가 추가로 향상됩니다 (예: 평균 71.49% 에서 72.48% 로). 이는 노이즈 수준 전반에 걸친 증거를 집계하는 것의 이점을 검증합니다.

• 선호도 정렬 (ReFL):

  • SD3.5-M 에서 ReFL(보상 피드백 학습) 을 사용한 후 훈련 정렬 실험에서, DiNa-LRM 은 HPSv3 에 비해 프록시 점수에서 더 빠른 수렴을 보였으며 "골든" 지표 (홀드아웃 PickScore) 가 개선되었습니다.
  • 효율성: DiNa-LRM 은 VLM 기반 보상에 비해 리소스 비용을 크게 절감합니다. 피크 VRAM 사용량은 51.4% 감소하고, 보상 계산 FLOPs 는 71.1% 감소하며, 최적화 단계 FLOPs 는 46.4% 감소합니다.

• 일반화:

  • 이 공식화는 서로 다른 확산 백본 (SDXL, FLUX.1-Dev, Z-Image-Turbo) 전반에 걸쳐 일반화되어, 이전 확산 기반 베이스라인을 일관되게 능가합니다.
  • 전이성 실험은 SD3.5-M 에서 훈련된 보상 모델이 동일한 잠재 공간 내에서 SD3.5-L 을 효과적으로 정렬할 수 있음을 보여줍니다.

5. 중요성 및 주장

이 논문은 DiNa-LRM이 선호도 정렬을 위한 픽셀 공간 VLM 보상에 비해 실용적이고 효율적인 대안을 제공한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

• 구조적 우위: 잠재 공간에서 네이티브로 작동함으로써 잠재 - 픽셀 불일치와 디코딩의 계산 오버헤드를 제거하여, 보상 - 기울기 정렬 방법과 높은 호환성을 가집니다.
• 최적화 친화적: 정렬 품질을 희생하지 않으면서 훈련 역학을 개선하고 리소스 요구 사항을 줄이는 미분 가능하고 효율적인 신호를 제공합니다.
• 확장성: 노이즈 앙상블 메커니즘은 추론 시간 성능을 확장하는 독창적이고 확산 기반의 방법을 제공합니다.

저자들은 한계를 인정하며, 보상이 특정 잠재 공간에 종속되어 재훈련 없이 백본 간 전이가 제한된다는 점과, 잠재 공간 모델링이 특정 픽셀 수준 아티팩트를 과소평가하여 장기 최적화 중 보상 해킹 행동 (예: 부당한 인간 개입 또는 스타일적 편향) 을 초래할 수 있음을 지적합니다. 그들은 향후 연구가 더 강력한 통합 백본에 집중하고 경량 픽셀 공간 정규화를 통합해야 한다고 제안합니다.