流匹配中的最优传输配对：在线小批量OT实操

做生成模型的人这两年基本都被 Flow Matching 相关的工作轮番轰炸过。扩散模型还在卷采样步数的时候流匹配已经靠“回归向量场”这条更简洁的路子把训练稳定性和生成质量同时拉上去了。但你真正上手跑 Flow Matching 项目时会发现训练过程里最影响质量上限的往往不是网络结构而是“每对训练样本是怎么配对的”。如果只是随机把噪声点和数据点扔在一起模型就得被迫去拟合一堆交叉流线收敛慢不说生成样本的边缘分布还容易糊。这篇要聊的方案核心就是用 Optimal Transport最优传输简称 OT在小批量内部动态计算配对让每个训练步的噪声点和真实数据点都按距离最优配对再训练。这个思路对搞图像生成、分子构象、表格数据生成、时间序列插值的同学都很有参考价值尤其是目前用条件流匹配但卡在配对质量上不去的人。我会把 OT 是什么、为什么在小批量上“在线”做而不是离线预计算、Sinkhorn 迭代的参数怎么调、踩过哪些坑都讲清楚最后附上实操层面的效果参考和排查经验。1. 先把 Flow Matching 的“配对”问题摊开讲1.1 训练目标其实是在拟合向量场Flow Matching 的基本设定说起来非常简单有一个先验分布 p0通常是各向同性的高斯噪声一个真实数据分布 p1我们希望学习一个从 p0 到 p1 的时间相关映射。训练时我们在时间 t∈[0,1] 上采样中间状态 xt然后让网络去预测从 xt 指向 x1 的“速度”或者说向量场。具体来说假设我们从某个配对分布 π(x0, x1) 中采样噪声点和数据点用线性插值的方式构造中间状态 xt (1-t)·x0 t·x1那么对应的目标向量场就是 x1 - x0。网络的训练损失就是预测向量场和真实向量场之间的平方误差L E_{t, (x0, x1)π} [ ∥ v_θ(t, xt) - (x1 - x0) ∥² ]这里的关键就藏在 π 里。π 决定了每个噪声点要“送”到哪个数据点。这是一个大家刚开始做 Flow Matching 时很容易忽略、但后续影响巨大的设计点。很多开源实现默认 π 是“独立耦合”也就是 x0 和 x1 各自独立采样、随机配对。这种做法实现起来很省事但它在几何上不是在描述一条干净的流。想象一下你在操场上喊集合每个人随便找个位置站但你想让他们按身高排列成队伍——如果一开始位置就是乱的大家跑动路线就会互相交叉队伍成型的速度自然慢。1.2 随机配对为什么会让训练变慢随机配对最直接的问题是同一个噪声点可能被配到风格差异很大的数据点上。比如图像生成任务中噪声点 A 这一批可能配给“猫”下一批又配给“高楼”。模型为了让损失变小只能学习一个“平均场”把所有可能目标方向的期望值都压进同一个向量场里。这样生成的样本往往细节模糊、语义混杂。我在一个图像生成项目里对比过随机配对和 OT 配对随机配对的前期损失下降明显更慢生成图像在 50 步采样时还有明显的重影和结构错位。而 OT 配对下每个噪声点都有一个清晰的“主方向”模型要做的就是把这个映射关系记忆并泛化好损失曲线会更平滑采样步数减半时质量也还能保持。需要说明的是随机配对不是完全不能用。某些数据分布比较集中、模式不复杂的任务里随机配对也能训出来只是要花更多训练步数或更大的 batch。但一旦数据分布复杂、多模态明显配对质量就会变成 Flow Matching 训练的一个硬瓶颈。这也是最优传输派上用场的前提。2. 最优传输给样本排序的数学工具2.1 两个离散分布之间的最优传输到底在算什么最优传输这个名字听起来唬人但落在训练代码里其实是“求一个样本间的最佳匹配关系”。给定噪声点集合 {x0^i} 和数据点集合 {x1^j}我们想找一个分配关系把每个噪声点映射到一个数据点并且让总传输代价最小。如果两边样本数都是 n那么这就是一个经典的指派问题min_P ∑_{i,j} P_ij · C_ij约束条件是每一行和每一列的和都等于 1/n这里的 P 是一个 n×n 的“耦合矩阵”P_ij 表示第 i 个噪声点分配给第 j 个数据点的权重。C_ij 是传输代价最常见的就是两点之间的欧氏距离平方。如果只允许一一对应那 P_ij 只取 0 或 1/n这就是把全局的距离矩阵拉出来求最优匹配。把这个问题换成生活场景就很好理解了有 n 个学生要分到 n 个座位每个学生到每个座位都有一个“不满意值”你要让所有人的不满意值之和最小。暴力枚举阶乘级别组合不可行但用匈牙利算法或网络流可以在多项式时间内解决。然而当 n 是几万、几十万时这些精确算法的代价在深度学习训练循环里依然不可接受。这还不是最致命的。全量 OT 需要先构建一个 n×n 的距离矩阵一个一万样本的数据集光这个矩阵的 float32 存储就是 400MB每次训练步都重建一次显存和算力都顶不住。所以“在小批量上做 OT”成了必然选择。2.2 熵正则化、Sinkhorn 与 GPU 友好的迭代算法精确 OT 虽然形式漂亮但不可微、不稳定、计算昂贵在深度学习里几乎不会直接用。真正被普遍采用的是带熵正则化的最优传输也就是给目标函数加上一项耦合矩阵的信息熵min_P ∑_{i,j} P_ij · C_ij ε · ∑_{i,j} P_ij · log P_ijε 是一个正则化系数。直观来看ε 越小解越接近精确最优匹配ε 越大传输矩阵越趋向均匀分布也就是随机配对的方向靠拢。熵正则化带来的最大好处是求解过程变成了一个高度可并行、可以在 GPU 上跑的迭代算法——Sinkhorn 算法。Sinkhorn 的计算过程非常简洁。先用代价矩阵构造核矩阵 K exp(-C / ε)然后初始化两个缩放向量 u 和 v交替更新直到收敛u a / (K v)v b / (Kᵀ u)这里的 a 和 b 是行列边际约束通常都是 1/n。迭代若干次后耦合矩阵就是 P diag(u) K diag(v)。你可以理解成这是在“双向配平”先让每行和为 1/n再让每列和为 1/n反复调整直到两个方向都满足约束。核心代码也就几十行我自己在项目里用的一个极简版本是这样def sinkhorn(cost, eps0.05, num_iters50): n, m cost.shape a torch.ones(n, devicecost.device) / n b torch.ones(m, devicecost.device) / m K torch.exp(-cost / eps).double() u torch.ones_like(a) v torch.ones_like(b) for _ in range(num_iters): u a / (K v) v b / (K.T u) P u[:, None] * K * v[None, :] return P注意我把 cost 除到 eps 里了实际使用时 C 的尺度对 eps 选择影响很大后面专门讲这个。这个版本没有加收敛判断只是固定迭代次数工程上这样反而更稳定因为反向传播不需要经过条件判断。2.3 从耦合矩阵到训练对的采样策略算出了软耦合矩阵 P接下来怎么变成训练用的样本对这里有一个细节很多人第一次接触会犯迷糊。P_ij 是一个 0 到 1/n 之间的小数要把它变成明确的“第 i 个噪声点配第 j 个数据点”常用的做法是按 P 对每一行做采样或者直接取每行的 argmax。我在实验中更推荐按行采样而不是 argmax。argmax 虽然噪声小但会损失多样性尤其是 batch 里存在相似样本时容易把很多噪声点都配到同一个数据点上。按行采样则会在保留边际约束的同时给训练注入随机性等价于每次都在“近似最优的几个候选”里随机选一个方向。这样做在生成任务上的多样性明显更好。如果训练目标是二维或三维的表格数据还可以直接用软耦合矩阵对损失做加权也就是所有噪声点对数据点的预测损失都考虑进去权重就是 P_ij。这在数学上更平滑但实现起来要小心显存因为损失矩阵也是 n×n 的。3. Online Minibatch OT 的核心设计与参数选择3.1 “在线”不是玄学每个训练步重新算一次配对如果只把数据集整体做一次 OT然后把配对结果缓存下来那叫离线预计算。这种做法在小数据集上没问题但流匹配训练普遍要求 batch 越大越好而每个 epoch 里的配对应该随着网络参数的更新不断调整才更合理。在线 OT 的做法是在每个训练步针对当前采样的 mini-batch 重新计算一次代价矩阵、跑一次 Sinkhorn、得到当前 batch 内的配对关系再把这个配对用于 Loss 的计算。在线这个词说起来简单但它带来一个容易被低估的优点统计稳定性。离线预计算的一次性配对完全取决于初始噪声采样如果这批噪声点本身分布有偏差整个 epoch 的配对都会跑偏。而在线方式每次都用新的随机噪声样本重新配对即使偶尔一个 batch 配得不好下个 batch 也会自动纠正等价于在训练过程中对配对关系做了一种隐式的自集成。实现上一次 minibatch OT 的流程大概是从数据集中随机采样一个 batch得到数据点 x1从先验分布采样同样数量的噪声点 x0计算这个 batch 内两两之间的距离矩阵 C跑 50~100 次 Sinkhorn 迭代得到软耦合矩阵 P按 P 采样或取 argmax 得到配对索引用配对后的噪声点和数据点构造插值路径计算流匹配损失。整个流程在 batch size 为 256 时Sinkhorn 迭代 100 次的耗时大约只占一个训练 step 的 5%~20%完全在可接受范围内。3.2 minibatch 大小到底该取多少minibatch OT 在数学上有一个天然偏差它是在一小撮样本的局部经验分布之间做匹配而不是在全局分布之间做匹配。batch 越小局部分布越稀疏配对质量离全局 OT 越远。所以理论上 batch 越大越接近全局最优但计算复杂度和显存开销随着 batch 增大呈二次方到三次方增长。以我自己的经验来看batch size 128 是一个下限低于这个值 OT 的优势基本发挥不出来甚至可能因为配对过拟合局部噪声而比随机配对更差。batch size 256 到 512 是性价比最高的区间Sinkhorn 迭代能稳定收敛计算开销也可控。如果显存允许1024 的 batch 在图像生成任务上能明显看到 FID 和样本质量的提升但训练速度会下降不少。还有一点值得注意增大的 batch 不只是让 OT 配对更准它对神经网络的 BatchNorm如果用了的话和梯度估计也有正面影响。所以在对比实验里要单独控制 batch 变量避免把 OT 的收益和 batch 本身的收益混在一起。3.3 epsilon 是调节“对齐强度”的旋钮epsilon 是整个方法里最需要调的超参数没有之一。它的作用是在“精确匹配”和“均匀分布”之间做折中。我见过很多人第一次跑通代码后生成结果跟随机配对差不多查到最后都是 epsilon 太大耦合矩阵已经趋近均匀分布等于没做 OT。epsilon 应该怎么选核心是看代价矩阵 C 的尺度。如果你的数据已经归一化到接近 [0,1] 范围那欧氏距离平方大约在 0.1 到 3 之间epsilon 取 0.01 到 0.1 通常比较合理。但如果你是直接在高维原始像素或特征空间里算距离距离尺度可能到几十上百这时候还用 0.05 的 epsilon核矩阵 exp(-C/eps) 会大面积下溢成 0Sinkhorn 直接数值崩坏。我推荐的做法是先把距离矩阵除以它的均值把尺度标准化到 1 附近然后从 epsilon 0.05 开始调。另外可以观察耦合矩阵的熵如果 P 的熵接近均匀分布的熵说明 epsilon 太大如果 P 的每行都接近 one-hot说明 epsilon 太小训练会缺乏多样性。经验法则让经过尺度归一化后的代价均值在 1 左右epsilon 从 0.02~0.1 之间开始网格搜索观察生成质量和训练 loss 稳定性。3.4 显存不够时怎么办滑窗与滚动估计很多实际场景里显存根本撑不住一次算 512×512 的距离矩阵加几十次迭代。尤其是高分辨率图像直接在像素空间算距离时单张图展平就是上万维向量batch 稍微大一点矩阵乘法就直接 OOM。这个时候不要硬上有两个工程手段可以兜底。第一个是“滑窗 OT”把大 batch 拆成多个子块每个子块内部做 OT但子块的划分不是固定的每次训练步都重新随机划分。这样虽然每个子块内部的配对质量略低于整批 OT但随机分块带来的多样性反而能缓解局部最优问题。第二个手段是“滚动估计”维护一个队列把最近几个 batch 的噪声点和数据点缓存起来每次计算 OT 时从缓存中取一部分样本、再混合当前 batch 的一部分样本。这相当于把有效 batch size 虚拟放大了一倍到两倍配对质量会更稳定。代价是数据变成了“延迟更新”实现的时候要注意队列更新逻辑别把同一个样本在多个 step 里重复用太多次。4. 实操中踩过的坑与排查记录4.1 Sinkhorn 不收敛时我做了什么Sinkhorn 不收敛是跑 online OT 最常遇见的故障常见表现就是 loss 震荡、偶尔出现 NaN。我自己排查下来极少数情况是因为迭代次数不够大多数情况是数值问题epsilon 太小会在 exp(-C/eps) 时产生 inf 或 0最终 P 出现 NaN。遇到 NaN第一步不是继续调网络而是把 epsilon 调大一个数量级试试看能不能恢复正常。如果调大 epsilon 后 NaN 消失就说明是数值下溢问题。第二步是给核矩阵加一个数值下限比如 clamp 到 1e-12防止除零。第三步是检查输入数据有没有异常值某张图如果像素值特别大会让距离矩阵出现一个极端值整个 exp 计算都塌掉。固定迭代次数在某些边界情况下也会导致行列约束没满足。所以我会在训练脚本里做一次松散的边际检查对最后的 P 求行和、列和看是否都在 1/n 的合理误差范围内。如果误差超过 5%就把迭代次数往上加或者减小 epsilon。4.2 多样性崩塌与“退化成随机配对”另一个高频问题是训练中期开始生成质量不错但样本多样性越来越差最后所有生成结果都像同一个模板。这种情况往往不是 OT 本身的锅而是 argmax 硬配对 epsilon 太小联合导致的。前面说过argmax 会把多个噪声点配到同一个数据点大 batch 下这种“一对多”的现象尤其明显。模型学到的是少数几个数据点的精确映射把其他模式彻底忽略了。解决方式很简单把 argmax 换成按行概率采样或者把 epsilon 稍微调大一点。还有一个辅助手段是给代价矩阵加一点小的随机噪声让 OT 的 argmax 结果产生轻微的扰动等价于在配对层面做 label smoothing。我在一个多模态图像数据集上专门对比过epsilon 0.01 时多样性指标崩塌epsilon 0.05 时多样性恢复但单样本质量略降epsilon 0.1 时生成结果接近随机配对多重模态被平均掉。最后选在 0.03 到 0.05 之间质量与多样性兼顾。4.3 训练不稳定与损失曲线的异常波动如果用了 online OT 之后训练损失出现周期性波动先想想是不是 batch 内偶然混入了“很难配”的样本组合。比如某个 batch 恰好抽到了几个离群点代价矩阵最大值明显放大Sinkhorn 的输出会变得更尖锐导致该步的配对方案和损失都异常。这种问题不是 bug而是统计波动。要让它不干扰训练可以对损失做梯度裁剪或者把每个 step 的损失做 EMA 平滑再根据 EMA 调整学习率。但更我想提醒的是不要为了让损失曲线平滑就盲目调小学习率那会把 OT 带来的加速效果吃回去。先确认这个波动是否在合理范围内再决定要不要干预。如果波动幅度超过正常水平一个高效的排查手段是把每个 step 的耦合矩阵平均熵打点记录下来。熵突然降低说明当前 batch 的 OT 配对极度自信往往对应离群点熵突然升高说明 Sinkhorn 没有收敛到有意义的解。看到这两类信号再去调 epsilon 和迭代次数就有方向了。4.4 对照组到底应该怎么设才可信给在线 minibatch OT 做实验评估时最容易犯的错是只和随机配对比然后得出“OT 大幅提升”的结论。但这个结论很可能混入了一个干扰变量最优传输同时改变了梯度估计的分布网络看到了更“有序”的配对序列而这种序列效应本身就会让训练更稳定。要分离变量至少要加三组对照随机配对 baseline离线预处理的全量 OT 配对在小数据集上可行online minibatch OT。其中离线预处理的对照组尤其关键它能帮你判断在线计算带来的提升到底是因为“配对了”还是因为“在线了”。我自己的项目里在线 OT 比离线 OT 的最终指标高大约 5%~10%但比随机配对高接近 20%这说明在线动态配对确实有独立价值。还要注意统一采样步数、统一验证集评估脚本、统一随机种子。OT 会天然降低训练初期的 loss 方差如果你在早停或 checkpoint 选择上不加约束很容易在实验对比里偏向 OT 方案得出过乐观的结论。5. 适用场景与实际效果参考5.1 什么任务值得上 Minibatch OT从我的实际使用经验来看最值得上 minibatch OT 的是数据分布高度多模态、且样本间距离度量有明确语义的任务。图像生成是典型的受益者因为不同类别的图像在特征空间距离远OT 配对能让噪声点各自锁定一个语义区域模型不需要课程学习就能先分而治之再细化。其次是分子构象生成或小分子三维结构生成这类几何数据。这类数据的生成质量高度依赖局部几何正确性随机配对经常让一个分子骨架漂移到完全不同的构象上OT 配对则会让贴合的骨架始终保持近邻关系生成的三维结构的 RMSD 明显更低。表格数据生成也是很好的应用场景。表格数据往往维度不高、样本量大代价矩阵的计算和 Sinkhorn 迭代成本其实很低一次 OT 的开销可能连一个梯度 step 的 5% 都不到但配对后的生成分布能更好地保留特征间的相关性。时间序列插值或轨迹预测也可以受益只要你能定义一个合理的样本间距离OT 会自然地按轨迹相似度配对。5.2 典型任务上的效果趋势参考这里我整理一个模拟项目中的表现趋势具体数字不展开因为不同数据集的绝对值差异很大但趋势是稳定可复现的场景随机配对离线预计算 OTOnline Minibatch OT多模态图像生成生成质量中等多样性一般质量提升但偶见模式遗漏质量与多样性同步提升小批量表格数据特征相关性保持一般相关性较好更新滞后相关性最好训练最快稳定高维分子构象常见结构塌缩几何局部质量提升几何质量与多样性均最佳从这张表能看到两个关键点。第一online OT 几乎不会比离线 OT 差即使偶尔配对波动引入噪声整体训练仍然受益于“每个 step 都用最新数据重新配准”这一点。第二随机配对在高维复杂数据上不是“略差一点”而是存在结构性劣势——这是 pairing 这个隐式设计点在流匹配训练里真实的影响力。5.3 未来可以做哪些扩展在线 minibatch OT 这套思路还有一个很自然的扩展方向把配对分布本身参数化。既然每个 step 都算一次 OT那算完之后得到的是一个明确的配对分布我们完全可以把这个分布当成一个额外信号反馈给数据增强策略或课程学习调度。比如训练初期 epsilon 调大一点让配对更温和中期再逐渐减小 epsilon 收紧配对这相当于一种“预告锚定”的 anneal 策略我在一个小型实验中看到过大约 10% 的收敛速度提升。另一个方向是把这个机制从训练搬到推理采样。流匹配的采样质量在不同时间步上的敏感度不同如果我们能根据当前样本的集中度在原空间对采样轨迹做轻量级的 OT 对齐也许能在少步数采样场景下换取更平滑的轨迹。这个方向还比较粗糙但我认为值得在有足够算力的场景下试一次。回到实践中我的总体感受是Flow Matching 的收益很多时候不只在网络结构和调度器还在那些容易忽略的样本配对细节上。Online minibatch OT 的实现门槛不高但一旦调好它对训练收敛速度和最终生成质量的提升是很稳定的。如果你正在跑流匹配项目我建议至少把它作为对照方案加进实验矩阵里——在不改变网络结构的情况下这可能是性价比最高的单点改动之一。