Aligning Few-Step Diffusion Models with Dense Reward Difference Learning

이 논문은 제한된 단계의 확산 모델이 특정 하위 목표와 잘 정렬되도록 하기 위해, 밀집 보상 차이를 학습하여 저분산 혼합 단계 최적화를 가능하게 하는 새로운 강화학습 프레임워크인 단계별 확산 정책 최적화 (SDPO) 를 제안합니다.

핵심

이 논문은 **"몇 걸음 만에 그림을 그리는 AI(확산 모델) 가 더 똑똑하고 아름다운 그림을 그릴 수 있도록 가르치는 새로운 방법"**에 대한 이야기입니다.

기존의 AI 는 그림을 그릴 때 50100 번이나 반복해서 수정하는 (단계가 많은) 방식이 일반적이었습니다. 하지만 최근에는 14 번만 반복해도 고화질 그림을 그리는 'few-step(소단계)' AI가 등장했습니다. 문제는 이 빠른 AI 를 특정 목적 (예: "더 예쁘게", "사람이 좋아할 만한 스타일로") 에 맞게 수정할 때, 기존의 학습 방법이 잘 먹히지 않는다는 점입니다.

저자팀은 이를 해결하기 위해 **SDPO(Stepwise Diffusion Policy Optimization)**라는 새로운 방법을 개발했습니다. 이를 일상적인 비유로 설명해 드리겠습니다.

---

🎨 비유: "빠르게 그림 그리는 화가와 교정 선생님"

1. 기존 방법의 문제점: "완성된 작품만 보고 점수 매기기"

기존의 학습 방법 (RL) 은 화가가 그림을 완성했을 때만 점수를 매겨주었습니다.

• 상황: 화가가 100 번의 붓질로 그림을 그립니다. 마지막에 "아주 잘 그렸네!"라고 점수를 줍니다.
• 문제: 하지만 '소단계' 화가는 1~4 번의 붓질로 그림을 완성합니다.
  • 붓질이 너무 적어서 중간 과정이 불분명합니다.
  • 마지막에 점수만 받으면, 화가는 "어디를 어떻게 고쳐야 할지" 전혀 알 수 없습니다. (신호 부족)
  • 그래서 화가는 엉뚱한 방향으로 그림을 그리거나, 그림이 흐릿해집니다.

2. SDPO 의 핵심 아이디어: "중간중간 피드백을 주는 '밀도 높은' 지도"

이 논문은 화가에게 그림을 그리는 과정 전체에 걸쳐 세밀하게 피드백을 주는 시스템을 만들었습니다.

① '예상된 완성도'를 실시간으로 확인하기 (Dual-State Sampling)

• 비유: 화가가 붓질을 할 때마다, "지금 이 붓질로 완성하면 어떨까?"라고 가상의 완성된 그림을 미리 그려봅니다.
• 효과: 실제 그림이 완성되지 않았더라도, AI 는 "지금 단계에서 이 정도면 점수가 얼마나 나올까?"를 예측할 수 있게 됩니다. 마치 그림을 그리는 도중에도 "지금 이 부분은 색감이 좋네, 저 부분은 어둡네"라고 중간중간 점수를 매길 수 있게 되는 것입니다.

② "모든 단계에 점수를 다 매기지 않아도 돼" (잠재적 유사성 기반 예측)

• 문제: 매번 붓질할 때마다 전문가 (점수판) 를 불러와 점수를 매기면 비용이 너무 많이 듭니다.
• 해결: SDPO 는 **가장 중요한 3 곳 (시작, 중간, 끝)**에서만 점수를 매깁니다. 그리고 나머지 단계들은 "이 단계는 시작과 비슷하니까 점수도 비슷할 거야"라고 유사도를 계산해 자동으로 점수를 추정합니다.
• 효과: 비용은 줄이면서, 화가는 마치 모든 단계에 점수를 받은 것처럼 세밀한 피드백을 받게 됩니다.

③ "과거의 실수를 반복하지 않기" (단계별 뒤섞기 학습)

• 비유: 화가가 그림을 그릴 때, 항상 "1 번 붓질 -> 2 번 붓질 -> 3 번 붓질" 순서대로만 배우면, 1 번 붓질만 잘하고 3 번은 못 할 수도 있습니다.
• 해결: SDPO 는 학습할 때 붓질 순서를 무작위로 섞어서 가르칩니다. (예: 3 번 -> 1 번 -> 2 번)
• 효과: 화가는 어떤 순서로 그림을 그려도 모든 단계에서 똑똑하게 대응할 수 있게 됩니다.

---

🚀 왜 이 방법이 특별한가요?

1. 빠른 학습 (Few-Step): 기존 방법들은 1~4 단계처럼 단계가 적은 AI 를 가르칠 때 실패했습니다. 하지만 SDPO 는 중간 과정의 '예상 완성도'를 활용해서, 단계가 적어도 확실하게 가르칩니다.
2. 비용 절감: 매번 전문가에게 점수를 매기게 하지 않고, 스마트하게 추정해서 계산 비용을 아끼면서도 높은 성능을 냅니다.
3. 안정성: 기존 방법들은 학습 중 그림이 망가지거나 (불안정), 특정 스타일만 따라 하는 (과적합) 문제가 있었지만, SDPO 는 단계별 뒤섞기와 정교한 점수 예측으로 이를 해결했습니다.

💡 결론

이 논문은 **"빠르게 그림을 그리는 AI 가, 완성된 결과물뿐만 아니라 그리는 과정 하나하나에서도 올바른 방향으로 나아가도록, 중간중간 세밀하고 똑똑한 피드백을 주는 새로운 지도법"**을 제안한 것입니다.

이 방법을 사용하면 AI 는 적은 단계로도 훨씬 더 아름답고, 우리가 원하는 스타일의 그림을 정확하게 그려낼 수 있게 됩니다. 마치 초고속으로 그림을 그리는 천재 화가가, 교정 선생님의 세심한 지도를 받아 더욱 완벽해진 것과 같습니다.

기술 요약

1. 문제 제기 (Problem)

최근 텍스트-to-이미지 생성을 위한 표준 확산 모델 (Diffusion Models) 을 소수의 단계 (Few-Step, 예: 1~4 단계) 로 압축하는 연구가 활발합니다. 이는 고해상도 이미지 생성 속도를 획기적으로 높여줍니다. 그러나 이러한 Few-Step Diffusion Models은 기존 강화학습 (RL) 기반 정렬 방법론을 적용할 때 다음과 같은 심각한 한계에 직면합니다.

• 희소 보상 (Sparse Reward) 의 한계: 기존 RL 방법 (예: DDPO) 은 주로 전체 경로 (Trajectory) 의 최종 결과물에만 보상을 부여합니다. Few-Step 환경에서는 상태 공간 (State Space) 이 매우 제한적이고 샘플 품질이 낮아, 최종 보상만으로는 효과적인 정책 업데이트를 위한 충분한 신호 다양성을 확보하기 어렵습니다.
• 훈련 불안정성: Few-Step 경로를 그대로 RL 에 적용하면 훈련이 불안정해지고 수렴이 어렵습니다.
• 혼합 단계 (Mixed-Step) 최적화의 문제: 다양한 길이의 경로를 섞어 학습하면 상태 공간은 넓어지지만, 서로 다른 역학 (Dynamics) 으로 인해 분산 (Variance) 이 커져 정책 최적화가 불안정해집니다. 또한 중간 단계에 대한 밀집 보상 (Dense Reward) 이 없어 중간 예측 품질을 개선하기 어렵습니다.

2. 제안 방법론: SDPO (Stepwise Diffusion Policy Optimization)

저자들은 위 문제들을 해결하기 위해 SDPO라는 새로운 RL 미세 조정 (Finetuning) 프레임워크를 제안합니다. SDPO 의 핵심 구성 요소는 다음과 같습니다.

A. 이중 상태 경로 샘플링 (Dual-State Trajectory Sampling)

• 개념: 기존 확산 모델의 단일 노이즈 상태 ($x_t$) 추적 대신, 각 단계에서 **노이즈 상태 ($x_t$)**와 **예측된 깨끗한 상태 ($\hat{x}_0$)**를 동시에 추적합니다.
• 동작: Few-Step 모델의 강력한 단일 단계 제거 (Denoising) 능력을 활용하여, 각 단계마다 중간 단계에서 생성된 '예측된 깨끗한 이미지'를 실제 최종 출력의 대리자 (Surrogate) 로 사용합니다.
• 효과: 이를 통해 서로 다른 길이의 경로에서 얻은 최종 출력을 일관된 중간 '깨끗한 상태' 시퀀스로 매핑할 수 있어, **모든 단계에 대해 밀집 보상 (Dense Reward)**을 부여하면서도 경로 간 일관된 역학을 유지하고 분산을 줄일 수 있습니다.

B. 잠재 유사도 기반 밀집 보상 예측 (Latent Similarity-based Dense Reward Prediction)

• 문제: 모든 단계에서 보상 함수 (Reward Function) 를 호출하는 것은 계산 비용이 매우 큽니다.
• 해결: 보상 호출을 최소화하기 위해 3 번의 선택적 쿼리 (시작, 끝, 그리고 적응적으로 선택된 '앵커' 단계) 만 수행합니다.
• 기법:
  1. 적응형 앵커 선택: 시작과 끝 단계의 잠재 표현 (Latent Representation) 과의 코사인 유사도를 계산하여, 두 극단과 가장 다른 (정보량이 많은) 중간 단계를 '앵커'로 선택합니다.
  2. 유사도 가중 보간: Lipschitz 연속성 (Lipschitz Continuity) 가정을 바탕으로, 쿼리된 3 지점의 보상을 잠재 공간 유사도를 가중치로 하여 나머지 단계의 보상을 예측합니다.
• **효과:**昂贵的인 보상 호출 횟수를 획기적으로 줄이면서도 밀집 보상의 이점을 유지합니다.

C. 밀집 보상 차이 학습 (Dense Reward Difference Learning)

• 목표: 전체 경로의 누적 보상이 아닌, 단계별 (Stepwise) 보상 차이를 최적화합니다.
• 수식: 두 개의 독립적인 이중 상태 경로에 대해, 각 단계 $t$에서의 로그 가능도 비율 (Log-likelihood ratio) 의 차이와 예측된 밀집 보상의 차이를 매칭하는 MSE 손실 함수를 정의합니다.
• 장점: 경로 수준의 업데이트보다 더 빈번하고 세밀한 정책 업데이트가 가능해집니다.

D. 추가 최적화 기법

• 단계별 이점 차이 학습 (Stepwise Advantage Difference Learning): 시간적 의존성을 고려하여 미래 단계의 할인된 반환 (Discounted Return) 을 현재 단계의 이점 (Advantage) 추정에 반영합니다.
• 시간적 중요도 가중 (Temporal Importance Weighting): Few-Step 환경에서 초기 단계 (Low-step) 의 최적화를 우선시하기 위해, 단계가 앞설수록 가중치를 높이는 지수적 감쇠 가중치를 적용합니다.
• 단계 셔플링 그라디언트 업데이트 (Step-Shuffled Gradient Updates): 고정된 단계 순서로 인한 과적합을 방지하기 위해 미니배치 내 단계 순서를 무작위로 섞어 그라디언트를 업데이트합니다.

3. 주요 기여 (Key Contributions)

1. 이중 상태 샘플링 메커니즘: Few-Step 확산 모델을 위해 설계된 새로운 샘플링 방식으로, 밀집 보상 피드백과 낮은 분산의 혼합 단계 최적화를 가능하게 합니다.
2. 잠재 유사도 기반 보상 예측 전략: 비용이 많이 드는 보상 쿼리를 최소화하면서도 효율적인 밀집 보상을 제공하는 새로운 전략을 개발했습니다.
3. 밀집 보상 차이 학습 목적 함수: 경로 전체가 아닌 개별 단계에서의 보상 차이를 최적화하여 Few-Step 모델의 정밀한 학습을 가능하게 합니다.
4. 통합 SDPO 프레임워크: 단계별 이점 학습, 시간적 가중치, 단계 셔플링 등을 통합하여 극단적으로 낮은 단계 (1~4 단계) 에서도 견고하고 효율적인 최적화를 달성합니다.

4. 실험 결과 (Results)

• 비교 대상: DDPO, PRDP/REBEL, D3PO 등 기존 RL 및 선호도 기반 정렬 방법과 비교했습니다.
• 성능:
  • 샘플 효율성: SDPO 는 적은 학습 샘플로 더 높은 보상 점수 (Aesthetic Score, ImageReward 등) 를 달성했습니다. 특히 1~4 단계의 극단적 Few-Step 설정에서 기존 방법들이 훈련 불안정성이나 성능 저하를 보인 반면, SDPO 는 일관된 성능 향상을 보였습니다.
  • 일반화 능력: 학습에 사용되지 않은 새로운 프롬프트 (Unseen Prompts) 에 대해서도 우수한 생성 품질을 유지했습니다.
  • 확장성: 잠재 일관성 모델 (Latent Consistency Models) 과 텍스트-to-멀티뷰 (Multiview) 생성 모델에서도 SDPO 가 기존 방법 (DDPO, REBEL) 을 능가하는 성능을 입증했습니다.
• 정성적 결과: DDPO 는 Few-Step 설정에서 흐릿하거나 품질이 낮은 이미지를 생성하는 반면, SDPO 는 선명하고 보상과 정렬된 고품질 이미지를 생성했습니다.

5. 의의 및 결론 (Significance)

이 논문은 Few-Step Diffusion Models의 실용적 적용을 위한 핵심 병목 현상이었던 '보상 정렬의 어려움'을 해결했습니다. 기존 RL 방법론이 가진 희소 보상과 높은 분산의 문제를 이중 상태 샘플링과 밀집 보상 예측을 통해 우아하게 해결함으로써, 초고속 이미지 생성 모델이 사용자 선호도나 미적 기준과 완벽하게 정렬될 수 있는 길을 열었습니다. 이는 실시간 이미지 생성 애플리케이션 및 효율적인 AI 에이전트 개발에 중요한 기여를 할 것으로 기대됩니다.