[论文解读] Sinkhorn Natural Gradient for Generative Models

Zebang Shen, Zhenfu Wang+2 more|arXiv (Cornell University)|Nov 9, 2020

Constraint Satisfaction and Optimization参考文献 22被引用 4

一句话总结

本文提出Sinkhorn自然梯度（SiNG），一种新颖的优化方法，通过在概率空间中使用Sinkhorn散度作为度量执行最速下降，以训练生成模型。与现有自然梯度方法不同，SiNG以对数时间复杂度（相对于精度）计算Sinkhorn信息矩阵（SIM），实现无需梯度函数的精确且高效的更新，并在蒙特卡洛积分下提供稳定的实证估计。

ABSTRACT

We consider the problem of minimizing a functional over a parametric family of probability measures, where the parameterization is characterized via a push-forward structure. An important application of this problem is in training generative adversarial networks. In this regard, we propose a novel Sinkhorn Natural Gradient (SiNG) algorithm which acts as a steepest descent method on the probability space endowed with the Sinkhorn divergence. We show that the Sinkhorn information matrix (SIM), a key component of SiNG, has an explicit expression and can be evaluated accurately in complexity that scales logarithmically with respect to the desired accuracy. This is in sharp contrast to existing natural gradient methods that can only be carried out approximately. Moreover, in practical applications when only Monte-Carlo type integration is available, we design an empirical estimator for SIM and provide the stability analysis. In our experiments, we quantitatively compare SiNG with state-of-the-art SGD-type solvers on generative tasks to demonstrate its efficiency and efficacy of our method.

研究动机与目标

• 为解决现有自然梯度方法在生成建模中的局限性，特别是梯度函数计算不可行以及基于Wasserstein距离的替代方法计算成本过高的问题。
• 开发一种直接在概率空间中使用Sinkhorn散度的自然梯度方法，确保重参数化不变性。
• 推导出Sinkhorn信息矩阵（SIM）的显式表达式，其计算复杂度随目标精度呈对数增长。
• 设计一种在仅有蒙特卡洛样本可用时仍保持稳定的SIM实证估计器（eSIM），确保实际训练中的鲁棒性。

提出的方法

• 将SiNG提出为在配备Sinkhorn散度的概率空间上的最速下降方法，其推导基于在半径为ε的Sinkhorn球内最小化泛函的极限形式。
• 将Sinkhorn信息矩阵（SIM）定义为Sinkhorn散度关于参数θ的Hessian矩阵，并证明其可通过O(log(1/ε))次操作完成计算。
• 证明SIM的评估仅需从生成测度αθ中采样，无需计算梯度函数，这与Fisher-Rao自然梯度不同。
• 提出基于经验Sinkhorn散度Hessian的SIM实证估计器（eSIM），并证明其在蒙特卡洛采样下的稳定性。
• 利用共轭梯度法与PyTorch的autograd，高效计算涉及eSIM的矩阵-向量乘积，而无需显式构造矩阵。
• 使用可微的Sinkhorn求解器（如geomloss）隐式构建Sinkhorn势函数，从而支持计算图中的反向传播。

实验结果

研究问题

• RQ1 能否为生成模型设计一种自然梯度方法，避免对梯度函数的依赖，同时保持重参数化不变性？
• RQ2 Sinkhorn散度的Hessian矩阵（即SIM）是否可显式表达且具有低计算复杂度？
• RQ3 能否构造一种在从生成分布中进行蒙特卡洛采样时仍保持稳定的SIM实证估计器？
• RQ4 在生成建模任务中，SiNG相较于SGD类求解器在收敛速度和样本效率方面的表现如何？
• RQ5 SiNG的连续时间极限是否保持重参数化不变性，以满足深度生成模型中稳健优化的要求？

主要发现

• Sinkhorn信息矩阵（SIM）具有显式表达式，且可基于所需精度ε以O(log(1/ε))次操作完成计算，从而实现高效且精确的更新。
• 所提出的SiNG算法在连续时间极限下对重参数化保持不变，确保对非线性、非单射映射Tθ的鲁棒性。
• SIM的实证估计器eSIM在蒙特卡洛采样下保持稳定，其收敛性基于Sinkhorn势函数Fréchet导数的连续性得到保证。
• 实验结果表明，SiNG在生成建模基准测试中，无论是训练效率还是样本质量，均优于当前最先进的SGD类求解器（如Adam、RMSProp）。
• 在使用DC-GAN架构训练CIFAR-10和CelebA数据集时，该方法实现了更快的收敛速度和更低的FID分数，优于基线优化器。
• 实现中利用PyTorch的autograd与geomloss库，高效计算梯度和矩阵-向量乘积，而无需显式构造SIM。