如果你在生成模型上摸爬滚打过一段时间大概率会撞上这样的场景明明每次迭代的 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 TransportOT做的事情很朴素给定两组点找到一种一对一的对应方式让总的传输代价最小。放进流匹配训练里就是对每个 batch 的噪声点和数据点做一次优化配对使得所有配对的欧氏距离总和最小。这样带来的改变是插值路径不再交叉向量场在不同区域内的方向变得一致回归目标的方差大幅下降。更进一步因为 OT 配对天然对应着两点之间的最短路线学习到的路径也变得更直采样速度随之提升。那为什么不直接在全量数据集上做 OT答案很残酷计算量撑不住。在有 $N$ 个数据点的情况下精确求最优传输的时间复杂度接近 $O(N^3)$在几万条数据上都跑不动更别提真实场景里动辄百万、千万级别的数据集。小批量最优传输Minibatch OT就是在这里登场的——只在每个训练 batch 的内部做一次 OT 配对把问题规模从全量 $N$ 降到 batch 大小 $n$ 的范围计算成本从不可行变成完全可以接受。2. 在线还是离线这一刀切在哪里2.1 离线 OT 的甜头和苦果很多人第一次接触 Minibatch OT 时会想到一个省事的方案训练前先跑一次批量 OT把整个数据集所有样本的配对算好之后每次训练迭代都直接读取固定结果。这个方案叫离线 OTOffline/Precomputed OT。乍看之下它能白赚 OT 的收益还不占用训练时的计算资源何乐而不为我最初也是这么干的结果踩了大坑。离线 OT 的问题藏在一个非常隐蔽的地方固定配对意味着训练数据的多样性被锁死了。打个比方你给一个学生发了一套固定的练习题他反复做同样那几道题很快就能把答案背下来。考试时碰到稍微变形的题目立刻傻眼。离线 OT 就是这种感觉——模型每次迭代看到的都是同样的 $x_0 \to x_1$ 对应关系它不是在学会如何从任意噪声点映射到数据分布而是在死记硬背一套固定的映射。训练 loss 下降得异常快但真正采样时生成的样本要么重复要么分布不完整。我还有过一个更直观的观察在 CIFAR 级别的图像生成任务上离线 OT 的模型在训练早期指标追得很紧但到了训练后期完全停滞。原因就在于那些固定的配对让模型过拟合到了特定的传输路径上失去了探索其他合理映射的机会。2.2 在线 OT 的核心区别每次 step 重新分配在线 OTOnline 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 大小。当 $n64$ 或 $n128$ 时这个计算量完全不在话下即便放在 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, devicex0.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, epsilon0.05, iterations50): # 代价矩阵 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 看到的配对只基于它本地的小 batchOT 效果的全局性会减弱。汇总成一个大 batch 统一做 OT配对质量更高但需要先把所有特征收集到主节点上算完再分发回各卡通信成本上升。我的个人建议是如果你的单卡 batch 已经大于等于 256就各做各的。这个规模下的配对质量已经足够好了没有必要为了追求更优的全局配对去付出通信代价。只有当单卡 batch 特别小比如 64而总 batch 又很大时才值得考虑跨卡汇总的方案。4. 实际训练中我观察到的效果与几个关键参数4.1 一个可复现的玩具实验双月形分布为了把在线 OT 和离线 OT、独立采样分个高下我设计了一套非常干净的基础实验。数据分布用一个双月形二维 toy setbatch 大小取 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 size512 的性价比最高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 配对本身也纳入可微训练、或者用更高效的近似算法去逼近大规模数据集上的最优传输。如果你手头正好也在折腾流匹配强烈建议先从这个小模块下手试试水。