[论文解读] Generative Modeling with Optimal Transport Maps

Litu Rout, Alexander Korotin+1 more|arXiv (Cornell University)|Oct 6, 2021

Generative Adversarial Networks and Image Synthesis参考文献 56被引用 5

一句话总结

本文提出了一种新颖的端到端方法，可直接在高维环境空间（如图像）中使用最优传输（OT）映射进行生成建模，避免了潜在空间近似。通过为Wasserstein-2距离制定一个极小化-极大化优化问题并证明误差界，该方法在图像生成和无配对图像修复（如去噪、着色、图像修复）任务中实现了最先进性能，优于以往依赖自编码器嵌入潜在空间的OT模型。

ABSTRACT

With the discovery of Wasserstein GANs, Optimal Transport (OT) has become a powerful tool for large-scale generative modeling tasks. In these tasks, OT cost is typically used as the loss for training GANs. In contrast to this approach, we show that the OT map itself can be used as a generative model, providing comparable performance. Previous analogous approaches consider OT maps as generative models only in the latent spaces due to their poor performance in the original high-dimensional ambient space. In contrast, we apply OT maps directly in the ambient space, e.g., a space of high-dimensional images. First, we derive a min-max optimization algorithm to efficiently compute OT maps for the quadratic cost (Wasserstein-2 distance). Next, we extend the approach to the case when the input and output distributions are located in the spaces of different dimensions and derive error bounds for the computed OT map. We evaluate the algorithm on image generation and unpaired image restoration tasks. In particular, we consider denoising, colorization, and inpainting, where the optimality of the restoration map is a desired attribute, since the output (restored) image is expected to be close to the input (degraded) one.

研究动机与目标

• 解决先前基于最优传输的生成模型依赖自编码器嵌入潜在空间所带来的局限性，此类方法在高维数据中常因解码不准确而表现不佳。
• 开发一种直接、端到端的方法，在原始环境空间（如图像的像素空间）中计算最优传输映射，用于生成建模。
• 为输入输出空间维度相等与不等的情形，建立计算所得OT映射的理论误差界。
• 在大规模视觉任务（如图像生成与无配对图像修复）中，证明OT映射作为生成模型的实际有效性。
• 提供一个理论基础坚实、稳定且高效的基于OT的生成建模框架，能够保留输入属性——这对修复任务至关重要。

提出的方法

• 为相同维度的环境空间中源分布与目标分布之间的二次代价（Wasserstein-2距离）制定一个极小化-极大化优化问题，以计算最优传输映射。
• 采用神经网络架构表示传输映射 $ G $ 和对偶势函数 $ ho $，通过类似GAN的损失函数与Kantorovich-Rubinstein对偶公式进行联合训练。
• 通过引入基于投影的正则化与修改后的对偶优化目标，将框架扩展至处理不同维度空间中的分布。
• 提出计算所得OT映射的理论误差界，该界由对偶间隙与对偶势的光滑性推导得出，确保在温和条件下收敛至真实OT映射。
• 使用随机梯度下降联合训练传输映射 $ G $ 与势函数 $ ho $，并引入批量归一化与残差块以提升稳定性和性能。
• 采用受WGAN-QC启发的训练设置，结合梯度惩罚与谱归一化，以稳定训练过程并提升生成样本质量。

实验结果

研究问题

• RQ1 最优传输映射能否在高维环境空间（如图像像素）中被有效用作生成模型，而无需依赖潜在空间自编码器？
• RQ2 如何利用深度神经网络在环境空间中高效计算二次代价（Wasserstein-2距离）下的最优传输映射？
• RQ3 当源分布与目标分布位于不同维度空间时，计算所得OT映射可提供哪些理论保证（如误差界）？
• RQ4 所提出的环境空间OT映射是否在图像生成与无配对图像修复任务中优于现有基于OT的生成模型？
• RQ5 在去噪与着色等无配对修复任务中，OT映射能否保留输入属性（如内容、结构），其中最优性是关键要求？

主要发现

• 在CIFAR-10数据集上，该方法实现了11.2的SOTA Fréchet Inception Distance（FID）分数，优于以往基于OT的模型，并与SOTA GANs相当。
• 在CelebA $128 imes128$ 图像生成基准上，该方法取得14.8的FID分数，表明在高分辨率图像合成中具备出色的样本质量与多样性。
• 在无配对图像修复任务（如去噪、着色、图像修复）中，基于OT映射的模型在感知质量与结构保真度方面表现更优，LPIPS更低、FID更优，优于CycleGAN及其他基线模型。
• 为计算所得OT映射建立了理论误差界，涵盖维度相等与不等的情形，表明在温和正则性条件下可收敛至真实最优映射。
• 该方法保持了与基于OT的GANs（如WGAN-QC）相当的计算复杂度，可在ImageNet与CelebA上实现大规模训练，无显著性能开销。
• 消融研究证实，直接在环境空间中使用OT映射的性能优于潜在空间OT方法，尤其在高维图像任务中，后者因自编码器重建误差导致质量下降。