[论文解读] Amortized Projection Optimization for Sliced Wasserstein Generative Models

Khai Nguyen, Nhat Ho|arXiv (Cornell University)|Mar 25, 2022

Human Pose and Action Recognition被引用 4

一句话总结

该论文提出了一种近似投影优化方法，通过学习预测最优投影方向，加速切片Wasserstein生成模型，从而消除迭代优化循环。该方法引入了三种近似模型——线性、广义线性和非线性，分别生成新的小批量损失函数（LA-SW、GA-SW、NA-SW），在训练速度和生成性能方面优于Max-SW和大L值下的标准SW。

ABSTRACT

Seeking informative projecting directions has been an important task in utilizing sliced Wasserstein distance in applications. However, finding these directions usually requires an iterative optimization procedure over the space of projecting directions, which is computationally expensive. Moreover, the computational issue is even more severe in deep learning applications, where computing the distance between two mini-batch probability measures is repeated several times. This nested loop has been one of the main challenges that prevent the usage of sliced Wasserstein distances based on good projections in practice. To address this challenge, we propose to utilize the learning-to-optimize technique or amortized optimization to predict the informative direction of any given two mini-batch probability measures. To the best of our knowledge, this is the first work that bridges amortized optimization and sliced Wasserstein generative models. In particular, we derive linear amortized models, generalized linear amortized models, and non-linear amortized models which are corresponding to three types of novel mini-batch losses, named amortized sliced Wasserstein. We demonstrate the favorable performance of the proposed sliced losses in deep generative modeling on standard benchmark datasets.

研究动机与目标

• 解决切片Wasserstein GAN中迭代投影优化带来的高计算成本问题。
• 减少嵌套优化循环（全局模型训练 + 局部投影优化），从而加快训练速度。
• 在不牺牲样本质量的前提下，提升训练效率和生成性能。
• 通过学习具有信息量的投影方向，弥合近似优化与最优传输之间的差距。
• 开发可扩展、可微分的模型，以预测小批量分布的最优投影。

提出的方法

• 提出近似模型，用于预测小批量概率测度之间的最优投影方向，替代迭代优化。
• 设计三种类型的近似模型：线性、广义线性和非线性，每种对应一种新型小批量损失函数（LA-SW、GA-SW、NA-SW）。
• 训练神经网络，将输入小批量映射到最优投影方向，实现端到端反向传播。
• 在切片Wasserstein距离计算中使用预测方向，替代蒙特卡洛采样或迭代最大化。
• 将近似损失集成到SNGAN等深度生成模型中，实现模型参数与投影方向的联合优化。
• 通过参数化投影预测器并使用随机梯度下降，确保可微性和可扩展性。

实验结果

研究问题

• RQ1 近似优化能否替代切片Wasserstein GAN中迭代投影搜索，以减少训练时间？
• RQ2 不同近似模型架构（线性、广义线性、非线性）在预测信息性投影方向方面的表现如何比较？
• RQ3 与Max-SW和大L值下的标准SW相比，所提出的近似切片Wasserstein损失是否能提升训练收敛性和样本质量？
• RQ4 与现有SW变体相比，所提出的近似模型在计算和内存开销方面如何？
• RQ5 所学习的投影方向是否能在不同数据分布和小批量配置下实现泛化？

主要发现

• 所提出的LA-SW损失在训练速度上与L=1的SW和Max-SW相当，同时保持了高质量的样本生成。
• 与L ≥ 100的SW和T₂ ≥ 10的Max-SW相比，LA-SW在速度和内存使用上均更优，是效率最高的变体。
• GA-SW和NA-SW的训练速度略慢，内存占用更高（最高达6.16 GB），但FID和IS得分优于Max-SW和大L值下的SW。
• LA-SW生成的图像在定性上更优且一致，与定量指标结果相符，可视化未进行选择性展示。
• 近似模型显著减轻了嵌套优化负担，实现了切片Wasserstein GAN的可扩展训练。
• 实证结果表明，LA-SW、GA-SW和NA-SW在CIFAR10、CelebA、STL10和CelebA-HQ上均优于Max-SW和标准SW，FID和Inception Score表现更优。