前置知识
本文涉及到的两个核心公式来自Flow Matching与MeanFlow，分别是边缘速度场$\\\textit{u}_t(\\\textit{x})$的形式以及MeanFlow恒等式。在阅读本文前，需要对Flow Matching的基本原理，以及这两篇文章的内容有比较清晰的了解。
引言
在Mean Flows for One-step Generative Modeling一文中，作者提出了基于一种物理一致性的Consistency Model，相比于传统的CMs或是sCM，MeanFlow，无需ema之类的操作，也不需要蒸馏，而是可以直接从头开始训练，得到一个具有少步生成能力的模型——随着步数的增加，MeanFlow的生成效果也会变好。但是MeanFlow的定义式存在对$\\\textit{u}_\theta(\\\textit{x},\\\textit{t},\\\textit{r})$求关于$\\\textit{x}_r$和$\\\textit{t}$的偏导的项，训练的过程会极其复杂，计算量非常大。因此MeanFlow经常被诟病的问题就是难以训练，在复现时我也遇到了同样的问题。甚至Kaiming在github上传的项目中，足足在128*8的batch_size上训练了800k个iter。
但是我们回顾MeanFlow，其实另一个与模型参数$\\\mathit{\theta}$相关的项$\\\textit{u}_\theta(\\\textit{x},\\\textit{t},\\\textit{t})$，作者并没有采用模型来估计，而是直接选取了条件速度场(Conditional Flow Matching)的形式，即认为$\\\textit{u}_\theta(\\\textit{x},\\\textit{t},\\\textit{t})=\\\textit{x}_1-\\\textit{x}_0$。也就是说，作者认为边缘速度场与条件速度场在瞬时速度上的差异是不大的，这就给了我们启发：如果我们能将MeanFlow的定义式中的平均速度$\\\textit{u}_\theta(\\\textit{x},\\\textit{t}, \\\textit{r})$转化为某个确知的瞬时速度，是否就能够用条件速度场替代呢？这样，我们就能够直接求出偏导数的显式表达式了，而无需再通过雅各比之类的方法求导，能够节省相当多时间成本。
边缘速度场的一般形式
我们的想法是在边缘速度场里面寻找平均速度的替代形式。那么我们就要先求出一个解析形式的边缘速度场，在Flow Matching一文中，我们显式定义了边缘速度场的形式，它等于： $$ \begin{equation} u_t(x)=\int v_t(x|x_1) \frac{p_t(x|x_1)q(x_1)}{p_t(x)}dx_1, \\\label{marginalvelocity} \tag{1} \end{equation} $$ 在这个式子中，按照EDM的定义，条件速度场就应该等于： $$ \begin{equation} v_t(x|x_1)=\dot{f}(t)x_1 + \dot{g}(t)x_0 \\\label{conditionalvelocity} \tag{2} \end{equation} $$ 它们满足边界条件： $$ \begin{equation} \begin{cases} f(0) = 0, \quad f(1) = 1\\ g(0) = 1, \quad g(1) = 0 \end{cases} \label{edgecondition} \tag{3} \end{equation} $$ 条件概率流的形式是一个高斯分布，满足： $$ \begin{equation} p_t(x|x_1)\\\sim\mathcal{N}(x|f(t)x_1,g^2(t)\boldsymbol{I}), \label{conditionalpath} \tag{4} \end{equation} $$ 而两个边缘分布$\\\textit{q}(\\\textit{x}_1)$和$\\\textit{p}_t(\\\textit{x})$都是未知的。因此我们几乎无法显式计算得到$\\\textit{u}_t(\\\textit{x})$的具体形式。连表达式都无法求解，我们还能怎么从中找到平均速度与边缘速度的关系呢？
不妨先从简单的分布出发，假设$\\\textit{q}(\\\textit{x}_1),\\\textit{q}(\\\textit{x}_0)\sim \mathcal{N}(0,1)$， 重新推导公式$\eqref{marginalvelocity}$，显然概率的部分是一个标准的贝叶斯公式，即：
$$ \begin{equation} \frac{p_t(x|x_1)q(x_1)}{p_t(x)}=p_t(x_1|x), \label{bayes} \tag{5} \end{equation} $$
那么整个积分式变为：
$$ \begin{equation} \begin{aligned} \int v_t(x|x_1) \frac{p_t(x|x_1)q(x_1)}{p_t(x)} dx_1 &= \int v_t(x|x_1)p_t(x_1|x)dx_1 \\ &= \int (\dot{f}(t)x_1 + \dot{g}(t)x_0)p_t(x_1|x)dx_1, \end{aligned} \label{int1} \tag{6} \end{equation} $$
由于$\\\textit{x}=\\\textit{f}(\\\textit{t})\\\textit{x}_1+\\\textit{g}(\\\textit{t})\\\textit{x}_0$，那么公式$\eqref{int1}$有 $$ \begin{equation} \begin{aligned} u_t(x) &= \int (\dot{f}(t)x_1 + \dot{g}(t)x_0)p_t(x_1|x)dx_1\\ &= \int \left(\dot{f}(t)x_1 + \frac{\dot{g}(t)(x-f(t)x_1)}{g(t)}\right)p_t(x_1|x)dx_1\\ &= \int \left(\left(\dot{f}(t)-\frac{f(t)\dot{g}(t)}{g(t)}\right)x_1 + \frac{\dot{g}(t)}{g(t)}x\right)p_t(x_1|x)dx_1\\ &= \left(\dot{f}(t)-\frac{f(t)\dot{g}(t)}{g(t)}\right) \underbrace{\int x_1 p_t(x_1|x)dx_1}_{\mathbb{E}[x_1|x]} + \frac{\dot{g}(t)}{g(t)}x \underbrace{\int p_t(x_1|x)dx_1}_{1}, \end{aligned} \label{int2} \tag{7} \end{equation} $$
注意到第一项积分等于$\\\textit{x}_1$在$\\\textit{x}$条件下的均值，第二项是概率密度函数的积分，显然等于$\\\textit{1}$，由于$\\\textit{p}(\\\textit{x}_0)$和$\\\textit{p}(x_1)$都是高斯分布，它们的线性混合$\\\textit{p}_t(\\\textit{x})$也是高斯分布。此时，后验均值$\mathbb{E}[\\\textit{x}_1|\\\textit{x}]$具有解析形式：
$$ \begin{equation} \mathbb{E}[x_1|x] = \frac{f(t)}{f(t)^2+g(t)^2}x, \label{conditionalestimation} \tag{8} \end{equation} $$
于是经过一些暴力的计算，我们就得到了边缘速度场$\\\textit{u}_t(\\\textit{x})$的一般形式
$$ \begin{equation} \begin{aligned} u_t(x) &= \left(\dot{f}(t)-\frac{f(t)\dot{g}(t)}{g(t)}\right) \left(\frac{f(t)}{f(t)^2+g(t)^2}x\right) + \frac{\dot{g}(t)}{g(t)}x \\ &= \left[ \frac{f(t)\dot{f}(t)}{f(t)^2+g(t)^2} - \frac{f(t)^2\dot{g}(t)}{g(t)(f(t)^2+g(t)^2)} + \frac{\dot{g}(t)}{g(t)} \right] x \\ &= \left[ \frac{f(t)\dot{f}(t)}{f(t)^2+g(t)^2} + \frac{\dot{g}(t)(-f(t)^2 + (f(t)^2+g(t)^2))}{g(t)(f(t)^2+g(t)^2)} \right] x \\ &= \left[ \frac{f(t)\dot{f}(t)}{f(t)^2+g(t)^2} + \frac{g(t)^2\dot{g}(t)}{g(t)(f(t)^2+g(t)^2)} \right] x \\ &= \frac{f(t)\dot{f}(t)+g(t)\dot{g}(t)}{f(t)^2+g(t)^2} x, \end{aligned} \label{marginalvelocityexpression} \tag{9} \end{equation} $$
Accelerated Rectified Flow
上述的推导中，$\\\textit{u}, \\\textit{x}$都是向量，我们暂时不要纠结那么多符号上的问题，只来关注这个形式。对于ReFlow而言，我们采用的瞬时速度$\\\textit{u}_t(\\\textit{x}|\\\textit{x}_1)$是一个匀速直线运动，这意味着，我们希望$\\\textit{u}_t(\\\textit{x})$关于$\\\textit{t}$是一个常数，这能否实现呢。
非常显然，这是无法实现的，当我们采用ReFlow的条件，令
$$ \begin{equation} f(t)=t, \\\quad g(t)=1-t, \label{reflow} \tag{10} \end{equation} $$
代入边缘速度场的公式$\eqref{marginalvelocityexpression}$，我们能够得到
$$ \begin{equation} u_t=\frac{2t-1}{2t^2-2t+1}x, \label{reflowmarginalvelocity} \tag{11} \end{equation} $$
显然无法得到一个常数。
退而求其次，我们看看这个$\\\textit{u}_t$是否能够是个关于$\\\textit{t}$的线性函数，也就是考虑
$$ \begin{equation} \frac{f(t)\dot{f}(t)+g(t)\dot{g}(t)}{f(t)^2+g(t)^2} x = (kt+b)x, \tag{12} \end{equation} $$
我们又知道，在Flow Matching中，由Fokker-Planck方程定义了一个连续性条件，它在这里的直观表现就是$\\\textit{f}(\\\textit{t})+\\\textit{g}(\\\textit{t})=1$，代入之后，经过复杂的微分方程计算，我们能够得到$\\\textit{f}(\\\textit{t})$的一个具体形式：
$$ \begin{equation} f(t)=\frac{1}{2}(1-\sqrt{2e^{kt(1-t)}-1}), \\\quad t \\\in [0,0.5] \\\quad or \\\quad f(t)=\frac{1}{2}(1+\sqrt{2e^{kt(1-t)}-1}), \\\quad t \\\in [0.5,0.1] \tag{13} \end{equation} $$
我们取第一个解，然后将$\\\textit{t}$替换为$\\\frac{\\\textit{t}}{2}$，这样，取值范围可以回到$[\\\textit{0},\\\textit{1}]$。于是有
$$ f(t)=-\frac{1}{2}(1-\sqrt{2e^{Ct(2-t)}-1}), \tag{14} $$
求它的二阶泰勒展开，我们能够得到
$$ f(t)=\frac{1}{2}Ct(2-t), \\\label{arflow} \tag{15} $$
结果非常惊人，它代表着，我们的条件速度场的定义式是一个匀加速直线运动！类似地，它的一阶泰勒公式等于$-\frac{\\\textit{C}}{\\\textit{2}}\\\textit{t}$，刚好是ReFlow的形式。但是非常显然，一阶泰勒展开的精度并不如二阶，且我们只要求边缘速度场是线性，而并不能保证它是常数，这就意味着ReFlow得到的结果与边缘速度场的误差永远大于二阶泰勒的定义。我们称这个二阶泰勒为Accelerated Rectified Flow。
公式$\eqref{arflow}$告诉我们，如果让条件速度场是一个匀加速场，那么得到的边缘速度场也可以是一个线性匀加速场，且速度变化率与条件速度场完全一致！在实验中我们也验证了这个结果，在经典的8gaussian2moons问题中，我们得到的条件速度场与边缘速度场非常吻合！
事实上，这个结果也是我在实验中偶然发现的，在采样的过程中，原有的flow的速度大小总是在最终阶段突然下降。而将ReFlow改成了减速运动后（加噪过程加速，去噪过程就是减速），得到的速度大小竟然出奇的吻合。这也很好理解，在生成初期，模型主要开始构建轮廓特征等，因此速度是比较大的，后面精细化纹理时，速度理应是比较小的。这和我们理论推导的结果也是不谋而合。
值得一提的是，以上这些都是假设初末分布都是标准高斯的情况，在实际的训练过程中，这是没法满足的。在低维数据中，其实这种影响并不大，因为模型的参数量尚且足够拟合出一个比较好的路径。然而只要到达了高维，模型就没法拟合得十分精确了，生成的图像会很模糊。因此在高维生成的时候，现有的方法通常会加入VAE。早期的一些解释都认为是VAE的特征压缩功能使得模型能够学习到一个比较好的结果。但是这其实是不准确的，已经有工作指出VAE的压缩性能十分有限，并不具备质变的潜力。实际上，根据我们得到的结果，加入VAE的效果好的原因其实在于VAE强制让源数据分布趋近于了标准高斯，这样就满足了我们的条件，这个时候边缘速度场和条件速度场之间就产生了显式的关联，生成结果自然就变好了。
ARCFlow: Accelerated Rectified Consistency Flow
这时候可能有人会有一些疑问，讲了这么多，一开始说的如何在瞬时速度场里面找到平均速度场的解析形式还是没有讲啊？现在我们已经得到了一个吻合相当好的匀变速边缘速度场。高中物理告诉我们，匀变速直线运动的质点，在任意时间段内，平均速度等于中间时刻的瞬时速度。那么这是不是意味着我们如果取中间时刻的瞬时速度作为平均速度的值，是否就能够直接得到解析形式的MeanFlow恒等式了呢？
现在，让我们考虑MeanFlow恒等式，根据论文中的表述，它就等于： $$ \begin{equation} u_{\\\text{tgt}} = v(z_t, t) - (t - r) \\\left(v(z_t, t) \\\partial_{z_t} u_{\theta}(z_t, t, r) + \\\partial_t u_{\theta}(z_t, t, r)\\\right), \tag{16} \label{meanflowidentity} \end{equation} $$
$\\\textit{u}_{\theta}(\\\textit{z}_t,\\\textit{r})$就是$[\\\textit{t},\\\textit{r}]$时间段内的平均速度。这里的$\\\textit{v}(\\\textit{z}_t,\\\textit{t})$，原论文中采用了一种折中方案，认为$\\\textit{v}(\\\textit{z}_t,\\\textit{t})$就等于条件速度场，也就是$\dot{\\\textit{f}}(\\\textit{t})-\dot{\\\textit{g}}(\\\textit{t})$，在ReFlow的形式下，就等于$\\\textit{x}_1-\\\textit{x}_0$。实际上，这种设计就是边缘速度场的一阶泰勒。而我们的方案也认为瞬时的条件速度场等于边缘速度场，但是我们采用了二阶泰勒，显然误差是小于一阶的。 此外，虽然 MeanFlow做出的这个约束有着十分严谨的定义，而不是CMs的很工程化的思维，但是它有一个最致命的问题，就是没办法计算右侧的两项偏导，只能通过雅各比计算梯度，在模型训练的时候又要计算一次梯度，即便原论文用了stop_gradient来减小计算量，雅各比矩阵的计算仍然是一笔不小的开销。这其中最重要的原因就是没办法认为条件平均速度场等于边缘平均速度场，只能认为边缘瞬时速度场等于条件瞬时速度场，而这两者之间根据ReFlow的定义是无法建立起联系的。那么这时候，ARFlow就派上用场了，我们可以用它构造一个新的一致性条件，前文已经证明了，我们的边缘瞬时速度场的线性性质足够好，因此可以认为一段时间内的边缘平均速度就等于中间时刻的瞬时速度： $$ \begin{equation} u_{\theta}(z_t, t, r) = v_{\theta}(z_{\frac{t+r}{2}}, \frac{t+r}{2}) = \frac{1}{2}(v_{\theta}(z_t, t) + v_{\theta}(z_r, r)), \tag{17} \label{arcflowidentity} \end{equation} $$
那么这个速度应该等于$\frac{1}{2}(\dot{\\\textit{f}}(\\\textit{t})-\dot{\\\textit{g}}(\\\textit{t}) + \dot{\\\textit{f}}(\\\textit{r})-\dot{\\\textit{g}}(\\\textit{r}))$，将$\eqref{arflow}$代入，我们就能够得到这个平均速度的表达式：
$$ \begin{equation} u_{\theta}(z_t, t, r) = (x_1-x_0)(2-t-r) \tag{18} \\\label{arcmeanvelocity} \end{equation} $$
由于$\\\textit{x}_t$关于$\\\textit{x}_1,\\\textit{x}_0,\\\textit{t}$的表达式是已知的，因此我们最后就能够求出整个$\eqref{meanflowidentity}$的右侧的两个偏微分的解析式，MeanFlow恒等式就可以转化为为：
$$ v_{\theta}(z_r, r) = \frac{ u_{\theta}(z_r, t, r) + (t - r) (x_1 - x_r) \left(\frac{2t + r - 3}{(r - 1)^3}\right) }{ 1 - (t - r) \left(\frac{t + r - 2}{(r - 1)^2}\right) } \tag{19} \\\label{arcflowtarget} $$
这里我们为了方便解释，将求导的对象改成了$\\\textit{r},\\\textit{r}<\\\textit{t}$，这样$\\\textit{u}_{\theta}(\\\textit{z}_r,\\\textit{t},\\\textit{r})$的物理意义就是从$\\\textit{z}_r$开始走到时间$\\\textit{t}$的平均速度。等式的右边所有变量已知，等式的左边就等于条件速度场，也是已知。这样就能够直接进行训练，而不需要雅各比、 stop_gradient、ema等一系列操作了。以上这个式子我们称为Accelerated Rectified Consistency Flow (ARCFlow)，它通过加速运动内在地整流了边缘速度场，使得其有较好的线性性质，换言之，更加接近了条件速度场。同时，线性性质又使得平均速度能够被计算，从而降低了计算量，减轻训练负担。后续的实验中，我们会逐步跟进实验结果，目前这个研究还停留在理论阶段，仍然存在一些训练不稳定，考虑在后续的实验中继续完善。Alea jacta est!，静待花开。