Neural Monge Map estimation and its applications

TL;DR

论文提出Neural Monge Map,以无配对样本学习一般成本的最优传输映射,并在CelebA上达成FID 3.7109。

cs.LG 🔴 高级 2021-06-08 17 次浏览
Jiaojiao Fan Shu Liu Shaojun Ma Haomin Zhou Yongxin Chen
最优传输 Monge映射 神经网络 无配对学习 生成模型

核心发现

方法论

作者将Monge问题中的约束T#ρa=ρb通过拉格朗日乘子f显式处理,构造鞍点目标L(T,f)=E[c(X,T(X))-f(T(X))+f(Y)]。以普通神经网络Tθ和fη分别表示传输映射与对偶势,交替执行K1步映射最小化和K2步对偶最大化。算法只需分别采样ρa、ρb,不要求密度表达式、配对样本或熵正则化,并支持不同维度与一般成本函数。

关键结果

  • 在无配对文本到图像任务中,模型使用CLIP ViT-L/14的文本和图像嵌入,并以负余弦相似度为成本。在Laion art上生成结果与配对训练的DALL·E2-Laion具有竞争力;在其未见过配对数据的CC-3M上仍能生成合理图像,且对无关文本的相似度很低。
  • 在CelebA 64×64无配对修复中,源域和目标域各含80k图像。以未遮挡区域MSE为成本时,α=1000取得FID 4.7621,α=10,000取得3.7109,均优于WGAN-GP的18.7942和离散OT基线的视觉质量。
  • 理论上,定理3表明即使学得的映射暂未严格满足推送约束,鞍点值仍等于一般OT距离;定理4给出||T−T*||L2(βρa)≤√[2(E1+E2)]的后验误差界,将训练中的对偶间隙转化为映射可靠性指标。

研究意义

该工作把连续空间中的最优传输从经验性的分布匹配推进到可泛化的确定性映射学习。相比依赖经验点对的EMD或Sinkhorn,所得Tθ可作用于未见样本;相比WGAN只追求输出分布接近,本文同时优化运输成本并估计OT距离。无配对文本、图像数据即可完成跨模态转换,降低了昂贵配对标注的需求,为生成、修复、超分辨率和跨域生物数据对齐提供统一框架。

技术贡献

核心技术是直接对Monge约束施加拉格朗日乘子,而非先离散化Kantorovich问题或加入熵正则。该形式将连续期望转化为GAN式双网络鞍点优化,复杂度为O(K(K1+K2)B)。理论上,作者证明了鞍点与Monge解、Kantorovich对偶及一般OT值之间的一致性,并在c(x,y)=||x−y||²/2或−x·y等成本下建立基于E1、E2的加权L2误差估计。

新颖性

相较多数基于样本的Monge估计器仅适用于二次成本、相同维度或凸势网络,本文直接支持一般成本、不同维度和无配对边缘样本。方法并非首次使用神经最优传输,但其“普通映射网络+对偶势网络+一般成本”的统一鞍点形式,以及可计算的对偶间隙误差证书,是主要创新。

局限性

  • 误差定理要求紧支撑、源分布绝对连续、成本满足可逆混合Hessian等正则条件,因而不能自动覆盖离散分布、奇异流形或非光滑成本。
  • 鞍点训练可能不稳定,且理论误差依赖对偶间隙和c-凹性;论文没有给出大规模文本到图像的完整数值表,也未系统比较训练时间与更多现代扩散基线。
  • 球面实验需将测地距离线性化并额外投影到陆地,说明复杂几何约束仍需要任务特定工程处理。

未来方向

后续可研究更稳定的双层优化、自动满足c-凹性的网络参数化,以及非紧支撑和奇异分布下的统计误差。还可扩展到随机或多对一传输、流形上的内禀成本、在线数据和大规模扩散模型,并报告样本复杂度、计算成本及与最新生成模型的系统比较。

AI 总览摘要

最优传输提供了把一个概率分布“搬运”为另一个分布的原则性方法,但传统离散算法依赖大量样本,难以处理连续高维空间,也无法自然作用于新样本。WGAN能学习分布匹配映射,却不保证运输成本最优;Sinkhorn则通常输出经验耦合而非可泛化的确定性映射。

Neural Monge Map把Monge问题改写为拉格朗日鞍点问题:映射网络Tθ负责低成本搬运,势函数网络fη负责检验输出是否遵守目标分布。训练目标为E[c(X,Tθ(X))−fη(Tθ(X))+fη(Y)],只需分别从两个边缘分布采样。方法支持一般成本和不同维度,并通过对偶间隙E1、E2给出||T−T*||L2(βρa)≤√[2(E1+E2)]的后验误差界。

实验显示,该框架不仅适用于合成分布,也能处理真实任务。无配对CLIP文本到图像实验在Laion art和CC-3M上生成合理图像;CelebA 64×64无配对修复中,α=1000和10,000分别达到FID 4.7621和3.7109,优于WGAN-GP的18.7942。代价是训练稳定性和理论假设仍有限,复杂几何还需额外处理;但该工作展示了无需配对数据即可学习可泛化最优映射的可行路径。

深度分析

研究背景

最优传输从Kantorovich线性规划和Sinkhorn算法发展为神经OT。离散EMD/Sinkhorn适合经验分布,却难以泛化到未见样本;熵正则化还会引入偏差。Makkuva等人的ICNN方法主要针对二次成本,且ICNN在大规模任务中表达力受限。本文回到确定性Monge映射,试图在连续、高维、无配对场景中直接学习运输函数。

核心问题

给定ρa∈Rn和ρb∈Rm,需要求T使T#ρa=ρb,并最小化E[c(X,T(X))]。困难包括高维离散化的维数灾难、边缘密度通常未知、样本不配对、输入输出维度可能不同,以及成本不一定是平方欧氏距离。算法还必须提供对未见样本的映射,而非仅返回有限样本上的耦合矩阵。

核心创新

  • �� 将推送约束写入拉格朗日函数,得到一般成本下的maxη minθ鞍点训练。• 用不受凸性限制的神经网络直接表示Tθ和fη,避免ICNN表达瓶颈。• 只使用独立边缘样本,支持不同维度和连续分布。• 定理2、3证明鞍点值与Monge/一般OT问题的一致性。• 定理4把两个对偶间隙转化为加权L2映射误差证书。

方法详解

  • �� 输入:分别采样X~ρa、Y~ρb,指定可微成本c(x,y)。
  • �� 建模:Tθ:Rn→Rm产生确定性运输;fη:Rm→R是拉格朗日对偶势。
  • �� 目标:L=1/BΣ[c(Xk,Tθ(Xk))−fη(Tθ(Xk))+fη(Yk)]。
  • �� 优化:每轮先更新θ,使L下降K1步,再更新η,使L上升K2步;最终输出Tθ。
  • �� 理论:若存在一致鞍点,满足推送约束时得到Monge最优解;即使未严格满足,鞍点值仍为C(ρa,ρb)。
  • �� 诊断:用E1=L(T,f)−infT'L(T',f)和E2=supf'infT'L(T',f')−infT'L(T',f)估计映射误差。

实验设计

实验覆盖文本到图像、图像修复和球面人口运输。文本任务使用Laion art、Conceptual Captions 3M、CLIP ViT-L/14及DALL·E2扩散解码器,成本为负余弦相似度。修复任务使用CelebA 64×64和128×128,源、目标各80k图像,成本为未遮挡区域MSE,比较WGAN-GP和Perrot等离散OT。球面实验使用Doxsey-Whitfield等人的人口数据,并比较线性映射基线。

结果分析

Laion art上,模型生成嵌入与真实图像的余弦相似度优于基线;CC-3M虽存在更强配对噪声,仍生成合理图像且无关文本过拟合较低。CelebA 64×64中,WGAN-GP FID为18.7942,本文α=0、1000、10,000分别为4.7621、3.7109、6.7479,说明成本设计显著影响真实性。图像质量优于使用1000个样本训练的Perrot离散OT映射。

应用场景

可用于无配对文本到图像、遮挡图像修复、超分辨率、风格迁移和跨域数据对齐。应用者只需获得两个边缘数据集和任务相关成本,例如CLIP负余弦相似度或未遮挡区域MSE。训练后Tθ能处理新输入,适合配对数据昂贵、持续变化或隐私受限的场景。

局限与展望

理论结果依赖紧支撑、绝对连续源分布、混合Hessian可逆和c-凹对偶势等条件;实际神经训练还可能陷入不稳定鞍点。普通网络虽灵活,却缺少自动保证推送约束和结构保持的机制。球面实验用线性化成本并额外把海上点投回陆地,表明复杂几何、非光滑成本和多峰随机映射仍需专门设计。

通俗解读 非专业人士也能看懂

把源分布想成一座仓库里的货物,把目标分布想成另一座仓库。我们想设计一套搬运规则:每件货物都能找到新位置,搬完后整体摆放形状像目标仓库,而且总路程或总费用尽可能低。

传统方法会把所有货物两两配对,画出巨大的运输表。这在货物很多时非常慢,而且新来一件货物时,表格无法直接告诉它该去哪。本文训练一个“搬运员”网络T:输入货物的位置,直接给出目的地;同时训练一个“检查员”网络f,奖励符合目标仓库布局、惩罚不合理搬运。两者一边竞争一边改进。

重要的是,训练时不需要知道每件源货物对应哪件目标货物,只要分别看到两边的样本。费用也不必只是直线距离:文本到图像可以用语义相似度,修复照片可以只计算没有被遮住的区域。于是模型学到的是一条可重复使用的搬运规则,而不是一次性的配对清单。

简单解释 像给14岁少年讲一样

想象你在玩一个“整理房间”的游戏。左边房间里有很多乱放的积木,右边房间展示了最终应该有的积木分布。你的任务不是把每一块积木和右边某一块提前配好,而是训练一个机器人:看到左边的积木,就决定它应该搬到哪里,同时尽量少走冤枉路。

论文里的Tθ就是这个机器人,fη像裁判。机器人负责搬东西,裁判检查搬完后整体是不是越来越像目标房间。机器人和裁判轮流升级:机器人降低搬运费用,裁判提高检查标准。这样,即使训练资料没有“这块对应那块”的答案,机器人也能从两组独立样本中学会规则。

它还能处理很酷的任务:把文字变成图像条件、补全被方块挡住的脸,甚至把地球人口重新分布到陆地上。CelebA修复实验的FID最低达到3.7109,而WGAN-GP是18.7942,数字越低通常代表图像更像真实照片。

当然,机器人不一定每次都训练稳定,复杂形状还需要额外规则。它也不是凭空创造答案,而是按照研究者指定的“什么叫好搬运”来学习。

术语表

Monge map(Monge映射)

把源分布中的每个点确定性地映射到目标分布的函数,并使总运输成本最小。它不同于允许拆分质量的随机运输计划。

论文学习的核心对象Tθ。

Optimal transport(最优传输)

在保持两个边缘分布不变的条件下,寻找成本最低的分配或映射。其最优值衡量两个分布的几何差异。

论文用它统一定义距离、映射与对偶问题。

Pushforward(推送分布)

经函数T作用后,源分布变成的分布,记为T#ρa。要求T#ρa=ρb意味着映射后的样本服从目标分布。

Monge约束的核心。

Kantorovich duality(Kantorovich对偶)

将运输最小化问题改写为对偶势函数的最大化问题。对偶势提供检验运输可行性和成本的方式。

fη来自该思想,并构成鞍点目标。

Duality gap(对偶间隙)

当前解与固定对偶或全局鞍点之间的目标差距。间隙越小,通常意味着解越接近最优。

E1、E2被用于定理4的误差证书。

FID(Fréchet Inception Distance)

比较生成图像与真实图像特征分布的指标,通常数值越低越好。它不能完全代表人类感知质量。

CelebA修复的主要定量指标。

开放问题 这项研究留下的未解疑问

  • 1 神经鞍点优化在超高维、非光滑或多峰分布上何时稳定,论文尚未给出普适样本复杂度与训练收敛率。
  • 2 如何在不牺牲表达力的情况下强制满足推送约束、c-凹性和流形几何,仍缺少统一架构。
  • 3 文本到图像实验主要展示定性结果,未来需要与最新扩散模型进行大规模、可复现的定量评测。

应用场景

近期应用

无配对文本到图像

拥有独立文本集合和图像集合的团队可提取CLIP嵌入,以负余弦相似度训练Tθ,再将输出嵌入输入冻结的扩散解码器。该流程减少文本—图像逐项配对的标注成本。

无配对图像修复

医疗或人脸数据可分别收集遮挡图像与完整图像,用未遮挡区域MSE约束运输结果。CelebA实验表明α=1000可达到FID 4.7621,α=10,000达到3.7109。

远期愿景

跨域科学数据对齐

未来可将该映射用于不同实验平台、组织或模态之间的细胞、医学影像和传感器数据转换,前提是设计可信成本并处理批次差异、隐私和流形结构。

原文摘要

Monge map refers to the optimal transport map between two probability distributions and provides a principled approach to transform one distribution to another. Neural network based optimal transport map solver has gained great attention in recent years. Along this line, we present a scalable algorithm for computing the neural Monge map between two probability distributions. Our algorithm is based on a weak form of the optimal transport problem, thus it only requires samples from the marginals instead of their analytic expressions, and can accommodate optimal transport between two distributions with different dimensions. Our algorithm is suitable for general cost functions, compared with other existing methods for estimating Monge maps using samples, which are usually for quadratic costs. The performance of our algorithms is demonstrated through a series of experiments with both synthetic and realistic data, including text-to-image generation and image inpainting tasks.

cs.LG math.OC