超越直线性：通过分位数对齐树耦合的非交叉流匹配

# 超越直线性：基于分位数对齐树耦合的非交叉流匹配

来源：https://arxiv.org/html/2609.03443

**Junyi Lin**
单位：中国人民大学统计与大数据研究院，中国北京
邮箱：[junyilin@ruc\.edu\.cn](mailto:[email protected])

**Mengyu Li**
单位：清华大学统计与数据科学系，中国北京
邮箱：[2025104242@ruc\.edu\.cn](mailto:[email protected])

**Jingxuan Hu**
单位：中国人民大学统计与大数据研究院，中国北京
邮箱：[chengmeng@ruc\.edu\.cn](mailto:[email protected])

**Kejun He**††thanks:通讯作者
单位：中国人民大学统计与大数据研究院，中国北京
邮箱：[kejunhe@ruc\.edu\.cn](mailto:[email protected])

**Cheng Meng**††thanks:通讯作者
单位：中国人民大学统计与大数据研究院，中国北京
邮箱：[mengyuli@tsinghua\.edu\.cn](mailto:)

###### 摘要
流匹配的性能在很大程度上取决于源分布与目标分布之间耦合的质量。然而，独立耦合常常导致路径交叉和局部速度模糊，而基于最优传输的耦合通常构建成本高昂。为解决此挑战，我们提出**分位数对齐树流匹配（QAT-FM）**，这是一种高效的结构化耦合策略。它通过分位数对齐树结构，在高斯先验与目标数据分布之间构建层次化耦合。QAT-FM在$\mathcal{O}(Nd\log N)$时间内完成耦合构建，并支持复杂度为$\mathcal{O}(d)$的逐对源采样，从而能够为大规模高维生成任务进行可扩展训练。理论上，我们证明QAT耦合满足边缘一致性，诱导出非交叉线性插值路径，并且在中间时刻相比独立耦合一致性地改善了路径分离度，从而缓解了局部速度模糊问题。QAT-FM进一步自然地扩展到条件生成，在保持全局高斯对齐的同时实现了结构化条件耦合。在多个基准数据集上的实验表明，QAT-FM在显著降低耦合构建成本的同时，取得了具有竞争力的生成性能。

## 1 引言
近年来，基于连续时间流的生成模型在生成建模领域日益凸显其重要性（Chen et al., 2018; Grathwohl et al., 2019; Lipman et al., 2023; Liu et al., 2023）。一个代表性的类别是连续归一化流（CNFs），它将连续变换定义为常微分方程（ODEs）的解流，其时间依赖的速度场由神经网络参数化。采样时，通过数值积分学习到的ODE动力学来生成新数据。为提高训练效率，流匹配（Lipman et al., 2023）直接使用从预指定的条件概率路径中导出的目标来监督速度场，避免了传统CNF训练中所需的重复数值ODE积分，从而实现了生成模型的可扩展训练。流匹配已在包括图像生成（Ma et al., 2024; Esser et al., 2024）、语音合成（Le et al., 2023; Liu et al., 2024）和文本生成（Hu et al., 2024; Hu et al., 2026）在内的多个领域取得了具有竞争力的性能。

尽管流匹配提高了训练连续时间生成模型的可扩展性，但其性能对训练期间构建的条件概率路径和源-目标耦合仍然敏感（Lipman et al., 2023; Tong et al., 2024; Pooladian et al., 2023）。在常见的线性路径公式中，流匹配通过回归噪声与数据样本之间线性插值路径上的目标速度来训练速度场；因此，不同的耦合选择直接改变了监督信号的结构。当训练路径交叉，或在具有显著不同目标速度的同一中间区域过于接近地通过时，会产生局部的*速度模糊*；在均方误差训练下，这些不一致的速度方向被平均，增加了拟合速度场的难度（Guo and Schwing, 2025）。相比之下，具有*非交叉*特性或良好*路径分离度*的训练路径减少了这种局部模糊性，提供了从中间状态到目标速度的更具确定性的映射，从而提高了训练稳定性。

在条件生成中，耦合策略还必须保持样本与条件之间的对齐，因为忽略条件信息的耦合可能产生几何上有利但条件不一致的监督信号（Chemseddine et al., 2025; Cheng and Schwing, 2025）。几种现有方法试图通过提高耦合质量来缓解路径交叉和速度模糊。整流流（Rectified Flow）中的回流（reflow）过程使用学习到的模型重新构建噪声与数据之间的配对，产生更确定、更直的生成轨迹；但这需要额外的采样和再训练步骤，带来巨大的计算开销。基于最优传输（OT）的方法通过最小化传输成本来构建几何最优的配对；在二次成本下，诱导的位移插值表现出有利的非交叉结构，从而缓解速度模糊。然而，精确的经验OT计算成本高昂且难以扩展到大规模数据集；在条件生成中，保持条件一致性通常需要条件感知的约束或逐条件耦合，进一步增加了算法的复杂性。小批量OT（Pooladian et al., 2023; Tong et al., 2024）通过在每个小批量内求解OT来大幅降低此成本，但所得到的耦合主要是局部的，无法在全数据集级别严格保持非交叉结构。

上述大多数方法仅在有限的高斯样本上构建耦合；SD-FM（Mousavi-Hosseini et al., 2026）通过基于梯度的优化进一步改善了连续高斯信息的对齐，但其迭代优化带来了巨大的采样或优化成本，并且不能产生显式的低成本全局耦合构建。

表1：流匹配耦合策略对比。
- 1 符号说明。N：数据集大小；B：小批量大小；d：数据维度；K：SD优化步骤；E：回流轮数；$\Theta$：一次网络前向传播。报告的复杂度忽略常数因子。
- 2 ✓ 表示该特性在条件生成中也得以保持。

表1从五个维度比较了代表性的流匹配耦合策略：耦合构建成本（Build Cost）、采样成本（Sample Cost）、耦合是否在全数据集级别诱导非交叉路径（Non-Cross）、是否改善路径分离度（Path Sep.）、以及源边缘分布是否与连续高斯先验对齐（Gaussian Align.）。如图所示，先前的方法通常在计算效率、路径几何和高斯先验对齐之间进行权衡。

为解决这些限制，我们提出**分位数对齐树流匹配（QAT-FM）**，这是一种低成本的结构化耦合策略，能够实现非交叉路径、改善路径分离度和连续高斯先验对齐，其中部分特性自然地扩展到条件生成。

**贡献。** 我们的主要贡献如下：
- 1）我们提出**分位数对齐树（QAT）**，一种基于二叉树结构分位数对齐的高效半离散传输构建方法。QAT在经验数据分布和连续高斯先验之间诱导出精确耦合。其构建复杂度为$\mathcal{O}(Nd\log N)$，逐对采样复杂度为$\mathcal{O}(d)$，使QAT能够扩展到大规模、高维的生成建模任务。
- 2）QAT耦合支持多种树分裂策略，并能自然地将条件信息融入树构建和耦合过程中。
- 3）我们证明QAT诱导出有利的路径结构：所得的线性插值路径满足*非交叉*特性，且在中间时刻的路径分离度一致大于标准独立耦合流匹配的分离度，从而缓解了局部速度模糊。
- 4）我们在低维合成数据、无条件生成、离散类条件生成、连续文本条件生成以及大规模图像合成上进行了系统性实验。结果表明，QAT-FM在显著低于基于OT方法的耦合构建复杂度下，实现了具有竞争力的生成质量。

## 2 预备知识
#### 符号说明。令$\mathcal{P}(\mathbb{R}^{d})$表示$\mathbb{R}^{d}$上的博雷尔概率测度集。对于$\alpha\in\mathcal{P}(\mathbb{R}^{d})$和$\beta\in\mathcal{P}(\mathbb{R}^{d'})$，一个联合分布$\pi\in\mathcal{P}(\mathbb{R}^{d}\times\mathbb{R}^{d'})$若其边缘分布分别为$\alpha$和$\beta$，则称为$\alpha$和$\beta$的一个*耦合*；我们用$\Pi(\alpha,\beta)$表示所有此类耦合的集合。对于向量$\bm{x}\in\mathbb{R}^{d}$，令$\bm{x}^{(j)}$表示其第$j$个坐标。对于$\rho\in\mathcal{P}(\mathbb{R}^{d})$，令$\rho^{(j)}$表示其第$j$个边缘分布，$F_{\rho}^{(j)}$为$\rho^{(j)}$的累积分布函数（CDF），$\bigl(F_{\rho}^{(j)}\bigr)^{-1}$为其广义逆。对于可测映射$T:\mathbb{R}^{d}\to\mathbb{R}^{k}$，我们用$T_{\#}\rho$表示$\rho$在$T$下的推前（pushforward）。

#### 基于流的生成建模。给定一个时间依赖的速度场$v:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$，记$v_{t}(x):=v(t,x)$。由$v_{t}$诱导的*流映射*$\phi_{t}:\mathbb{R}^{d}\to\mathbb{R}^{d}$定义为以下ODE的解：
$$
\frac{d}{dt}\phi_{t}(x)=v_{t}(\phi_{t}(x)),\qquad\phi_{0}(x)=x.
$$
令$\rho_{t}=(\phi_{t})_{\#}\gamma$，其中$\gamma$表示源分布，则对$(\rho_{t},v_{t})$满足连续性方程：
$$
\partial_{t}\rho_{t}+\nabla\cdot(\rho_{t}v_{t})=0,
$$
且$\rho_{0}=\gamma$。当$(\phi_{1})_{\#}\gamma=\mu$时，流将源$\gamma$传输到目标分布$\mu$。因此，基于流的生成建模的目标是学习一个参数化速度场$v_{\theta}$，其诱导的流映射满足$(\phi_{1})_{\#}\gamma=\mu$。

流匹配（Lipman et al., 2023）为此速度场提供了一种易于处理的训练过程。它指定了一个连接$\gamma$到$\mu$的概率路径$(\rho_{t})_{t\in[0,1]}$以及一个相容的目标速度场$u_{t}$。然后通过回归目标学习参数化速度场：
$$
\mathcal{L}_{\mathrm{FM}}(\theta)=\mathbb{E}_{\begin{subarray}{c}t\sim\mathrm{Unif}[0,1]\\ x_{t}\sim\rho_{t}\end{subarray}}\left\|v_{\theta}(t,x_{t})-u_{t}(x_{t})\right\|^{2}.
$$
这种公式避免了在训练过程中通过数值ODE求解器进行反向传播。

#### 条件流匹配。实际上，边缘速度场$u_{t}$通常难以直接构建。条件流匹配¹通过配对端点条件来解决此问题：给定$\pi\in\Pi(\gamma,\mu)$，绘制$(x_{0},x_{1})\sim\pi$，指定一个条件路径$p_{t}(\cdot\mid x_{0},x_{1})$及其相容速度$u_{t}(\cdot\mid x_{0},x_{1})$，并训练$v_{\theta}(t,x_{t})$以匹配$u_{t}(x_{t}\mid x_{0},x_{1})$，其中$x_{t}\sim p_{t}(\cdot\mid x_{0},x_{1})$。

¹ 此处“条件”指的是以配对端点为条件。为避免与条件生成混淆，除特别说明外，本文其余部分将此公式简称为流匹配（FM）。

一个规范且广泛使用的选择是线性插值路径：
$$
x_{t}=(1-t)x_{0}+tx_{1},\qquad u_{t}(x_{t}\mid x_{0},x_{1})=x_{1}-x_{0},
$$
在此路径下，训练目标简化为：
$$
\mathcal{L}_{\mathrm{CFM}}(\theta)=\mathbb{E}_{\begin{subarray}{c}t\sim\mathrm{Unif}[0,1]\\ (x_{0},x_{1})\sim\pi\end{subarray}}\left\|v_{\theta}\left(t,(1-t)x_{0}+tx_{1}\right)-(x_{1}-x_{0})\right\|^{2}.
$$
这表明端点配对由耦合$\pi\in\Pi(\gamma,\mu)$控制。最简单的选择是独立耦合$\pi=\gamma\otimes\mu$。一个更有原则的替代方案是*OT耦合*：
$$
\pi_{\mathrm{OT}}\in\arg\min_{\pi\in\Pi(\gamma,\mu)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|x_{0}-x_{1}\|^{2}\,d\pi(x_{0},x_{1}),
$$
在线性路径下，该耦合倾向于产生更短、几何上更自然的配对轨迹（Pooladian et al., 2023; Tong et al., 2024）。然而，计算精确的全局OT耦合在大规模上成本高昂；实践中，使用小批量OT作为近似，以牺牲耦合精确性和跨批量一致性为代价降低计算成本。

对于特殊情况$\gamma=\mathcal{N}(0,I_{d})$，Mousavi-Hosseini et al. (2026)提出在(4)中求解一个半离散OT问题，以更好地对齐连续的高斯信息。