Flow-GRPO：通过在线 RL 训练 Flow matching 模型

1 前言

本期内容，我们讲Flow-GRPO，他将基于强化学习的GRPO用于Flow matching，并在多个测试指标上获得了巨大的突破，下面让我们来看一下
视频：Flow-GRPO：通过在线 RL 训练 Flow matching 模型

参考论文：Flow-GRPO: Training Flow Matching Models via Online RL

参考代码：Flow-GRPO:
Training Flow Matching Models via Online RL

2 引入

在Flow matching当中，已经可以取得相当不错的效果了，一些基于此开发的模型，如SD3.5的生成质量也相当不错。然而，与最先进的模型相比，SD3.5的指标质量仍然有待提高。比如，与GPT-4o相比，SD3.5显然落后一大截。

在NLP领域，将基于RL（强化学习）的方法引入其中已经证明可以取得相当不错的效果，该方法可以让模型的生成结果更加的趋近于人类的偏好，比如DPO、GRPO等等

强化学习除了应用于NLP领域，在CV领域中也逐渐大放异彩，而Flow-GRPO，就是将GRPO用于Flow matching当中。

3 Flow matching

先回顾一下Flow matching，假定存在$x_0\sim X_0$为真实的数据样本，$x_1\sim X_1$为噪声样本，以Rectified flow为例，任意时刻的状态可以表示为

$$x_t=(1-t)x_0+t{x}_1$$
其中$t\in [0,1]$，我们可通过训练得到一个近似向量场$v_{\theta }(x_t,t)$
$$\mathcal{L}(\theta )={\mathbb{E}}_{t,x_0,x_1}\left[\Vert v-v_{\theta }(x_t,t){\Vert }^2\right]$$
其中，向量场$v=x_1-x_0$

4 方法

论文以SD3.5为例，将Flow-GRPO应用于T2I（文生图）当中。熟悉GRPO的小伙伴都知道，要使用GRPO的方法对Flow进行训练，要先解决ODE的问题：

• ODE无法在同一条件下生成多个样本，因此需要进行ODE到SDE的转化

4.1 GRPO

RL的优化目标一般为
$$\underset{\theta }{\max }{\mathbb{E}}_{(s_0,a_0,\cdots ,s_T,a_T)\sim {\pi }_0}\left[\sum _{t=0}^T\left(R(s_t,a_t)-\beta D_{KL}({\pi }_{\theta }(\cdot |s_t)||{\pi }_{ref}(\cdot |s_t))\right)\right]$$
去噪过程可以表示为一个MDP，给定提示词c，Flow可以得到一组图像 { x 0 i } i = 1 G \{ x_0^i \}_{i=1}^G {x0i​}i=1G​，还有对应的一个采样轨迹$\{(x_T^i,x_{T-1}^i,\cdots ,x_0^i){\}}_{i=1}^G$，我们可通过组归一化来计算第i张图形的优势，即
A ^ t i = R ( x 0 i , c ) − mean ( { R ( x 0 i , c ) } i = i G ) std ( { R ( x 0 i , c ) } i = 1 G ) \hat A_t^i=\frac{R(x_0^i,c)-\text{mean}(\{R(x_0^i,c)\}_{i=i}^G)}{ \text{std}(\{R(x_0^i,c)\}_{i=1}^G)} A^ti​=std({R(x0i​,c)}i=1G​)R(x0i​,c)−mean({R(x0i​,c)}i=iG​)​
最大化GRPO的优化目标
J Flow-GRPO ( θ ) = E c ∼ C ， { x i } i = 1 G ∼ π θ old ( ⋅ ∣ c ) f ( r , A ^ , θ , ε , β ) \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) JFlow-GRPO​(θ)=Ec∼C，{xi}i=1G​∼πθold​​(⋅∣c)​f(r,A^,θ,ε,β)
其中
$$\begin{gathered} f(r,\hat{A},\theta ,\epsilon ,\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,\text{clip}(1-\epsilon ,1+\epsilon ){\hat{A}}_t^i\right)-\beta D_{KL}({\pi }_{\theta }||{\pi }_{ref})\right), \\ \text{and}{r}_t^i(\theta )=\frac{p_{\theta }(x_{t-1}^i|x_t^i,c)}{p_{{\theta }_{old}}(x_{t-1}^i|x_t^i,c)} \end{gathered}$$

4.2 从 ODE 到 SDE

如上式可见，无论是计算优势函数，还是优化目标当中，都依赖于随机采样来得到不同的轨迹。而基于ODE的去噪过程显然是不满足这一要求的，为此，我们需要把去噪过程从ODE转变为SDE，这样就有了随机性。

那么如何将ODE转化为SDE呢？其实，我们可以得到下面的等式（稍后证明）
$$d{x}_t=\left[v_t(x_t)+\frac{{\sigma }_t^2}{2t}(x_t+(1-t)v_t(x_t))\right]dt+{\sigma }_{t}d\bar{w}\text{\hspace{1em}(1)}$$
$d\bar{w}$表示维纳过程增量，${\sigma }_t$是用于控制稳定程度的

可以看到，Eq.(1)仅仅依赖于向量场$v$，我们完全可以使用学习到的近似向量场$v_{\theta }$去表示他。我们可以使用任意一个数值求解器，来得到生成轨迹

如欧拉-丸山法

去噪过程为
$$x_{t+\Delta t}=x_t+\left[v_{\theta }(x_t,t)+\frac{{\sigma }_t^2}{2t}(x_t+(1-t)v_{\theta }(x_t,t))\right]\Delta t+{\sigma }_t\sqrt{\Delta t}\epsilon \text{\hspace{1em}(2)}$$
其中$\epsilon \sim \mathcal{N}(0,I)，{\sigma }_t=a\sqrt{\frac{t}{1-t}}$，$a$是控制噪声水平的超参数。

依据正态分布的性质可知，Eq.(2)，也就是${\pi }_{\theta }(x_{t-1}|x_t,c)$服从正态分布，那么很显然，我们可以直接KL散度为
$$D_{KL}({\pi }_0||{\pi }_{ref})=\frac{\Vert {\bar{x}}_{t+\Delta t,\theta }-{\bar{x}}_{t+\Delta t,ref}{\Vert }^2}{2{\sigma }_t^2\Delta t}=\frac{\Delta t}{2}{\left(\frac{{\sigma }_t(1-t)}{2t}+\frac{1}{{\sigma }_t}\right)}^2\Vert v_{\theta }(x_t,t)-v_{ref}(x_t,t){\Vert }^2\text{\hspace{1em}(3)}$$

这里直接代入KL散度公式来计算即可

5 Denoising Reduction

为了生成高质量的图像，Flow matching通常需要很多的去噪步骤，这使得RL训练的数据收集成本非常高。

论文发现，在进行RL训练的时候，是不需要太多的采样步数的，而在推理的时候保持原始的采样步依然能够获取高质量的样本。

为此，以SD3.5为例，在进行RL训练的时候，令采样时间步T=10；而在推理的时候，保持SD3.5默认设置T=40。

6 模型图

模型的训练流程见下图：

首先，分别采样5个高斯白噪声$s_0$，将提示词“A photo of four cups”作为条件，使用SDE数值求解器采样（T=10）得到$s_T$。然后将$s_T$送进奖励模型，得到$R^1,R^2,R^3,R^4,\cdots ,R^G$作为奖励。用这些奖励根据上面的优势函数计算优势得到${\hat{A}}^1,{\hat{A}}^2,{\hat{A}}^3,{\hat{A}}^4,\cdots ,{\hat{A}}^G$，最后送进Flow-GRPO的损失函数计算损失即可。

7 数学证明

7.1 Eq.(1)证明

要将ODE转换成对应的SDE，就要先从ODE开始，我们有
$$d{x}_t=v_{t}dt\text{\hspace{1em}(4)}$$
依据先前讲过的SDE，我们有对应的方程
$$d{x}_x=f_{\text{SDE}}(x_t,t)dt+{\sigma }_{t}dw\text{\hspace{1em}(5)}$$
我们需要求出$f_{\text{SDE}}$和$v_t$的关系式

依据Flow matching所提到的FP方程，Eq.(4)和Eq.(5)都有一个对应的连续性方程来表达概率密度路径$p_t$。对于Eq.(5)，就是对应的FP方程（证明过程见什么是Fokker-Planck方程），即
KaTeX parse error: Undefined control sequence: \part at position 2: \̲p̲a̲r̲t̲ ̲_tp_t(x) = -\na…
而Eq.(4)对应的连续性方程为：
$${\partial }_t{p}_t(x)=-\nabla \cdot [v_t(x_t,t)p_t(x)]\text{\hspace{1em}(7)}$$
当$p_t$和$v_t$的关系满足Eq.(7)，则我们说向量场$v$能够生成对应的路径$p_t$。Eq.(6)同理。

那么接下来就简单了，联立Eq.(6)和Eq.(7)
$$-\nabla \cdot [f_{\text{SDE}}(x_t,t)p_t(x)]+\frac{1}{2}{\nabla }^2[{\sigma }_t^2{p}_t(x)]=-\nabla \cdot [v_t(x_t,t)p_t(x)]\text{\hspace{1em}(8)}$$
因为
$$\nabla \log p_t(x)=\frac{1}{p_t(x)}\cdot \nabla p_t(x)\to \nabla p_t(x)=p_t(x)\cdot \nabla \log p_t(x)$$
对Eq.(8)左侧第二项进行一下变化
$$\begin{array}{rl}{\nabla }^2[{\sigma }_t^2{p}_t(x)]=& {\sigma }_t^2{\nabla }^2{p}_t(x)\\ =& {\sigma }_t^2\nabla \cdot (\nabla p_t(x))\\ =& {\sigma }_t^2\nabla \cdot (p_t(x)\nabla \log p_t(x))\end{array}$$
所以Eq.(8)等于：
$$\begin{array}{rl}-\nabla \cdot [f_{\text{SDE}}(x_t,t)p_t(x)]+\frac{1}{2}{\sigma }_t^2\nabla \cdot (p_t(x)\nabla \log p_t(x))& =-\nabla \cdot [v_t(x_t,t)p_t(x)]\\ -f_{\text{SDE}}(x_t,t)p_t(x)+\frac{1}{2}{\sigma }_t^2{p}_t(x)\nabla \log p_t(x)& =-v_t(x_t,t)p_t(x)\\ f_{\text{SDE}}(x_t,t)p_t(x)& =v_t(x_t,t)p_t(x)+\frac{1}{2}{\sigma }_t^2{p}_t(x)\nabla \log p_t(x)\\ f_{\text{SDE}}(x_t,t)& =v_t(x_t,t)+\frac{1}{2}{\sigma }_t^2\nabla \log p_t(x)\end{array}\text{\hspace{1em}(9)}$$
这样的话，我们就得到了$f_{\text{SDE}}$和$v_t$的关系式了

依据Score-Based Generative Modeling through Stochastic Differential Equations，正向过程Eq.(5)有对应的反向过程为
$$d{x}_t=[f(x_t,t)-g^2(t)\nabla \log p_t(x_t)]dt+g(t)d\bar{w}\text{\hspace{1em}(10)}$$
其中，在本篇文章中，我们是让$g(t)={\sigma }_t$，将Eq.(9)代入至Eq.(10)
$$\begin{array}{rl}d{x}_t=& \left[v_t(x_t,t)+\frac{1}{2}{\sigma }_t^2\nabla \log p_t(x_t)-{\sigma }_t^2\nabla \log p_t(x_t)\right]dt+{\sigma }_{t}d\bar{w}\\ d{x}_t=& \left[v_t(x_t,t)-\frac{{\sigma }_t^2}{2}\nabla \log p_t(x_t)\right]dt+{\sigma }_{t}d\bar{w}\end{array}\text{\hspace{1em}(11)}$$
对于Eq.(11)，已知$v_t$，一旦$\nabla \log p_t(x_t)$也是已知的，那么就没有未知变量了，也就可以使用数值求解器生成样本了。因此我们还需要求解$\nabla \log p_t(x_t)$。

对于前向加噪过程，我们有$x_t={\alpha }_t{x}_0+{\beta }_t{x}_1$，在本期的Flow中，我们将加噪过程定义为${\alpha }_t=1-t;\beta =t$，$x_t$服从的概率分布为（假设一维的情况）
p t ∣ 0 ( x t ∣ x 0 ) = N ( x t ∣ a t x 0 , β t 2 I ) = 1 β t 2 π exp ⁡ { − ( x t − a t x 0 ) 2 2 β t 2 } p_{t|0}(x_t|x_0) = \mathcal{N}(x_t|a_tx_0,\beta_t^2I) = \frac{1}{\beta_t\sqrt{2\pi}}\exp\{-\frac{(x_t-a_tx_0)^2}{2\beta_t^2}\} pt∣0​(xt​∣x0​)=N(xt​∣at​x0​,βt2​I)=βt​2π ​1​exp{−2βt2​(xt​−at​x0​)2​}
其对数结果为
log ⁡ p t ∣ 0 ( x t ∣ x 0 ) = log ⁡ ( 1 β t 2 π exp ⁡ { − ( x t − a t x 0 ) 2 2 β t 2 } ) = log ⁡ 1 β t 2 π − ( x t − a t x 0 ) 2 2 β t 2 \begin{aligned} \log p_{t|0}(x_t|x_0) = &\log \left( \frac{1}{\beta_t\sqrt{2\pi}}\exp\{-\frac{(x_t-a_tx_0)^2}{2\beta_t^2}\} \right) \\= &\log \frac{1}{\beta_t\sqrt{2\pi}} -\frac{(x_t-a_tx_0)^2}{2\beta_t^2} \end{aligned} logpt∣0​(xt​∣x0​)==​log(βt​2π ​1​exp{−2βt2​(xt​−at​x0​)2​})logβt​2π ​1​−2βt2​(xt​−at​x0​)2​​
所以
$$\nabla \log p_{t|0}(x_t|x_0)=-\frac{x_t-{\alpha }_t{x}_0}{{\beta }_t^2}=\frac{{\beta }_t{x}_1}{{\beta }_t^2}=-\frac{x_1}{{\beta }_t}$$
因此
$$\begin{array}{rl}\nabla \log p_t(x_t)=& \frac{1}{p_t(x_t)}\nabla p_t(x_t)\\ =& \frac{1}{p_t(x_t)}\nabla \int p_{t,0}(x_t,x_0)d{x}_0\\ =& \frac{1}{p_t(x_t)}\int \nabla p_{t,0}(x_t,x_0)d{x}_0\\ =& \frac{1}{p_t(x_t)}\int \nabla \left[p_{t|0}(x_t|x_0)p_0(x_0)\right]d{x}_0\\ =& \frac{1}{p_t(x_t)}\int p_0(x_0)\nabla p_{t|0}(x_t|x_0)d{x}_0\\ =& \frac{1}{p_t(x_t)}\int p_0(x_0)\cdot p_{t|0}(x_t|x_0)\nabla \log p_{t|0}(x_t|x_0)d{x}_0\\ =& \frac{1}{p_t(x_t)}\int p_{t,0}(x_t,x_0)\nabla \log p_{t|0}(x_t|x_0)d{x}_0\\ =& \frac{1}{p_t(x_t)}\int p_{0|t}(x_0|x_t)p_t(x_t)\nabla \log p_{t|0}(x_t|x_0)d{x}_0\\ =& \int p_{0|t}(x_0|x_t)\nabla \log p_{t|0}(x_t|x_0)d{x}_0\\ =& {\int }_{x_0}{\int }_{x_1}{p}_{0|t}(x_0,x_1|x_t)d{x}_1\nabla \log p_{t|0}(x_t|x_0)d{x}_0\\ =& {\int }_{x_0}{\int }_{x_1}{p}_{0|t}(x_0,x_1|x_t)\nabla \log p_{t|0}(x_t|x_0)d{x}_{1}d{x}_0\\ =& \mathbb{E}\left[\nabla \log p_{t|0}(x_t|x_0)|x_t\right]\\ =& \mathbb{E}\left[-\frac{x_1}{{\beta }_t}|x_t\right]\\ =& -\frac{1}{{\beta }_t}\mathbb{E}\left[x_1|x_t\right]\end{array}\text{\hspace{1em}(12)}$$
对于向量场$v$，在我们之前的表达式中，是有$v_t=x_1-x_0$。然而，由于路径存在交叉点，所以我们之前说过，我们学习到的$v_{\theta }$其实并不等于$v_t$，而是$v_t$的数学期望。我们可以通过以下来证明：
$$\begin{array}{rl}\mathcal{L}=& {\int }_0^1{\mathbb{E}}_{x_0,x_1}\left[\Vert x_1-x_0-v_{\theta }(x_t,t){\Vert }^2\right]\\ =& {\int }_0^1{\mathbb{E}}_{x_0,x_1}\left[\Vert x_1-x_0{\Vert }^2+||v_{\theta }(x_t,t)|{|}^2-2(x_1-x_0{)}^T{v}_{\theta }(x_t,t)\right]dt\\ =& {\int }_0^1\left\{{\mathbb{E}}_{x_0,x_1}\left[||v_{\theta }(x_t,t)|{|}^2\right]-2{\mathbb{E}}_{x_0,x_1}\left[(x_1-x_0{)}^T{v}_{\theta }(x_t,t)\right]\right\}dt+C\\ =& {\int }_0^1\left\{{\mathbb{E}}_{x_t}\left[||v_{\theta }(x_t,t)|{|}^2\right]-2{\mathbb{E}}_{x_0,x_1}\left[(x_1-x_0{)}^T{v}_{\theta }(x_t,t)\right]\right\}dt+C\end{array}\text{\hspace{1em}(13)}$$
第一项是因为给定$x_0,x_1$，有$x_t=t{x}_1+(1-t)x_0$，所以可以直接写成关于$x_t$的数学期望。

第二项我们可以继续变化，由全期望公式：$\mathbb{E}Y={\mathbb{E}}_X[{\mathbb{E}}_Y(Y|X)]$，可得
$$\begin{array}{rl}{\mathbb{E}}_{x_0,x_1}\left[(x_1-x_0{)}^T{v}_{\theta }(x_t,t)\right]=& {\mathbb{E}}_{x_t}[{\mathbb{E}}_{x_0,x_1}\left[(x_1-x_0{)}^T{v}_{\theta }(x_t,t)|x_t\right]]\\ =& {\mathbb{E}}_{x_t}[{\mathbb{E}}_{x_0,x_1}{\left[(x_1-x_0)|x_t\right]}^T{v}_{\theta }(x_t,t)]\end{array}$$
所以Eq.(13)为
$$\begin{array}{rl}\mathcal{L}=& {\int }_0^1\left\{{\mathbb{E}}_{x_t}\left[||v_{\theta }(x_t,t)|{|}^2\right]-2{\mathbb{E}}_{x_t}[{\mathbb{E}}_{x_0,x_1}{\left[(x_1-x_0)|x_t\right]}^T{v}_{\theta }(x_t,t)]\right\}dt+C\\ =& {\int }_0^1{\mathbb{E}}_{x_t}\left[\Vert {\mathbb{E}}_{x_0,x_1}[x_1-x_0|x_t]-v_{\theta }(x_t,t){\Vert }^2\right]dt+C^{\prime }\end{array}\text{\hspace{1em}(14)}$$
此时我们不难看出，我们所学习到的$v_{\theta }(x_t,t)={\mathbb{E}}_{x_0,x_1}[x_1-x_0|x_t]$

我们继续转化
$$\begin{array}{rl}{v}_{\theta }(x_t,t)=& {\mathbb{E}}_{x_0,x_1}[x_1-x_0|x_t]\\ =& {\mathbb{E}}_{x_0,x_1}[x_1|x_t]-{\mathbb{E}}_{x_0,x_1}[x_0|x_t]\\ =& {\mathbb{E}}_{x_0,x_1}[x_1|x_t]-{\mathbb{E}}_{x_0,x_1}\left[\frac{x_t-t{x}_1}{1-t}|x_t\right]\\ =& {\mathbb{E}}_{x_0,x_1}[x_1|x_t]-{\mathbb{E}}_{x_0,x_1}\left[\frac{x_t}{1-t}|x_t\right]+{\mathbb{E}}_{x_0,x_1}\left[\frac{t{x}_1}{1-t}|x_t\right]\\ =& {\mathbb{E}}_{x_0,x_1}[x_1|x_t]-\frac{x_t}{1-t}+\frac{t}{1-t}{\mathbb{E}}_{x_0,x_1}\left[x_1|x_t\right]\\ =& -\frac{x_t}{1-t}+\frac{1}{1-t}{\mathbb{E}}_{x_0,x_1}\left[x_1|x_t\right]\\ =& -\frac{x_t}{1-t}+\frac{1}{1-t}\cdot (-{\beta }_t\nabla \log p_t(x_t))\\ =& -\frac{x_t}{1-t}-\frac{t}{1-t}\cdot \nabla \log p_t(x_t)\end{array}$$
把$\nabla$单独放等式左侧可得
$$\nabla \log p_t(x_t)=-\frac{x}{t}-\frac{1-t}{t}{v}_{\theta }(x_t,t)$$
把它代入到Eq.(11)可得最终的表达式
$$\begin{array}{rl}d{x}_t=& \left[v_t(x_t,t)-\frac{{\sigma }_t^2}{2}\left(-\frac{x}{t}-\frac{1-t}{t}{v}_{\theta }(x_t,t)\right)\right]dt+{\sigma }_{t}d\bar{w}\\ d{x}_t=& \left[v_t(x_t,t)+\frac{{\sigma }_t^2}{2t}\left(x+(1-t)v_{\theta }(x_t,t)\right)\right]dt+{\sigma }_{t}d\bar{w}\end{array}\text{\hspace{1em}(15)}$$
至此得证

8 参考

[1] 深入理解Rectified Flow，完善统一扩散框架 - 知乎

9 结束

好了，本期内容到此为止了，如有问题，还望指出，阿里嘎多！