기술 요약: 확산 보상 (Diffused Reward) 을 활용한 Diff-Instruct (Didr)

문제 제기

최근 한 단계 텍스트-이미지 (T2I) 생성기의 발전은 잠재 노이즈를 단일 순전파로 직접 이미지로 매핑함으로써 실시간 합성을 가능하게 했습니다. 그러나 강화 학습 (RL) 을 통해 이러한 생성기를 인간의 선호도와 정렬하는 것은 근본적인 구조적 과제를 제시합니다.

이전 한 단계 RLHF 방법들 (예: Diff-Instruct*, Diff-Instruct++) 은 이미지 공간의 보상 최적화와 잠재 공간의 분포 매칭을 결합합니다. 이러한 접근법들은 일반적으로 선호도 신호 (보상) 를 깨끗한 이미지 종점 ($x_0$) 에서만 적용하는 반면, 전체 노이즈 확산 궤적에 걸쳐 참조 분포에 대한 Kullback-Leibler (KL) 정규화를 부과합니다. 저자들은 이 패러다임에서 결정적인 구조적 불일치를 식별합니다: 보상은 종단 상태에서만 작용하는 반면, 정규화자는 모든 노이즈 레벨에 걸쳐 적용된다는 점입니다.

이러한 불일치는 **종단 보상 지배 (terminal reward domination)**라고 명명된 실패 모드로 이어집니다. 전방 노이즈가 자연스럽게 이미지 세부 정보를 가리기 때문에, 고노이즈 레벨에서 모드 발산에 대한 KL 정규화 패널티는 현저히 약화됩니다. 결과적으로 최적화 알고리즘은 이러한 궤적 취약점을 이용하여 참조 분포를 포기하고 충분한 KL 패널티를 부과받지 않은 채 고보상 영역으로 붕괴시킵니다. 이는 보상 점수의 향상으로 이어지지만, 종종 이미지 충실도와 분포적 안정성을 희생합니다.

방법론: 확산 보상 (Diffused Reward) 을 활용한 Diff-Instruct (Didr)

이러한 불일치를 해결하기 위해 저자들은 적분 KL (IKL) 최소화에서 파생된 데이터 프리, 궤적 레벨 정렬 프레임워크인 Didr을 제안합니다. 핵심 통찰은 RLHF 최적의 보상 기울어진 깨끗한 이미지 분포를 확산 궤적의 모든 노이즈 레벨로 체계적으로 전파하는 것입니다.

1. 원칙적인 궤적 목표

이 방법은 KL 정규화 RLHF 하에서 깨끗한 이미지에 대한 올바른 정렬 목표를 설정하는 것으로 시작합니다. 최적의 깨끗한 이미지 분포 $q^*(x_0|c)$는 보상 기울어진 밀도입니다:
$$q^*(x_0|c) \propto q_0(x_0|c) \exp(r(x_0, c)/\tau)$$
여기서 $q_0$는 참조 분포, $r$은 보상 함수, $\tau$는 온도 매개변수입니다.

보상을 $x_0$에서만 적용하는 대신, Didr 은 이 최적 목표를 참조 전방 과정을 통해 $q^*$를 확산시킴으로써 전체 궤적으로 "들어올립니다". 이는 모든 노이즈 레벨에서 정의된 보상 기울어진 주변 분포 (reward-tilted marginals) $\{q^*_t\}_{t \in [0, T]}$의 집합을 산출합니다:
$$q^*_t(x_t|c) = \int q_t(x_t|x_0) q^*(x_0|c) dx_0$$
저자들은 생성기의 주변 분포와 이러한 보상 기울어진 주변 분포 간의 적분 KL 발산을 최소화하는 것이 표준 KL 정규화 RLHF 목적 함수와 동일한 최소화자를 공유함을 증명합니다. 이는 모든 노이즈 레벨에서 보상과 정규화가 균형을 이루도록 하여 종단 보상 지배를 방지합니다.

2. 확산 보상 점수 (Diffused Reward Score, DRS)

IKL 목적 함수를 미분함으로써 저자들은 목표 점수 함수를 유도합니다. 최적 점수는 두 가지 구성 요소로 분해됩니다:
$$s_{target} = s_{ref} + s_r$$
여기서 $s_{ref}$는 참조 점수 함수이고, $s_r$은 **확산 보상 점수 (Diffused Reward Score, DRS)**입니다. DRS 는 모든 노이즈 레벨에서 정의된 보상 주도 수정 항으로 작용합니다:
$$s_r(x_t, t, c) = \nabla_{x_t} \log \mathbb{E}_{x_0 \sim q(x_0|x_t, c)} \left[ \exp\left(\frac{r(x_0, c)}{\tau}\right) \right]$$
분류기 가이드 (1 차 근사) 와 달리, DRS 는 전체 로그 지수 기울기를 유지하며 노이즈 관측치로부터 깨끗한 이미지를 식별할 수 있는 곳에만 선호도 정보를 전파합니다.

3. 확산 보상 프록시 (Diffused Reward Proxy, DRP)

DRS 는 처리 불가능한 사후 기대값 $q(x_0|x_t, c)$를 필요로 하므로, 저자들은 이 방법을 실용적으로 만들기 위해 **확산 보상 프록시 (DRP)**를 도입합니다. DRP 는 고정된 참조 모델의 미분 가능한 짧은 단계 탈노이즈 체인을 사용하여 사후 기대값을 근사합니다.

• 메커니즘: 노이즈 상태 $x_t$에서, 이 방법은 고정된 참조 모델을 사용하여 $K$개의 독립적인 $S$-단계 탈노이즈 체인을 실행하여 근사적인 깨끗한 샘플 $\{\hat{x}^{(k)}_0\}$을 생성합니다.
• 기울기 추정: DRS 는 경로별 보상 기울기의 소프트맥스 가중 평균으로 근사됩니다:
  $$\hat{s}_r(x_t, t, c) = \frac{1}{\tau} \sum_{k=1}^K \omega^{(k)} \nabla_{x_t} r(\hat{x}^{(k)}_0, c)$$
  여기서 가중치 $\omega^{(k)}$는 $\exp(r(\hat{x}^{(k)}_0, c)/\tau)$에 비례합니다.
• 효율성: 이 접근법은 참조 모델의 활성화 상태를 저장하거나 이미지 학습 데이터가 필요 없이 보상 기울기가 궤적을 통해 안정적으로 역전파되도록 합니다.

4. 학습 프레임워크

학습 과정은 두 단계 사이를 오가며 진행됩니다:

1. 조교 (Teaching Assistant, TA) 업데이트: TA 점수 모델 ($s_\psi$) 은 탈노이즈 점수 매칭 (DSM) 을 통해 훈련되어 생성기의 진화하는 주변 분포를 추적하고, 암시적 생성기 점수를 대리합니다.
2. 생성기 업데이트: 생성기 ($g_\theta$) 는 CFG 보정된 참조 점수와 DRP 추정의 합에 대한 TA 점수를 비교하는 IKL 목적 함수의 기울기를 사용하여 업데이트됩니다.

주요 기여

1. 종단 보상 지배의 식별: 이 논문은 이전 한 단계 RLHF 방법들의 구조적 실패 모드를 공식적으로 식별하고 분석합니다. 여기서 최적화 알고리즘은 노이즈 레벨을 이용하여 고보상 영역으로 붕괴시켜 의도된 RLHF 균형을 위반합니다.
2. 원칙적인 궤적 레벨 정렬: 저자들은 RLHF 최적의 깨끗한 이미지 밀도를 전방으로 전파함으로써 수학적으로 건전한 궤적 레벨 목표를 도출하는 프레임워크인 Didr 을 제안합니다. 이는 모든 노이즈 레벨에서 보상과 정규화가 균형을 이루도록 보장합니다.
3. DRS 의 분석적 유도: 이 논문은 확산 보상 점수를 분석적으로 유도하여 궤적 목표를 참조 점수와 보상 유도 수정 항으로 정확히 분해함을 보여줍니다.
4. 실용적 추정기 (DRP): 저자들은 이미지 데이터 없이 안정적인 훈련을 가능하게 하는 미분 가능한 짧은 단계 탈노이즈를 통해 DRS 를 계산하는 제로 데이터 추정기인 확산 보상 프록시를 도입합니다.

경험적 결과

광범위한 실험은 Didr 이 기존 한 단계 베이스라인보다 일관되게 우수함을 보여줍니다:

• 파레토 우세: SDXL 백본 (1024×1024) 에서 표준 Didr 변형은 PickScore–FID 트레이드오프에서 기존 모든 한 단계 정렬 베이스라인을 파레토 우세하며, FLUX-dev 및 SDXL-DPO 와 같은 다단계 모델보다 두 축 모두에서 동시에 뛰어납니다.
• 선호도 지표: 확장된 훈련 변형 (Didr_longer) 은 평가된 한 단계 모델 중 가장 높은 PickScore(23.9) 를 달성했으며, 가장 강력한 이전 베이스라인 (Diff-Instruct*) 대비 ImageReward 를 8.9% 향상시켰습니다. 또한 HPSv2.1 및 미적 점수 (Aesthetic Score) 에서도 선도합니다.
• 일반화: Didr 은 60 억 파라미터 Z-Image 백본으로 성공적으로 이전됩니다. 단일 생성 단계에서 선호도 지표에 대해 50 단계 교사 모델을 능가하여 아키텍처 전반에 걸친 확장성을 입증합니다.
• 정성적 개선: 시각적 비교는 Didr 이 이전 한 단계 방법들보다 더 선명한 세부 사항과 자연스러운 색상을 생성하며, 보상 최적화 베이스라인에서 흔히 보이는 과포화 및 화가 같은 아티팩트를 피함을 보여줍니다.

중요성 및 주장

이 논문은 원칙적인 궤적 레벨 보상 전파가 종단 전용 보상 방법의 고유한 불일치를 해결하는 효과적인 정렬 전략임을 주장합니다. RLHF 최적 분포를 확산 과정을 통해 엄격하게 전파함으로써 Didr 은 한 단계 생성기 간 선호도 정렬과 이미지 충실도 사이의 최상의 트레이드오프를 달성합니다.

저자들은 그들의 방법이 이미지 학습 데이터가 필요하지 않으며, 텍스트 프롬프트와 고정된 참조 모델에만 의존한다고 강조합니다. 그들은 Didr 을 실시간 추론 속도를 유지하면서 다단계 교사 모델의 선호도 정렬을 일치시키거나 능가할 수 있게 하는 한 단계 생성기를 위한 솔루션으로 위치시킵니다. 이 연구는 이전 방법에서 식별된 구조적 불일치가 종단 상태뿐만 아니라 전체 궤적을 정렬함으로써만 해결될 수 있는 근본적인 한계임을 시사합니다.