[论文解读] Entropic estimation of optimal transport maps

Aram-Alexandre Pooladian, Jonathan Niles‐Weed|arXiv (Cornell University)|Sep 24, 2021

Markov Chains and Monte Carlo Methods参考文献 49被引用 6

一句话总结

本文提出了一种计算高效的最优传输映射估计器，通过熵正则化和Sinkhorn算法实现。通过熵耦合的巴氏投影，该方法在有限样本下实现了统计一致性，提供了比以往在高维或大规模数据集下计算缓慢的方法更具可扩展性的替代方案。

ABSTRACT

We develop a computationally tractable method for estimating the optimal map between two distributions over $\mathbb{R}^d$ with rigorous finite-sample guarantees. Leveraging an entropic version of Brenier's theorem, we show that our estimator -- the \emph{barycentric projection} of the optimal entropic plan -- is easy to compute using Sinkhorn's algorithm. As a result, unlike current approaches for map estimation, which are slow to evaluate when the dimension or number of samples is large, our approach is parallelizable and extremely efficient even for massive data sets. Under smoothness assumptions on the optimal map, we show that our estimator enjoys comparable statistical performance to other estimators in the literature, but with much lower computational cost. We showcase the efficacy of our proposed estimator through numerical examples, even ones not explicitly covered by our assumptions. By virtue of Lepski's method, we propose a modified version of our estimator that is adaptive to the smoothness of the underlying optimal transport map. Our proofs are based on a modified duality principle for entropic optimal transport and on a method for approximating optimal entropic plans due to Pal (2019).

研究动机与目标

• 解决现有最优传输映射估计器在高维或大规模样本下计算不可行的问题。
• 开发一种利用熵正则化和Sinkhorn算法实现可处理且可扩展的最优传输映射估计方法。
• 在真实映射满足光滑性假设的前提下，为所提出的估计器提供有限样本下的统计收敛保证。
• 通过在大规模数据上实现高效且可并行计算，弥合理论最优传输与实际机器学习应用之间的差距。
• 建立熵耦合的巴氏投影与Brenier定理之间的规范联系，从而支持样本外扩展。

提出的方法

• 利用最优传输问题的熵正则化，将原始Monge问题转化为计算上可处理的形式。
• 使用Sinkhorn算法高效计算从i.i.d.样本中获得的经验测度之间的熵正则化耦合。
• 对熵正则化耦合应用巴氏投影，得到将输入点映射到耦合下条件期望的映射估计器。
• 建立巴氏投影与熵正则化最优传输中对偶势函数梯度之间的理论联系，推广Brenier定理。
• 采用修改的对偶原理及Pal（2019）提出的近似技术，分析估计器的统计性质。
• 利用链式法和覆盖数论证，界定经验过程并推导出有限样本风险界。

实验结果

研究问题

• RQ1 我们能否构建一个在高维或大规模数据集下计算高效且统计一致的最优传输映射估计器？
• RQ2 熵耦合的巴氏投影是否能提供真实最优传输映射的有效且一致的估计？
• RQ3 在真实映射满足光滑性假设的前提下，此类估计器的有限样本收敛速率如何？
• RQ4 与现有最先进估计器相比，该方法在时间复杂度和可并行性方面表现出怎样的计算效率？
• RQ5 巴氏投影能否被解释为在熵正则化下经典最优传输理论的自然延伸？

主要发现

• 在真实映射的逆映射 $T_0^{-1}$ 具备 $\mathcal{C}^\alpha$ 光滑性且 $\alpha \in (1,3]$ 的条件下，估计器的收敛速率为 $\mathbb{E}\|\hat{T}-T_0\|_{L^2(P)}^2 \lesssim n^{-\frac{\alpha+1}{2(d'+\alpha+1)}}\log n$，其中 $d' = 2\lceil d/2 \rceil$。
• 该估计器计算高效，由于使用了Sinkhorn算法，时间复杂度与样本数的平方成正比，且在GPU上高度可并行化。
• 首次为熵耦合的巴氏投影建立了有限样本收敛保证，该方法此前仅被启发式使用。
• 在数值基准测试中，该方法在统计精度和计算速度方面均优于现有估计器，尽管收敛速率略慢于极小化最大风险的估计器。
• 巴氏投影被证明对应于熵正则化最优传输中对偶势函数的梯度，提供了理论基础并支持样本外预测。
• 分析依赖于修改的对偶原理和Hölder空间的覆盖数界，常数依赖于 $\varepsilon^{-d/2}$，反映了正则化与近似误差之间的权衡。