Beyond Trajectory Matching: ReFlow with Marginal Distribution Alignment¶

会议: NeurIPS 2026
arXiv: 2606.29287
OpenReview: https://openreview.net/forum?id=5a2Wi08KqE
代码: https://github.com/hansyan/PeRFlow
领域: 图像生成
关键词: 连续流生成模型, 轨迹蒸馏, 边际分布对齐, 少步采样, 概率流ODE

一句话总结¶

针对传统 ReFlow 仅依赖轨迹匹配导致端点边际分布欠定（under-determined）的理论缺陷，本文提出一种免额外训练网络的边际分布对齐正则项（MA-Reflow），通过伴随 ODE 追踪对数密度变化并结合冻结教师模型的 score 函数，严格约束流模型端点分布，实现高质量少步（4步）图像生成。

研究背景与动机¶

连续正规化流（Continuous Normalizing Flows, CNF）与扩散模型的概率流常微分方程（PF-ODE）为高质量图像合成提供了确定性的动力学框架。然而，由于速度场在时空维度的高度非线性与复杂曲率，标准数值求解器往往需要数十甚至上百步离散迭代才能生成高质量样本。为了加速推理，以 Rectified Flow（ReFlow）和分段流（Piecewise ReFlow, PeRFlow）为代表的蒸馏技术，通过分段直线化教师模型的轨迹，使学生网络能在极少步数（如 4 步 Euler 积分）下逼近目标采样分布，展现出极具潜力的蒸馏范式。

尽管轨迹匹配损失在实践中极大简化了采样路径，但其底层目标潜藏着一个基础性的理论漏洞：轨迹匹配损失能否完全决定学生模型诱导的生成分布？ 本文指出，在有限模型容量和有限优化预算的现实设定下，速度场的点对点匹配必然存在残差。流模型的概率密度演化由连续性方程（Liouville 方程）主导，本质上取决于速度场的空间散度（即一阶空间导数），而均方误差（MSE）形式的轨迹匹配损失仅约束速度向量本身的 $L_2$ 范数，无法对散度施加直接监督。作者通过数学构造严格证明，存在两个速度场具有完全相同的轨迹匹配损失，但诱导出的端点边际分布截然不同（一个导致均值偏移，另一个导致协方差膨胀）。这表明单纯优化轨迹匹配存在分布欠定性，极易造成生成质量与分布漂移。

为破解这一分布级欠定困境，必须在不牺牲 ODE 连续动力学与求解器灵活性的前提下引入分布层面的正则化约束。现有的对抗蒸馏或分布匹配蒸馏（如 DMD/DMD2）要么引入不稳定的判别器，要么需要额外训练动态学生 score 估计器，带来严重的训练震荡与算力开销。本文的核心 idea 是：在分段 ReFlow 的每个蒸馏时间窗口端点引入显式的边际分布对齐（Marginal Alignment）正则项，利用伴随增广 ODE 追踪学生动力学的前向对数密度变化，并直接由冻结教师的 score 函数计算 KL 散度梯度，无需引入任何判别器或额外可训练网络，从理论与算法上实现对端点分布总变差界（Total Variation Bound）的严格级联控制。

方法详解¶

整体框架¶

MA-Reflow 建立在分段 ReFlow（PeRFlow）的架构之上，将总采样时间区间 $[0, 1]$ 划分为 $K$ 个子区间 $[t_{k-1}, t_k]$。在每个子区间内，学生模型不仅要拟合连接教师起点与终点的直线位移速度（保持轨迹笔直易解），还要在时间端点 $t_k$ 处对其诱导的边际分布与教师边际分布进行精准对齐。

整体训练流程采用双目标解耦优化机制：对于每个批次数据，首先基于前向加噪边际分布抽取起点 $X \sim \pi_{t_{k-1}}$，并利用冻结的教师模型求解 ODE 获得窗口终点 $Y$；主干支路计算标准的 PeRFlow 轨迹匹配损失；与此同时，系统周期性激活边际对齐支路，通过增广学生速度场前向积分输出状态与累积散度变化，并结合冻结教师的 score 函数计算无需梯度回传教师的正则项损失，联合更新学生速度场网络参数。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["参考加噪数据<br/>X ~ π(t_{k-1})"] --> B["教师模型终点采样<br/>Y = Φ_g(X)"]
    B --> C["分段轨迹匹配<br/>直线常数速度场拟合"]
    A --> D["增广学生动力学积分<br/>状态演化与散度追踪"]
    D --> E["冻结教师 Score 引导<br/>计算端点 KL 梯度"]
    C --> F["端点边际分布对齐<br/>联合梯度反向传播更新"]
    E --> F
    F --> G["高质量 4 步少步生成器"]

关键设计¶

1. 轨迹匹配与边际分布欠定性机理分析：揭示纯速度损失对分布几何约束的失效

在 CNF 动力学中，设初始噪声分布为 $\pi_0 = \mathcal{N}(0, I_d)$，学生速度场 $\vec{f}_\theta(x, \tau)$ 驱动状态演化。根据 Liouville 连续性方程，边际密度沿轨迹的对数变化率严格遵循散度变换： $$\frac{\mathrm{d}}{\mathrm{d}\tau}\log q_{\theta,\tau}(X_\tau) = -\mathrm{div}_x \vec{f}_\theta(X_\tau, \tau)$$ 然而，常规 PeRFlow 优化的轨迹匹配损失仅计算速度差值的 $L_2$ 模长： $$\mathcal{L}_{\mathrm{TM}}^{(k)}(\theta) = \mathbb{E}_{X\sim\pi_{t_{k-1}}}\int_{t_{k-1}}^{t_k} \left\|\vec{f}_\theta\left(X+\frac{\tau-t_{k-1}}{t_k-t_{k-1}}(Y-X), \tau\right) - \frac{Y-X}{t_k-t_{k-1}}\right\|^2 \mathrm{d}\tau$$ 由于 $L_2$ 范数并不限制空间导数 $\nabla_x \vec{f}_\theta$，作者给出了理论构造（Theorem 1）：对于常数教师场 $\vec{g}_\psi \equiv \mu$，构造两个扰动速度场 $\vec{f}_1(x, t) = \mu + \varepsilon e$ 与 $\vec{f}_2(x, t) = \mu + \varepsilon \langle x, e \rangle e$（其中 $\langle e, \mu \rangle = 0, \|e\|=1$）。两者在所有窗口上的轨迹匹配损失完全相同（均为 $\varepsilon^2(t_k - t_{k-1})$），但推导出的端点分布截然不同：$\vec{f}_1$ 导致均值偏移 $\mathcal{N}(\mu + \varepsilon e, I_d)$，而 $\vec{f}_2$ 导致特定轴向的协方差扩张 $\mathcal{N}(\mu, I_d + (e^{2\varepsilon}-1)ee^\top)$，其相对真实分布的 KL 散度显著不等。这严密证明了仅靠轨迹匹配无法唯一定义正确的终点生成分布。

2. 训练协议边际分布对齐正则项：解耦递归依赖并推导级联 TV 误差上界

直接最小化端到端采样分布的 KL 散度 $\mathrm{KL}(p_{\theta, t_k} \| q_{\psi, t_k})$ 需要在整个链路上展开全部前向采样步，导致巨大的显存开销与梯度爆炸。MA-Reflow 提出基于“训练协议边际”（Training-Protocol Marginal）的局部化目标：在第 $k$ 个窗口中，输入直接采用基准前向扩散边际 $\pi_{t_{k-1}}$（与参数 $\theta$ 无关），定义推进后的分布为 $\tilde{p}_{\theta^{(k)}} = (\Phi_{\vec{f}_\theta, t_{k-1}\to t_k})_\# \pi_{t_{k-1}}$，局部的边际对齐损失定义为向教师端点边际 $q_{\psi, t_k}$ 的 KL 散度： $$\mathcal{L}_{\mathrm{MA}}^{(k)}(\theta) = \mathrm{KL}\left(\tilde{p}_{\theta^{(k)}} \,\big\|\, q_{\psi, t_k}\right)$$ 作者在 Theorem 2 中证明了级联总变差（Total Variation, TV）误差界：记第 $k$ 步教师诱导边缘与参考加噪边缘的分布偏差为 $\delta_k = \mathrm{TV}(q_{\psi, t_k}, \pi_{t_k})$，则推理阶段最终时刻累积误差满足级联不等式： $$\mathrm{TV}\left(p_{\theta, t_K}, q_{\psi, t_K}\right) \le \sum_{k=1}^K \sqrt{\frac{1}{2}\mathcal{L}_{\mathrm{MA}}^{(k)}(\theta)} + \sum_{k=1}^{K-1}\delta_k$$ 这一理论结果证明，尽管训练时将输入解耦为各窗口独立的参考边际分布，最小化局部正则项 $\mathcal{L}_{\mathrm{MA}}^{(k)}$ 仍能严格压低最终生成样本与真实教师目标之间的全局分布距离。

3. 基于增广 ODE 与冻结教师 Score 的免判别器梯度求解机制

为使 KL 散度具备高效可微性，作者基于变量代换原理推导了 Liouville 积分形式（Lemma 3），将分布层面的积分转化为在基准分布 $X \sim \pi_{t_{k-1}}$ 上的期望求值： $$\mathcal{L}_{\mathrm{MA}}^{(k)}(\theta) = \mathbb{E}_{X\sim\pi_{t_{k-1}}}\left[\log \pi_{t_{k-1}}(X) + \Delta_{t_{k-1}\to t_k}^{\mathrm{Stu}}(X; \theta) - \log q_{\psi, t_k}\left(X_{t_{k-1}\to t_k}^{\mathrm{Stu}}(X; \theta)\right)\right]$$ 其中对数密度变化量 $\Delta_{t_{k-1}\to t_k}^{\mathrm{Stu}}(X; \theta) = -\int_{t_{k-1}}^{t_k} \mathrm{div}_x \vec{f}_\theta(X_\tau, \tau)\mathrm{d}\tau$。对参数 $\theta$ 求导后，基准先验项梯度为零，而教师边际项的梯度恰好通过链式法则转化为教师模型的对数梯度——即 score 函数 $s_\psi(x, t_k) = \nabla_x \log q_{\psi, t_k}(x)$。由此推导出无需对教师反向传播的 stop-gradient 训练目标（Theorem 4）： $$\nabla_\theta \mathcal{L}_{\mathrm{MA}}^{(k)}(\theta) = \nabla_\theta \mathbb{E}_{X\sim\pi_{t_{k-1}}}\left[\Delta_k^{\mathrm{Stu}} - \mathrm{stopgrad}\left(s_\psi(X_k^{\mathrm{Stu}}, t_k)\right)^\top X_k^{\mathrm{Stu}}\right]$$ 在工程实现中，将状态演化与散度积分打包为增广速度场向量 $\vec{f}_\theta^{\mathrm{aug}}((x, c), t) = (\vec{f}_\theta(x, t), -\mathrm{div}_x \vec{f}_\theta(x, t))^\top$，采用 4 阶 Runge-Kutta（RK4）数值积分器向前积分。对高维散度 $\mathrm{div}_x \vec{f}_\theta$，采用 Hutchinson 随机迹估计配合有限差分逼近：$\nabla_x \vec{f}_\theta(x, t)\epsilon \approx [\vec{f}_\theta(x+\sigma\epsilon, t) - \vec{f}_\theta(x, t)] / \sigma$，仅需单次扰动即可无偏估计散度，全程无需任何辅助网络或对抗判别器。

损失函数 / 训练策略¶

MA-Reflow 的完整联合优化目标为： $$\mathcal{L}_{\mathrm{total}}(\theta) = \mathcal{L}_{\mathrm{TM}}(\theta) + \lambda \mathcal{L}_{\mathrm{MA}}(\theta)$$ 在训练策略与超参数设置上： - 分段配置：统一采用 $K=4$ 个时间窗口。在各窗口内，教师模型采用 8 步 DDIM 采样（等价于重标定坐标系下的 Euler 步）生成目标终点 $Y$。 - 交替计算加速：由于增广 ODE 的 RK4 求解相对昂贵，算法设置交替比率 $m=9$（即每 9 次标准 PeRFlow 步骤计算 1 次边际对齐正则项），使正则化计算仅占总训练耗时约 30%。 - 权重与数值尺度：正则项权重设为 $\lambda = 10^{-6}$（基于两项损失的梯度模长量级匹配）；有限差分扰动步长 $\sigma = \sigma_0 / \sqrt{d}$，其中 SD 1.5 设 $\sigma_0 = 10$，SDXL 设 $\sigma_0 = 1$。 - 算力与收敛：基于 2 块 NVIDIA H200 GPU 采用 BF16 混合精度训练，在预训练 PeRFlow 权重上微调 5K 步即可收敛（约 1 天）。

实验关键数据¶

主实验¶

评估在标准 MS-COCO 验证集上展开，对比主流训练无关求解器与各前沿 4 步少步蒸馏方法（包括一致性模型 LCM/PCM、对抗蒸馏 SDXL-Lightning/Turbo 以及分布匹配模型 DMD2）。

表 1: Stable Diffusion v1.5 在 MS-COCO 上的生成质量评估

类别 / 方法	采样步数	COCO-2014 FID (↓)	COCO-2017 FID (↓)	COCO-2017 CLIP (↑)
训练无关求解器
DPMSolver	25	9.78	20.10	31.80
DPMSolver	8	22.44	-	-
DPMSolver++	4	22.36	-	-
DDIM (Teacher)	32	10.05	21.98	31.22
蒸馏少步生成模型
LCM-LoRA	4	23.62	36.46	29.10
2-ReFlow	4	15.32	32.97	28.93
Flash Diffusion	4	12.41	-	-
LCM	4	-	24.39	30.50
PeRFlow (基线)	4	12.01	23.81	30.24
PCM	4	11.70	-	-
MA-Reflow (本文, K=4)	4	10.66	22.76	31.03

在 SD 1.5 主实验中，MA-Reflow 在 4 步采样下将 COCO-2014 FID 从 PeRFlow 的 12.01 显著降低至 10.66（提升 1.35），超越了 PCM 与 Flash Diffusion，接近 32 步 DDIM 教师模型的表现（10.05）；在 COCO-2017 上 FID 从 23.81 降至 22.76，CLIP 图文对齐得分提升至 31.03。

表 2: SDXL 在 MS-COCO-2014-10K 与 COCO-2017 上的 4 步采样性能对比

模型与蒸馏方法	采样步数	COCO-2014-10K FID (↓)	COCO-2017 FID (↓)
SDXL-Lightning	4	24.56	30.16
SDXL-Turbo	4	23.19	30.92
HyperSD	4	22.95	30.58
LCM	4	22.16	27.09
PCM	4	21.04	-
PeRFlow	4	20.99	27.06
Rectified Diffusion	4	19.71	25.81
DMD2	4	19.32	26.19
MA-Reflow (本文, K=4)	4	19.17	25.94

在更大规模的 SDXL 架构上，MA-Reflow 展现出优异扩展性：在 4 步推理下取得 19.17 的 COCO-2014-10K FID，位列所有 4 步蒸馏方案第一，优于需要复杂对抗/分布近似的 DMD2（19.32）与 SDXL-Lightning（24.56）。

消融实验¶

为验证边际对齐正则项的有效性，作者在 CelebA-HQ $256 \times 256$ 数据集上使用潜空间流匹配（LFM）模型进行了不同初始化条件下的消融对比。

表 3: CelebA-HQ 潜空间流匹配（LFM）初始化与持续微调消融实验（FID ↓）

设置模式	方法配置	4 步采样 FID	8 步采样 FID	核心结论与机制说明
设置 A 从教师模型权重初始化（训练 100K 步）	LFM (Direct 无蒸馏)	54.41	21.29	零训练下直接少步采样失真严重
	PeRFlow	18.17	12.87	纯轨迹匹配使 4 步生成可用
	MA-Reflow	17.25	13.04	加入边际对齐后 4 步 FID 明显改善
设置 B 从收敛的 PeRFlow 权重继续微调（额外 100K 步）	PeRFlow (继续微调)	23.87	13.85	轨迹匹配已过拟合，指标显著恶化（+5.70）
	MA-Reflow (继续微调)	16.78	11.98	打破瓶颈，4 步与 8 步指标持续刷新最优

关键发现¶

纯轨迹匹配的性能饱和与退化：在设置 B 中，在 PeRFlow 收敛检查点上继续执行轨迹匹配优化，4 步 FID 不降反升（从 18.17 退化至 23.87），验证了模型容量受限时，轨迹位移误差无法再给生成分布提供有效增益；而 MA-Reflow 成功打破该瓶颈，进一步压低 FID 至 16.78。
超参敏感性鲁棒：在 SD 1.5 上的参数扫描表明，正则项系数 $\lambda \in [10^{-7}, 10^{-5}]$ 区间内均稳定超越 PeRFlow 基线，默认值 $\lambda = 10^{-6}$ 表现最佳；同时该正则项在分类器自由引导（CFG）系数 $w \in [2, 7]$ 的大范围扫描中保持极佳稳定性。

亮点与洞察¶

抓住轨迹匹配的本质短板：直击 ReFlow 依赖速度 $L_2$ 损失无法约束空间散度与连续性方程的理论漏洞，构造了同等轨迹误差下边缘分布显著发散的严格反例，数学直觉深刻。
巧妙规避辅助网络：无需像对抗蒸馏那样训练易崩溃的判别器，也无需像 DMD 那样训练滞后的学生 score 估计网络，直接借助增广 ODE 与冻结教师的 score 函数实现闭式梯度回传，兼顾了理论自洽与工程优雅性。
保持连续流动力学特性：区别于离散跃迁的一致性模型（Consistency Models），MA-Reflow 的产物依然是连续时间速度场，不仅保留了与任意可变步长 ODE 求解器的兼容性，还能自然享用未来更高阶数值算法的红利。

局限与展望¶

计算代价与积分步骤依赖：训练期间为了追踪对数密度，必须利用 RK4 数值积分器求解增广 ODE，即使用了交替更新（每 9 步调 1 次），依然占用了约 30% 的额外训练时长。
极端单步（1-Step）生成的未涉足：目前主实验集中在 4 步与 8 步的分段设置（$K=4$），对于 $K=1$ 的极限单步生成尚未进行大规模实验验证。
对非 ODE 架构的不可迁移性：该框架高度依赖概率流 ODE 的连续变化率公式（Liouville 散度定理），无法直接推广至不可微跃迁的跳步或离散采样架构。

评分¶

新颖性: ⭐⭐⭐⭐☆ [揭示 ReFlow 轨迹匹配的分布欠定性反例，理论机制清晰且具有启发性]
实验充分度: ⭐⭐⭐⭐☆ [覆盖 SD 1.5、SDXL 及 CelebA-HQ，涵盖不同初始化机制消融与敏感度分析]
写作质量: ⭐⭐⭐⭐⭐ [数学推导严丝合缝，定理与反例构造清晰，整体逻辑环环相扣]
价值: ⭐⭐⭐⭐☆ [为扩散模型与流模型的高效少步蒸馏提供了一种免外挂网络的新范式，实用性强]