跳转至

DuoFlow: JVP-Free Finite-Difference Mean Flows for One-Step Image Generation

会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/stepfun-ai/DuoFlow
领域: 图像生成
关键词: 单步生成、MeanFlow、有限差分、无JVP、轨迹一致性

一句话总结

针对 MeanFlow 依赖高开销雅可比-向量积(JVP)导致训练难以扩展的瓶颈,DuoFlow 提出误差驱动的无 JVP 有限差分框架,通过随机符号单边差分控制步长截断误差、同状态联合估计与渐进自引导对齐轨迹速度误差,在 ImageNet \(256 \times 256\) 从头训练下将单步 FID 降低约 30%。

研究背景与动机

单步图像生成是扩散与流匹配模型迈向实时部署的关键方向。相比需要几十乃至数百步迭代数值采样的传统方案,单步模型能极大压缩推理延迟与算力开销。在这一探索中,MeanFlow 通过建立耦合区间均值速度与瞬时速度的恒等式,为无教师蒸馏、从头训练大步长传输网络提供了严格的数学支撑。然而在实际训练中,MeanFlow 必须沿着模型预测的轨迹速度计算方向导数以强制满足该恒等式,主流实现全面依赖自动微分中的雅可比-向量积(JVP)。在现代分布式训练与编译器技术栈中,JVP 引入了显著的额外前向/反向图开销,严重阻碍算子融合(如 FlashAttention/Fused Kernels)与编译器优化,使其极难扩展到更大规模模型与更高分辨率。

摆脱对 JVP 的依赖看似可以直接采用经典数值分析中的有限差分法,但简单的算子替换在 MeanFlow 中极易导致训练发散或生成质量严重坍塌。根本原因在于 MeanFlow 的方向导数并非沿着固定外生场计算,而是沿着模型自身正在学习的动态轨迹速度前进——监督信号的准确性同时受制于数值差分的步长截断以及轨迹速度本身的预测漂移。这两个误差通道在训练中相互交织,导致确定性差分带来的系统偏差随迭代累积,将参数推向损失曲面的高曲率恶化区域。

本文的切入角度是摒弃将差分视作黑盒算子替换的朴素思路,转向显式的误差驱动机制设计。作者将总差分监督误差正交解耦为步长截断误差与轨迹速度误差两大通道,并据此设计针对性的算法消除机制。核心 idea:通过仅需一次额外前向的随机带符号单边差分在期望上消除一阶截断偏差,并在同一采样状态下联合预测均值与瞬时速度配合渐进自引导对齐,构建高效、稳定且无 JVP 的自洽均值流训练框架。

方法详解

整体框架

DuoFlow 的核心在于保持 MeanFlow 恒等式监督形式不变的前提下,用分析驱动的差分估计器彻底替代 JVP。其训练流程在每一个优化批次仅需两次标准的前向传播,完全保留在常规的前向/反向图计算范式中,天然兼容硬件算子融合。模型在采样状态 \(z_t\) 下同时输出瞬时速度 \(v_\theta\) 与区间均值速度 \(u_\theta\),利用扰动后的状态进行单次额外前向评估方向导数,再借助余弦对齐驱动的动态混合因子 \(\kappa\) 构造自引导监督目标。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入潜变量状态 z_t 与时间 (r, t)"] --> B["同状态联合速度估计<br/>单网络共享参数输出 u_θ 与 v_θ"]
    B --> C["随机符号单边差分<br/>构造扰动状态并前向评估方向导数"]
    C --> D["渐进自引导对齐<br/>基于批次余弦相似度计算权重 κ 融合目标"]
    D --> E["分离损失归一化优化<br/>更新均值与瞬时速度参数"]

关键设计

1. 误差驱动正交分解:显化差分监督的失真源头 在理想的 MeanFlow 恒等式监督下,模型应当沿着真实最优轨迹速度 \(v^\star\) 计算区间均值速度 \(u\) 的方向全导数 \(\mathcal{D}^{v^\star}u\)。而在无 JVP 的数值差分设定中,监督信号面临两层误差逼近: $$ \hat{\mathcal{D}} u - \mathcal{D}^{v^\star} u = \underbrace{(\hat{\mathcal{D}} u - \mathcal{D}^{v_\theta} u)}{E}}} + \underbrace{(\mathcal{D}^{v_\theta} u - \mathcal{D}^{v^\star} u){E $$ 其中 }}\(E_{\text{trunc}}\) 为有限差分步长带来的步长截断误差,而 \(E_{\text{traj}}\) 为采用未成熟的模型自身速度 \(v_\theta\) 替代最优速度 \(v^\star\) 引发的轨迹速度误差。这一数学分解指明了算法设计的唯一抓手:若单纯压小步长 \(h\) 降低截断误差,若轨迹速度 \(v_\theta\) 自身漂移,总体导数依然严重偏离;反之,若差分带有确定性偏置,累积误差会摧毁轨迹学习。

2. 随机符号单边差分:以单次前向换取期望二阶截断收敛 传统二阶中心差分需要两次额外前向评估,计算开销较大;而常规确定性前向差分虽仅需一次前向,其截断误差中带有不可消除的一阶系统性偏置 \(\frac{h}{2}(\mathcal{D}^{v_\theta})^2 u_0\)。这种相干偏差会在长程训练迭代中持续累积,导致模型跑偏甚至发散。DuoFlow 引入随机符号单边差分,按均匀分布采样随机符号 \(\sigma \sim \text{Unif}\{+1, -1\}\),并根据机器精度自适应确定扰动基础步长 \(h_{\text{base}}\)(如 BF16 下设为 \(\sqrt{\varepsilon_{\text{mach}}}\))以及尺度归一化因子,构造状态扰动 \(z_{t+\sigma h} = z_t + \sigma h v_\theta(z_t, t)\)。其方向导数估计式为: $$ \hat{\mathcal{D}}h^\sigma u\theta = \frac{u_\theta(z_{t+\sigma h}, r, t+\sigma h) - u_\theta(z_t, r, t)}{\sigma h} = \mathcal{D}^{v_\theta}u_0 + \frac{\sigma h}{2}(\mathcal{D}^{v_\theta})^2 u_0 + \mathcal{O}(h^2) $$ 在对随机符号 \(\sigma\) 求期望后,\(\mathcal{O}(h)\) 阶线性偏差精确抵消,使估计器在期望意义下享有 \(\mathcal{O}(h^2)\) 的二阶截断精度: $$ \mathbb{E}\sigma \left[ \hat{\mathcal{D}}_h^\sigma u\theta \right] = \mathcal{D}^{v_\theta} u_\theta + \mathcal{O}(h^2) $$ 零均值的随机扰动方差在 mini-batch 训练中会被自然平滑平均,以极低代价消除了长期累积偏差。

3. 同状态联合速度估计:以共享几何约束收敛轨迹误差 为了压制 \(E_{\text{traj}}\),DuoFlow 抛弃了将瞬时速度场 \(v\) 与区间均值流场 \(u\) 割裂训练的常规方式,通过共享主干网络参数,仅通过时间区间条件调节(\(r=t\) 时退化为瞬时速度,一般 \(r<t\) 时为区间均值速度),使两者在同一个输入采样点 \(z_t\) 下完成联合推断。由于两类速度预测受同一套局部几何特征与隐层表征约束,瞬时速度的优化能够即时反馈并改善有限差分扰动所依赖的切线方向,显著抑制了速度场不一致带来的轨迹漂移。

4. 渐进自引导对齐与分离损失归一化:构建自洽不动点动力学 在训练初期,模型自预测的速度场充满噪声,完全依赖自预测会导致误差雪崩;而若始终依赖外生目标(如真实数据差值 \(\dot{z}_t = \epsilon - x_0\)),轨迹动力学又无法与模型自洽。DuoFlow 设计了基于余弦相似度自适应门控的混合机制:计算批次内模型预测速度 \(v_\theta\) 与外生/CFG目标速度 \(v_{\text{ext}}\) 之间的余弦相似度,动态确定权重 \(\kappa = 0.5 + 0.5 \cdot \frac{\bar{d}}{\sqrt{\max(\bar{q}_v \bar{q}_t, \epsilon)}}\),从而得到平滑演进的目标速度 \(v_{\text{target}} = \kappa \, \text{sg}[v_\theta] + (1-\kappa) v_{\text{ext}}\)。随着训练进行,\(\kappa\) 单调上升并逐渐饱和,驱动轨迹朝向自洽不动点收敛。为避免 \(\kappa\) 变化引起损失尺度失衡,优化采用分离损失归一化(Detached-loss Normalization),平衡瞬时速度与均值速度的梯度更新。

损失函数 / 训练策略

定义瞬时速度与均值速度的监督目标与损失: $$ u_{\text{target}} = v_{\text{target}} - \text{sg}[\Delta_\theta], \quad \Delta_\theta = (t-r)\hat{\mathcal{D}}h^\sigma u\theta $$ $$ \mathcal{L}u = |u\theta - u_{\text{target}}|2^2, \quad \mathcal{L}_v = |v\theta - v_{\text{target}}|_2^2 $$ 引入分离损失归一化 \(\bar{\mathcal{L}}_u = \frac{\mathcal{L}_u}{\text{sg}[\mathcal{L}_u] + \epsilon}\)\(\bar{\mathcal{L}}_v = \frac{\mathcal{L}_v}{\text{sg}[\mathcal{L}_v] + \epsilon}\),最终优化总损失为: $$ \mathcal{L} = \bar{\mathcal{L}}_u + w_v \bar{\mathcal{L}}_v, \quad w_v = \frac{1}{(1 - \bar{\kappa})^2} $$ 其中 \(\text{sg}[\cdot]\) 表示梯度截断(stop-gradient)。全流程仅需常规 Adam 优化器,学习率设为 \(2 \times 10^{-4}\),配合 EMA(衰减率 0.9999)完成训练。

实验关键数据

主实验

在 ImageNet \(256 \times 256\) 类别条件图像生成任务上,DuoFlow 与多类从头训练的扩散/流匹配基线进行了严格匹配的性能对比(无额外教师蒸馏,统一使用相同的预训练 VAE 分词器,主线模型训练 240 epochs)。

模型类别 / 方法 参数量 采样步数 (NFE) FID (↓) 相比 MeanFlow 相对提升
iCT-XL/2 (从头一致性训练) 675M 1 34.24 -
Shortcut-XL/2 675M 1 10.60 -
MeanFlow-B/2 (JVP基线) 131M 1 6.17 基线
DuoFlow-B/2 (本文) 131M 1 5.19 -15.9%
MeanFlow-M/2 (JVP基线) 308M 1 5.01 基线
DuoFlow-M/2 (本文) 308M 1 4.72 -5.8%
MeanFlow-L/2 (JVP基线) 459M 1 3.84 基线
DuoFlow-L/2 (本文) 459M 1 3.14 -18.2%
MeanFlow-XL/2 (JVP基线) 676M 1 3.43 基线
DuoFlow-XL/2 (本文) 676M 1 2.48 -27.7%
DuoFlow-XL/2 (640 epochs) 676M 1 1.99 -
MeanFlow-XL/2 (2-NFE) 676M 2 2.93 基线
DuoFlow-XL/2 (2-NFE) 676M 2 1.96 -33.1%

注:数据摘自原论文 Table 1。在 XL/2 架构下,DuoFlow 1-NFE 达到 2.48 FID,大幅超越同等算力下 MeanFlow 的 3.43 FID;当扩展至 640 epochs 时更是达到 1.99 FID。

消融实验

在 80 epochs、学习率 \(4 \times 10^{-4}\) 的匹配快速消融基准下,系统拆解各模块对 ImageNet \(256 \times 256\) (1-NFE) 的贡献:

配置 / 变体 随机符号差分 联合速度估计 渐进自引导 FID (↓) 相对变化说明
MeanFlow (JVP 基线复现) - - - 14.13 原始 JVP 监督
DuoFlow 确定性单边差分 19.06 存在 \(\mathcal{O}(h)\) 累积偏差,性能剧烈恶化 (-34.9%)
DuoFlow 随机带符号差分 14.11 消除线性偏差,追平 JVP 基线 (+0.0%)
DuoFlow + 联合速度估计 11.30 显著压低轨迹速度漂移 (+20.0%)
DuoFlow (完整模型) 9.17 形成自洽不动点,相对 JVP 提升 +35.1%

注:数据摘自原论文 Table 3。

此外,论文在 Table 2 中对方向导数计算算子进行了微观评测:在 Batch Size 为 1 至 64 时,单次额外前向的差分代理比 PyTorch 的 torch.func.jvp 加速 \(7.17\times\)\(34.83\times\),显存峰值降低 \(1.71\times\)\(2.68\times\)。在 8 卡 H800 的端到端训练中,DuoFlow 训练迭代耗时由 MeanFlow 的 0.89s/iter 缩减到 0.35s/iter,实现 \(2.54\times\) 端到端挂钟加速,峰值显存由 97GB 降至 73GB。

关键发现

  • 随机符号差分是数值稳定性的定海神针:确定性单边差分由于固有的 \(\mathcal{O}(h)\) 相干截断误差,在训练中期会导致参数冲入高曲率危险区乃至发散(FID 仅能达到 19.06);随机符号变换在数学上将期望误差压至 \(\mathcal{O}(h^2)\),通过 batch 随机抵消使得差分逼近在长时间训练中保持稳定。
  • 早期方差敏感 vs 后期偏差敏感:实验发现,在训练极早期确定性差分与 JVP 的余弦相似度反而略高(因为此时速度场自身误差大、确定性差分方差较小);但随着训练深入,系统性截断偏置成为主导瓶颈,随机符号差分全面反超并持续逼近 JVP 真值。
  • 联合估计与自引导的协同互补:联合网络架构将均值与瞬时速度锁定在同一局部流形上,使导数评估与 JVP 的余弦相似度全阶段显著高于独立网络;而自适应 \(\kappa\) 避免了固定 \(\kappa=0.9\) 导致的早期自欺欺人崩溃(固定 0.9 会导致 FID 坍塌至 35.33)。

亮点与洞察

  • 分析驱动而非黑盒替换:没有将有限差分当成一个简单的“无梯度近似算子”,而是深入 MeanFlow 导数耦合特质,建立 \(E_{\text{trunc}} + E_{\text{traj}}\) 显式误差预算,指引出了极简却精妙的算法组合。
  • 期望抵消一阶项的巧思:利用 mini-batch 梯度下降天然平滑零均值方差的特点,以单次额外前向加符号随机扰动 \(\sigma \in \{+1, -1\}\),用 \(\mathcal{O}(1)\) 的计算增量实现了本需双倍前向的二阶精度特性。
  • 软硬件系统层面的深度解锁:将算法彻底拉回常规的前向/反向执行图,彻底释放了 PyTorch 2.x 编译器及 Fused FlashAttention 等现代底层加速利器的全部潜力。

局限与展望

  • 适用场景边界:当前验证主要集中在 ImageNet \(256 \times 256\) 类别条件潜空间生成,尚未扩展至更高分辨率(如 \(512 \times 512\) / \(1024 \times 1024\))、大规模文生图(Text-to-Image)或视频生成等复杂多模态骨干中。
  • 对局部流形光滑性的隐式依赖:泰勒展开的一阶项消除建立在速度场充分光滑的理论假设之上,在极端多模态冲突或强非线性边界处,截断高阶项可能带来瞬时数值震荡。
  • 未来方向:探索结合 REPA 表征对齐与 LightningDiT 结构优化,以及将误差驱动差分原则推广至其他高阶流匹配(如 iMF / Consistency Models)中。

相关工作与启发

  • vs MeanFlow [Geng et al., 2025]: MeanFlow 奠定了均值流与瞬时速度恒等式的理论基础,但使用 JVP 作为黑盒求导工具,显存占用大且阻碍编译器优化。DuoFlow 证明通过精细的误差分解与随机符号差分,完全可以抛弃 JVP 并取得比原始 JVP 更佳的生成质量(FID 从 3.43 降至 2.48)与 \(2.54\times\) 端到端加速。
  • vs iCT / Shortcut Models: 传统从头单步生成模型(如 Shortcut Models)在单步下 FID 往往在 10 以上,而 DuoFlow 在 1-NFE 下达到 2.48,显著拉近了从头单步生成与多步 ODE 采样的质量鸿沟。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ [首次从截断与轨迹双误差通道解构 MeanFlow 方向导数,提出期望二阶收敛的随机符号单边差分]
  • 实验充分度: ⭐⭐⭐⭐⭐ [覆盖算子微观测试、端到端训练吞吐、消融与敏感性分析,数据详实严谨]
  • 写作质量: ⭐⭐⭐⭐⭐ [数学推导清晰优美,机理解释透彻深入,逻辑层层递进]
  • 价值: ⭐⭐⭐⭐⭐ [彻底移除大模型训练中恶心棘手的 JVP 依赖,为大规模单步流生成模型的基础设施扩展扫清了工程障碍]