[论文解读] Optimal transport natural gradient for statistical manifolds with continuous sample space

Yifan Chen, Wuchen Li|arXiv (Cornell University)|May 22, 2018

Geometric Analysis and Curvature Flows参考文献 47被引用 11

一句话总结

本文通过将连续样本空间的参数模型的密度空间中的 $L^2$-Wasserstein 度量张量拉回到参数空间，提出了 Wasserstein 统计流形，从而实现一种自然梯度下降方法，该方法在最小化 Wasserstein 距离时优于标准梯度下降。该方法渐近表现为牛顿法，并在高斯分布、混合分布、伽马分布和拉普拉斯分布中表现出鲁棒性和高效性。

ABSTRACT

We study the Wasserstein natural gradient in parametric statistical models with continuous sample spaces. Our approach is to pull back the $L^2$-Wasserstein metric tensor in the probability density space to a parameter space, equipping the latter with a positive definite metric tensor, under which it becomes a Riemannian manifold, named the Wasserstein statistical manifold. In general, it is not a totally geodesic sub-manifold of the density space, and therefore its geodesics will differ from the Wasserstein geodesics, except for the well-known Gaussian distribution case, a fact which can also be validated under our framework. We use the sub-manifold geometry to derive a gradient flow and natural gradient descent method in the parameter space. When parametrized densities lie in $\bR$, the induced metric tensor establishes an explicit formula. In optimization problems, we observe that the natural gradient descent outperforms the standard gradient descent when the Wasserstein distance is the objective function. In such a case, we prove that the resulting algorithm behaves similarly to the Newton method in the asymptotic regime. The proof calculates the exact Hessian formula for the Wasserstein distance, which further motivates another preconditioner for the optimization process. To the end, we present examples to illustrate the effectiveness of the natural gradient in several parametric statistical models, including the Gaussian measure, Gaussian mixture, Gamma distribution, and Laplace distribution.

研究动机与目标

• 开发基于最优传输的参数统计模型在连续样本空间上的黎曼几何框架。
• 通过将 $L^2$-Wasserstein 度量张量拉回到参数空间，定义 Wasserstein 自然梯度。
• 建立诱导度量张量为正定且由此形成的流形定义良好的条件。
• 证明 Wasserstein 自然梯度下降在渐近意义上表现如牛顿法，用于最小化 Wasserstein 距离。
• 在多个参数族（包括高斯混合和拉普拉斯分布）中实证验证该方法的有效性。

提出的方法

• 将无限维密度空间中的 $L^2$-Wasserstein 度量张量拉回到有限维参数空间。
• 利用椭圆型偏微分方程（通过假设1）在参数空间中定义诱导度量张量，确保其正定性。
• 推导出一维样本空间中度量张量的显式公式，从而实现自然梯度的解析计算。
• 基于 Wasserstein 统计流形实现自然梯度下降算法，采用前向欧拉法对梯度流进行离散化。
• 通过精确计算 Wasserstein 距离的 Hessian 矩阵，证明 Wasserstein 自然梯度下降在渐近下表现如牛顿法。
• 基于 Hessian 结构设计预条件器，以提升优化性能。

实验结果

研究问题

• RQ1 能否一致地将 $L^2$-Wasserstein 度量张量拉回到有限维参数空间，从而构成黎曼流形？
• RQ2 当最小化 Wasserstein 距离时，所得的 Wasserstein 自然梯度下降是否优于标准梯度下降？
• RQ3 在参数空间中，诱导度量张量在何种条件下为正定且定义良好？
• RQ4 在渐近状态下，Wasserstein 自然梯度与牛顿法在 Wasserstein 损失最小化中的关系如何？
• RQ5 当真实分布位于参数族之外时，Wasserstein 自然梯度是否仍具鲁棒性？

主要发现

• 通过精确计算 Wasserstein 距离的 Hessian 矩阵，证明 Wasserstein 自然梯度下降在最小化 Wasserstein 距离时，其渐近收敛行为与牛顿法相似。
• 在模型正确设定的情况下，Wasserstein GD 平均仅需 6.29 次迭代即可达到 0.2490 的平均目标值，显著优于标准 GD（56.36 次迭代）的 Wasserstein 损失。
• 在最大似然估计中，Fisher-Rao 自然梯度下降优于标准 GD 和 Wasserstein GD，表明最优性依赖于几何结构。
• 在模型设定错误的情况下（目标分布为拉普拉斯，模型为高斯混合），Wasserstein GD 在 Wasserstein 损失下仍保持高效（平均迭代次数：6.29），而 Fisher-Rao GD 在 MLE 下表现更优（平均迭代次数：4.24）。
• 该方法对模型设定错误具有鲁棒性，这在多种参数族与真实分布的实验中得到验证。
• 在一维样本空间中推导出度量张量的显式公式，从而实现解析计算与理论分析。