Stepwise Credit Assignment for GRPO on Flow-Matching Models

1. 문제 정의 (Problem Statement)

텍스트-to-이미지 생성을 위한 흐름 매칭 (Flow-Matching) 모델에 강화 학습 (RL) 을 적용할 때, 기존 방법인 Flow-GRPO는 다음과 같은 근본적인 한계를 가지고 있습니다.

• 균일한 신용 할당 (Uniform Credit Assignment): Flow-GRPO 는 생성 과정의 모든 단계 (denoising steps) 에 최종 이미지 하나만 받은 보상 (reward) 을 동일하게 할당합니다. 즉, 최종 이미지의 품질이 좋으면 모든 단계가 '잘했다'고 간주하고, 나쁘면 모두 '나빴다'고 간주합니다.
• 시간적 구조 무시: 확산 (diffusion) 생성 과정은 시간적 계층 구조를 가집니다.
  • 초기 단계: 저주파수 정보 (구도, 레이아웃, 객체 배치) 를 결정합니다.
  • 후기 단계: 고주파수 정보 (세부 묘사, 질감) 를 정교화합니다.
  • 균일한 신용 할당은 이러한 구조를 무시하여, 초기 단계의 나쁜 결정이 나중에 수정되더라도 최종 보상이 좋으면 그 나쁜 단계를 강화시켜 버리는 문제가 발생합니다.
• 중간 상태의 보상 신호 부재: 중간 단계의 노이즈 상태 ($x_t$) 에서는 직접 보상을 평가할 수 없어, 최종 결과만 보고 학습하므로 학습 효율이 낮고 수렴 속도가 느립니다.

2. 제안 방법 (Methodology)

저자들은 Stepwise-Flow-GRPO를 제안하여 각 디노이징 단계별로 신용 (credit) 을 할당하는 방식을 도입했습니다. 주요 구성 요소는 다음과 같습니다.

가. 단계별 보상 추정 (Stepwise Rewards via Tweedie's Formula)

• 중간 노이즈 상태 $x_t$에서 직접 보상을 계산할 수 없으므로, Tweedie 공식을 활용하여 $x_t$로부터 깨끗한 이미지 $\hat{x}_0(t) = E[x_0|x_t]$를 추정합니다.
• 이 추정된 이미지 $\hat{x}_0(t)$에 대해 보상 모델 (Reward Model) 을 실행하여 각 단계 $t$에서의 보상 $r_t$를 계산합니다.
• 이를 위해 $T'$개의 ODE 서브스텝을 사용하여 고품질의 중간 이미지를 추정합니다.

나. 이득 기반 신용 할당 (Gain-Based Credit Assignment)

• 최종 보상만 사용하는 대신, 각 단계가 보상을 얼마나 개선했는지를 기준으로 학습합니다.
• 이득 (Gain): $g_t = r_{t-1} - r_t$로 정의합니다. (보통 $t$가 1 에서 0 으로 감소하므로, $r_{t-1}$이 더 최근의 상태입니다. 즉, 이전 단계 대비 보상 증가분을 의미합니다.)
• 그루프 상대적 이득 (Group-Relative Gains): 모든 트래젝토리와 모든 단계에 걸쳐 이득을 정규화하여 상대적 이득 (Advantage) $\tilde{A}_t$를 계산합니다.
  • 초기 단계의 이득이 일반적으로 크므로, 이를 정규화할 때 단계별로 나누지 않고 전체적으로 (Jointly) 정규화하여 시간적 구조를 보존합니다.
• 이 방식은 서브모듈러 최적화 (Submodular Optimization) 이론과 연결되어, 각 단계의 한계적 개선 (marginal improvement) 을 극대화함으로써 전역 최적에 가까운 정책을 학습할 수 있음을 시사합니다.

다. 개선된 샘플링 SDE (DDIM-Inspired SDE)

• 기존 Flow-GRPO 의 SDE 는 탐색을 위해 노이즈를 주입하지만, 이로 인해 중간 생성 이미지가 노이즈가 많아 보상 모델의 평가 정확도를 떨어뜨립니다.
• 저자들은 DDIM 스타일의 업데이트 규칙을 도입하여, 확률적 성질 (탐색) 을 유지하면서도 더 깨끗한 이미지를 생성하는 새로운 SDE 를 제안했습니다.
• 이 SDE 는 Flow-ODE 의 마진 (marginal) 분포를 보존하면서도 시각적 품질을 높여 보상 신호의 신뢰성을 개선합니다.

3. 주요 기여 (Key Contributions)

1. Stepwise-Flow-GRPO 도입: 최종 이미지 보상이 아닌, 각 확산 단계의 **상대적 이득 (relative gain)**을 기반으로 신용을 할당하는 새로운 프레임워크를 제안했습니다. 이는 별도의 크리틱 (critic) 모델 없이도 세밀한 신용 할당을 가능하게 합니다.
2. DDIM 기반 SDE 개선: RL 탐색을 위한 확률적 샘플링을 유지하면서 이미지 품질을 높여 보상 신호를 개선하는 새로운 SDE 를 설계했습니다.
3. 이론적 및 실험적 검증:
   • 적응형 서브모듈러 최적화 이론과 연결하여 방법론의 타당성을 입증했습니다.
   • 다양한 보상 모델 (PickScore, ImageReward, UnifiedReward) 과 데이터셋 (GenEval) 에서 기존 Flow-GRPO 대비 **샘플 효율성 (sample efficiency)**과 수렴 속도가 월등히 우수함을 보였습니다.

4. 실험 결과 (Results)

• 샘플 효율성 및 수렴 속도:
  • 모든 설정에서 Stepwise-Flow-GRPO 는 Flow-GRPO 보다 더 적은 학습 단계 (steps) 로 더 높은 보상을 달성했습니다.
  • 실제 시간 (wall-clock time) 기준으로도 중간 이미지 생성에 따른 오버헤드가 있음에도 불구하고, 학습 효율이 높아 전체적으로 20~40% 빠른 수렴을 보였습니다.
• GenEval 벤치마크 성능:
  • 공간 추론 (Spatial Reasoning): 객체의 위치 관계 (예: "오른쪽", "위쪽") 를 더 정확하게 생성했습니다.
  • 속성 바인딩 (Attribute Binding): 색상과 객체의 연결 (예: "분홍색 오븐", "초록색 오토바이") 이 더 정확했습니다.
  • 카운팅 (Counting): 객체 개수 (예: "4 개의 프리스비") 를 더 정확하게 표현했습니다.
  • 특히 CFG(Classifier-Free Guidance) 가 4.5 일 때, 기존 방법보다 모든 하위 카테고리에서 성능이 우수했습니다.
• OCR 및 복잡한 작업:
  • 텍스트 렌더링 (OCR) 과 같은 계층적 구조가 중요한 작업에서 Flow-GRPO 는 학습 중 발산하거나 성능이 저하되는 반면, 제안 방법은 안정적으로 학습하고 높은 성능을 유지했습니다.
• 동시 작업 (Concurrent Works) 비교:
  • TempFlow-GRPO 나 Granular-GRPO 와 달리, 수동으로 노이즈 레벨을 가중치하거나 특정 단계만 학습하는 방식이 아니라, **데이터 기반의 자동 이득 (gain)**을 통해 모든 단계를 최적화하며 더 효율적입니다.

5. 의의 및 결론 (Significance)

이 논문은 강화 학습을 흐름 매칭 (Flow-Matching) 기반 이미지 생성 모델에 적용할 때 발생하는 **신용 할당 문제 (Credit Assignment Problem)**를 해결했습니다.

• 생성 과정의 시간적 구조 존중: 확산 모델이 초기에는 구도를, 후기에는 디테일을 학습한다는 점을 고려하여, 각 단계의 기여도를 정밀하게 평가함으로써 학습 효율을 극대화했습니다.
• RL 기반 생성 모델의 새로운 표준: 기존에 최종 결과만 평가받던 RL 기반 생성 모델이, 중간 과정의 품질을 평가하고 개선할 수 있는 패러다임을 제시했습니다.
• 실용적 가치: 더 적은 계산 비용으로 더 높은 품질의 이미지를 생성할 수 있게 되어, 텍스트-to-이미지 생성 모델의 실용성과 안정성을 크게 향상시켰습니다.

결론적으로, Stepwise-Flow-GRPO는 단순한 보상 최적화를 넘어, 생성 과정의 역동적인 구조를 이해하고 이를 학습 신호로 활용함으로써 차세대 고품질 이미지 생성 모델의 핵심 기술로 자리 잡을 것으로 기대됩니다.