技术摘要：通过边界条件改进整流流（Rectified Flow）

问题陈述

整流流（Rectified Flow, RF），也称为流匹配（Flow Matching），已成为一种强大的生成模型框架，它通过学习一个速度场（velocity field）来将噪声分布转化为目标数据分布，并通过常微分方程（ODE）实现这一过程。尽管在文本生成图像和视频生成等任务中取得了成功，但作者指出标准 RF 实现中存在一个关键局限性：学习到的速度场往往无法满足理论推导的边界条件。

具体而言，在噪声（$t=0$）与数据（$t=1$）之间的线性插值下，最优速度场应满足：

1. 左边界 ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. 右边界 ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

标准的神经网络参数化速度场并不能内在强制执行这些约束。正如论文中所展示的可视化结果，在 $t=1$ 处的预测速度往往偏离预期的恒等映射。这种违反行为在随机采样（例如使用 Langevin 动力学）过程中尤为有害。由于得分函数估计（通过 Tweedie 公式）涉及对 $(1-t)$ 的除法，靠近 $t=1$ 时不准确的速度会导致得分估计发散。这会导致采样不稳定，产生过平滑的图像或“卡通化”伪影，尤其是在使用较大的步长或噪声尺度时。

方法论

为了解决这个问题，作者提出了边界强制整流流模型（Boundary-enforced Rectified Flow Model，简称 Boundary RF Model）。该方法通过两种不同的变体，将边界约束显式地整合到速度场的参数化中：

1. 基于掩码的边界 RF 模型 (Mask-based Boundary RF Model)

该变体通过将速度场构建为边界项与可学习神经网络输出的线性组合，来同时强制执行两个边界条件（$t=0$ 和 $t=1$）：
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
其中：

• $m_\theta(\mathbf{x}, t)$ 是标准的神经网络（如 U-Net 或 DiT）。
• $C = \mathbb{E}[\mathbf{X}_1]$。
• $f(t), g(t), h(t)$ 是旨在在特定边界处消失或激活的标量函数（例如 $g(1)=f(0)=h(0)=h(1)=0$ 且 $g(0)=f(1)=1$）。
• 作者建议使用“标准余弦”（Standard Cosine）选择：$g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, 以及 $h(t) = \sin(\pi t)$。

这种构造保证了无论神经网络的输出如何，边界条件都能通过设计得到满足。

2. 基于减法的边界 RF 模型 (Subtraction-based Boundary RF Model)

受“仅强制执行 $t=1$ 条件即可获得显著经验收益”这一观察的启发，作者提出了一个更简单的单边界变体：
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
该公式通过使 $m_\theta(\mathbf{x}, 1)$ 项在 $t=1$ 时抵消，从而内在满足了 $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$。虽然这种方法提供了设计的简洁性和灵活性，但它也带来了计算成本，因为每次评估都需要进行额外的正向传播来计算 $m_\theta(\mathbf{x}, 1)$。

对随机采样的影响

通过强制执行 $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$，边界 RF 模型确保了 Tweedie 公式中的分子（$t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$）在 $t \to 1$ 时趋于零。这防止了得分函数 $\nabla \log \rho_t$ 的发散，从而稳定了随机采样器（SDE 采样器），并允许在使用较大的噪声尺度或步长时进行稳健的生成。

核心贡献

1. 识别局限性： 论文强调了未受约束的速度场违反了理论边界条件，从而导致不准确的 ODE 轨迹和不稳定的随机采样。
2. 创新的参数化方案： 引入了 Boundary RF 模型，提供了两种实用的变体（基于掩码和基于减法），能够以极小的代码修改量强制执行这些条件。
3. 稳定得分函数： 证明了边界强制执行能直接稳定 $t=1$ 附近的得分函数，缓解了困扰原生 RF 的发散问题。
4. 经验验证： 通过综合实验，展示了在不同数据集和模型规模下，生成质量和采样稳定性的一致提升。

实验结果

作者使用 Diffusion Transformer (DiT) 架构在 CIFAR-10 和 ImageNet (256$\times$256 和 512$\times$512) 上评估了他们的模型。

• 定量改进：

  • 在 ImageNet 256$\times$256 上，基于减法的 Boundary RF 模型实现了 6.32 (ODE 采样) 和 6.38 (SDE 采样) 的 FID，而 vanilla RF 基线分别为 6.87 和 6.79。
  • 基于掩码的模型实现了 6.63 (ODE) 和 6.18 (SDE) 的 FID。
  • 在两个数据集上，IS（Inception Score）、精度（Precision）和召回率（Recall）也都观察到了提升。
  • 基于减法的变体一致获得了最高的 IS 分数。

• 采样鲁棒性：

  • 在使用具有不同噪声尺度 ($c$) 的超调采样器（Overshooting Sampler）的实验中，Boundary RF 模型表现出显著更稳定的性能，并且随着随机性的增加，其样本质量的退化程度远低于 vanilla RF。
  • 定性结果显示，Boundary RF 模型生成的图像更清晰，伪影更少，物体保真度更高，特别是在 vanilla RF 容易出现过平滑现象的随机采样设置中。

• 可扩展性：

  • 当扩展到更大规模的模型（DiT-L, DiT-XL）和更高分辨率（512$\times$512）时，Boundary RF 模型的优势依然存在。

意义与主张

该论文声称，Boundary RF 模型为 Rectified Flow 的一个基本局限性提供了一个简单、易于实现的解决方案。通过显式强制执行边界条件，该方法：

1. 提升了生成质量： 在 FID 和其他指标上一致优于 vanilla RF。
2. 实现了稳定的随机采样： 解决了终端时间附近的得分函数发散问题，从而支持更稳健的采样策略。
3. 易于适配： 可以通过极小的架构改动集成到现有的 RF 模型（如使用 U-Net 或 DiT 的模型）中。

作者将这项工作定位为对 Rectified Flow 进行实际部署时必要的精细化改进，特别是在需要高质量随机采样的场景下。他们指出，虽然本文重点在于 RF，但强制执行边界条件的核心概念可以通过微调扩展到扩散模型。未来的工作拟探索自适应边界强制执行以及在文本生成视频中的应用。