1. 项目概述与核心思路最近在折腾一些图像生成和风格迁移的项目发现一个绕不开的核心问题如何高效、准确地度量并迁移两个数据分布之间的差异。比如你想让一张风景照拥有另一张画的色调或者让一个生成模型学会MNIST手写数字的分布。传统方法里最优传输Optimal Transport, OT是个数学上很优雅的框架它把这个问题看作一个“搬沙子”的最优方案——用最小的成本把一堆沙子源分布搬成另一堆沙子目标分布的形状。但玩过的人都知道这玩意儿计算起来是真要命尤其是面对高维图像数据动辄就是O(n³)的复杂度直接劝退。所以社区里一直在找更快的路子。切片最优传输Sliced OT是个聪明的思路与其在高维空间里硬算不如把分布投影到一堆随机的一维方向上在一维空间里算瓦瑟斯坦距离这个有闭式解快得很然后再把这些距离平均起来。这大大降低了计算成本。但每次遇到新的一对分布你还是得重新做一遍投影和计算对于需要频繁计算OT的任务比如训练生成模型开销依然不小。这就引出了“摊销最优传输”Amortized OT的概念。它的核心思想是“一次学习多次使用”。我们训练一个神经网络让它学会根据输入的源分布和目标分布直接预测出它们之间的最优传输计划或相关的耦合矩阵。一旦模型训练好再遇到新的分布对直接前向传播一下就能得到结果省去了每次迭代求解的昂贵过程。这有点像你背熟了九九乘法表以后算乘法就不用再掰手指头了。我这次重点实验和剖析的是两种具体的摊销OT方法基于回归的摊销OTRA-OT和基于目标的摊销OTOA-OT。它们的目标一致但实现路径和哲学有所不同。RA-OT更直接它让神经网络去回归一个预先用经典OT算法比如Sinkhorn计算好的“标准答案”传输计划。OA-OT则更“端到端”它不依赖外部标签而是定义一个可微的OT目标函数比如切片瓦瑟斯坦距离然后让神经网络通过优化这个目标来学习如何预测传输计划。简单说RA-OT是“模仿学习”OA-OT是“强化学习”自监督学习。这篇文章我就结合在MNIST灰度传输、球形供需匹配和颜色迁移这三个任务上的实测来拆解这两种方法。特别是我们会深入一个关键的“旋钮”切片投影的数量L。它就像相机的分辨率太少看不清细节太多可能又浪费。通过系统的消融实验我们来看看这个参数到底怎么影响效果和速度以及在实际项目中该如何设置。最后我们还会看到它们在CIFAR-10图像生成这种高维任务上的威力证明其可扩展性。2. 核心原理与方案选型解析2.1 从最优传输到切片与摊销为什么是它们要理解RA-OT和OA-OT我们得先捋清楚几个关键概念。最优传输OT的经典形式可以表述为给定两个概率分布比如两堆点集寻找一个传输计划一个矩阵表示从源点i到目标点j运多少“质量”使得总传输成本最小。这个成本通常是点对之间的欧氏距离的p次方。当p2时就是著名的2-Wasserstein距离。直接求解OT是个线性规划问题计算复杂度极高。熵正则化与Sinkhorn算法为了可计算我们引入一个熵正则项把问题变成凸优化然后用Sinkhorn迭代算法求解。这成了目前最流行的近似OT方法但它仍然需要迭代每次面对新分布对都要重算。切片最优传输Sliced OT提供了一种降维打击的思路。其核心是**拉东变换Radon Transform**的一个特例。对于任何高维分布我们随机选取L个方向单位球面上的随机向量。对于每个方向我们将分布中所有点投影到这个方向的一维直线上。神奇的是两个高维分布之间的切片瓦瑟斯坦距离等于所有一维投影之间的瓦瑟斯坦距离的期望。而一维瓦瑟斯坦距离有解析解——只需要对两个一维投影后的样本进行排序然后计算对应序号的样本差值的绝对值之和即可。计算复杂度从O(n³)降到了O(L n log n)其中排序是主要开销。L决定了我们“观察”分布的视角数量。摊销最优传输Amortized OT则引入了学习的思想。我们训练一个参数化的模型比如一个神经网络f_θ。输入是源分布和目标分布的某种表示比如点云或它们的统计特征输出是预测的传输计划P_θ。训练完成后对于新的分布对我们只需计算f_θ的一次前向传播就能得到近似的OT方案避免了运行Sinkhorn或多次切片投影计算。这本质上是将在线计算成本转移到了离线的模型训练上。2.2 RA-OT vs. OA-OT两条不同的学习路径现在来看我们实验的两位主角RA-OT (Regression-based Amortized OT)核心思想监督学习。我们需要一个“教师”来生成训练数据。这个教师通常是一个精确但较慢的OT求解器比如收敛后的Sinkhorn算法。对于每一对训练分布源和目标我们用教师算法计算出“真实”的最优传输计划P_true。学习目标训练神经网络f_θ使其预测的传输计划P_θ尽可能接近P_true。损失函数通常采用均方误差MSE或Frobenius范数L_RA ||P_θ - P_true||^2。优点训练目标清晰、稳定因为P_true提供了一个明确的监督信号。如果教师算法足够精确RA-OT可以学到一个很好的近似。挑战依赖高质量的“真值”标签。生成这些标签本身计算昂贵需要运行很多次Sinkhorn。而且如果教师算法在某些情况下本身就有误差或偏差学生模型也会继承这些偏差。OA-OT (Objective-based Amortized OT)核心思想自监督学习。它不依赖于外部标签。相反它直接优化一个与OT相关的、可微的代理目标函数。学习目标通常这个代理目标是切片瓦瑟斯坦距离Sliced Wasserstein Distance, SWD。神经网络f_θ预测一个传输计划P_θ然后我们基于P_θ构造出传输后的样本或计算耦合成本再与目标分布计算SWD。损失函数就是L_OA SWD(Transported_Source_via_Pθ, Target)。工作原理模型通过反向传播直接学习如何生成一个传输计划使得按照这个计划移动源样本后得到的分布与目标分布在切片距离上尽可能接近。它不需要知道“真正”的OT计划长什么样只关心结果好不好。优点无需昂贵的预计算标签训练数据准备简单。理论上可以探索到不同于传统算法的新解。挑战优化可能更困难因为SWD作为损失函数可能存在平坦区域或梯度问题。需要仔细设计网络结构和训练技巧。为什么同时研究它们因为它们代表了摊销学习的两种主流范式。RA-OT告诉我们在有高质量监督信号时模型能学到多好。OA-OT则探索了在弱监督甚至无监督下模型的潜力。在实际应用中如果你有计算资源预生成标签RA-OT可能更稳如果你追求纯粹的端到端和便捷OA-OT更有吸引力。2.3 关键超参数切片投影数L的角色在切片方法中投影方向的数量L是一个至关重要的超参数。你可以把它理解为模型“理解”高维分布的“词汇量”。L太小例如35模型只能从极少数几个随机方向“看”数据。这就像只用三五个角度去观察一个复杂的雕塑很可能会错过大量的结构信息。对于MNIST数字可能无法区分“3”和“8”的细微弯曲对于颜色迁移无法捕捉色彩分布中复杂的多模态特性。结果就是模型性能较弱预测的传输计划粗糙误差大。L适中例如102050随着视角增多模型能捕获到更多的分布结构信息。性能会显著提升误差快速下降。我们的实验表明在大多数任务上当L从3增加到20左右时提升最为明显。L很大例如100甚至更多性能提升会进入平台期。因为增加到一定程度后新增加的随机方向带来的信息增量变得有限可能和已有方向高度冗余。此时再增加L对精度的改善微乎其微。最棒的一点是在我们的实现中增加L几乎不增加额外的计算时间。这是因为切片计算是高度并行的——所有L个投影方向的计算可以同时进行。主要的开销在于对每个投影后的样本进行排序O(n log n)而排序操作对于每个L是独立的且可以批处理。所以训练和推理时间在L从3到100的变化中基本保持恒定。这意味着在实践中我们可以放心地设置一个较大的L比如50或100以确保性能饱和而无需担心效率惩罚。这为参数选择提供了极大的便利。3. 实验设置与任务拆解为了全面评估RA-OT和OA-OT我们设计了三个具有代表性的任务覆盖了从简单到复杂、从低维到高维特性的不同场景。3.1 任务一MNIST灰度传输这是一个经典的分布对齐基准测试。源分布从MNIST测试集中随机采样的一批手写数字图像灰度展平为784维向量并归一化为概率质量。目标分布另一批随机采样的MNIST数字图像。目标学习一个传输计划将源批次的数字“转换”为目标批次的数字同时尽可能保持数字的语义结构例如把“7”变成“1”是合理的但把“7”变成“8”可能就不对。评估指标我们计算预测的传输计划与Sinkhorn算法求得的“地面真值”计划之间的均方根误差Plan RMSE。这个误差越小说明预测的计划越接近最优解。为什么选它MNIST数据具有明确的类别结构和相对较低的维度784维适合作为初始验证看方法是否能捕捉到离散的、结构化的分布差异。3.2 任务二球形供需匹配这个任务模拟了一个更几何化、连续的空间。源分布需求模拟全球人口分布点密集地集中在几个大洲区域。目标分布供给均匀散布在全球球面上的点代表资源供给点。挑战数据位于球面2维流形上标准的欧氏切片投影需要调整为球形切片使用球面上的随机大圆进行投影。这测试了方法在非欧几里得空间上的适应性。目标找到一个传输计划将密集的人口需求点匹配到均匀的供给点上理想情况下应产生平滑的、全局连贯的匹配模式而不是产生局部畸变。评估指标同样使用Plan RMSE并与基于球形切片的Sinkhorn结果对比。为什么选它测试方法在流形数据上的泛化能力以及处理具有明确几何约束的全局匹配问题的能力。3.3 任务三颜色迁移这是计算机视觉中的一个实用任务复杂度最高。源与目标两张不同的自然彩色图像。处理将每张图像的所有像素的RGB颜色值三维视为一个3维点云构成一个经验分布。目标学习一个传输计划将源图像的颜色分布映射到目标图像的颜色分布上。执行传输后源图像应具有目标图像的色调和色彩风格但同时保留自身的纹理和内容信息。挑战颜色分布通常是复杂、多模态的比如一张图同时有蓝天、绿树、肤色。传输计划需要精细地处理这些模式避免产生色块、伪影或内容扭曲。评估指标除了Plan RMSE我们更注重定性视觉评估。好的结果应该颜色迁移自然内容保真度高。为什么选它这是高维3维虽然不高但分布复杂、多模态分布对齐的典型代表非常考验模型的表达能力。3.4 模型实现与训练细节对于RA-OT和OA-OT我们采用了一个轻量级的多层感知机MLP作为摊销网络f_θ。输入对于点云数据我们先将源点和目标点分别通过一个共享权重的点云编码器如PointNet的小型变体提取特征然后将两个特征向量拼接后送入MLP。对于图像数据如颜色迁移我们使用RGB值直接作为点坐标。输出一个batch_size x batch_size的矩阵表示预测的传输耦合经过softmax或sinkhorn归一化以确保是双随机矩阵。训练使用Adam优化器学习率设为1e-3。RA-OT使用MSE损失OA-OT使用切片瓦瑟斯坦距离作为损失。对于所有任务我们固定训练迭代次数为5000次并使用一个包含50个小批量M50的经验分布来训练摊销网络以增加其泛化能力。注意这里M50是一个关键技巧。我们不是只用一对分布来训练而是从数据集中采样50个不同的源-目标对小批量用这些对的聚合经验分布来训练网络。这迫使网络学习的是“如何为这一类分布对计算OT”而不是记忆某一个特定的配对从而大大提升了泛化能力。4. 核心实验结果投影数L的消融分析这是我们本次探索的重头戏。我们系统地改变了切片投影的数量L ∈ {3, 5, 10, 20, 50, 100}在三个任务上分别测试了RA-OT和OA-OT的性能。下表汇总了关键结果基于提供的Table 9数据整理L值任务方法Plan RMSE (×10⁻⁶, ↓)训练时间 (秒)推理时间 (毫秒)3MNISTRA-OT9.25 ± 3.733.2237.47 ± 3.91OA-OT8.50 ± 3.5515.3537.38 ± 2.65SphericalRA-OT0.71 ± 0.172.1730.51 ± 4.43OA-OT0.41 ± 0.1618.6234.20 ± 5.70ColorRA-OT25.60 ± 10.426.6816.48 ± 0.96OA-OT23.96 ± 6.7916.8617.29 ± 1.1520MNISTRA-OT7.62 ± 3.113.1838.11 ± 2.45OA-OT6.05 ± 2.5315.3537.83 ± 2.23SphericalRA-OT0.74 ± 0.192.3332.06 ± 4.97OA-OT0.38 ± 0.1618.2330.89 ± 4.24ColorRA-OT9.39 ± 5.507.0317.64 ± 1.28OA-OT9.11 ± 5.1117.3117.48 ± 1.37100MNISTRA-OT7.77 ± 3.063.0339.36 ± 3.61OA-OT6.02 ± 2.5215.7838.92 ± 2.23SphericalRA-OT0.78 ± 0.192.5341.96 ± 6.29OA-OT0.39 ± 0.1919.3741.03 ± 2.76ColorRA-OT9.99 ± 5.617.4017.40 ± 0.83OA-OT9.00 ± 5.0218.1317.76 ± 0.79核心发现解读性能随L增加而提升但存在收益递减点在所有任务上当L从极小的3或5开始增加时RMSE误差显著下降。特别是在颜色迁移任务上RA-OT的误差从25.6降至9.39L20OA-OT从23.96降至9.11提升超过50%。这印证了L太小时模型“看不清”复杂分布结构的假设。大约在L20左右性能提升曲线开始变得平缓。增加到50或100时改善已经非常微小有时甚至因为随机性而有轻微波动。计算成本几乎恒定这是切片方法一个非常强大的优势。观察“训练时间”和“推理时间”两列无论是3个投影还是100个投影时间开销几乎没有系统性增长。训练时间主要取决于网络前向/反向传播和优化步骤推理时间主要取决于网络计算。切片投影的计算是高度并行化的增加L只是增加了并行计算的任务量在GPU等硬件上几乎不增加额外延迟。这意味着在实践中我们可以毫无负担地选择一个较大的L如50或100来确保性能饱和而不用担心效率损失。OA-OT与RA-OT的比较精度在MNIST和球形任务上OA-OT的RMSE普遍低于RA-OT表明其自监督学习的目标函数能够引导网络找到更接近真实OT的解。在颜色迁移上两者在L较大时表现相当。训练时间OA-OT的训练时间显著长于RA-OT约5-6倍。这是因为OA-OT的损失函数SWD计算涉及每个投影方向的排序操作并且在反向传播时需要计算SWD对网络参数的梯度这比RA-OT简单的MSE损失要复杂。适用场景如果你追求极致的精度且有足够的训练时间预算OA-OT是更好的选择。如果你需要快速训练和部署或者无法获得高质量的OT真值标签RA-OT是更实用的选择。任务难度差异颜色迁移任务的RMSE值远高于其他两个任务数量级在10⁻⁵ vs 10⁻⁶。这反映了颜色分布对齐本身是一个更困难的问题存在更多模糊性和可能的解。同时也说明即使对于复杂任务摊销OT方法也能给出有意义的近似。实操心得基于这些结果我的经验是将L设置为20到50之间是一个很好的起点。这个范围在大多数任务上已经能捕获足够的信息达到接近饱和的性能。如果你对性能有极致要求可以尝试增加到100。几乎不用担心L增大会拖慢你的程序这是切片方法送给我们的“免费午餐”。在训练OA-OT时需要对训练时间有更多耐心但其最终精度往往值得等待。5. 高维图像生成实战CIFAR-10上的流匹配微调理论分析和低维实验固然重要但真正的考验在于高维现实任务。我们选择在CIFAR-10图像生成上测试摊销OT方法如何赋能基于流匹配Flow Matching的生成模型。5.1 背景流匹配与最优传输耦合流匹配是一种新兴的生成模型框架它通过学习一个将简单先验分布如高斯噪声映射到复杂数据分布如图像的常微分方程ODE来生成数据。其中条件流匹配Conditional Flow Matching, CFM需要为每个噪声样本配对一个数据样本这个配对策略直接影响生成路径的“直”度straightness进而影响采样效率。最朴素的是独立耦合I-CFM即随机配对噪声和数据。这通常会导致弯曲、交叉的流轨迹需要更多的数值积分步骤NFE来生成高质量样本。最优传输耦合OT-CFM使用OT计划来配对噪声和数据可以产生更直、不交叉的流轨迹从而降低NFE提升采样效率。但问题是每次训练迭代都计算OT耦合即使在小批量上开销巨大。5.2 实验设置公平的微调竞技场为了进行公平比较我们设计了一个严谨的微调实验基模型我们首先用一个标准的I-CFM方法在CIFAR-10上从头训练一个U-Net模型约3500万参数直到收敛40万步。这个模型作为我们所有后续实验的起点。微调协议我们加载这个相同的预训练模型 checkpoint。然后仅改变小批量内的配对策略进行为期10个epoch的微调。比较四种策略I-CFM基线独立随机配对。OT-CFM使用精确的Sinkhorn算法计算每个小批量的OT配对。OA-OT使用我们训练的OA-OT摊销网络来预测OT配对。RA-OT使用我们训练的RA-OT摊销网络来预测OT配对。关键细节使用超大批次B2048来模拟数据密集型训练场景这对OT-CFM的计算是巨大挑战。对于摊销方法OA-OT/RA-OT我们使用M50个小批量训练好的预测器L100。微调时采用保守的超参数更低的学习率5e-5、梯度裁剪0.5、更短的warmup以防止灾难性遗忘。评估使用FID弗雷歇起始距离衡量生成图像的质量使用NFE函数评估次数衡量采样效率。NFE越低说明流轨迹越直采样越快。5.3 结果与分析实验结果基于提供的Table 10令人振奋方法FID (↓)NFE/样本 (↓)训练时间/epoch (秒, ↓)预计算时间 (秒, ↓)I-CFM3.638146.61481.90.0OT-CFM3.630146.86744.20.0OA-OT3.575146.00607.012.0RA-OT3.543147.10618.116.4核心结论效率的巨大优势OT-CFM由于要在每个大批次2048上实时计算OT训练时间比I-CFM基线增加了超过50%744.2秒 vs 481.9秒。而我们的摊销方法OA-OT和RA-OT虽然也有预计算摊销网络的时间约12-16秒这是一次性开销但在训练时的每个epoch上它们只比I-CFM慢约25-30%远低于OT-CFM。这证明了摊销策略成功地将在线计算瓶颈转移到了离线阶段。性能不降反升更令人惊喜的是摊销OT不仅仅是为了快它们还带来了更好的生成质量RA-OT取得了最好的FID分数3.543超越了所有其他方法包括精确的OT-CFM。OA-OT则取得了最低的NFE146.00意味着它学习到了更直的流轨迹理论上采样更快。这表明摊销网络并非简单地“模仿”OT在优化过程中它可能学习到了某种更有利于生成任务的规则化或平滑特性。可视化验证从提供的流轨迹图如从高斯先验到8-Gaussians、Moons、S-Curve分布可以清晰看到I-CFM的轨迹黑色点移动的橄榄色路径是弯曲、交叉的。OT-CFM的轨迹则变得笔直、不交叉。而我们的OA-OT和RA-OT几乎完美地复现了OT-CFM这种笔直、无纠缠的轨迹显著降低了“碰撞能量”这与它们优秀的NFE指标相符。避坑技巧在高维流匹配微调中使用摊销OT时一个关键点是微调策略要温和。预训练模型已经学到了丰富的视觉特征我们只想调整其“运动路径”由配对策略决定。因此必须使用极低的学习率、梯度裁剪和warmup否则容易破坏已学到的特征导致生成质量崩溃。我们的实验设置学习率5e-5梯度裁剪0.5是一个有效的起点。6. 常见问题与实战排查指南在实际实现和应用摊销OT时我踩过不少坑。这里总结几个最常见的问题和解决思路。6.1 训练不稳定或发散问题描述特别是OA-OT损失值震荡剧烈甚至变成NaN。可能原因与解决学习率过高OT损失曲面可能很崎岖。尝试将学习率从1e-3降低到5e-4或1e-4。梯度爆炸SWD损失计算中涉及排序其梯度在某些情况下可能很大。引入梯度裁剪Gradient Clipping例如设置范数阈值为1.0或5.0。投影方向L太少在训练初期如果L太小比如3损失的估计方差会很大导致梯度噪声大。确保使用足够大的L≥20进行训练。网络输出未归一化网络直接输出的矩阵可能不是有效的耦合矩阵行和与列和不为1。需要在网络末端添加Sinkhorn归一化层即使只迭代几次或简单的行/列softmax以确保输出是双随机矩阵。6.2 泛化能力差问题描述在训练集上表现很好但在测试集或新数据对上性能骤降。可能原因与解决过拟合摊销网络可能记住了训练用的特定分布对。增加M用于训练的经验分布小批量数量是提升泛化的最有效手段。不要只用一对分布训练用几十甚至上百对。网络容量不足或过度MLP太浅可能学不到复杂映射太深可能过拟合。对于点云输入一个包含残差连接的3-5层MLP通常是个不错的起点。可以尝试在输入前加入一个简单的点云特征提取层如通过几个线性层和最大池化获取全局特征。训练数据分布不具代表性确保用于训练摊销网络的小批量分布覆盖了测试时可能遇到的分布多样性。对于图像颜色迁移应从大量不同主题和色调的图像中采样训练对。6.3 颜色迁移出现色斑或伪影问题描述迁移后的图像颜色不自然出现大块均匀色斑或局部颜色失真。可能原因与解决传输计划过于“硬”如果预测的耦合矩阵非常稀疏接近one-hot会导致每个源像素只匹配到极少数目标像素容易产生色块。在OA-OT的SWD计算中或网络输出后可以加入一个小的熵正则项鼓励更平滑的传输计划。忽略了空间信息我们的方法只操作颜色直方图完全丢弃了像素位置信息。对于要求空间一致性的风格迁移这可能不够。可以考虑将像素的坐标x, y或深度特征与颜色值拼接作为5维或更高维的点让OT同时考虑颜色和空间的相似性。但这会提高计算维度。后处理执行OT传输后可以对结果图像进行轻微的双边滤波或引导滤波在平滑颜色的同时保持边缘这能有效减少伪影。6.4 与基线方法对比时效果不明显问题描述在简单任务如2D高斯上摊销OT的效果和精确OT差不多但没显示出明显优势。正确认识在低维、简单分布上精确OT如Sinkhorn本身已经很快很准。摊销OT的优势在于高维和重复计算的场景。它的价值体现在批量推理速度一旦训练好对成千上万对新分布进行推理的速度是常数级的而Sinkhorn需要线性或更差的增长。嵌入工作流可以作为一个可微层嵌入到更大的端到端模型中如图像生成器进行联合训练这是迭代式OT求解器难以做到的。处理超大批次如上文的CIFAR-10实验当批次大到Sinkhorn都算不动时摊销OT是唯一可行的选择。6.5 如何选择RA-OT还是OA-OT这是一个实践中的关键决策。我的建议如下选择RA-OT如果你拥有充足的计算资源来预生成高质量的OT“真值”标签例如可以用Sinkhorn在较小的子集上精细计算。追求更快的训练速度和更稳定的训练过程。你的任务中精确模仿某个特定OT求解器的行为很重要。选择OA-OT如果你无法获得或生成OT真值标签例如在无监督或自监督学习框架中。追求极致的最终性能并且愿意投入更长的训练时间。希望探索可能超越传统OT解法的、由目标函数引导的新型传输计划。需要将OT模块完全无缝地集成到一个可微分的管道中。最后别忘了利用投影数L的免费午餐特性。在你的硬件允许的范围内尽可能设置一个较大的L比如50或100这几乎总能带来性能提升或至少更稳定的训练而成本微乎其微。
摊销最优传输实战：RA-OT与OA-OT在分布对齐中的性能对比与调优
1. 项目概述与核心思路最近在折腾一些图像生成和风格迁移的项目发现一个绕不开的核心问题如何高效、准确地度量并迁移两个数据分布之间的差异。比如你想让一张风景照拥有另一张画的色调或者让一个生成模型学会MNIST手写数字的分布。传统方法里最优传输Optimal Transport, OT是个数学上很优雅的框架它把这个问题看作一个“搬沙子”的最优方案——用最小的成本把一堆沙子源分布搬成另一堆沙子目标分布的形状。但玩过的人都知道这玩意儿计算起来是真要命尤其是面对高维图像数据动辄就是O(n³)的复杂度直接劝退。所以社区里一直在找更快的路子。切片最优传输Sliced OT是个聪明的思路与其在高维空间里硬算不如把分布投影到一堆随机的一维方向上在一维空间里算瓦瑟斯坦距离这个有闭式解快得很然后再把这些距离平均起来。这大大降低了计算成本。但每次遇到新的一对分布你还是得重新做一遍投影和计算对于需要频繁计算OT的任务比如训练生成模型开销依然不小。这就引出了“摊销最优传输”Amortized OT的概念。它的核心思想是“一次学习多次使用”。我们训练一个神经网络让它学会根据输入的源分布和目标分布直接预测出它们之间的最优传输计划或相关的耦合矩阵。一旦模型训练好再遇到新的分布对直接前向传播一下就能得到结果省去了每次迭代求解的昂贵过程。这有点像你背熟了九九乘法表以后算乘法就不用再掰手指头了。我这次重点实验和剖析的是两种具体的摊销OT方法基于回归的摊销OTRA-OT和基于目标的摊销OTOA-OT。它们的目标一致但实现路径和哲学有所不同。RA-OT更直接它让神经网络去回归一个预先用经典OT算法比如Sinkhorn计算好的“标准答案”传输计划。OA-OT则更“端到端”它不依赖外部标签而是定义一个可微的OT目标函数比如切片瓦瑟斯坦距离然后让神经网络通过优化这个目标来学习如何预测传输计划。简单说RA-OT是“模仿学习”OA-OT是“强化学习”自监督学习。这篇文章我就结合在MNIST灰度传输、球形供需匹配和颜色迁移这三个任务上的实测来拆解这两种方法。特别是我们会深入一个关键的“旋钮”切片投影的数量L。它就像相机的分辨率太少看不清细节太多可能又浪费。通过系统的消融实验我们来看看这个参数到底怎么影响效果和速度以及在实际项目中该如何设置。最后我们还会看到它们在CIFAR-10图像生成这种高维任务上的威力证明其可扩展性。2. 核心原理与方案选型解析2.1 从最优传输到切片与摊销为什么是它们要理解RA-OT和OA-OT我们得先捋清楚几个关键概念。最优传输OT的经典形式可以表述为给定两个概率分布比如两堆点集寻找一个传输计划一个矩阵表示从源点i到目标点j运多少“质量”使得总传输成本最小。这个成本通常是点对之间的欧氏距离的p次方。当p2时就是著名的2-Wasserstein距离。直接求解OT是个线性规划问题计算复杂度极高。熵正则化与Sinkhorn算法为了可计算我们引入一个熵正则项把问题变成凸优化然后用Sinkhorn迭代算法求解。这成了目前最流行的近似OT方法但它仍然需要迭代每次面对新分布对都要重算。切片最优传输Sliced OT提供了一种降维打击的思路。其核心是**拉东变换Radon Transform**的一个特例。对于任何高维分布我们随机选取L个方向单位球面上的随机向量。对于每个方向我们将分布中所有点投影到这个方向的一维直线上。神奇的是两个高维分布之间的切片瓦瑟斯坦距离等于所有一维投影之间的瓦瑟斯坦距离的期望。而一维瓦瑟斯坦距离有解析解——只需要对两个一维投影后的样本进行排序然后计算对应序号的样本差值的绝对值之和即可。计算复杂度从O(n³)降到了O(L n log n)其中排序是主要开销。L决定了我们“观察”分布的视角数量。摊销最优传输Amortized OT则引入了学习的思想。我们训练一个参数化的模型比如一个神经网络f_θ。输入是源分布和目标分布的某种表示比如点云或它们的统计特征输出是预测的传输计划P_θ。训练完成后对于新的分布对我们只需计算f_θ的一次前向传播就能得到近似的OT方案避免了运行Sinkhorn或多次切片投影计算。这本质上是将在线计算成本转移到了离线的模型训练上。2.2 RA-OT vs. OA-OT两条不同的学习路径现在来看我们实验的两位主角RA-OT (Regression-based Amortized OT)核心思想监督学习。我们需要一个“教师”来生成训练数据。这个教师通常是一个精确但较慢的OT求解器比如收敛后的Sinkhorn算法。对于每一对训练分布源和目标我们用教师算法计算出“真实”的最优传输计划P_true。学习目标训练神经网络f_θ使其预测的传输计划P_θ尽可能接近P_true。损失函数通常采用均方误差MSE或Frobenius范数L_RA ||P_θ - P_true||^2。优点训练目标清晰、稳定因为P_true提供了一个明确的监督信号。如果教师算法足够精确RA-OT可以学到一个很好的近似。挑战依赖高质量的“真值”标签。生成这些标签本身计算昂贵需要运行很多次Sinkhorn。而且如果教师算法在某些情况下本身就有误差或偏差学生模型也会继承这些偏差。OA-OT (Objective-based Amortized OT)核心思想自监督学习。它不依赖于外部标签。相反它直接优化一个与OT相关的、可微的代理目标函数。学习目标通常这个代理目标是切片瓦瑟斯坦距离Sliced Wasserstein Distance, SWD。神经网络f_θ预测一个传输计划P_θ然后我们基于P_θ构造出传输后的样本或计算耦合成本再与目标分布计算SWD。损失函数就是L_OA SWD(Transported_Source_via_Pθ, Target)。工作原理模型通过反向传播直接学习如何生成一个传输计划使得按照这个计划移动源样本后得到的分布与目标分布在切片距离上尽可能接近。它不需要知道“真正”的OT计划长什么样只关心结果好不好。优点无需昂贵的预计算标签训练数据准备简单。理论上可以探索到不同于传统算法的新解。挑战优化可能更困难因为SWD作为损失函数可能存在平坦区域或梯度问题。需要仔细设计网络结构和训练技巧。为什么同时研究它们因为它们代表了摊销学习的两种主流范式。RA-OT告诉我们在有高质量监督信号时模型能学到多好。OA-OT则探索了在弱监督甚至无监督下模型的潜力。在实际应用中如果你有计算资源预生成标签RA-OT可能更稳如果你追求纯粹的端到端和便捷OA-OT更有吸引力。2.3 关键超参数切片投影数L的角色在切片方法中投影方向的数量L是一个至关重要的超参数。你可以把它理解为模型“理解”高维分布的“词汇量”。L太小例如35模型只能从极少数几个随机方向“看”数据。这就像只用三五个角度去观察一个复杂的雕塑很可能会错过大量的结构信息。对于MNIST数字可能无法区分“3”和“8”的细微弯曲对于颜色迁移无法捕捉色彩分布中复杂的多模态特性。结果就是模型性能较弱预测的传输计划粗糙误差大。L适中例如102050随着视角增多模型能捕获到更多的分布结构信息。性能会显著提升误差快速下降。我们的实验表明在大多数任务上当L从3增加到20左右时提升最为明显。L很大例如100甚至更多性能提升会进入平台期。因为增加到一定程度后新增加的随机方向带来的信息增量变得有限可能和已有方向高度冗余。此时再增加L对精度的改善微乎其微。最棒的一点是在我们的实现中增加L几乎不增加额外的计算时间。这是因为切片计算是高度并行的——所有L个投影方向的计算可以同时进行。主要的开销在于对每个投影后的样本进行排序O(n log n)而排序操作对于每个L是独立的且可以批处理。所以训练和推理时间在L从3到100的变化中基本保持恒定。这意味着在实践中我们可以放心地设置一个较大的L比如50或100以确保性能饱和而无需担心效率惩罚。这为参数选择提供了极大的便利。3. 实验设置与任务拆解为了全面评估RA-OT和OA-OT我们设计了三个具有代表性的任务覆盖了从简单到复杂、从低维到高维特性的不同场景。3.1 任务一MNIST灰度传输这是一个经典的分布对齐基准测试。源分布从MNIST测试集中随机采样的一批手写数字图像灰度展平为784维向量并归一化为概率质量。目标分布另一批随机采样的MNIST数字图像。目标学习一个传输计划将源批次的数字“转换”为目标批次的数字同时尽可能保持数字的语义结构例如把“7”变成“1”是合理的但把“7”变成“8”可能就不对。评估指标我们计算预测的传输计划与Sinkhorn算法求得的“地面真值”计划之间的均方根误差Plan RMSE。这个误差越小说明预测的计划越接近最优解。为什么选它MNIST数据具有明确的类别结构和相对较低的维度784维适合作为初始验证看方法是否能捕捉到离散的、结构化的分布差异。3.2 任务二球形供需匹配这个任务模拟了一个更几何化、连续的空间。源分布需求模拟全球人口分布点密集地集中在几个大洲区域。目标分布供给均匀散布在全球球面上的点代表资源供给点。挑战数据位于球面2维流形上标准的欧氏切片投影需要调整为球形切片使用球面上的随机大圆进行投影。这测试了方法在非欧几里得空间上的适应性。目标找到一个传输计划将密集的人口需求点匹配到均匀的供给点上理想情况下应产生平滑的、全局连贯的匹配模式而不是产生局部畸变。评估指标同样使用Plan RMSE并与基于球形切片的Sinkhorn结果对比。为什么选它测试方法在流形数据上的泛化能力以及处理具有明确几何约束的全局匹配问题的能力。3.3 任务三颜色迁移这是计算机视觉中的一个实用任务复杂度最高。源与目标两张不同的自然彩色图像。处理将每张图像的所有像素的RGB颜色值三维视为一个3维点云构成一个经验分布。目标学习一个传输计划将源图像的颜色分布映射到目标图像的颜色分布上。执行传输后源图像应具有目标图像的色调和色彩风格但同时保留自身的纹理和内容信息。挑战颜色分布通常是复杂、多模态的比如一张图同时有蓝天、绿树、肤色。传输计划需要精细地处理这些模式避免产生色块、伪影或内容扭曲。评估指标除了Plan RMSE我们更注重定性视觉评估。好的结果应该颜色迁移自然内容保真度高。为什么选它这是高维3维虽然不高但分布复杂、多模态分布对齐的典型代表非常考验模型的表达能力。3.4 模型实现与训练细节对于RA-OT和OA-OT我们采用了一个轻量级的多层感知机MLP作为摊销网络f_θ。输入对于点云数据我们先将源点和目标点分别通过一个共享权重的点云编码器如PointNet的小型变体提取特征然后将两个特征向量拼接后送入MLP。对于图像数据如颜色迁移我们使用RGB值直接作为点坐标。输出一个batch_size x batch_size的矩阵表示预测的传输耦合经过softmax或sinkhorn归一化以确保是双随机矩阵。训练使用Adam优化器学习率设为1e-3。RA-OT使用MSE损失OA-OT使用切片瓦瑟斯坦距离作为损失。对于所有任务我们固定训练迭代次数为5000次并使用一个包含50个小批量M50的经验分布来训练摊销网络以增加其泛化能力。注意这里M50是一个关键技巧。我们不是只用一对分布来训练而是从数据集中采样50个不同的源-目标对小批量用这些对的聚合经验分布来训练网络。这迫使网络学习的是“如何为这一类分布对计算OT”而不是记忆某一个特定的配对从而大大提升了泛化能力。4. 核心实验结果投影数L的消融分析这是我们本次探索的重头戏。我们系统地改变了切片投影的数量L ∈ {3, 5, 10, 20, 50, 100}在三个任务上分别测试了RA-OT和OA-OT的性能。下表汇总了关键结果基于提供的Table 9数据整理L值任务方法Plan RMSE (×10⁻⁶, ↓)训练时间 (秒)推理时间 (毫秒)3MNISTRA-OT9.25 ± 3.733.2237.47 ± 3.91OA-OT8.50 ± 3.5515.3537.38 ± 2.65SphericalRA-OT0.71 ± 0.172.1730.51 ± 4.43OA-OT0.41 ± 0.1618.6234.20 ± 5.70ColorRA-OT25.60 ± 10.426.6816.48 ± 0.96OA-OT23.96 ± 6.7916.8617.29 ± 1.1520MNISTRA-OT7.62 ± 3.113.1838.11 ± 2.45OA-OT6.05 ± 2.5315.3537.83 ± 2.23SphericalRA-OT0.74 ± 0.192.3332.06 ± 4.97OA-OT0.38 ± 0.1618.2330.89 ± 4.24ColorRA-OT9.39 ± 5.507.0317.64 ± 1.28OA-OT9.11 ± 5.1117.3117.48 ± 1.37100MNISTRA-OT7.77 ± 3.063.0339.36 ± 3.61OA-OT6.02 ± 2.5215.7838.92 ± 2.23SphericalRA-OT0.78 ± 0.192.5341.96 ± 6.29OA-OT0.39 ± 0.1919.3741.03 ± 2.76ColorRA-OT9.99 ± 5.617.4017.40 ± 0.83OA-OT9.00 ± 5.0218.1317.76 ± 0.79核心发现解读性能随L增加而提升但存在收益递减点在所有任务上当L从极小的3或5开始增加时RMSE误差显著下降。特别是在颜色迁移任务上RA-OT的误差从25.6降至9.39L20OA-OT从23.96降至9.11提升超过50%。这印证了L太小时模型“看不清”复杂分布结构的假设。大约在L20左右性能提升曲线开始变得平缓。增加到50或100时改善已经非常微小有时甚至因为随机性而有轻微波动。计算成本几乎恒定这是切片方法一个非常强大的优势。观察“训练时间”和“推理时间”两列无论是3个投影还是100个投影时间开销几乎没有系统性增长。训练时间主要取决于网络前向/反向传播和优化步骤推理时间主要取决于网络计算。切片投影的计算是高度并行化的增加L只是增加了并行计算的任务量在GPU等硬件上几乎不增加额外延迟。这意味着在实践中我们可以毫无负担地选择一个较大的L如50或100来确保性能饱和而不用担心效率损失。OA-OT与RA-OT的比较精度在MNIST和球形任务上OA-OT的RMSE普遍低于RA-OT表明其自监督学习的目标函数能够引导网络找到更接近真实OT的解。在颜色迁移上两者在L较大时表现相当。训练时间OA-OT的训练时间显著长于RA-OT约5-6倍。这是因为OA-OT的损失函数SWD计算涉及每个投影方向的排序操作并且在反向传播时需要计算SWD对网络参数的梯度这比RA-OT简单的MSE损失要复杂。适用场景如果你追求极致的精度且有足够的训练时间预算OA-OT是更好的选择。如果你需要快速训练和部署或者无法获得高质量的OT真值标签RA-OT是更实用的选择。任务难度差异颜色迁移任务的RMSE值远高于其他两个任务数量级在10⁻⁵ vs 10⁻⁶。这反映了颜色分布对齐本身是一个更困难的问题存在更多模糊性和可能的解。同时也说明即使对于复杂任务摊销OT方法也能给出有意义的近似。实操心得基于这些结果我的经验是将L设置为20到50之间是一个很好的起点。这个范围在大多数任务上已经能捕获足够的信息达到接近饱和的性能。如果你对性能有极致要求可以尝试增加到100。几乎不用担心L增大会拖慢你的程序这是切片方法送给我们的“免费午餐”。在训练OA-OT时需要对训练时间有更多耐心但其最终精度往往值得等待。5. 高维图像生成实战CIFAR-10上的流匹配微调理论分析和低维实验固然重要但真正的考验在于高维现实任务。我们选择在CIFAR-10图像生成上测试摊销OT方法如何赋能基于流匹配Flow Matching的生成模型。5.1 背景流匹配与最优传输耦合流匹配是一种新兴的生成模型框架它通过学习一个将简单先验分布如高斯噪声映射到复杂数据分布如图像的常微分方程ODE来生成数据。其中条件流匹配Conditional Flow Matching, CFM需要为每个噪声样本配对一个数据样本这个配对策略直接影响生成路径的“直”度straightness进而影响采样效率。最朴素的是独立耦合I-CFM即随机配对噪声和数据。这通常会导致弯曲、交叉的流轨迹需要更多的数值积分步骤NFE来生成高质量样本。最优传输耦合OT-CFM使用OT计划来配对噪声和数据可以产生更直、不交叉的流轨迹从而降低NFE提升采样效率。但问题是每次训练迭代都计算OT耦合即使在小批量上开销巨大。5.2 实验设置公平的微调竞技场为了进行公平比较我们设计了一个严谨的微调实验基模型我们首先用一个标准的I-CFM方法在CIFAR-10上从头训练一个U-Net模型约3500万参数直到收敛40万步。这个模型作为我们所有后续实验的起点。微调协议我们加载这个相同的预训练模型 checkpoint。然后仅改变小批量内的配对策略进行为期10个epoch的微调。比较四种策略I-CFM基线独立随机配对。OT-CFM使用精确的Sinkhorn算法计算每个小批量的OT配对。OA-OT使用我们训练的OA-OT摊销网络来预测OT配对。RA-OT使用我们训练的RA-OT摊销网络来预测OT配对。关键细节使用超大批次B2048来模拟数据密集型训练场景这对OT-CFM的计算是巨大挑战。对于摊销方法OA-OT/RA-OT我们使用M50个小批量训练好的预测器L100。微调时采用保守的超参数更低的学习率5e-5、梯度裁剪0.5、更短的warmup以防止灾难性遗忘。评估使用FID弗雷歇起始距离衡量生成图像的质量使用NFE函数评估次数衡量采样效率。NFE越低说明流轨迹越直采样越快。5.3 结果与分析实验结果基于提供的Table 10令人振奋方法FID (↓)NFE/样本 (↓)训练时间/epoch (秒, ↓)预计算时间 (秒, ↓)I-CFM3.638146.61481.90.0OT-CFM3.630146.86744.20.0OA-OT3.575146.00607.012.0RA-OT3.543147.10618.116.4核心结论效率的巨大优势OT-CFM由于要在每个大批次2048上实时计算OT训练时间比I-CFM基线增加了超过50%744.2秒 vs 481.9秒。而我们的摊销方法OA-OT和RA-OT虽然也有预计算摊销网络的时间约12-16秒这是一次性开销但在训练时的每个epoch上它们只比I-CFM慢约25-30%远低于OT-CFM。这证明了摊销策略成功地将在线计算瓶颈转移到了离线阶段。性能不降反升更令人惊喜的是摊销OT不仅仅是为了快它们还带来了更好的生成质量RA-OT取得了最好的FID分数3.543超越了所有其他方法包括精确的OT-CFM。OA-OT则取得了最低的NFE146.00意味着它学习到了更直的流轨迹理论上采样更快。这表明摊销网络并非简单地“模仿”OT在优化过程中它可能学习到了某种更有利于生成任务的规则化或平滑特性。可视化验证从提供的流轨迹图如从高斯先验到8-Gaussians、Moons、S-Curve分布可以清晰看到I-CFM的轨迹黑色点移动的橄榄色路径是弯曲、交叉的。OT-CFM的轨迹则变得笔直、不交叉。而我们的OA-OT和RA-OT几乎完美地复现了OT-CFM这种笔直、无纠缠的轨迹显著降低了“碰撞能量”这与它们优秀的NFE指标相符。避坑技巧在高维流匹配微调中使用摊销OT时一个关键点是微调策略要温和。预训练模型已经学到了丰富的视觉特征我们只想调整其“运动路径”由配对策略决定。因此必须使用极低的学习率、梯度裁剪和warmup否则容易破坏已学到的特征导致生成质量崩溃。我们的实验设置学习率5e-5梯度裁剪0.5是一个有效的起点。6. 常见问题与实战排查指南在实际实现和应用摊销OT时我踩过不少坑。这里总结几个最常见的问题和解决思路。6.1 训练不稳定或发散问题描述特别是OA-OT损失值震荡剧烈甚至变成NaN。可能原因与解决学习率过高OT损失曲面可能很崎岖。尝试将学习率从1e-3降低到5e-4或1e-4。梯度爆炸SWD损失计算中涉及排序其梯度在某些情况下可能很大。引入梯度裁剪Gradient Clipping例如设置范数阈值为1.0或5.0。投影方向L太少在训练初期如果L太小比如3损失的估计方差会很大导致梯度噪声大。确保使用足够大的L≥20进行训练。网络输出未归一化网络直接输出的矩阵可能不是有效的耦合矩阵行和与列和不为1。需要在网络末端添加Sinkhorn归一化层即使只迭代几次或简单的行/列softmax以确保输出是双随机矩阵。6.2 泛化能力差问题描述在训练集上表现很好但在测试集或新数据对上性能骤降。可能原因与解决过拟合摊销网络可能记住了训练用的特定分布对。增加M用于训练的经验分布小批量数量是提升泛化的最有效手段。不要只用一对分布训练用几十甚至上百对。网络容量不足或过度MLP太浅可能学不到复杂映射太深可能过拟合。对于点云输入一个包含残差连接的3-5层MLP通常是个不错的起点。可以尝试在输入前加入一个简单的点云特征提取层如通过几个线性层和最大池化获取全局特征。训练数据分布不具代表性确保用于训练摊销网络的小批量分布覆盖了测试时可能遇到的分布多样性。对于图像颜色迁移应从大量不同主题和色调的图像中采样训练对。6.3 颜色迁移出现色斑或伪影问题描述迁移后的图像颜色不自然出现大块均匀色斑或局部颜色失真。可能原因与解决传输计划过于“硬”如果预测的耦合矩阵非常稀疏接近one-hot会导致每个源像素只匹配到极少数目标像素容易产生色块。在OA-OT的SWD计算中或网络输出后可以加入一个小的熵正则项鼓励更平滑的传输计划。忽略了空间信息我们的方法只操作颜色直方图完全丢弃了像素位置信息。对于要求空间一致性的风格迁移这可能不够。可以考虑将像素的坐标x, y或深度特征与颜色值拼接作为5维或更高维的点让OT同时考虑颜色和空间的相似性。但这会提高计算维度。后处理执行OT传输后可以对结果图像进行轻微的双边滤波或引导滤波在平滑颜色的同时保持边缘这能有效减少伪影。6.4 与基线方法对比时效果不明显问题描述在简单任务如2D高斯上摊销OT的效果和精确OT差不多但没显示出明显优势。正确认识在低维、简单分布上精确OT如Sinkhorn本身已经很快很准。摊销OT的优势在于高维和重复计算的场景。它的价值体现在批量推理速度一旦训练好对成千上万对新分布进行推理的速度是常数级的而Sinkhorn需要线性或更差的增长。嵌入工作流可以作为一个可微层嵌入到更大的端到端模型中如图像生成器进行联合训练这是迭代式OT求解器难以做到的。处理超大批次如上文的CIFAR-10实验当批次大到Sinkhorn都算不动时摊销OT是唯一可行的选择。6.5 如何选择RA-OT还是OA-OT这是一个实践中的关键决策。我的建议如下选择RA-OT如果你拥有充足的计算资源来预生成高质量的OT“真值”标签例如可以用Sinkhorn在较小的子集上精细计算。追求更快的训练速度和更稳定的训练过程。你的任务中精确模仿某个特定OT求解器的行为很重要。选择OA-OT如果你无法获得或生成OT真值标签例如在无监督或自监督学习框架中。追求极致的最终性能并且愿意投入更长的训练时间。希望探索可能超越传统OT解法的、由目标函数引导的新型传输计划。需要将OT模块完全无缝地集成到一个可微分的管道中。最后别忘了利用投影数L的免费午餐特性。在你的硬件允许的范围内尽可能设置一个较大的L比如50或100这几乎总能带来性能提升或至少更稳定的训练而成本微乎其微。
