Few‑Step 확산 모델을 위한 밀집 보상 차이 학습 기반 정책 최적화

초록

본 논문은 1~4 단계의 초저스텝 확산 모델이 원하는 이미지 품질·선호도와 정렬되지 않는 문제를 해결하고자, 노이즈 상태와 예측된 클린 상태를 동시에 추적하는 듀얼‑스테이트 샘플링과 잠재 유사도 기반 밀집 보상 예측을 도입한다. 이를 바탕으로 단계별 보상 차이를 최소화하는 Dense Reward Difference Learning을 설계하고, 단계별 어드밴티지, 시간 가중치, 스텝 셔플링 업데이트 등을 결합한 Stepwise Diffusion Policy Optimization(SDPO) 프레임워크를 제안한다. 실험 결과, SDPO는 기존 DDPO 대비 극히 낮은 스텝에서도 높은 품질·보상 정렬 이미지를 지속적으로 생성한다.

상세 분석

SDPO는 기존 확산 모델을 MDP로 보는 강화학습 접근을 근본적으로 재구성한다. 첫 번째 핵심은 ‘듀얼‑스테이트’ 샘플링이다. 일반적인 확산에서는 매 스텝마다 노이즈 상태 xₜ만을 관찰하지만, Few‑Step 모델은 단일 스텝에서 강력한 디노이징 능력을 갖고 있어, 해당 스텝의 예측 클린 이미지 \hat{x}ₜ⁰를 동시에 추출한다. 이렇게 두 상태를 동시에 저장하면, 각 스텝마다 \hat{x}ₜ⁰에 직접 보상 함수 R을 적용해 ‘밀집 보상’ r(sₜ,aₜ)=R(\hat{x}ₜ⁰,c)를 얻을 수 있다. 결과적으로 스텝 수와 무관하게 일정한 보상 신호가 제공되어, 기존의 최종‑스텝만 보상을 주는 희소 보상 방식에서 발생하던 과적합과 신호 부족 문제를 해소한다.

두 번째 기여는 ‘잠재 유사도 기반 보상 예측’이다. 모든 스텝에서 실제 보상을 쿼리하면 비용이 크게 증가한다. 저자들은 잠재 공간에서 \hat{x}ₜ⁰와 앵커 스텝의 \hat{x}_{anchor}⁰ 사이의 코사인 유사도를 이용해 보상을 선형 보간한다. 이때 Lipschitz 연속성을 가정해 보상 변화가 부드럽다고 전제함으로써, 적은 수의 실제 쿼리만으로도 고품질의 밀집 보상 시퀀스를 재구성한다.

세 번째 핵심은 ‘밀집 보상 차이 학습(Dense Reward Difference Learning)’이다. 기존 DPO‑계열 방법은 전체 트래젝터리의 로그우도 비율 차이를 보상 차이와 매핑했지만, 단계별 차이를 활용하지 못했다. SDPO는 각 스텝 t에 대해 \Delta\hat{A}_t = \hat{R}t - \hat{R}{t-1}와 로그우도 비율 차이 \Delta\tilde{\rho}_t를 MSE 손실로 매칭한다. 이렇게 하면 정책 업데이트가 매 스텝마다 이루어져 학습 효율이 크게 향상된다.

추가적인 안정화 기법으로는 (1) ‘단계별 어드밴티지 추정’—시간 할인 없이 순수 차이를 사용해 장기 의존성을 보존, (2) ‘시간 가중치’—초기 스텝에 더 큰 가중치를 부여해 Few‑Step 상황에서 중요한 초기 디노이징을 강조, (3) ‘스텝 셔플링 그라디언트 업데이트’—각 배치 내 스텝 순서를 무작위로 섞어 상관관계에 의한 편향을 감소시킨다.

실험에서는 SDPO를 기존 DDPO와 동일한 데이터·프롬프트 조건에서 1,2,4,8 스텝 모델에 적용하였다. 평가 지표는 PickScore 기반 보상, CLIP‑Score, FID 등이다. 모든 스텝에서 SDPO는 DDPO 대비 평균 12%~25% 높은 보상 점수를 기록했고, 특히 1‑2 스텝에서는 이미지 흐림 현상이 현저히 감소했다. 시각적 예시에서도 SDPO가 복잡한 프롬프트(예: “cyberpunk cat wearing black leather jacket”)를 정확히 반영한 고해상도 이미지를 생성함을 확인할 수 있다.

요약하면, SDPO는 (1) 듀얼‑스테이트를 통한 밀집 보상 제공, (2) 효율적인 보상 예측, (3) 단계별 차이 학습, (4) 다양한 안정화 기법을 결합해 Few‑Step 확산 모델의 RL 파인튜닝을 실용적인 수준으로 끌어올렸다. 이는 고해상도 이미지 생성에서 비용·시간을 크게 절감하면서도 사용자 정의 목표와 정렬된 결과를 얻고자 하는 실무·연구 현장에 큰 파급 효과를 기대한다.