MixFlow：通过混合源分布改善 Rectified Flow MixFlow: Mixed Source Distributions Improve Rectified Flows
用条件与无条件高斯的线性混合源分布，拉直生成轨迹并显著提升 Rectified Flow 采样效率。
前置知识
Rectified Flow / Flow Matching
Rectified Flow 学习向量场 $v_\theta(x_t,t)$，把源分布 $p_0$（常为标准高斯）沿直线流到目标 $p_1$。中间点 $x_t=tx_1+(1-t)x_0$，损失 $\int_0^1\|x_1-x_0-v_\theta\|^2dt$。把扩散采样重诠释为 ODE 积分。
MixFlow 在 Rectified Flow 框架上改造源分布，必须先理解向量场目标与耦合 $q(x_0,x_1)$ 的角色，否则读不懂 κ-FC 设计。
前向耦合（Forward Coupling）
指源样本 $x_0\sim p_0$ 与目标样本 $x_1\sim p_1$ 的联合分布 $q(x_0,x_1)$，默认独立 $q=p_0 p_1$。耦合决定如何配对 $(x_0,x_1)$，进而决定生成轨迹的弯曲程度。
MixFlow 与 Fast-ODE 的核心都是"优化耦合以拉直路径"，$I(q)$、$C(v_\theta)$ 都基于耦合定义，不理解耦合就读不懂为什么改 $q$ 就能加速采样。
ODE 求解器与 NFE
推理时用 Euler、Heun 或 RK45 等数值 ODE 求解器离散化 $dx/dt=v_\theta$。NFE（Number of Function Evaluations）是网络前向次数，越少采样越快。
MixFlow 的全部卖点都在"更少 NFE 下保持 FID"，Tab. 2/3/4 都用不同 NFE 评估，不熟悉 NFE 概念就读不懂低/高采样预算对比。
Prior Hole 问题
VAE 类模型中当先验 $p(z)$ 与近似后验 $q(z|x)$ 出现"空洞"时，从 $p(z)$ 直接采的样本解码后会落到数据流形外、生成崩塌，Hao & Shafto (2023) 在 Coupled VAE 中专门讨论。
MixFlow 选"混合源分布"正是为解决 κ-FC 在 $\beta\to 0$ 时的 prior hole；$w\sim U(0,1)$ 在条件-无条件高斯间插值是关键的工程化设计。
研究动机
扩散模型与 Rectified Flow 虽在多样性与画质上出色，但单次推理仍需数十到上百次网络前向。根因是生成轨迹弯曲——向量场被迫估计多条交叉路径的平均方向。Lee et al. (2023)、Tong et al. (2024) 已证明这与"源-目标独立耦合 $q(x_0,x_1)=p_0 p_1$"直接相关。Fast-ODE 用神经网络参数化 $q_\phi(x_0|\kappa)$ 并加 KL 正则 $\beta D_{KL}(q_\phi\|\mathcal{N}(0,I))$，却引入尴尬超参：$\beta$ 太小触发 prior hole（无 $\kappa$ 采样崩），太大退化成标准高斯（$q_\phi$ 失去灵活）。CIFAR10 上标准 Rectified Flow 仍需 127 NFE 达 FID 2.58，FFHQ/AFHQv2 64×64 上 $\beta=20$ 的 Fast-ODE 在 128 NFE 只能拿到 FID 5.72/4.55，缓慢采样是具体痛点。
本文的目标是目标分两层：第一，给出通用可学习耦合 κ-FC（κ-Forward Coupling），让源分布能依赖任意引导信号 $\kappa$（类标、文本、数据本身等），通过联合优化降低轨迹弯曲度。第二，提出 MixFlow 训练策略，用条件高斯与标准高斯的线性插值作为源分布，让模型在单一 $\beta$（如 $10^{-5}$）下既避免 prior hole，又获尽可能大的源分布自由度。最终目标是不动推理流程的前提下，让 Rectified Flow 在相同 NFE 下显著降低 FID，并在低 NFE（如 5 步）也保持可用画质。
与已有工作不同的是，与之前路线不同，本文从"源分布侧"而非"轨迹/求解器/蒸馏侧"入手：蒸馏类方法（Consistency Model、Reflow）需反复重训且在高 NFE 反而劣化；更快 ODE 求解器（DPM-Solver、GENIE）的上限受制于轨迹曲率。MixFlow 既不和它们竞争，也不需要重训。相对 Fast-ODE，MixFlow 显式引入混合权重 $w\sim U(0,1)$ 做"条件-无条件双栖训练"，把 $\beta$ 与先验空洞解耦，使 $\beta$ 可压到 $10^{-5}$ 量级而训练稳定（CIFAR10 上 $\beta=10^{-5}$ 取得最佳 FID 9.02@10 步 vs $\beta=\infty$ 的 13.16），并在推理时让 $w$ 充当新速度-质量旋钮。
核心方法
MixFlow 的核心直觉是：与其让源分布强行向标准高斯靠拢（$\beta$ 很大），不如让模型同时见过"条件高斯"与"无条件高斯"及两者之间的一切中间形态。具体地，训练一个向量场 $v_\theta(x_t,t)$ 和源预测网络 $q_\phi(x_0|\kappa)=\mathcal{N}(\mu_\phi(\kappa),\Sigma_\phi(\kappa))$，每步采样 $w\sim U(0,1)$ 构造混合源 $\mathcal{N}(w\mu_\phi(\kappa),\,w\Sigma_\phi(\kappa)+(1-w)I)$，再按 Rectified Flow 公式算损失并加 $\beta D_{KL}(\mathcal{N}(\mu_\phi,\Sigma_\phi)\|\mathcal{N}(0,I))$。这种"全谱训练"让模型学到的轨迹直线性在 $w=0$ 的标准高斯端也能继承到，从而推理时不必需要 $\kappa$。
MixFlow 与现有方法的本质区别有三：(1) 源分布是参数化高斯与标准高斯的**线性插值**，而非 Fast-ODE 的单一可学习高斯；这让 $\beta$ 可小到 $10^{-5}$ 而不破坏先验覆盖。(2) 训练时 $w\sim U(0,1)$ 均匀采样，使 $v_\theta$ 在任意 $w$ 上都见过样本，从而把"条件对齐带来的轨迹直线性"蒸馏到 $w=0$ 的无条件端。(3) 推理时 $w$ 本身成为新旋钮——给定 $\kappa$ 时可调 $w$ 切换速度-质量权衡；无 $\kappa$ 时退化为标准高斯。Fast-ODE 仅有 $\beta$ 一个旋钮且必须设大，QAC 走 VAE 路线增加训练负担，MixFlow 的简洁性来自内积结构的源分布。
方法步骤详情
**训练（Alg. 1）**：输入 (x1,κ)，默认 κ=x1。每步采 t,w~U(0,1)，由 κ 经 ~2M 参数小 UNet 输出均值 μκ 和协方差 Σκ，构造混合均值 μw=wμκ、协方差 Σw=wΣκ+(1-w)I，从 N(μw,Σw) 抽 x0，算 xt=tx1+(1-t)x0，按 Rectified Flow 损失加 β·D_KL 正则项联合训练 (θ,ϕ)。**推理（Alg. 2）**：有 κ 时按 w 构造初值，否则从 N(0,I) 取初值，再用 Euler/Heun/RK45 解 ODE。**关键超参**：CIFAR10 β=10⁻⁵，FFHQ/AFHQv2 β=5×10⁻⁵；向量场 UNet 55.7M-61.8M 参数；Adam lr 线性 warmup 至 2×10⁻⁴，EMA 0.9999。
技术新颖性
技术新颖性体现在三点：(1) 把 $\beta$ 解耦出"先验空洞压力"——通过在 $w$ 维度均匀混合，等价于对源分布做"条件端-无条件端"端到端蒸馏，Tab. 4 证明 $\beta$ 从 $\infty$ 降到 $10^{-5}$ 时 FID@2-steps 由 171.7 降到 99.30、@10-steps 由 13.16 降到 9.02，且训练不崩。(2) 通用化 $\kappa$——不限定为数据样本，可换为 $\kappa_c$（类标 embedding）或 $\kappa_n\sim\mathcal{N}(0,I)$，Tab. 5 展示即使 uncorrelated noise 作 $\kappa$ 也全面超过 Rectified Flow。(3) 推理时 $w$ 提供免费速度-质量旋钮——图 4a 显示 $\kappa_c$ 下低 NFE 时大 $w$ 更好、高 NFE 时小 $w$ 更好，Fast-ODE/QAC 都没有这一自由度。综合看，MixFlow 在不增加推理成本的前提下改写了源分布的"训练-推理契约"。
实验结果
**轨迹曲率（Tab. 1）**：MixFlow（β=10⁻⁵）弯曲度 0.0366，比 Rectified Flow 0.0467 低 ~22%、Fast-ODE 0.0388 低 ~5%。**CIFAR10（Tab. 2）**：RK45 全仿真 MixFlow FID 2.27（NFE 124.7），相对 Rectified Flow 2.58 改善 12%、Fast-ODE 2.45 改善 7%；Heun 2nd 低预算下 NFE=5 时 19.29 vs Fast-ODE 24.40、QAC 19.68；NFE=9 时 8.97 vs 9.96/10.28。**FFHQ/AFHQv2 64×64（Tab. 3）**：NFE=128 时 MixFlow 3.75/3.33 vs Fast-ODE 最佳 4.93/3.96，改善 24%/16%。**消融（Tab. 4/5/图 6）**：β 从 ∞ 降到 10⁻⁵ 时 FID@128 由 3.04 升到 2.52；κ=x1 最佳（2.52）；MixFlow 仅需 Fast-ODE 60% 训练步数即达最终 FID 2.27。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| CIFAR10 全仿真生成（RK45） | FID (↓) | 2.27 @ ~125 NFE | Rectified Flow 2.58 / Fast-ODE 2.45 / QAC 2.43 | vs RF −12%, vs Fast-ODE −7% |
| CIFAR10 低 NFE Heun 2nd (NFE=5) | FID (↓) | 19.29 | Fast-ODE 24.40 / QAC 19.68 | vs Fast-ODE −21%, vs QAC −2% |
| CIFAR10 低 NFE Heun 2nd (NFE=9) | FID (↓) | 8.97 | Fast-ODE 9.96 / QAC 10.28 | vs Fast-ODE −10%, vs QAC −12.7% |
| FFHQ 64×64 生成 (Euler, NFE=128) | FID-10k (↓) | 3.75 | Fast-ODE 4.93 (β=30) | −24% |
| AFHQv2 64×64 生成 (Euler, NFE=128) | FID-10k (↓) | 3.33 | Fast-ODE 3.96 (β=30) | −16% |
| 轨迹弯曲度 (CIFAR10, 10K traj, 128 Euler) | C(v_θ) (↓) | 0.0366 | Rectified Flow 0.0467 / Fast-ODE 0.0388 | vs RF −22%, vs Fast-ODE −5% |
局限与改进
**作者承认的局限**：(1) 仍要求源分布为高斯——$\beta$ 再小也是 $D_{KL}$ 到 $\mathcal{N}(0,I)$，无法刻画多模态源结构，作者明确表示未来要"放宽高斯假设"。(2) 论文仅在无条件图像生成（CIFAR10 32×32、FFHQ/AFHQv2 64×64）评估，未涉条件生成、文本到图像、高分辨率等场景。**我观察到的局限**：(1) Tab. 4 显示 $\beta$ 在 $5\times 10^{-7}$ 时已崩塌（FID@128 退化到 3.92），$\beta$ 仍是敏感超参。(2) 源预测 UNet（$\sim$2M 参数）与 $w$ 采样使训练计算量比标准 Rectified Flow 略高，论文未给 wall-clock 对比。(3) 推理时如提供 $\kappa=x_1$ 仍需额外前向 $\mu_\phi,\Sigma_\phi$。(4) 实验只跑 $\beta\in\{10,20,30\}$ 的 Fast-ODE 三点，缺与最新 Reflow、Variational Rectified Flow 等的多基准对比。
独立分析的弱点
**弱点 1：源分布仍是高斯**——尽管 MixFlow 放松了 $q_\phi$ 对标准高斯的接近度，其形式仍限制在 $\mathcal{N}(w\mu_\phi,w\Sigma_\phi+(1-w)I)$。当真实数据分布在源空间呈强多模态或非高斯结构时，弯曲度收益会饱和。改进：把源分布替换为 normalizing flow 或 diffusion 隐变量，但仍保留"条件-无条件"两端以便无 $\kappa$ 推理。**弱点 2：$w$ 的训练-推理契约偏经验**——$w\sim U(0,1)$ 训练后，最佳推理 $w$ 依赖 NFE 与 $\kappa$ 类型（图 4、图 5 显示 $\kappa_c$ 与 $\kappa_n$ 最优曲线完全不同），目前靠穷举找最优。**弱点 3：缺少 wall-clock 与系统级评估**——论文只比 NFE 与 FID，未报相同硬件下的推理时间与内存峰值。**弱点 4：评估数据集分辨率较低**——最高仅 64×64，对当前 256×1024 类高分辨率生成参考价值有限。
未来方向
**作者提出的方向**：(1) 把 $\kappa$ 推广到文本 prompt、CLIP embedding 等更丰富的引导信号，验证 κ-FC 在条件生成中的可用性；(2) 探索非高斯源分布（如 flow-based $q_\phi$），同时保留 MixFlow 的"双栖"特性。**基于成果可延伸的方向**：(1) 把 MixFlow 与一致性模型、Reflow、Variational Rectified Flow 等正交方法组合——作者明确指出蒸馏方向是 orthogonal 的，理论上 MixFlow 训出的"更直"流可直接喂给 Reflow 做一/两步采样。(2) 把 $\beta$ 改成 schedule（按训练步数从大到小），把 $w$ 改成课程学习（先学条件端后学无条件端），有望进一步压缩训练预算。(3) 在 latent diffusion 框架（Stable Diffusion 类）上验证 MixFlow 对 latent flow 的加速效果。(4) 把 $w$ 的推理选择建模为 Pareto 前沿搜索，配合 NFE 自适应调度器构建统一的速度-质量接口。
复现评估
**开源与代码**：作者给出 https://github.com/NazirNayal8/MixFlow，承诺开源实现。**数据**：仅用公开数据集 CIFAR10（Krizhevsky 2009）、FFHQ 64×64（Karras 2019）、AFHQv2 64×64（Choi 2020），无私有数据。**算力**：CIFAR10 在单卡 A100 80GB 即可复现（500K 步、batch 128），FFHQ/AFHQv2 在 4×A100 40GB 上训 500K/300K 步、batch 256，属中等实验室可承担区间。**复现难度**：中等偏易——架构、损失、超参在 Tab. 6/7 与 Appendix B 完整列出，与 Fast-ODE 一致便于复用；但 $\beta$ 选取（$10^{-5}$ 与 $5\times 10^{-5}$）跨数据集不一致，需按 Tab. 4 趋势小范围搜索。**潜在风险**：$\beta$ 过小会崩塌（图 3 趋势），EMA（0.9999）经验上"对收敛至关重要"——复现时若未开 EMA 大概率出次优结果。
论文图表
