[论文解读] Kernelized Wasserstein Natural Gradient

Michael Arbel, Arthur Gretton+2 more|arXiv (Cornell University)|Oct 21, 2019

Advanced Neuroimaging Techniques and Applications参考文献 50被引用 5

一句话总结

本文提出核化Wasserstein自然梯度（KWNG），通过在再生核希尔伯特空间（Reproducing Kernel Hilbert Space）中采用对偶公式，提出一种可扩展的方法来近似Wasserstein度量的自然梯度。该方法实现了高效、曲率不变的优化，并具有理论保证，在CIFAR-10和CIFAR-100上的实验表明，其性能优于标准SGD和自适应方法，尤其在病态参数化设置下表现更优。

ABSTRACT

Many machine learning problems can be expressed as the optimization of some cost functional over a parametric family of probability distributions. It is often beneficial to solve such optimization problems using natural gradient methods. These methods are invariant to the parametrization of the family, and thus can yield more effective optimization. Unfortunately, computing the natural gradient is challenging as it requires inverting a high dimensional matrix at each iteration. We propose a general framework to approximate the natural gradient for the Wasserstein metric, by leveraging a dual formulation of the metric restricted to a Reproducing Kernel Hilbert Space. Our approach leads to an estimator for gradient direction that can trade-off accuracy and computational cost, with theoretical guarantees. We verify its accuracy on simple examples, and show the advantage of using such an estimator in classification tasks on Cifar10 and Cifar100 empirically.

研究动机与目标

• 为解决高维深度学习模型中精确自然梯度计算的计算不可行性问题。
• 开发一种基于Wasserstein度量的可扩展、参数化不变的优化方法，以在病态参数化设置下实现更优的收敛性。
• 提供一种灵活的估计器，通过再生核希尔伯特空间中的核近似，实现精度与计算成本之间的权衡。
• 通过实证验证该方法在具有挑战性优化景观的监督学习任务中，其精度与有效性。

提出的方法

• 该方法利用受限于再生核希尔伯特空间（RKHS）的Wasserstein度量对偶公式，推导出自然梯度的闭式表达式。
• 通过在RKHS中求解一个凸泛函优化问题来近似自然梯度方向，从而实现高效计算。
• 估计器利用核方法和基函数，在精度与计算成本之间实现权衡，并具有理论收敛保证。
• 通过Levenberg-Marquardt动力学自适应更新阻尼参数，以稳定训练过程。
• 该方法可通过相同的对偶公式框架推广至其他度量，包括Fisher-Rao度量。
• 通过在CIFAR-10和CIFAR-100上使用残差网络进行端到端训练，将该方法应用于监督学习。

实验结果

研究问题

• RQ1 基于核的Wasserstein自然梯度近似是否能在保持计算可行性的同时实现高精度，适用于深度学习模型？
• RQ2 在病态参数化设置下，所提出的KWNG估计器与标准SGD及自适应方法（如Adam）相比，在收敛速度和泛化能力方面表现如何？
• RQ3 核带宽和基函数数量对自然梯度近似精度有何影响？
• RQ4 自适应阻尼方案在实际训练中如何影响训练稳定性和性能？

主要发现

• 与基线方法相比，KWNG估计器在自然梯度近似中表现出显著更低的相对误差，尤其在高维设置下。
• 在CIFAR-10和CIFAR-100上，KWNG在训练和测试精度方面均优于SGD和Adam，尤其在病态情况下表现更优。
• 该方法收敛速度优于标准SGD，并在Hessian条件数较高时，测试精度高于对角预条件和KFAC。
• 随着样本数量和基函数数量的增加，估计器的相对误差减小，且在d=1至d=10的维度范围内表现稳定。
• 自适应阻尼可提升训练稳定性，且核函数选择（高斯核与有理二次核）影响性能，其中高斯核在大多数情况下表现更优。