[论文解读] Vector Quantized Wasserstein Auto-Encoder

Tung-Long Vuong, Trung Le+5 more|arXiv (Cornell University)|Feb 12, 2023

Generative Adversarial Networks and Image Synthesis被引用 5

一句话总结

本文提出向量量化Wasserstein自编码器（VQ-WAE），一种通过最小化离散码字序列分布与数据分布之间Wasserstein距离来学习深度离散潜在表征的新方法。通过利用Wasserstein距离的理论特性并将目标函数重构为可训练形式，VQ-WAE在代码本利用率、图像重建与生成质量，以及潜在表征的解耦与聚类性方面，均优于VQ-VAE和SQ-VAE变体。

ABSTRACT

Learning deep discrete latent presentations offers a promise of better symbolic and summarized abstractions that are more useful to subsequent downstream tasks. Inspired by the seminal Vector Quantized Variational Auto-Encoder (VQ-VAE), most of work in learning deep discrete representations has mainly focused on improving the original VQ-VAE form and none of them has studied learning deep discrete representations from the generative viewpoint. In this work, we study learning deep discrete representations from the generative viewpoint. Specifically, we endow discrete distributions over sequences of codewords and learn a deterministic decoder that transports the distribution over the sequences of codewords to the data distribution via minimizing a WS distance between them. We develop further theories to connect it with the clustering viewpoint of WS distance, allowing us to have a better and more controllable clustering solution. Finally, we empirically evaluate our method on several well-known benchmarks, where it achieves better qualitative and quantitative performances than the other VQ-VAE variants in terms of the codebook utilization and image reconstruction/generation.

研究动机与目标

• 解决现有VQ-VAE变体在代码本利用率和表征质量方面的局限性。
• 从最优传输理论的生成视角探索深度离散表征的学习方法。
• 构建一个理论基础坚实的框架，将Wasserstein距离与向量量化相联系，以实现更好的聚类与解耦。
• 通过确保代码本中码字的更均匀、更有意义的使用，提升所学表征的质量。
• 证明最小化Wasserstein距离可带来更优的重建效果、生成性能及下游表征表现。

提出的方法

• 构建一个生成模型，其中解码器通过最小化Wasserstein距离，将码字序列的离散分布映射到数据分布。
• 推导出一个等价目标函数，包含重建项以及潜在表征与码字分布之间的Wasserstein距离。
• 引入一个理论框架，将Wasserstein距离与聚类相联系，从而实现对潜在空间结构的更好控制。
• 采用确定性编码器与解码器，并结合向量量化，其中量化通过潜在空间中的最近邻分配定义。
• 应用梯度截断操作以稳定训练，确保梯度通过编码器流动，同时在反向传播过程中保持代码本固定。
• 采用简单的离散潜在序列先验（如自回归模型），以实现从学习到的码字生成图像。

Figure 1: Reconstruction results for the FFHQ dataset.

实验结果

研究问题

• RQ1 Wasserstein距离是否可在类似VAE的框架中有效用于改进离散表征学习？
• RQ2 最小化潜在码分布与数据分布之间的Wasserstein距离是否能带来更好的代码本利用率与表征质量？
• RQ3 Wasserstein距离的理论特性是否可用于实现更具结构化与聚类性的潜在表征？
• RQ4 在图像重建、生成与下游分类方面，VQ-WAE相较于VQ-VAE与SQ-VAE的表现如何？
• RQ5 所提出的框架是否能通过改进优化过程，自然地生成更多样化且更有意义的码字？

主要发现

• 在所有数据集上，VQ-WAE在潜在空间的单层分类准确率最高，CIFAR10为50.19%，MNIST为95.62%，SVHN为38.38%，优于VQ-VAE与SQ-VAE。
• 在图像生成方面，VQ-WAE的FID得分最低：CIFAR10上无条件生成为87.73，有条件生成为88.51；MNIST上无条件生成为8.21，有条件生成为3.88，表明其图像质量更优。
• VQ-WAE的代码本利用率显著提升，更多码字被有效使用，死码或使用不足的码字数量明显少于VQ-VAE与SQ-VAE。
• 潜在空间可视化显示，VQ-WAE的表征在各类别间形成分离良好、紧凑的聚类，而VQ-VAE与SQ-VAE的表征则分布零散且类别间重叠明显。
• 消融实验确认，VQ-WAE能最优地利用代码本，从而生成更多样化且更有意义的码字，并构建更结构化的潜在流形。
• 理论分析表明，Wasserstein距离公式中的最优耦合（定理2.3中的OP）提示了一条学习码字序列联合分布的路径，有望进一步提升生成建模性能。

Figure 2: Latent distribution over the codebook on test-set.