Neurips 2025
Jie Liu, Gongye Liu, Jiajun Liang, Yangguang Li, Jiaheng Liu, Xintao Wang, Pengfei Wan, Di Zhang, Wanli Ouyang
8 May 2025
[paper]

Background

Flow Matching 모델은 이미지 생성 분야에서 굉장히 강한 성능을 보였지만, 모델이 인간 선호도에 맞추는 것은 아직도 어려운 일이다. LLM 분야에서 GRPO는 Deepseek R1 등에서 critic 모델 없이도 RLHF를 하여 굉장히 강력한 퍼포먼스를 보인 선례가 존재한다. 해당 논문에서는 Flow Matching 모델에서도 GRPO를 도입할 수 있는 방법을 제안한다.

preliminary

Flow Matching

\[x_t = (1 - t) x_0 + t x_1\]

$x_0$를 true 분포에서 샘플링한 데이터 샘플, $x_1$을 노이즈 분포에서 샘플링한 데이터 샘플이라 할때 최근 Rectified Flow 모델에서 $x_t$는 위와 같이 표기할 수 있다. 이때 Rectified Flow는 데이터 분포와 노이즈 분포 사이를 가장 직선에 가까운 경로로 연결하는 ODE모델을 학습시키는 방법론을 의미한다. (논문 링크)

\[\mathcal{L}(\theta) = \mathbb{E}_{t, \mathbf{x}_0 \sim X_0, \mathbf{x}_1 \sim X_1} \left[ \left\| \mathbf{v} - \mathbf{v}_{\theta}(\mathbf{x}_t, t) \right\|^2 \right]\]

이때 velocity field $v = x_1 - x_0$ 일때 위 flow matching 목적함수를 최소화함으로서 트랜스포머 모델이 velocity field $v_\theta(x_t,t)$를 예측할 수 있게 된다.

GRPO(Group Relative Policy Optimization)

GRPO는 DeepSeek R1 모델의 핵심 알고리즘으로 기존의 PPO 모델의 actor critic 방식에서 벗어나 하나의 질문에 대해 여러개의 답변을 생성 한 이후 그들간의 상대적 우열을 비교해 학습하는 방식이다.

기존 PPO에서는 Actor(행동을 하는 모델)과 더불어 현재 상태에서 value를 계산하는 critic 모델을 학습시에 GPU에 올려야한다. Critic 모델은 actor 모델만큼 용량이 커 GPU의 VRAM을 많이 차지해 LLM의 RLHF를 어렵게하는 요인이 되었다.

\[A_i = \frac{r_i - \text{mean}(\{r_1, ..., r_G\})}{\text{std}(\{r_1, ..., r_G\}) + \epsilon}\]

GRPO에서는 프롬프트 하나를 이용해 얻은 여러 답변의 평균을 baseline으로 삼아, 각 group 내에서 상대적인 advantage를 계산한다. 이를 통해 별도의 critic모델이 필요가 없어진다.

\[\mathcal{J}_{GRPO}(\theta) = \mathbb{E}\left[q \sim P(Q), \{o_i\}_{i=1}^G \sim \pi_{\theta_{old}}(O|q)\right] \frac{1}{G} \sum_{i=1}^G \frac{1}{|o_i|} \sum_{t=1}^{|o_i|} \left\{ \min \left[ \frac{\pi_\theta(o_{i,t}|q, o_{i,<t})}{\pi_{\theta_{old}}(o_{i,t}|q, o_{i,<t})} \hat{A}_{i,t}, \text{clip}\left( \frac{\pi_\theta(o_{i,t}|q, o_{i,<t})}{\pi_{\theta_{old}}(o_{i,t}|q, o_{i,<t})}, 1-\epsilon, 1+\epsilon \right) \hat{A}_{i,t} \right] - \beta D_{KL}[\pi_\theta || \pi_{ref}] \right\},\]

좀 길긴하지만 GRPO의 목적함수는 다음과 같다. $\frac{\pi_\theta}{\pi_{\theta_{old}}}$는 PPO와 마찬가지로 이전 정책 대비 현재 정책의 변화비율을 의미한다. 다음으 정책이 너무 급격하게 변하지 않도록 $[1-\epsilon, 1+\epsilon]$ 범위에서 clipping한다. $A_i$는 위에서 구한 advantage로 평균보다 잘하면 +, 못했다면 -를 주어 방향을 정하고, KLD항을 두어 원본 모델과 너무 멀어지지 않도록 규제를 걸어준다.

Denoising as an MDP(Markov Decision Process)

\[Markov Decision Process (MDP) : (S, A, ρ_0, P, R)\]

Flow Matching 모델에서 디노이징하는 과정을 MDP로 표현 가능하다.

State $s_t \triangleq (c, t, x_t)$ : 현재 모델이 처한 상태를 의미, 이때 c는 프롬프트, t는 타임스텝, $x_t$는 노이즈가 낀 현재 상태의 이미지 즉 latent를 의미한다.
Action $a_t \triangleq x_{t-1}$ : 모델이 취할 행동으로, 여기서는 한 단계 다음의 깨끗해진 이미지를 의미한다.
Policy $\pi(a_t | s_t) \triangleq p_\theta(x_{t-1} | x_t, c)$ : 현재 state에서 어떤 행동을 취할지 결정하는 함수로, 이는 우리가 현재 학습시키려는 모델($\theta$)를 의미한다.
Transition $P(s_{t+1} | s_t, a_t) \triangleq (\delta_c, \delta_{t-1}, \delta_{x_{t-1}})$ : 내가 행동을 취했을때 환경이 어떻게 변하는가를 의미하며, Dirac delta $\delta$ 를 통해 표기됐다. 이는 프롬프트, 다음 타임 스텝, 다음 상태의 이미지는 결정론적이며 바뀌지 않는다는 것을 의미한다.
Initial State $\rho_0(s_0) \triangleq (p(c), \delta_T, \mathcal{N}(0, I))$ : 초기상태는 유저가 입력한 프롬프트, 완전한 노이즈인 타임스텝 T, 그리고 완전한 가우시안 노이즈에서 시작된다.
Reward $R(s_t, a_t)$ : Reward는 잘했는지 못했는지를 점수로서 매기는 것이며, 중간 상태일 때는 점수를 주지 않고 t=0일때, 즉 최종 이미지를 보고 평가를 하게 된다.

Methodology

Flow GRPO

\[\mathcal{J}_{\text{Flow-GRPO}}(\theta) = \mathbb{E}_{c \sim \mathcal{C}, \{x^i\}_{i=1}^G \sim \pi_{\theta_{\text{old}}}(\cdot|c)} f(r, \hat{A}, \theta, \varepsilon, \beta)\]

Flow GRPO의 학습목표는 목적함수를 최대화하는 $\theta$를 찾는 것이다. 여기서 다양한 프롬프트의 평균을 계산하고, 하나의 프롬프트에 대해 G개의 이미지 그룹을 생성한다. 이때 업데이트 직전의 정책을 사용하게 된다.

\[\begin{align*} f(r, \hat{A}, \theta, \varepsilon, \beta) &= \frac{1}{G} \sum_{i=1}^{G} \frac{1}{T} \sum_{t=0}^{T-1} \left( \min \left( r_t^i(\theta) \hat{A}_t^i, \operatorname{clip}(r_t^i(\theta), 1 - \varepsilon, 1 + \varepsilon) \hat{A}_t^i \right) - \beta D_{\text{KL}}(\pi_\theta || \pi_{\text{ref}}) \right), \\ r_t^i(\theta) &= \frac{p_\theta(\mathbf{x}_{t-1}^i | \mathbf{x}_t^i, \mathbf{c})}{p_{\theta_{\text{old}}}(\mathbf{x}_{t-1}^i | \mathbf{x}_t^i, \mathbf{c})}. \end{align*}\]

여기서 해당 수식은 위에 GRPO의 목적함수와 유사하게 PPO의 loss function을 확장한 것이다. GRPO의 특성으로 그룹 내 G개의 샘플에 대해 평균을 내고, 시간(step)에 대해서도 평균을 내어 모든 step에서 잘했는지를 평가한다. Clipping하는 부분과 KLD 패널티를 주는 부분은 GRPO와 유사하다.

그 아래 $r_t^i(\theta)$는 현재 스텝 t에서 t-1로 가는 전이확률의 비율을 의미한다. 이때 Diffusion 모델에서 이 모델을 가우시안 분포로 가정하므로 실제로 계산할 때는 평균과 분산의 차이를 통해 계산할 수 있게 된다.

ODE to SDE

GRPO는 샘플링을 할때 exploration을 위해 SDE를 사용하다. DDPM은 기본적으로 SDE를 사용하기 때문에 문제가 없지만 Flow Matching 모델은 forward 과정에서 ODE를 사용한다.

\[\mathrm{d}x_t = v_t \mathrm{d}t\]

그렇기에 해당 논문에서는 ODE를 다음과 같은 동치인 아래의 SDE로 변환하여 사용한다. (증명은 Appendix A)

\[d\boldsymbol{x}_t = \left[ \boldsymbol{v}_t(\boldsymbol{x}_t) + \frac{\sigma_t^2}{2t} \left(\boldsymbol{x}_t + (1 - t)\boldsymbol{v}_t(\boldsymbol{x}_t)\right) \right] dt + \sigma_t d\boldsymbol{w}.\]

여기서 Euler-Maruyama discretization를 사용하면 아래의 수식이 완성된다. Euler-Maruyama discretization는 SDE를 시간 간격으로 이산화하여 작성하는 방법을 의미한다.

\[\boldsymbol{x}_{t+\Delta t} = \boldsymbol{x}_t + \left[ \boldsymbol{v}_\theta(\boldsymbol{x}_t, t) + \frac{\sigma_t^2}{2t} \left(\boldsymbol{x}_t + (1-t)\boldsymbol{v}_\theta(\boldsymbol{x}_t, t)\right) \right] \Delta t + \sigma_t \sqrt{\Delta t} \epsilon\]

Denoising Reduction

해당 논문에서는 Online RL의 리소스를 절약하기 위해 Denoising때 Step 수를 10으로 (default는 40) 줄여서 실험했을때 성능은 크게 저하되지 않으면서 속도는 훨씬 빠르게 학습된다는 것을 발견했고, 이를 실험을 통해 보였다고 한다.

Overview Figure

Flow GRPO의 전체적인 맥락은 다음과 같다. ODE를 SDE로 바꿔서 샘플링을 하고, denoising step을 10으로 줄여서 GRPO를 적용하는 전체적인 하름을 보여준다.

Experiments

Experiment Setting

해당 논문에서는 총 4가지 부분에서 실험을 진행하게 된다.

Compositional Image Generation : 해당 부분에서 GenEval을 사용하게 되며, 크게 볼 때 프롬프트에 나와 있는 대로 카운팅은 잘되는지 그리고 위치와 색상은 정확하게 나오는지를 평가한다.
Visual Text Rendering : 말 그대로 텍스트를 잘 렌더링하는지를 평가하며 여기서는 OCR을 사용하였다.
Human Preference Alignment : 인간이 더 선호하는 이미지를 만들도록 하는 태스크로 여기서는 PickScore가 사용되었다. 해당 모델은 Stable Diffusion으로 만든 두 이미지중 어떤것을 사람이 더 선호하는지 데이터를 수집하여 CLIP로 학습시킨 모델을 이용한다.
Image Quality Evaluation Metric : 이미지 퀄리티 자체를 평가하는 지표로 Aesthetic Score, DeQA, ImageReward, UnifiedReward가 사용되었으며 모두 DrawBench를 통해 평가되었다.