基于改进VQGAN的向量量化图像建模

在大规模文本语料库上通过下一词预测进行预训练的语言模型，已在生成式与判别式语言任务中展现出卓越的零样本（zero-shot）、少样本（few-shot）、迁移学习及多任务处理能力。受此成功启发，本文提出一种向量量化图像建模（Vector-quantized Image Modeling, VIM）方法，该方法通过训练Transformer模型自回归地预测栅格化图像标记（rasterized image tokens），实现图像的高效建模。这些离散的图像标记由基于视觉Transformer架构的向量量化生成对抗网络（Vision-Transformer-based VQGAN，简称ViT-VQGAN）学习得到。我们首先在架构设计与码本学习（codebook learning）等方面对原始VQGAN进行了多项改进，显著提升了模型的效率与重建保真度。改进后的ViT-VQGAN进一步推动了向量量化图像建模任务的发展，包括无条件图像生成、类别条件图像生成以及无监督表征学习等。当在ImageNet数据集上以 $256 \times 256$ 分辨率进行训练时，我们的方法取得了175.1的Inception Score（IS）和4.17的Fréchet Inception Distance（FID），相较原始VQGAN的70.6（IS）和17.04（FID）实现了显著提升。在此基础上，基于ViT-VQGAN与无监督预训练的Transformer模型，我们采用类似于Image GPT（iGPT）的策略，通过平均中间层特征来评估预训练效果。在ImageNet上预训练的VIM-L模型，在相同模型规模下，线性探测（linear-probe）准确率从iGPT-L的60.3%提升至73.2%，表现显著优于iGPT-L。此外，VIM-L还超越了在更大模型规模和额外网络图像数据上训练的iGPT-XL模型，在多项评估指标上展现出更强的泛化能力。