[论文解读] Autoregressive Image Generation using Residual Quantization

Doyup Lee, Chiheon Kim+3 more|arXiv (Cornell University)|Mar 3, 2022

Advanced Vision and Imaging被引用 8

一句话总结

该论文提出了一种两阶段框架 RQ-VAE 和 RQ-Transformer，用于使用残差量化实现高分辨率自回归图像生成。通过采用固定码书大小的从粗到精的残差量化，RQ-VAE 将 256×256 图像压缩为 8×8×4 的码图，使 RQ-Transformer 能够以显著更低的计算成本和更快的采样速度生成高保真度图像，优于以往的自回归模型。

ABSTRACT

For autoregressive (AR) modeling of high-resolution images, vector quantization (VQ) represents an image as a sequence of discrete codes. A short sequence length is important for an AR model to reduce its computational costs to consider long-range interactions of codes. However, we postulate that previous VQ cannot shorten the code sequence and generate high-fidelity images together in terms of the rate-distortion trade-off. In this study, we propose the two-stage framework, which consists of Residual-Quantized VAE (RQ-VAE) and RQ-Transformer, to effectively generate high-resolution images. Given a fixed codebook size, RQ-VAE can precisely approximate a feature map of an image and represent the image as a stacked map of discrete codes. Then, RQ-Transformer learns to predict the quantized feature vector at the next position by predicting the next stack of codes. Thanks to the precise approximation of RQ-VAE, we can represent a 256$ imes$256 image as 8$ imes$8 resolution of the feature map, and RQ-Transformer can efficiently reduce the computational costs. Consequently, our framework outperforms the existing AR models on various benchmarks of unconditional and conditional image generation. Our approach also has a significantly faster sampling speed than previous AR models to generate high-quality images.

研究动机与目标

• 为解决自回归图像生成中的率失真权衡问题，其中缩短码序列长度对计算效率至关重要，但以往的 VQ 方法难以维持图像保真度。
• 通过精确逼近特征图以最小化离散码序列长度，实现在降低计算成本的同时实现高分辨率图像生成。
• 通过使用软标签和随机采样技术解决暴露偏差问题，提升自回归模型的训练稳定性和性能。
• 在保持或提升图像质量的同时，实现比现有自回归模型更快的采样速度。

提出的方法

• RQ-VAE 使用固定大小的码书进行残差量化，以从粗到精的方式迭代优化特征图的近似，将图像表示为堆叠的离散码。
• 通过应用 4 层残差量化，将 256×256 图像的空间分辨率降低至 8×8，同时在不增加码书大小的情况下保持高保真度重建。
• RQ-Transformer 将堆叠的码作为序列处理，并在每个位置自回归地预测下一个码堆，利用缩短的序列长度实现高效生成。
• 在 RQ-Transformer 训练过程中引入软标签和随机采样技术，以缓解暴露偏差并提升泛化能力。
• 框架采用栅格扫描顺序将 3D 码图转换为 1D 序列，以支持自回归建模。
• 重建保真度通过 rFID 评估，生成质量则在 ImageNet 和 LSUN 等基准上通过 FID 和 Inception Score 进行测量。

实验结果

研究问题

• RQ1 使用固定码书大小的残差量化能否在高分辨率图像生成中实现比标准 VQ-VAE 更好的重建质量与更低的分辨率？
• RQ2 通过 RQ-VAE 缩短码序列长度是否能在不损失图像质量的前提下，实现更快、更高效的自回归建模？
• RQ3 软标签和随机采样在多大程度上通过减少暴露偏差提升了 RQ-Transformer 的性能？
• RQ4 在无条件与条件图像生成任务中，该框架与先前自回归模型相比，在 FID、Inception Score 和采样速度方面表现如何？

主要发现

• RQ-VAE 使用 8×8×4 码图和码书大小 K=16,384 时，rFID 达到 10.2，显著优于使用 16×16 码图的 VQ-GAN，尽管空间分辨率降低了 4 倍。
• 在固定 K=16,384 的情况下，将量化深度 D 从 1 增加到 4，使 rFID 提升 3.1 分（从 13.3 提升至 10.2），证明了残差量化在不增加码书大小的情况下优于扩大码书规模。
• RQ-Transformer 在 ImageNet 256×256 无条件生成任务上实现了 SOTA 的 FID 得分 10.2，优于以往的自回归模型。
• 该框架在 LSUN、FFHQ 和 CC-3M 基准上实现了比以往自回归模型快 2.5 倍的采样速度，同时保持或提升了图像质量。
• 可视化结果表明，RQ-VAE 实现了从粗到精的重建：更深层的量化级别逐步添加更精细的图像细节，且在更高深度使用更低范数的码嵌入。
• 码使用分布显示各深度层级之间存在显著重叠，表明码书被高效复用，且共享码书被有效利用。