One-Sided Quantile Coupling for Flow Matching

技术摘要：用于流匹配（Flow Matching）的单侧分位数耦合

问题陈述
流匹配（Flow Matching, FM）通过回归连接简单源分布（通常为标准高斯分布）与目标数据分布的概率路径的速度场来训练连续时间生成模型。耦合（Coupling）的质量——即源样本与目标样本的配对方式——至该决定了训练轨迹的几何特性。标准的 FM 使用独立随机配对，这往往会导致长距离且冲突的路径，从而增加了目标速度场的方差并阻碍了训练效率。

虽然受最优传输（Optimal Transport, OT）启发的结构化耦合可以使这些路径变直并提高生成质量，但现有的方法（如 OT-CFM）依赖于解决小批量分配问题。这些方法需要构建成对代价矩阵并求解分配问题（例如通过匈牙利算法或 Sinkhorn 迭代），其计算成本随批量大小呈至少平方级（对于精确解则是立方级）增长。这限制了可扩展性，并可能引入不稳定性。本文探讨了是否可以在不匹配两个预采样批次的情况下，实现结构化耦合的几何收益，从而避免高昂的批量级传输成本。

方法论：分位数耦合流匹配 (QC-FM)
作者提出了分位数耦合流匹配 (QC-FM)，这是一种轻量级的单侧耦合方案。与传统方法采样噪声和数据批次后再进行重新配对不同，QC-FM 仅采样数据批次并直接构建配对的源样本。

1. 单侧构建：

   • 给定一个数据批次 $\mathcal{B}_x$，该方法将数据点投影到由矩阵 $U$ 定义的 $k$ 个随机正交方向（切片）上。
   • 对于每个方向，在批次内对数据点进行排序。这些经验秩（Empirical Ranks）被映射到高斯分位数（使用逆标准正态 CDF）。
   • 源潜码 $e^{(i)}$ 的构建需满足这些投影分位数约束。具体而言，源在切片方向上的投影与基于数据秩导出的高斯分位数相匹配。
   • 剩余维度（切片子空间的正交补）通过从条件高斯分布中采样来填充。这确保了源在未受约束的方向上保持随机性。

2. 高斯余项混合扩展：
   由于完全的 QC 耦合（即所有源点都通过秩进行构建）由于切片编码之间的依赖关系，对于 $k \ge 2$ 会偏离标准的 i.i.d. 高斯先验，因此作者引入了两种混合策略，以在保留 QC 收益的同时维持先验：

   • QC-FM-Mixture（混合型）： 对批次的“锚点”（Anchor）子集（例如 20%）应用 QC 构建，并将剩余的“余项”源样本与剩余的数据样本进行独立随机配对（基准方法）。这创建了结构化与非结构化耦合的混合体。
   • QC-FM-Adjacency（邻接型）： 对锚点子集应用 QC，但根据锚点诱导的局部几何结构对剩余的源样本进行配对。余项数据点被分配到其在结构化子空间中最接近的锚点，剩余的源样本则在这些组内进行配对。这在不复制目标的情况下传播了结构化偏差。

核心贡献

• QC-FM 算法： 一种新型的、批次局部的结构化耦合，它在随机切片上施加单调的一维传输。它消除了对成对代价矩阵和分配求解器的需求，将主要的计算成本从 $O(B^2)$ 或 $O(B^3)$ 降低到了 $O(kB \log B)$（排序）。
• 理论分析：
  • 方差缩减： 作者证明，对于固定的框架，QC 耦合消除了每个选定切片上的不可约回归方差，使得理想流在这些方向上是完全直线的。
  • 先验保持： 他们量化了当 $k=1$ 时，源边缘分布恰好为高斯分布。对于 $k \ge 2$，偏离高斯先验的情况被限制在切片编码的 Copula 中，且该偏差的传输成本是有界的。
• 混合训练方案： 引入了 Mixture 和 Adjacency 变体，以平衡注入几何偏差与保留来自 Baseline 的显式高斯信号之间的关系。

实验结果
该方法在匹配训练预算和架构的情况下，在 CIFAR-10、CelebA-64、FFHQ-64 和 ImageNet-64 上进行了评估。

• 性能： 在所有四个数据集上，QC-FM-Mixture 的表现始终优于独立耦合的 Baseline 以及精确的匈牙利 OT-CFM。在 FID-50K 指标上，相对于 Baseline 的提升范围从 6.3% (CIFAR-10) 到 12.9% (FFHQ) 不等。
• 效率： 在计算成本方面，QC-FM 比基于 OT 的方法显著更快。在批量大小为 2048 时，QC-FM-Mixture 比匈牙利 OT 求解器快 800 倍以上，比 Sinkhorn 求解器快 35 倍以上。
• 玩具数据集诊断： 在合成数据集（棋盘格和 8-高斯分布）上，QC-FM 变体产生了近乎直线的学习轨迹（路径长度比 $\approx 1$），并且与 Baseline 相比，目标速度方差显著降低，证实了耦合带来的几何收益。

意义与主张
论文声称，保留投影秩结构是向流匹配耦合中注入有用几何偏差的一种简单、可扩展且有效的方法，而无需解决小批量传输问题。

• 可扩展性： QC-FM 在非结构化 Baseline 与昂贵的批量级传输方法之间提供了一个极具吸引力的中间地带。它是计算轻量级的，并且可以通过投影数量 ($k$) 进行控制。
• 实用性： 该方法在推理阶段不需要任何学习组件；生成过程仍然从标准高斯开始。耦合是在训练期间即时构建的。
• 局限性： 作者谦虚地指出，QC-FM 是结构化耦合的一种批次局部替代方案，而非全局最优传输的近似。分位数是在小批量内估计的，理论控制的是每条切片的方差，而非保证全局轨迹的直线性或特定的 FID 提升。目前，切片数量 ($k$) 和锚点比例 ($p$) 的选择仍依赖于启发式方法和针对特定数据集的调优。

总之，这项工作表明，通过通过保持秩结构来实现单侧分位数耦合，可以优化流匹配中的样本质量和训练效率，为计算成本高昂的最优传输方法提供了一种实用的替代方案。