[论文解读] Sinkhorn-AutoDiff: Tractable Wasserstein Learning of Generative Models

Aude Genevay, Gabriel Peyré+1 more|arXiv (Cornell University)|Jun 1, 2017

Generative Adversarial Networks and Image Synthesis被引用 11

一句话总结

本文提出 Sinkhorn-AutoDiff，一种可微分、面向 GPU 优化的方法，用于使用最优传输（OT）损失训练大规模生成模型。通过结合熵正则化与 Sinkhorn 迭代的自动微分，该方法实现了稳定、可扩展的基于 OT 的训练，克服了高维分布中计算与梯度不稳定的难题。

ABSTRACT

The ability to compare two degenerate probability distributions (i.e. two probability distributions supported on two distinct low-dimensional manifolds living in a much higher-dimensional space) is a crucial problem arising in the estimation of generative models for high-dimensional observations such as those arising in computer vision or natural language. It is known that optimal transport metrics can represent a cure for this problem, since they were specifically designed as an alternative to information divergences to handle such problematic scenarios. Unfortunately, training generative machines using OT raises formidable computational and statistical challenges, because of (i) the computational burden of evaluating OT losses, (ii) the instability and lack of smoothness of these losses, (iii) the difficulty to estimate robustly these losses and their gradients in high dimension. This paper presents the first tractable computational method to train large scale generative models using an optimal transport loss, and tackles both these issues by relying on two key ideas: (a) entropic smoothing, which turns the original OT loss into one that can be computed using Sinkhorn fixed point iterations; (b) algorithmic (automatic) differentiation of these iterations. These two approximations result in a robust and differentiable approximation of the OT loss with streamlined GPU execution. The resulting computational architecture complements nicely standard deep network generative models by a stack of extra layers implementing the loss function.

研究动机与目标

• 解决在高维空间中支持于低维流形的退化概率分布之间的比较问题，此类问题在视觉与自然语言处理中普遍存在。
• 克服精确最优传输（OT）损失在深度生成建模中计算不可行与梯度不稳定的难题。
• 通过使损失可微分且适合 GPU 执行，实现基于 OT 目标函数的鲁棒、可扩展的生成模型训练。
• 提供一种实用、可即插即用的框架，将 OT 损失计算作为可微分层集成于深度生成架构中。

提出的方法

• 对原始最优传输问题应用熵正则化，将其转化为可通过 Sinkhorn 固定点迭代求解的正则化 OT 问题。
• 使用迭代 Sinkhorn 算法计算 Wasserstein 距离的可微分近似，从而支持梯度反向传播。
• 通过算法（自动）微分将 Sinkhorn 迭代集成到计算图中，以计算通过迭代过程的梯度。
• 设计一个可微分层堆叠结构，替代生成模型中的标准损失函数，从而实现基于 OT 目标函数的端到端训练。
• 利用 GPU 加速以优化 Sinkhorn 迭代及其梯度的计算，实现对大规模模型与高维数据的可扩展性。
• 通过在优化过程中控制熵正则化参数，确保损失的数值稳定性和平滑性。

实验结果

研究问题

• RQ1 熵正则化与自动微分能否结合使用，使最优传输损失在大规模深度生成模型中变得可行？
• RQ2 在高维数据中，可微分 Sinkhorn 近似与精确 OT 在训练稳定性与收敛性方面相比如何？
• RQ3 该方法在不修改网络架构的前提下，能在多大程度上集成到现有深度生成架构中？
• RQ4 熵正则化参数对近似精度与训练效率之间权衡的影响如何？
• RQ5 所提出的框架是否能在保持可微分性与 GPU 效率的同时，实现与生成建模中其他方法相当的性能？

主要发现

• 所提出的方法克服了精确 OT 在计算与可微分性方面的挑战，实现了生成模型使用 OT 损失的稳定高效训练。
• 熵正则化与 Sinkhorn 迭代的自动微分相结合，可产生适合深度学习的平滑、可微分的 Wasserstein 距离近似。
• 该方法在 GPU 硬件上实现了高度计算效率，使基于 OT 的大规模生成建模成为可能。
• 该框架可无缝作为可微分层集成于深度生成模型中，支持基于 OT 目标函数的端到端训练。
• 该方法在低维数据流形场景下，为传统基于散度的损失（如 GAN 或 VAE）提供了稳健的替代方案。
• 在标准 OT 与信息论散度失效的高维设置中，该方法展现出更优的稳定性和收敛性。