如果你在生成模型上摸爬滚打过一段时间，大概率会撞上这样的场景：明明每次迭代的 loss 都在稳步下降，生成出来的样本却还是糊成一团；或者训练到一半，发现不同类别的样本轨迹在插值空间里扭成一团，怎么调学习率都挽救不回来。今年我在折腾 Flow Matching（流匹配）方向的项目时，把大量精力花在了一个此前被我严重低估的细节上——训练时源样本和目标样本到底是怎么配对的。这个细节直指一个叫 Online Optimal Transport Minibatches 的技术方向，简单说就是在流匹配的训练过程中，每个 batch 内部实时计算一次最优传输配对，而不是提前算好一张固定的配对表。这套思路的收益非常直接：路径更直、收敛更快、生成质量明显提升，而且对计算资源的消耗远比想象中可控。
如果你正在做扩散模型、流匹配、或者任何基于插值路径的生成式训练，这篇文章应该能给你一些新的启发。我会把配对问题的底层逻辑、在线小批量最优传输的具体实现、扩展到大规模训练时的工程取舍，以及我实际跑实验撞见的一堆坑，一次讲清楚。
1. 先从问题的根上看：Flow Matching 到底在拼什么
1.1 三句话讲透训练范式
Flow Matching 的核心目标，是学习一个从噪声分布到数据分布的可微映射。它不像扩散模型那样需要模拟一个随机微分方程，而是直接构造一条确定性的概率路径，然后让神经网络去拟合这条路径上的向量场。
具体的训练目标函数长这样：
$$L = \mathbb{E}{t, x_0, x_1} \left[ | v\theta(x_t, t) - u_t |^2 \right]$$
其中 $x_0$ 从噪声分布采样得到，$x_1$ 从真实数据分布采样得到，中间点 $x_t$ 由两者插值而成：
$$x_t = (1 - t)x_0 + t x_1$$
而条件向量场 $u_t$ 在直线路径假设下就是 $x_1 - x_0$。这组公式看起来干净利落，但里面藏着一个容易被忽略的问题：$(x_0, x_1)$ 这个配对是谁决定的？
在最朴素的实现里，$x_0$ 和 $x_1$ 是独立抽取的。也就是说，从噪声分布取一个点，从数据分布取一个点，强行把它们配对、插值、训练。这种"随机盲配"的做法在数学上没问题，但在实际训练中会带来一系列连锁反应。
1.2 独立采样埋下的三处暗伤
第一处暗伤是路径交错。假设一个 batch 里有两组配对：$(n_1, d_1)$ 和 $(n_2, d_2)$，$n$ 表示噪声点，$d$ 表示数据点。线性插值本身是直的，但当 $d_1$ 恰好落在 $n_2$ 的延长线上、$d_2$ 落在 $n_1$ 的延长线上时，两条插值路径就会在中间某个位置交叉。
路径交叉的致命点在于：交叉点附近的空间被两个方向相反的向量场同时"拉扯"，神经网络在这里学到的梯度是相互矛盾的，优化器会陷入一个尴尬的平均状态。这就像两个人同时给你指路，一个让你往北，一个让你往南，最后你站在原地打转。
第二处暗伤是传输路径变长。独立采样配对后，两个分布之间的最短传输路径完全被忽略了。本来从 $n_1$ 到 $d_1$ 可能只需要走一小段直线，但因为配对错了对象，模型被迫学习一条绕远路的轨迹。到最后采样时，数值积分器需要更多步数才能完成从噪声到数据的映射，推理时间被白白拉长。
第三处暗伤是收敛效率下降。训练目标的回归目标 $u_t = x_1 - x_0$ 在独立配对下方差很大。同一时刻 $t$，不同样本对的向量场方向可能南辕北辙，神经网络很难从中提炼出稳定的模式。我在一个双月形分布的玩具实验里统计过，独立采样时向量场在不同区域的方差比最优传输配对大将近一个数量级。
1.3 最优传输为什么能救场
最优传输（Optimal Transport，OT）做的事情很朴素：给定两组点，找到一种一对一的对应方式，让总的传输代价最小。放进流匹配训练里，就是对每个 batch 的噪声点和数据点做一次优化配对，使得所有配对的欧氏距离总和最小。
这样带来的改变是：插值路径不再交叉，向量场在不同区域内的方向变得一致，回归目标的方差大幅下降。更进一步，因为 OT 配对天然对应着两点之间的"最短路线"，学习到的路径也变得更直，采样速度随之提升。
那为什么不直接在全量数据集上做 OT？答案很残酷：计算量撑不住。在有 $N$ 个数据点的情况下，精确求最优传输的时间复杂度接近 $O(N^3)$，在几万条数据上都跑不动，更别提真实场景里动辄百万、千万级别的数据集。小批量最优传输（Minibatch OT）就是在这里登场的——只在每个训练 batch 的内部做一次 OT 配对，把问题规模从全量 $N$ 降到 batch 大小 $n$ 的范围，计算成本从不可行变成完全可以接受。
2. 在线还是离线？这一刀切在哪里
2.1 离线 OT 的甜头和苦果
很多人第一次接触 Minibatch OT 时会想到一个"省事"的方案：训练前先跑一次批量 OT，把整个数据集所有样本的配对算好，之后每次训练迭代都直接读取固定结果。这个方案叫离线 OT（Offline/Precomputed OT）。乍看之下它能白赚 OT 的收益，还不占用训练时的计算资源，何乐而不为？
我最初也是这么干的，结果踩了大坑。离线 OT 的问题藏在一个非常隐蔽的地方：固定配对意味着训练数据的多样性被锁死了。
打个比方，你给一个学生发了一套固定的练习题，他反复做同样那几道题，很快就能把答案背下来。考试时碰到稍微变形的题目，立刻傻眼。离线 OT 就是这种感觉——模型每次迭代看到的都是同样的 $x_0 \to x_1$ 对应关系，它不是在学会如何从任意噪声点映射到数据分布，而是在死记硬背一套固定的映射。训练 loss 下降得异常快，但真正采样时，生成的样本要么重复，要么分布不完整。
我还有过一个更直观的观察：在 CIFAR 级别的图像生成任务上，离线 OT 的模型在训练早期指标追得很紧，但到了训练后期完全停滞。原因就在于那些固定的配对让模型过拟合到了特定的传输路径上，失去了探索其他合理映射的机会。
2.2 在线 OT 的核心区别：每次 step 重新分配
在线 OT（Online Minibatch OT）的思路完全不同：每个训练 step 都重新从数据集采样一批数据，重新从噪声分布采样一批噪声，然后在这个全新的 batch 内部计算一次最优传输配对。配对结果只对当前这次迭代生效，下一次迭代重新计算。
这样得到的收益是双重的。第一，每一轮训练都会看到不同的数据组合，模型被迫学会适应各种不同的配对方式，泛化能力明显更强。第二，因为每个 batch 的新鲜组合都能触发一次新的 OT 求解，配对本身始终是"当前最优"的，路径始终保持在最短状态附近。
我这里说的"在线"，在实现上其实就是把 OT 求解插进训练循环里。处理顺序如下：
- 从训练集里随机取一个 batch 的目标数据 $X_1$；
- 从标准正态分布里随机取同样数量的噪声点 $X_0$；
- 计算 $X_0$ 和 $X_1$ 之间的最优传输配对索引；
- 按照这个配对关系构造插值点 $X_t$ 和向量场目标 $U_t$；
- 正常做一次反向传播。
第 3 步是关键差异点。你会发现整个训练流程前前后后只多了这一个小小的组件，但效果的改变却非常直接。
2.3 在线 OT 的稳定性和 warm-start 技巧
在线 OT 有一个看起来吓人的问题：两个连续的 training step，由于 batch 的数据完全不同，算出来的 OT 配对也是完全不同的。有人担心这会引入额外的噪声，导致训练不稳定。
实际跑下来我发现，这种担忧属于"理论上成立，实践中可控"。因为 OT 配对是在低维空间（通常就是样本特征空间）上计算的，配对的变换是连续的。只要 batch 足够大、数据分布足够稳定，相邻两个 step 之间配对结果的变化幅度是有限的，模型完全适应得了。
有一个小技巧可以进一步消减这种波动：把上一个 step 的 Sinkhorn 迭代结果当作当前 step 的初始值。这个操作叫 warm-start。由于相邻 batch 的分布接近，warm-start 不仅让配对结果更平滑，还能大幅减少 Sinkhorn 的迭代次数，属于一句话能讲清楚但价值极高的工程优化。我在后面的代码实现里会再次提到。
3. 真正 Scale 起来：计算成本、算法选型与工程细节
3.1 成本拆解：精确算法和近似算法的分界线
在线 OT 要落地，必须先回答一个问题：每次训练迭代多出来这一步计算，到底贵不贵？答案取决于你用的求解算法。
经典的匈牙利算法（Hungarian Algorithm）可以算出精确的最优传输结果，时间复杂度是 $O(n^3)$，其中 $n$ 是 batch 大小。当 $n=64$ 或 $n=128$ 时，这个计算量完全不在话下，即便放在 CPU 上算也只要几毫秒。但当 $n$ 涨到 1024 或 2048，$O(n^3)$ 的增长曲线就开始变得狰狞，单次求解可能超过一个前向传播的时间。
替代方案是用熵正则化的 Sinkhorn 算法。它做的是近似求解，时间复杂度是 $O(k \cdot n^2)$，其中 $k$ 是迭代轮数。在同样的 batch 大小下，Sinkhorn 比匈牙利算法快一到两个数量级，而且天然适合在 GPU 上做矩阵运算加速。
我自己划定的分界线是 $n \leq 512$ 用匈牙利算法，$n > 512$ 换 Sinkhorn。这不是一个严格的业界标准，只是基于我的训练场景（单卡、图像分辨率偏低）总结出的经验值。如果你用的是多卡集群，batch 动辄几千，那直接默认 Sinkhorn 就好。
| 算法 | 时间复杂度 | 解的性质 | 适合场景 |
|---|---|---|---|
| 匈牙利算法 | $O(n^3)$ | 精确解 | batch 小于 512，追求最优配对 |
| Sinkhorn 算法 | $O(k n^2)$ | 近似解 | batch 大于 512，需要 GPU 加速 |
| Gumbel-Sinkhorn | $O(k n^2)$ | 可微分近似解 | 需要把配对纳入反向传播 |
3.2 最小工具集实现一次完整的在线 OT 训练循环
不引入任何复杂的框架，用 PyTorch 加 SciPy 就够了。下面这段代码展示了一个可行的训练循环框架。
import torch import numpy as np from scipy.optimize import linear_sum_assignment def assignment_cost(x0, x1): # x0: [n, d] 噪声样本，x1: [n, d] 数据样本 # 计算代价矩阵：每个噪声点到每个目标点的欧氏距离 x0 = x0.detach().cpu().numpy() x1 = x1.detach().cpu().numpy() cost_matrix = np.sum((x0[:, None, :] - x1[None, :, :]) ** 2, axis=-1) row_idx, col_idx = linear_sum_assignment(cost_matrix) return row_idx, col_idx # 返回配对索引 def train_step(model, optimizer, data_batch): # 准备噪声样本 n = data_batch.shape[0] x0 = torch.randn_like(data_batch) x1 = data_batch # 求当前 batch 的最优配对 row_idx, col_idx = assignment_cost(x0, x1) x0_perm = x0[torch.from_numpy(row_idx)] x1_perm = x1[torch.from_numpy(col_idx)] # 采样时间步，构造插值点和向量场目标 t = torch.rand(n, device=x0.device).unsqueeze(-1) xt = (1 - t) * x0_perm + t * x1_perm ut = x1_perm - x0_perm # 预测向量场，回归训练 vt = model(xt, t.squeeze(-1)) loss = torch.mean((vt - ut) ** 2) optimizer.zero_grad() loss.backward() optimizer.step() return loss.item()这段代码的核心只有assignment_cost这一步。用linear_sum_assignment求解欧氏代价矩阵下的精确最优配对，然后让噪声点和目标点按照配对结果对齐，后续的插值和训练完全不受影响。
需要注意的是detach()这一步。我把两个样本都从计算图里摘出来了，意味着 OT 配对的计算不参与反向传播。这是当前绝大多数流匹配训练的做法，因为配对本身只是一个"课程安排"，不需要优化。如果你在做的是可微分最优传输方向的探索，那另说。
3.3 用 Sinkhorn 把 batch 推到千级以上
如果你的 batch 已经超过了 1000，匈牙利算法就有点跟不上了。此时换成 Sinkhorn 近似，用矩阵运算解决问题。
Sinkhorn 的核心是一个迭代更新的过程，目的是让传输矩阵的行和列都归一化。在 PyTorch 里可以这样实现：
def sinkhorn_cost_matrix(x0, x1, epsilon=0.05, iterations=50): # 代价矩阵 C = torch.sum((x0[:, None, :] - x1[None, :, :]) ** 2, dim=-1) # 初始传输矩阵用 exp(-C/epsilon) 得到，模拟 KL 正则化 OT K = torch.exp(-C / epsilon) # 交替归一化：先归一行，再归一列 u = torch.ones_like(C[:, 0]) for _ in range(iterations): v = torch.ones_like(C[0, :]) u = 1.0 / (K @ v) v = 1.0 / (K.T @ u) P = u[:, None] * K * v[None, :] # 从传输矩阵里把离散配对抽出来：每一行取概率最大的列 col_idx = torch.argmax(P, dim=-1) row_idx = torch.arange(x0.shape[0]) return row_idx, col_idx这个实现里有几个值得反复调的点。第一，epsilon是熵正则的强度。epsilon越大，传输矩阵越平滑，配对越接近均匀分布，但会牺牲 OT 路径的"直"性；epsilon太小，矩阵接近 one-hot，失去了 Sinkhorn 的数值稳定性优势。我在图像任务里通常从 0.05 起步，用网格搜索的粗结果去微调。
第二，迭代次数。50 次往往是一个比较稳的起点。用 warm-start 策略后，第一次可能要 50 次收敛，后续迭代大概 15 到 20 次就能到同样精度。
第三，从传输矩阵P里抽取离散配对用到argmax，这一步不可导。但因为我们前面已经用detach隔离了梯度，这里完全不需要求导，所以没有任何问题。
3.4 多卡并行和通信开销的控制
把在线 OT 搬到多卡训练环境时，会多出来一个分布式的问题。每个 GPU 各自做一份 batch 的 OT，还是所有 GPU 汇总成一个大 batch 再统一做？两种做法各有取舍。
各做各的，实现最简单，几乎没有通信开销，但每个 GPU 看到的配对只基于它本地的小 batch，OT 效果的"全局性"会减弱。汇总成一个大 batch 统一做 OT，配对质量更高，但需要先把所有特征收集到主节点上，算完再分发回各卡，通信成本上升。
我的个人建议是：如果你的单卡 batch 已经大于等于 256，就各做各的。这个规模下的配对质量已经足够好了，没有必要为了追求更优的全局配对去付出通信代价。只有当单卡 batch 特别小（比如 64），而总 batch 又很大时，才值得考虑跨卡汇总的方案。
4. 实际训练中我观察到的效果与几个关键参数
4.1 一个可复现的玩具实验：双月形分布
为了把在线 OT 和离线 OT、独立采样分个高下，我设计了一套非常干净的基础实验。数据分布用一个双月形二维 toy set，batch 大小取 256，模型用一个 4 层 MLP，训练 50000 步。
跑下来的结论非常鲜明。独立采样的基线在训练早期 loss 降得很快，但 10000 步之后曲线就趋于平缓，最终的生成质量只能说是"能看"。离线 OT 的 loss 下降最快，但生成分布明显遗漏了双月形的一个"月牙"，说明固定配对过度牺牲了多样性。在线 OT 虽然在训练初期的 loss 下降不如离线 OT 快，但它的收敛曲线没有明显平台期，一路稳步向下，最后生成的分布完整度最高、边缘最清晰。
这个实验给我的启发是：loss 下降快不等于生成质量好。那些"绕近路"的训练策略，往往是把模型带进了一个局部最优，后期再想跳出来就很难了。
4.2 影响最大的几个超参数
在大量实验之后，我把影响在线 OT 训练效果的超参数排了个序。第一重要的是 batch size，第二是时间步采样策略，第三才是 Sinkhorn 的参数。
Batch size 的角色非常微妙。batch 越大，batch 内部做 OT 的结果越接近"全数据集 OT"的水平，路径越直。但是 batch 过大带来的问题是训练内存压力和 Sinkhorn 的计算时间上升。我在同样一个模型上对比过 128、512、1024 三个 batch size，512 的性价比最高，1024 比 512 提升不足 5%，但训练时间多了近一倍。
时间步采样策略是另一个经常被忽略的点。流匹配的插值在 $t$ 接近 0 和接近 1 的两端，对应的回归目标变化剧烈，模型更难以学会两端的向量场。我习惯用非均匀的时间采样，比如用 Beta 分布的密度把更多的采样点压在中段，或者采用端点加权的策略。这个策略对在线 OT 的增益尤其明显，因为 OT 配对的路径已经够直了，真正决定边界区域表现的就是时间步的覆盖质量。
Sinkhorn 的参数里，epsilon比迭代次数更重要。迭代次数不够只是导致归一化不到位，影响是轻微的；epsilon设错了则直接改变配对的性质。有一次我把epsilon设成 0.5，结果配对的发散程度接近完全随机的水平，在线 OT 的效果直接退化成和独立采样差不多。
4.3 一个不得不防的梯度尖刺问题
在线 OT 会给训练过程引入一种特有的不稳定现象，我称之为"配对切换尖刺"。现象是这样的：loss 在某个 step 突然冲高了几倍，但下一两步又回归正常。
原因很好理解。训练过程中模型在变，它对数据的表征在变，于是同一对 $(x_0, x_1)$ 在不同 step 计算出的代价矩阵也在变。某个 step 里，有一小部分样本的最优配对发生了跳变，换了新对象，于是这一部分的回归目标突然改变，loss 就爆一下。
尖刺本身并不可怕，怕的是优化器被它带偏。我的处理方案有三个：一是给 Sinkhorn 加大epsilon，让传输矩阵更平滑，减少配对跳变的幅度；二是降低学习率，尤其是在训练后期；三是给 loss 加一个小窗口的中值滤波，不从优化层面改变，只是让可视化不那么惊心动魄。
5. 常见问题排查手记
5.1 训练 loss 正常，生成样本却明显缺失某些类别
遇到这个问题先别急着调模型架构，很可能是离线 OT 的坑。检查一下训练代码，如果配对是提前算好的、固定不变的，那大概率就是死记硬背导致的过拟合。解决方案就是切成在线 OT，别省那点计算时间。
5.2 Sinkhorn 跑得比模型前向还慢怎么办
我会从三个方向同时排查。第一，看epsilon是否设得太小，过小的epsilon会让 Sinkhorn 陷入慢收敛。第二，看看是否开启了 warm-start，没有的话立即加上。第三，确认自己在用 GPU 上的矩阵运算，而不是把数据搬回 CPU 跑循环。
5.3 在线 OT 在训练后期 loss 反复震荡
通常是配对跳变幅度过大引起的。先把epsilon调大，再观察是否稳定。如果还不行，把 batch size 往上提一提。更大的 batch 意味着 OT 结果是基于更多样本的综合判断，配对跳变的发生频率会明显降低。
5.4 生成样本的边缘出现奇怪的分布塌缩
这个问题通常是时间步采样策略造成的。回归目标在 $t$ 靠近 1 时是直接从源跳向目标的最后一段，如果这段区域采样不足，边缘就会模糊或扭曲。我建议你把时间步分布画出来看一眼，看看是不是大量样本都集中在了 0.4 到 0.6 的中段。
一些实际操作后的体会
我个人在实际操作中的体会是，在线小批量最优传输并不是一个"锦上添花"的模块，而是直接决定了流匹配训练质量的上限。我在好几个项目里做了同样的切换：把固定配对换成在线 OT，不改变网络结构，不改变优化器，生成质量的提升都非常明显。代价只在训练阶段多出一点计算量，而这个计算量完全可以通过算法选型和工程优化压到可以接受的范围。
最后再分享一个小技巧。如果你做的任务允许，尽量在训练初期用较大的epsilon，让配对温和一点，等模型有了一定的先验能力之后，再逐步减小epsilon。我在图像生成实验里用这个退火策略，比全程固定参数的表现稳定不少。这个方向还有很多可以挖掘的点，比如把 OT 配对本身也纳入可微训练、或者用更高效的近似算法去逼近大规模数据集上的最优传输。如果你手头正好也在折腾流匹配，强烈建议先从这个小模块下手试试水。
