플로우 매칭 모델에서의 GRPO를 위한 단계별 신용 할당

Stepwise Credit Assignment for GRPO on Flow-Matching Models

Flow-GRPO는 강화 학습을 플로우 모델에 성공적으로 적용했지만, 모든 단계에 걸쳐 균일한 신용 할당을 사용합니다. 이는 확산 생성의 시간적 구조를 무시하는 것으로, 초기 단계는 전체적인 구성과 내용을 결정하며(저주파 구조), 후기 단계는 세부 사항과 질감을 결정합니다(고주파 세부 사항). 또한, 최종 이미지만을 기반으로 균일한 신용을 할당하면, 확산 경로의 후반부에 오류가 수정되는 경우, 비최적의 중간 단계를 의도치 않게 보상할 수 있습니다. 우리는 각 단계의 보상 개선에 따라 신용을 할당하는 Stepwise-Flow-GRPO를 제안합니다. Tweedie의 공식을 활용하여 중간 보상 추정치를 얻고, 이득 기반의 장점을 도입함으로써, 우리 방법은 더 높은 샘플 효율성과 빠른 수렴을 달성합니다. 또한, 정책 경사를 위한 확률성을 유지하면서 보상 품질을 향상시키는 DDIM에서 영감을 받은 SDE를 도입했습니다.

Flow-GRPO successfully applies reinforcement learning to flow models, but uses uniform credit assignment across all steps. This ignores the temporal structure of diffusion generation: early steps determine composition and content (low-frequency structure), while late steps resolve details and textures (high-frequency details). Moreover, assigning uniform credit based solely on the final image can inadvertently reward suboptimal intermediate steps, especially when errors are corrected later in the diffusion trajectory. We propose Stepwise-Flow-GRPO, which assigns credit based on each step's reward improvement. By leveraging Tweedie's formula to obtain intermediate reward estimates and introducing gain-based advantages, our method achieves superior sample efficiency and faster convergence. We also introduce a DDIM-inspired SDE that improves reward quality while preserving stochasticity for policy gradients.