🤖 machine learning

COT-FM: Cluster-wise Optimal Transport Flow Matching

本文提出了 COT-FM 框架，通过聚类目标样本并利用预训练流匹配模型反向构建专用源分布，将全局传输转化为更准确的局部直线性传输，从而在不改变模型架构的前提下显著加速采样并提升生成质量。

原作者： Chiensheng Chiang, Kuan-Hsun Tu, Jia-Wei Liao, Cheng-Fu Chou, Tsung-Wei Ke

发布于 2026-03-17

📖 1 分钟阅读☕ 轻松阅读

原作者： Chiensheng Chiang, Kuan-Hsun Tu, Jia-Wei Liao, Cheng-Fu Chou, Tsung-Wei Ke

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文介绍了一种名为 COT-FM 的新方法，旨在让 AI 生成图像、数据或机器人动作时更快、更准、更直。

为了让你轻松理解，我们可以把生成数据的过程想象成**“从起点（噪声）走到终点（真实数据）的旅程”**。

🌟 核心问题：为什么现在的 AI 走路总是“绕弯路”？

想象一下，你要教一群学生（AI 模型）从“一片混乱的迷雾”（噪声）走到“一个具体的目的地”（比如一张猫的照片）。

传统方法（随机配对）： 老师随便抓一个学生，让他随便指一个方向走。结果就是，有的学生往左，有的往右，大家挤在一起，路线互相交叉、打架。为了把大家带到终点，老师不得不画出一条弯弯曲曲、甚至打结的路线。
- 后果： 学生走得很累（需要很多步），而且容易走错地方，最后生成的猫可能长着狗耳朵（图像失真）。
现有的优化方法（批量最优）： 老师试图把学生分组，每组内部尽量走直线。但因为组太小，或者分组不够精准，路线还是会有点弯，或者不同组之间还是会撞车。

🚀 COT-FM 的解决方案：分而治之，各走各的“高速公路”

COT-FM 的核心思想是：“不要试图用一条路解决所有问题，我们要把大目标拆成小目标，给每个小目标定制专属路线。”

它主要做了三件聪明事：

1. 像“分班”一样聚类 (Clustering)

想象学校里有 1000 个学生，要让他们分别去“猫区”、“狗区”、“车区”。

旧方法： 让所有人混在一起，老师试图画一条万能路线。
COT-FM： 先把学生按“想去哪里”分成不同的班级（Cluster）。想画猫的归猫班，想画狗的归狗班。
- 比喻： 这样就把一个巨大的、混乱的交通网，拆成了几个清晰的小交通网。

2. 像“倒带”一样找起点 (Reverse ODE)

这是最精彩的一步。

旧方法： 老师假设所有学生都从同一个“迷雾中心”出发。但这不对，因为画猫的学生和画狗的学生，其实应该从迷雾里不同的位置出发，才能走直线。
COT-FM： 老师先让一个已经学会走路的“老学生”（预训练好的模型）带着大家走一遍。然后，老师把录像倒带（Reverse ODE），看着他们是怎么从终点走回起点的。
- 发现： 老师惊讶地发现，原来“猫班”的学生其实是从迷雾的左上角出发的，而“狗班”是从右下角出发的。
- 行动： 于是，老师为每个班级定制了专属的起点（比如：猫班从 A 点出发，狗班从 B 点出发）。

3. 像“修路”一样走直线 (Cluster-wise Optimal Transport)

现在，每个班级都有自己的起点和终点，而且人数变少了。

COT-FM： 老师在每个班级内部，重新规划了一条绝对笔直的高速公路，让“猫班”从 A 点直冲猫区，“狗班”从 B 点直冲狗区。
- 结果： 路线不再交叉，不再弯曲。学生只需要走很少的几步（甚至一步）就能到达终点，而且走得非常稳。

🎯 这带来了什么好处？

速度极快（少走路）： 因为路是直的，不需要绕弯。以前可能需要走 50 步才能画出一张图，现在走 1 步或 2 步就够了。
质量更高（不走样）： 因为路线不交叉、不混乱，生成的图像更清晰，机器人动作更精准（比如把杯子放进微波炉，不会卡住）。
通用性强（即插即用）： 这个方法不需要把 AI 的“大脑”（模型架构）换掉，只需要改变它“走路的方式”（训练策略）。就像给旧车换了更好的导航系统，车还是那辆车，但跑得更快了。

📊 实际效果如何？

论文在几个领域做了测试，效果都很惊人：

2D 图形： 以前生成的图形歪歪扭扭，现在像尺子画出来一样直。
图片生成 (CIFAR-10)： 以前画一张图要 10 步，现在 1 步就能画出很清晰的图，而且比以前的方法画得更好。
机器人控制： 以前机器人需要走 4 步才能把杯子放好，现在 COT-FM 让机器人1 步就能完成，而且成功率更高。

💡 总结

COT-FM 就像是一位高明的交通指挥官。它不再试图用一条拥挤、弯曲的大路把所有车（数据）都送到目的地，而是通过**“分班”和“倒带找路”，为每一类车规划了专属的直线高速**。

结果就是：路更直了，车跑得更快了，事故（生成错误）更少了。 这就是为什么它能用更少的步骤，生成更高质量的内容。

1. 研究背景与问题 (Problem)

Flow Matching (FM) 是一种强大的生成模型框架，通过学习一个确定性的向量场（Vector Field），将简单的源分布（如高斯分布）映射到复杂的目标数据分布。然而，现有的 FM 模型在实际应用中面临以下核心挑战：

弯曲的轨迹 (Curved Trajectories)： 大多数 FM 模型在训练时使用随机耦合 (Random Coupling) 或 批次最优传输 (Batch-wise Optimal Transport)。
- 随机耦合会导致源样本和目标样本之间的配对不一致，迫使模型学习一个“平均”的、弯曲的速度场。
- 批次最优传输虽然试图优化耦合，但由于批次大小有限，只能捕捉局部最优，导致全局路径依然弯曲。
离散化误差与采样效率： 弯曲的轨迹意味着在推理（采样）过程中，为了获得高质量样本，需要更多的积分步数（NFE, Number of Function Evaluations）。如果步数过少，离散化误差会导致样本偏离真实分布，生成质量下降。
现有解决方案的局限性：
- Rectified Flow： 通过迭代优化源和生成样本之间的耦合来拉直路径，但反复在自生成样本上训练容易导致模型崩溃 (Model Collapse)，降低生成质量。
- MeanFlow/Shortcut 方法： 试图学习平均向量场以跳过采样步骤，但这并未改变底层向量场的弯曲本质，仅减少了步数，却牺牲了生成质量。

核心痛点： 如何在保持模型架构不变的前提下，获得更直 (Straighter) 的传输路径，从而在极少的采样步数下（如 1 步或 few-step）实现高质量的生成？

2. 方法论 (Methodology)

作者提出了 COT-FM (Cluster-wise Optimal Transport Flow Matching)，这是一个通用的、即插即用的框架。其核心思想是**“分而治之” (Divide-and-Conquer)**：将全局最优传输问题分解为多个簇（Cluster）级别的局部最优传输问题。

核心流程：

数据聚类 (Clustering)：
- 将目标数据样本划分为 $K$ 个簇（ $C_1, ..., C_K$ ）。
- 对于无条件生成，使用无监督聚类（如 K-Means）；对于条件生成（如文本或类别），直接使用条件标签作为簇。
构建簇级源分布 (Cluster-wise Source Distributions)：
- 关键洞察： 同一簇内的数据样本通常由相似的源噪声生成。
- 逆向 ODE (Reverse ODE)： 利用一个预训练的 FM 模型，将每个簇内的数据样本 $x_1$ 逆向积分回源空间，得到对应的源样本估计 $\hat{x}_0$ 。
- 高斯拟合： 对每个簇的逆向源样本计算均值 ( $\mu_{0,k}$ ) 和协方差 ( $\Sigma_{0,k}$ )，从而为每个簇 $C_k$ 定义一个专属的高斯源分布 $p_{0,k} = \mathcal{N}(\mu_{0,k}, \Sigma_{0,k})$ 。
- 优势： 这种方法利用了预训练模型生成的非相交路径特性，确保了不同簇之间的源分布交叉较少。
簇内最优传输 (Cluster-wise OT)：
- 在每个簇内部，从估计的源分布 $p_{0,k}$ 采样源点，并与目标数据 $C_k$ 进行最优传输 (Optimal Transport, OT) 配对。
- 由于簇内样本量远小于全局，计算 OT 映射（或近似 OT）的复杂度大幅降低，且能更准确地捕捉局部结构。
交替优化 (Alternating Optimization)：
- 阶段 1： 基于当前的簇级源分布，计算簇内 OT 映射，构建目标向量场。
- 阶段 2： 使用构建好的向量场微调 FM 模型（回归条件流匹配损失）。
- 迭代： 重复上述过程（通常 2 次迭代即可收敛），不断更新源分布估计和模型参数。
推理采样 (Sampling)：
- 首先根据簇的大小比例随机选择一个簇索引 $k$ 。
- 从该簇对应的专属源分布 $p_{0,k}$ 中采样初始噪声 $x_0$ 。
- 通过微调后的 FM 模型进行积分，生成最终样本。

3. 关键贡献 (Key Contributions)

提出 COT-FM 框架： 首次将“簇级最优传输”引入 Flow Matching，通过为每个数据簇分配专属的源分布，显著拉直了传输轨迹。
即插即用 (Plug-and-Play)： 该方法不改变 FM 的模型架构或输入输出机制，仅修改目标概率路径的构建方式。可兼容 Rectified Flow、OT-CFM、MeanFlow 等多种基线模型。
无需模型崩溃的直化策略： 相比于 Rectified Flow 的迭代自训练，COT-FM 通过逆向预训练模型获取源分布，避免了在自生成数据上反复训练导致的模型崩溃问题。
理论洞察： 证明了将全局 OT 问题分解为局部簇级 OT 问题，不仅能降低计算复杂度，还能有效减少路径交叉，从而获得更直的向量场。

4. 实验结果 (Results)

COT-FM 在多个基准测试中均表现出显著优势，特别是在低采样步数 (Low NFE) 场景下：

2D 点云生成 (2D Benchmarks)：
- 在 5-Gaussians、Two Moons 和 Checkerboard 数据集上，COT-FM 的 Wasserstein 距离 和 轨迹曲率 (Curvature) 均优于所有基线（Rectified Flow, OT-CFM, MeanFlow）。
- 例如，在 5-Gaussians 上，Wasserstein 距离从 0.5421 降至 0.1995。
无条件图像生成 (CIFAR-10)：
- 1 步采样 (1-step)： FID 从 Rectified Flow 的 378.0 降至 205.0；MeanFlow 从 2.92 降至 2.60。
- 10 步采样： FID 从 12.6 降至 8.23。
- 证明了在极少步数下，COT-FM 能生成更清晰、结构更准确的图像。
条件图像生成 (ImageNet 256x256)：
- 基于 SiT 架构，COT-FM 在所有 NFE 设置下均降低了 FID。在低 NFE 区域（如 10 步），提升尤为明显（FID 从 8.25 降至 7.52）。
机器人操作 (LIBERO)：
- 在文本条件机器人策略任务中，COT-FM 仅需 1 步 (1 NFE) 即可达到 96.1% (Spatial) 和 94.5% (Long) 的成功率。
- 这一性能甚至超过了需要 4 步 的基线模型 FLOWER，且优于其他单步基线。

5. 意义与影响 (Significance)

突破采样效率瓶颈： COT-FM 证明了通过优化传输路径的几何结构（使其更直），可以极大地减少生成所需的积分步数，同时保持甚至提升生成质量。这对于实时性要求高的应用（如机器人控制、实时视频生成）至关重要。
解决“曲率”问题： 论文从理论上和实验上证实，生成质量低下的根本原因往往在于向量场的弯曲（由随机耦合引起），而不仅仅是模型容量不足。COT-FM 提供了一种优雅的解决方案。
通用性： 作为一种通用的概率路径重塑方法，COT-FM 为未来的 Flow Matching 研究提供了新的方向，即通过更智能的耦合策略（Coupling Strategy）来优化生成过程，而非单纯依赖更复杂的网络架构。

总结： COT-FM 通过“分簇”和“逆向源分布估计”的策略，成功地将 Flow Matching 的传输路径拉直，实现了在极低采样成本下的高质量生成，是生成模型领域在采样效率与质量平衡上的一项重要突破。