기술 요약: 확산 모델 정렬을 위한 명시적 비평가 안내

문제 제기

비미분 가능 목적 (예: 인간 선호도, 구성적 정확성, 텍스트 렌더링) 과 정렬되도록 확산 모델을 학습 후 조정하는 것은 사용자 중심 속성에서 사전 학습만으로는 종종 최적하지 않은 결과가 나오기 때문에 점점 더 중요해지고 있습니다. 이 분야의 핵심 과제는 **신용 할당 (credit assignment)**입니다. 확산 모델은 다단계 탈노이즈 궤적을 통해 샘플을 생성하지만, 보상은 일반적으로 최종 디코딩된 이미지만에 정의됩니다. 기존 솔루션은 다음과 같은 중대한 한계에 직면해 있습니다:

• 그룹 상대적 방법 (예: GRPO): 이들은 그룹 내 샘플을 비교함으로써 명시적 비평가를 우회합니다. 가치 함수 학습을 피할 수는 있지만, 결과적인 보상 신호는 거칠고 궤적 수준이며, 중간 탈노이즈 결정에 대한 세밀한 신용 할당이 부족합니다.
• 픽셀 공간 비평가: 기존 명시적 액터-비평가 방법은 종종 픽셀 공간에서 비평가를 정의하며, 중간 노이즈 잠재 변수로부터 재구성된 '클린' 이미지를 평가합니다. 이는 상태 불일치를 초래합니다. 정책은 노이즈 잠재 상태에서 작동하는 반면, 비평가는 표준 시각 인코더 (예: CLIP, BLIP) 에 대해 분포 밖 (out-of-distribution) 인 경우가 많은 프록시 상태를 평가합니다. 이러한 불일치는 취약한 가치 추정, 최적화 불안정성, 그리고 상당한 디코딩 오버헤드로 이어집니다.
• 보안 해킹 (Reward Hacking): 특정 검증 가능한 보상을 최적화하는 것은 종종 목표 지표를 극대화하면서 이미지 품질의 측정되지 않은 측면을 저하시키는 단축 솔루션 (예: 단순화된 배경, 반복적 패턴) 으로 이어집니다.

방법론

저자들은 확산 생성 과정에 비평가를 직접 통합함으로써 이러한 문제들을 해결하는 상태 정렬 잠재 액터-비평가 프레임워크를 제안합니다.

1. 상태 정렬 잠재 비평가

픽셀 공간에서 작동하는 대신, 이 프레임워크는 정책이 방문하는 노이즈 잠재 상태 $z_t$에 직접 가치 함수 $V_\phi(z_t, t, y)$를 학습합니다.

• 아키텍처: 비평가는 사전 학습된 확산 백본 (UNet 또는 DiT) 에서 초기화되며, 스칼라 값을 예측하기 위해 경량 아키텍처 변경이 가해집니다.
  • UNet: 특징 맵이 패치화되고, 학습 가능한 요약 토큰이 어텐션 헤드와 MLP 를 통해 정보를 집계합니다.
  • DiT: 학습 가능한 쿼리 토큰을 사용하여 중간 레이어 (예: 레이어 7, 15, 23) 에 경량 비평가 브랜치가 삽입되어 시각 정보를 집계합니다.
• 상태 정렬: 동일한 잠재 공간에서 작동하고 생성적 사전 지식을 공유함으로써, 비평가는 자연스럽게 타임스텝과 노이즈 수준을 인지하여 픽셀 공간 비평가의 분포 이동 문제를 피합니다.

2. 안정화 학습 기법

안정적인 액터-비평가 최적화를 보장하기 위해 두 가지 구체적인 기법이 도입됩니다:

• 타임스텝 인식 조건부: 가치 헤드는 타임스텝 $t$에 명시적으로 조건부 (AdaLN 스타일 변조를 통해) 가 되어, 궤적 전반에 걸친 가치 예측을 개선하기 위해 노이즈 수준에 직접 접근할 수 있게 합니다.
• 가치 사전 학습: 공동 액터-비평가 최적화 전에, 몬테카를로 반환 목표에 대한 회귀 손실을 사용하여 비평가만 업데이트하는 짧은 워밍업 단계가 수행됩니다. 이는 더 나은 초기 가치 추정을 제공하여 후속 PPO 업데이트를 안정화합니다.

3. 추론 시간 조종

비평가는 정책과 동일한 잠재 상태에 정의되므로, 추론 시간 안내 모듈로 재사용될 수 있습니다. 가치 함수의 기울기 $\nabla_{z_t} V_\phi$는 더 높은 최종 보상을 낼 것으로 예측되는 잠재 영역을 가리킵니다. 이 기울기는 샘플러 업데이트에 주입되어 (DDIM 의 노이즈 예측 또는 플로우 매칭의 속도 수정), 추가 학습 없이 탈노이즈 궤적을 효과적으로 조종합니다.

4. 다중 보상 최적화

보안 해킹을 완화하기 위해 프레임워크는 다중 보상 최적화로 확장됩니다. 공유 잠재 백본은 각 보상 유형별로 별도의 경량 가치 헤드를 사용합니다. 이 설계는 이질적인 보상 (예: 구성적 vs 미적) 이 비싼 특징 추출을 공유하면서 서로를 정규화할 수 있게 하여, 단일 보상 학습에서 흔히 발생하는 단축 솔루션으로의 붕괴를 방지합니다.

주요 기여

1. 상태 정렬 잠재 비평가: 사전 학습된 확산 모델을 타임스텝 조건부 가치 네트워크로 변환하는 프레임워크로, 노이즈 잠재 변수에 대한 가치 예측 및 추론 시간 조종을 가능하게 합니다.
2. 안정화 전략: 명시적 액터-비평가 확산 학습을 안정화하는 데 필수적인 기법으로 타임스텝 인식 조건부와 가치 사전 학습을 규명했습니다.
3. 다중 보상 확장: 공유 백본에 별도의 가치 헤드를 사용하여 상보적 보상을 공동으로 최적화하는 메커니즘으로, 프록시 과최적화와 보안 해킹을 완화합니다.
4. 실증적 우수성: UNet 및 DiT 백본에 걸쳐 단일 보상 및 다중 보상 벤치마크에서 그룹 상대적 (GRPO) 및 기존 액터-비평가 베이스라인보다 일관되게 우수한 성능을 발휘했습니다.

실험 결과

이 방법은 다양한 보상 유형에 걸쳐 Stable Diffusion 1.5 (UNet) 및 **Stable Diffusion 3.5M (DiT)**에서 평가되었습니다:

• 단일 보상 벤치마크: 이 방법은 CLIP, HPSv2.1, PickScore, GenEval, OCR 지표에서 최첨단 결과를 달성했습니다. 특히, 희소하고 이산적인 신호가 타임스텝별 가치 추정으로부터 가장 큰 혜택을 받는 검증 가능한 보상 (GenEval, OCR) 에서 유의미한 향상을 보였습니다. 또한 SD3.5M 에서 DDPO 가 수렴하지 못했던 상황에서 안정성을 입증했습니다.
• 다중 보상 벤치마크: 공동 최적화 작업 (예: CLIP + HPS + GenEval) 에서 이 방법은 GRPO 를 능가했으며, 특히 GenEval 에서 선호도 지표에서는 동등한 성능을 유지하면서 더 나은 결과를 보였습니다.
• 추론 시간 조종: 학습된 비평가를 샘플링 중 안내 신호로 적용하면 생성 품질이 추가로 향상되었습니다. SD1.5 에서는 비평가 안내만으로도 Best-of-N (BoN) 샘플링을 능가했으며, SD3.5M 에서는 안내와 BoN 을 결합했을 때 가장 좋은 결과를 얻었습니다.
• 효율성: 잠재 비평가는 가치 계산 (10 초 vs 2835 초) 및 학습 단계 (20 초 vs 5286 초) 모두에서 픽셀 공간 비평가 (CLIP/BLIP/HPS) 보다 훨씬 빨랐습니다.
• 애블레이션:
  • 가치 사전 학습을 제거하면 희소 보상에서 특히 성능이 저하되었습니다.
  • 다중 보상 작업용 별도 가치 헤드는 공유 헤드가 설계된 것보다 우수한 성능을 보였습니다.
  • Classifier-Free Guidance (CFG) 없이 학습하면 가장 안정적인 최적화와 가장 높은 최종 보상을 얻었으며, 높은 CFG 스케일과 관련된 불안정성과 아티팩트를 피할 수 있었습니다.

중요성 및 주장

이 논문은 현대 확산 모델에 강화 학습을 적용하기 위한 실용적인 레시피를 확립했다고 주장합니다. 그 중요성은 다음과 같습니다:

1. 상태 불일치 해결: 비평가를 정책의 잠재 상태 공간에 정렬함으로써 픽셀 공간 평가의 취약성과 오버헤드를 제거합니다.
2. 안정적인 액터-비평가 학습 가능: 잠재 정렬, 타임스텝 조건부, 가치 사전 학습의 조합은 확산 모델에 대한 명시적 액터-비평가 학습을 실현 가능하고 안정적으로 만들어, 그룹 상대적 방법보다 더 효율적인 샘플 효율 대안을 제공합니다.
3. 이중 유용성: 학습된 비평가는 학습 중 정책 최적화와 추론 시간 생성 과정 조종이라는 이중 목적을 수행하여, 학습 단계를 넘어 성능 향상을 제공합니다.
4. 보안 해킹 완화: 다중 보상 확장은 경쟁하는 목적을 균형 있게 조정하는 실용적인 메커니즘을 제공하여, 확산 RL 에서 단축 솔루션의 발생을 줄입니다.

저자들은 현재 이미지 확산 (비디오 아님) 에 초점을 맞추고 있으며, 단일 보상 학습에 비해 다중 보상 설정에서 잠재적 효율성 저하가 있을 수 있고, SDE 샘플링에 내재된 이산화 오류가 있음을 인정합니다.