[论文解读] Unleashing Transformers: Parallel Token Prediction with Discrete Absorbing Diffusion for Fast High-Resolution Image Generation from Vector-Quantized Codes

Sam Bond-Taylor, Peter Hessey+3 more|arXiv (Cornell University)|Nov 24, 2021

Generative Adversarial Networks and Image Synthesis参考文献 74被引用 6

一句话总结

该论文提出一种基于Transformer主干网络的并行非自回归扩散模型，用于快速生成高分辨率图像。通过随机掩码VQ（向量量化）标记并训练单一Transformer模型以重建这些标记，该方法实现了在超过训练数据分辨率下的快速、全局一致的图像生成，在保持较低计算成本的同时，实现了SOTA级别的FID、密度（Density）和覆盖率（Coverage）得分。

ABSTRACT

Whilst diffusion probabilistic models can generate high quality image content, key limitations remain in terms of both generating high-resolution imagery and their associated high computational requirements. Recent Vector-Quantized image models have overcome this limitation of image resolution but are prohibitively slow and unidirectional as they generate tokens via element-wise autoregressive sampling from the prior. By contrast, in this paper we propose a novel discrete diffusion probabilistic model prior which enables parallel prediction of Vector-Quantized tokens by using an unconstrained Transformer architecture as the backbone. During training, tokens are randomly masked in an order-agnostic manner and the Transformer learns to predict the original tokens. This parallelism of Vector-Quantized token prediction in turn facilitates unconditional generation of globally consistent high-resolution and diverse imagery at a fraction of the computational expense. In this manner, we can generate image resolutions exceeding that of the original training set samples whilst additionally provisioning per-image likelihood estimates (in a departure from generative adversarial approaches). Our approach achieves state-of-the-art results in terms of Density (LSUN Bedroom: 1.51; LSUN Churches: 1.12; FFHQ: 1.20) and Coverage (LSUN Bedroom: 0.83; LSUN Churches: 0.73; FFHQ: 0.80), and performs competitively on FID (LSUN Bedroom: 3.64; LSUN Churches: 4.07; FFHQ: 6.11) whilst offering advantages in terms of both computation and reduced training set requirements.

研究动机与目标

• 解决自回归VQ-VAE模型在高分辨率图像生成中采样速度慢、序列化的问题。
• 利用无约束的Transformer架构实现VQ标记的并行预测，避免顺序依赖偏差。
• 通过上下文窗口聚合技术，实现全局一致的高分辨率图像生成（甚至超过训练数据的分辨率）。
• 在降低计算与数据需求的前提下，实现似然性指标（密度、覆盖率）和FID的SOTA性能。
• 提供每张图像的似然估计，优于GAN基方法，从而支持更优的评估与潜在的条件生成。

提出的方法

• 采用离散吸收扩散过程，在训练期间以非自回归、无序的方式逐步掩码VQ标记。
• 训练单一非自回归Transformer模型，从掩码版本中重建原始VQ标记，实现并行预测。
• 模型端到端训练，以并行方式预测完整的VQ标记序列，利用自注意力机制捕获全局上下文信息。
• 推理阶段，模型从完全掩码状态去噪，通过聚合多个上下文窗口生成高分辨率图像。
• 该方法支持无条件与条件生成，包括通过针对性掩码与去噪实现的局部图像编辑。
• 该方法使用单一Transformer主干网络同时优化所有扩散步骤，相比逐步自回归扩散，显著减少参数量与推理时间。

实验结果

研究问题

• RQ1 非自回归、并行的VQ标记预测机制是否能在图像生成的速度与分辨率上超越自回归模型？
• RQ2 单一Transformer模型能否有效从随机掩码、无序输入中重建VQ标记，从而实现高质量图像生成？
• RQ3 该模型能否在保持全局一致性的前提下，生成超过原始训练分辨率的图像？
• RQ4 所提方法是否在似然性指标（密度、覆盖率）与FID上均达到SOTA性能，同时降低训练与推理成本？
• RQ5 该模型是否支持局部图像编辑与每张图像的似然估计，从而优于标准GAN基方法？

主要发现

• 在LSUN Bedroom数据集上，该方法实现了1.51的SOTA密度得分；在LSUN Churches上为1.12；在FFHQ上为1.20。
• 在LSUN Bedroom上实现了0.83的SOTA覆盖率得分；在LSUN Churches上为0.73；在FFHQ上为0.80。
• 在LSUN Bedroom上取得3.64的竞争力FID得分，在LSUN Churches上为4.07，在FFHQ上为6.11，优于大多数基线模型，仅略逊于StyleGAN2。
• 通过上下文窗口聚合与并行标记预测，模型可生成超过原始训练集分辨率的图像。
• 相比自回归VQ-VAE模型，该方法显著降低了计算成本与训练数据需求，同时实现快速采样与局部图像编辑能力。
• 该方法提供每张图像的似然估计，这是相较于GAN基模型的关键优势，支持更优评估与潜在的条件生成。