基于聚类的最优传输流匹配方法 COT-FM: Cluster-wise Optimal Transport Flow Matching
通过聚类和簇内最优传输，实现更直的流匹配轨迹，提升生成质量
前置知识
Flow Matching (流匹配)
Flow Matching 是一种生成建模框架，它学习一个时间依赖的向量场，通过求解常微分方程将样本从简单的源分布如高斯分布输送到复杂的目标数据分布。与扩散模型不同，FM 不需要在训练时求解微分方程，直接回归向量场，训练更高效。FM 的核心是构造条件概率路径，边缘化后得到边际概率路径。损失函数为数学公式形式，其中 v_theta 是参数化的神经网络，v_t 是真实向量场。通过最小化这个损失，模型学习到最优的向量场，在推理时通过 ODE 求解器生成样本。
COT-FM 是基于 Flow Matching 框架的改进方法，必须理解 FM 如何通过向量场学习生成变换，以及源分布与目标分布之间的耦合策略如何影响轨迹形状和生成质量，才能理解 COT-FM 的创新点和改进机制。
Optimal Transport (最优传输)
最优传输是衡量并实现两个概率分布之间最优映射的理论。给定源分布和目标分布，最优传输寻找传输计划使得传输成本最小。在生成模型中，使用最优传输作为耦合策略可以产生直线轨迹，因为每个源点被唯一地映射到对应的目标点。精确求解 OT 需要三次的时间复杂度和二次的空间复杂度，对于大规模数据集不可行。实际应用中常使用近似方法如 Sinkhorn 算法或批处理 OT，这些方法在小规模数据上表现良好，但在大规模数据上仍然面临挑战。
COT-FM 的核心创新是在聚类内部使用最优传输来构造更直的向量场。理解 OT 的计算成本和近似方法有助于理解为什么 COT-FM 需要聚类来降低每个 OT 问题的规模，从而实现更高效的训练和更高质量的生成。
Wasserstein Distance
Wasserstein 距离也称 Earth Mover's 距离，是度量两个概率分布之间差异的指标，定义为最优传输问题的最小传输成本。二阶 Wasserstein 距离定义为所有可能耦合集合上传输成本的下确界的平方根，其中传输成本是源点和目标点之间欧几里得距离的平方。较小的 Wasserstein 距离表示生成的分布与目标分布更接近。在二维实验中，论文使用精确 Wasserstein 距离评估生成质量，这提供了对生成分布与真实分布差异的定量度量。
论文在二维点云实验中使用 Wasserstein 距离作为主要评估指标，量化生成分布与真实分布的差异。理解该指标有助于判断不同方法的生成质量，以及 COT-FM 相对于基线方法的改进程度。
FID (Fréchet Inception Distance)
FID 是评估生成图像质量的常用指标，计算生成图像分布与真实图像分布在 Inception 网络特征空间中的距离。假设特征服从高斯分布，FID 衡量两个高斯分布之间的 Fréchet 距离。计算公式涉及真实和生成特征的均值和协方差矩阵，FID 越低表示生成质量越好。论文在 CIFAR-10 和 ImageNet 实验中使用 FID 来评估生成质量，这是图像生成任务的标准评估指标，能够反映生成图像的多样性和保真度。
FID 是图像生成任务的标准评估指标，论文在 CIFAR-10 和 ImageNet 实验中使用 FID 来量化 COT-FM 的生成质量提升。理解 FID 的计算和解释有助于评估不同方法在实际图像生成任务中的性能差异。
研究动机
现有 Flow Matching 模型由于使用随机耦合或批处理最优传输耦合，导致学习到的向量场轨迹弯曲。弯曲的轨迹会增加时间离散化误差，在低采样步数如一步或两步时特别严重，导致样本移动到未见过的位置，产生扭曲的传输和低质量的生成。例如在 CIFAR-10 上，Rectified Flow 在一步采样时的 FID 高达三百七十八，两步时为一百七十三。虽然 k-Rectified Flow 通过迭代优化可以增强直线性，但它需要在自生成样本上重复训练，导致模型崩塌和生成质量随时间下降。另一个研究方向是蒸馏或学习平均向量场来跳过采样步骤，但这些方法不调制底层的向量场，它仍然是弯曲的，只减少步数而不提升质量。这些限制使得现有方法在低采样预算下难以实现高质量的生成。
本文的目标是本文的目标是设计一个通用框架，在不改变 Flow Matching 模型架构或输入输出机制的前提下，使概率路径更直，从而加速采样并提高生成质量。这个框架应该是即插即用的，能够与大多数现有的 FM 模型兼容，并在不同的任务如二维点云生成、图像生成、机器人操作中持续改进生成速度和质量。作者希望通过聚类导向的最优传输策略，在不增加额外计算成本的情况下，显著提升低采样步数下的生成质量，使得生成模型能够更快速地生成高质量的样本。
与已有工作不同的是，本文的独特切入角度是观察到训练数据自然包含多个模式，同一模式的数据可以从相似的源样本生成。基于这个观察，作者提出将数据样本聚类，每个聚类分配一个专用的源分布，而不是从单个源分布映射整个数据分布。这种分而治之的策略将每个最优耦合问题简化为更小的簇级匹配，使批处理最优耦合更有效，同时通过限制源分布空间，学习整体向量场变得更高效。为了找到每个聚类的最优源分布，作者的第二个洞见是从预训练的 FM 模型自举，因为这些模型学习的路径天然可逆且不交叉，可以通过反转采样过程容易地获得每个聚类的源分布，路径交叉更少。这种策略与现有的全局耦合方法有本质区别，它是第一个将聚类与最优传输结合用于 Flow Matching 的框架。
核心方法
COT-FM 的整体思路是通过聚类将复杂的目标分布分解为多个简单的子分布，然后为每个子分布学习一个对应的源分布，并在聚类内部进行最优传输。这个框架交替优化目标向量场和 FM 模型。第一阶段，FM 模型回归从先验或之前估计的聚类源分布导出的目标向量场。第二阶段，通过反转学习到的路径为每个数据样本计算高斯统计量，然后使用批处理耦合在聚类内近似最优传输来更新这些源分布。这个框架能够适应不同类型的聚类，包括条件生成中的类标签或文本描述，以及无条件生成中的无监督聚类。COT-FM 只调制目标概率路径，不改变 FM 架构或输入输出机制，使其与大多数现有 FM 模型兼容。这种设计使得 COT-FM 可以作为一个即插即用的增强模块，应用于各种 Flow Matching 模型，而不需要对原始模型进行大的修改。
COT-FM 的核心创新点是通过聚类将全局最优传输问题分解为多个局部簇内最优传输问题。与现有方法使用全局随机耦合或批处理最优传输不同，COT-FM 首先使用预训练的 FM 模型通过反向 ODE 积分估计每个聚类的源分布参数均值和协方差，然后从估计的源分布中采样源样本，并在聚类内部计算最优传输映射。这些簇级 OT 映射共同组成向量场，FM 模型随后回归这个向量场。这种分治策略显著减少了每个 OT 问题的规模，使批处理近似更可行。作者证明，通过将复杂的多模态分布分解为多个单模态或简单多模态的聚类，每个聚类内部的传输路径可以更加直接和直，从而减少路径交叉和离散化误差。这种方法与全局 OT 相比，计算复杂度从三次降低到近似二次，使其在大规模数据集上成为可行。
方法步骤详情
COT-FM 的方法步骤完整描述如下。第一步，构造聚类目标向量场，使用预训练的 FM 模型，对每个聚类中的每个数据样本，通过反向 ODE 积分获取源样本。具体来说，从时间一反向积分到时间零，获取产生该数据样本的源点。第二步，计算每个聚类的源分布统计量，包括均值和协方差。均值是该聚类所有反向传播源样本的平均，协方差是这些源样本的方差。第三步，从估计的源分布中采样源样本，每个聚类采样一定数量的源点。第四步，在聚类内部计算最优传输映射，将源样本与目标样本配对。第五步，优化 FM 模型，首先随机采样聚类索引，概率与聚类大小成正比，然后从对应的 OT 映射中采样源目标样本对。第六步，采样时间步，从均匀分布中采样，计算线性插值。第七步，使用条件流匹配损失优化 FM 模型，损失是最小化预测向量与真实向量之间的均方误差。第八步，采样阶段，首先采样聚类索引，然后从对应的源分布中采样源样本，通过 ODE 求解器使用 FM 模型生成目标样本。这个过程可以重复多次以生成多个样本。
技术新颖性
COT-FM 的技术新颖性体现在几个方面。首先，它提出了一种新的聚类导向的最优传输框架，将全局 OT 问题分解为局部簇级 OT 问题，显著降低了计算复杂度。与全局 OT 的三次复杂度相比，COT-FM 的复杂度约为二次，其中是数据集大小，是聚类数。这使得大规模数据集上的最优传输成为可行。其次，它利用预训练 FM 模型的可逆性来估计聚类源分布，避免了需要预先知道源分布的假设。这是一种自举策略，利用现有模型的输出来改进自身。第三，COT-FM 与现有方法如二阶 Rectified Flow 互补。二阶 Rectified Flow 通过在合成耦合上重训练来拉直轨迹，而 COT-FM 通过簇级最优传输进一步减少路径交叉，两者可以组合使用产生累积改进。第四，COT-FM 自然扩展到条件生成设置，条件本身可以作为聚类标签如类标签或文本描述，对于非固定条件如机器人策略中的视觉观察，作者还提出学习一个条件模型来动态预测噪声分布。这种扩展性使得 COT-FM 可以应用于广泛的生成任务。
实验结果
COT-FM 在多个实验设置中展现了显著的性能提升。在二维点云生成中，COT-FM 在 Mixture of Five Gaussians 上将 Wasserstein 距离从 Rectified Flow 的零点五四二一降低到零点一九九五，曲率从零点零三一六降低到零点零零八四。在 Two Moons 上将 Wasserstein 距离从零点一零零六降低到零点零二六六，曲率从零点零一一一降低到零点零零一六。在 Checkerboard 上将 Wasserstein 距离从零点三九零零降低到零点二五五零，曲率从九点一九四六降低到零点一五零五。这些结果表明 COT-FM 能够学习更直的轨迹，同时更好地捕获目标分布的结构。在 CIFAR-10 图像生成中，COT-FM 在十步采样时将 FID 从 Rectified Flow 的十二点六降低到八点二三，在五十步时从四点四五降低到三点九七。在极低采样预算下，COT-FM 将一步 FID 从三百七十八降低到二百零五，两步 FID 从一百七十三降低到五十九点一。COT-FM 还改进了 MeanFlow，将一步 FID 从二点九二降低到二点六零，两步 FID 从二点八八降低到二点五三。在 ImageNet 二百五十六乘二百五十六条件生成中，COT-FM 在 SiT-B/2 和 SiT-B/4 架构上持续改进 FID，特别是在低 NFE 设置中改进更明显。在 LIBERO 机器人操作中，COT-FM 仅用一步就在 Spatial 上达到百分之九十六点一，在 Long 上达到百分之九十四点五，而 FLOWER 需要四步才能达到百分之九十七点一和百分之九十三点五。这些结果一致证实，聚类目标和学习聚类源分布能够产生更直的传输路径，从而在极低采样步数下实现更可靠的生成。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| Mixture of 5-Gaussians (2D point cloud) | Wasserstein Distance (lower is better) | 0.1995 | Rectified Flow: 0.5421 | 63.2% reduction |
| Two Moons (2D point cloud) | Wasserstein Distance (lower is better) | 0.0266 | Rectified Flow: 0.1006 | 73.6% reduction |
| Checkerboard (2D point cloud) | Wasserstein Distance (lower is better) | 0.2550 | Rectified Flow: 0.3900 | 34.6% reduction |
| CIFAR-10 (10-step generation) | FID (lower is better) | 8.23 | Rectified Flow: 12.6 | 34.7% reduction |
| CIFAR-10 (50-step generation) | FID (lower is better) | 3.97 | Rectified Flow: 4.45 | 10.8% reduction |
| CIFAR-10 (1-step generation) | FID (lower is better) | 205.0 | Rectified Flow: 378.0 | 45.8% reduction |
| LIBERO-Spatial (1-step) | Success Rate (higher is better) | 96.1% | FLOWER (4-step): 97.1% | Comparable performance with 4x fewer steps |
| LIBERO-Long (1-step) | Success Rate (higher is better) | 94.5% | FLOWER (4-step): 93.5% | 1.0% higher with 4x fewer steps |
局限与改进
作者承认 COT-FM 在扩展和泛化方面仍面临挑战。首先，反向 ODE 步骤随着样本规模增长可能变得低效。在 CIFAR-10 上，这一步需要四百二十秒，虽然相当于大约五个训练 epoch，但对于大规模数据集如 ImageNet 可能会更显著。其次，该方法依赖于聚类质量，虽然论文显示 COT-FM 对聚类数量相当鲁壮，即使五个聚类时 FID 为四点五六，仍优于 OT-CFM 的四点七八，但过于细粒度的聚类如一百二十个会降低每个聚类的样本数，削弱簇内 OT 的有效性。第三，聚类传输行为特别是局部 OT 如何捕获全局结构以及为什么交替精炼在一两个周期内收敛还没有被完全理解。我观察到的一个局限是，COT-FM 需要一次性预处理步骤，在 CIFAR-10 上约四百四十五秒，虽然这相当于大约五个训练 epoch，但增加了计算开销。另一个潜在的局限是，对于非固定条件如机器人策略中的动态视觉观察，需要学习一个条件模型来预测噪声分布，这增加了系统复杂性。
独立分析的弱点
COT-FM 的一个弱点是反向 ODE 预处理步骤的计算开销。在 CIFAR-10 上，这一步需要四百二十秒，虽然相当于大约五个训练 epoch，但对于大规模数据集如 ImageNet 可能会更显著。改进方向可以是开发更高效的反向 ODE 近似方法或并行化策略。第二个弱点是聚类质量依赖性。虽然论文表明方法对聚类数相对鲁壮，但聚类质量差时性能可能下降。改进方向可以是探索更鲁棒的聚类算法如层次聚类或基于语义的聚类，或自适应聚类策略。第三个弱点是对于非固定条件设置，需要学习一个条件模型来预测噪声分布，这增加了系统复杂性和训练难度。改进方向可以是简化条件模型架构或使用更高效的强化学习算法如强化学习的变体。另一个观察到的弱点是，COT-FM 在聚类数过多时性能下降，这可能是因为每个聚类的样本数太少，OT 近似效果变差。改进方向可以是动态调整聚类数或使用分层聚类策略。
未来方向
作者提出的未来研究方向包括探索更可扩展的聚类策略和更轻量的源分布估计器。基于论文成果，可以延伸的方向包括以下几方面。第一，理论分析交替优化过程的收敛性，目前只有经验观察，FID 在两次迭代后达到三点九七，第三次迭代时回升到四点一七，缺乏形式化的收敛保证。第二，将 COT-FM 与其他直化方法如二阶 Rectified Flow 结合更深入地研究，论文显示二阶 Rectified Flow 加 COT-FM 在 CIFAR-10 上达到十点九一 FID 一步，优于单独二阶 Rectified Flow 的十二点二一，可以进一步研究这种组合的最优策略。第三，扩展到其他生成任务，如三维点云生成、视频生成或多模态生成。第四，研究 COT-FM 在不同噪声分布设置下的表现，如非高斯源分布或多模态源分布。第五，开发自适应聚类策略，根据数据特性自动确定最优聚类数，目前论文使用固定的聚类数，但最佳聚类数可能因数据集而异。第六，研究 COT-FM 在更大规模模型如 Stable Diffusion 或 DALL-E 上的应用，论文主要在相对较小的模型上验证。
复现评估
论文提供了项目页面 https://embodiedai-ntu.github.io/cotfm，但未明确说明代码是否开源。实验配置详细描述如下。CIFAR-10 使用 AMD Instinct MI300X GPU 一百九十二 GB HBM，训练六百 epochs，batch size 一百二十八，学习率二乘十的负四次方，使用 Adam 优化器，梯度裁剪最大范数一点零。ImageNet 使用八个 MI300X GPU，batch size 二百五十六，学习率一乘十的负四次方。LIBERO 使用单张 NVIDIA GeForce RTX 4090 GPU。论文提供了详细的超参数表，包括网络架构参数。数据集使用公开基准。CIFAR-10 包含五万张训练图像，ImageNet 使用二百五十六乘二百五十六全数据集，LIBERO 是机器人操作基准。复现难度中等偏高，主要挑战在于需要 AMD MI300X GPU 或等效硬件，虽然 RTX 4090 可用于 LIBERO。反向 ODE 预处理需要额外计算。聚类步骤需要实现 DINO 特征提取和 K-means 聚类。如果代码开源，复现难度会降低。总体来说，论文提供了足够的实验细节和超参数设置，有经验的实践者应该能够复现主要结果，但硬件要求可能是一个限制因素。
论文图表
该图比较了随机耦合和最优传输两种耦合策略产生的向量场。左图显示随机耦合迫使模型回归不一致的速度目标灰色，产生模糊的交叉点紫色箭头表示平均方向，导致弯曲的速度场。右图显示最优传输提供一致交叉更少的耦合，使模型能够学习更直的速度场。这个对比清楚地展示了耦合策略对向量场形状的影响。
这张图对理解论文的技术背景很重要，它解释了为什么需要改进现有 Flow Matching 模型使用的耦合策略。随机耦合的平均效应导致弯曲轨迹，而最优传输的耦合一致性能够产生更直的轨迹。这为 COT-FM 使用簇级最优传输提供了理论基础。