단계별 보상 할당으로 강화학습 효율을 높인 Flow‑GRPO
본 논문은 기존 Flow‑GRPO가 최종 이미지에만 보상을 부여해 발생하는 시간적 신호 손실 문제를 해결한다. Tweedie 공식으로 중간 클린 이미지 추정치를 얻고, 각 디노이징 단계의 보상 증가량(gain)을 이용해 단계별 이득을 계산한다. 이를 GRPO에 적용해 샘플 효율성을 크게 향상시키고, DDIM‑영감을 받은 새로운 SDE로 노이즈를 감소시켜 정책 그라디언트 학습을 안정화한다.
저자: Yash Savani, Branislav Kveton, Yuchen Liu
본 논문은 텍스트‑투‑이미지 생성 모델에서 흐름 매칭(Flow‑matching) 방식을 강화학습으로 미세조정하는 기존 방법인 Flow‑GRPO의 근본적인 한계를 짚고, 이를 해결하기 위한 ‘Stepwise‑Flow‑GRPO’ 프레임워크를 제시한다. 먼저, 확산 기반 생성 과정은 시간에 따라 저주파(구성·배치)와 고주파(디테일·텍스처) 정보를 순차적으로 복원한다는 물리적·수학적 근거를 제시하고, 초기 단계에서의 오류가 후반에 교정될 경우 최종 보상만을 기준으로 전체 트래젝터리를 동일하게 강화하는 것이 비효율적임을 문제점으로 제시한다.
이를 해결하기 위해 두 가지 핵심 기술을 도입한다. 첫 번째는 Tweedie 공식에 기반한 중간 클린 이미지 추정이다. 흐름 모델이 매 단계 예측하는 노이즈 \( \hat{x}_1 \) 를 이용해 \( \hat{x}_0(t)=x_t - t\hat{x}_1 \) 를 계산함으로써, 각 디노이징 단계에서 기대되는 ‘깨끗한’ 이미지 \( \hat{x}_0(t) \) 를 거의 비용 없이 얻는다. 필요 시 5개의 서브스텝 ODE 통합을 통해 추정 품질을 강화하고, 이를 사전 학습된 이미지‑텍스트 정렬 보상 모델 \(R(\cdot, c)\) 에 입력해 단계별 보상 \(r_t\) 을 산출한다.
두 번째는 단계별 보상 차이 \(g_t = r_{t-1} - r_t\) 를 ‘gain’으로 정의하고, 이를 GRPO의 그룹‑상대 어드밴티지에 그대로 매핑하는 것이다. 이때 전체 gain은 \( \sum_t g_t = r_0 - r_T \) 로 텔레스코핑되므로, 개별 단계에서의 로컬 최적화가 전역 목표(최종 보상 최대화)와 동등하게 된다. 단계별 gain을 이용하면 보상이 증가한 단계는 긍정적으로 강화되고, 감소한 단계는 패널티를 받아, 비단순히 최종 이미지에만 의존하는 기존 방식보다 더 정교한 신호 전달이 가능해진다. 정규화는 모든 단계에 대해 공동으로 수행해 초기 단계의 큰 gain이 희석되지 않도록 설계하였다.
또한, 기존 Flow‑GRPO가 사용한 SDE는 마진을 정확히 맞추지만 샘플에 과도한 노이즈가 섞여 보상 모델이 학습하기 어려운 신호를 제공한다는 문제를 지적한다. 이를 보완하기 위해 DDIM‑영감을 받은 새로운 SDE를 제안한다. 이 SDE는 흐름‑ODE의 Fokker‑Planck 마진을 유지하면서, 디노이징 과정에서 노이즈를 감소시켜 중간 이미지의 품질을 높인다. 결과적으로 정책 그라디언트가 더 안정적인 보상 신호를 받게 되어 학습 효율이 크게 향상된다.
실험에서는 256개의 GenEval 프롬프트와 다양한 보상 모델(PickScore, R‑Score 등)을 사용해 Stepwise‑Flow‑GRPO와 기존 Flow‑GRPO를 비교하였다. 단계별 gain의 평균 절대값이 초기 단계에서 크게 나타나는 것을 확인했으며, 이는 초기 compositional 결정이 전체 보상 향상의 핵심임을 시사한다. 학습 곡선에서는 Stepwise‑Flow‑GRPO가 동일한 연산량 대비 30% 이상 빠르게 수렴하고, 최종 보상 점수가 0.05~0.08 정도 상승한다. 또한, 제안된 DDIM‑SDE와 결합했을 때는 특히 고해상도(512×512) 이미지에서 텍스트 정렬 정확도와 시각적 품질이 눈에 띄게 개선되었다.
본 연구는 별도 크리틱 네트워크 없이도 단계별 gain을 이용해 PPO‑형 크리틱의 역할을 수행함으로써 구현 복잡성을 크게 낮춘다. 또한, 흐름 매칭 모델에 대한 강화학습 적용을 보다 일반화 가능하고 효율적인 방법으로 확장한다는 점에서, 대규모 텍스트‑투‑이미지 시스템에 바로 적용할 수 있는 실용적인 기여를 제공한다.
원본 논문
고화질 논문을 불러오는 중입니다...
댓글 및 학술 토론
Loading comments...
의견 남기기