[论文解读] Image Compression with Product Quantized Masked Image Modeling

Alaaeldin El-Nouby, Matthew J. Muckley+4 more|arXiv (Cornell University)|Dec 14, 2022

Advanced Image Processing Techniques被引用 12

一句话总结

该论文提出了一种名为 PQ-MIM 的新型图像压缩框架，通过将产品量化（product quantization）与掩码图像建模（masked image modeling）相结合，实现了当前最优的率失真性能。通过用产品量化替代标量量化，并利用基于 Transformer 的条件熵模型来捕捉潜在码本中的空间依赖性，PQ-MIM 在 PSNR 方面表现具有竞争力，同时在感知质量（FID、KID）方面优于 HiFiC，且无需微调即可实现混合压缩-生成操作。

ABSTRACT

Recent neural compression methods have been based on the popular hyperprior framework. It relies on Scalar Quantization and offers a very strong compression performance. This contrasts from recent advances in image generation and representation learning, where Vector Quantization is more commonly employed. In this work, we attempt to bring these lines of research closer by revisiting vector quantization for image compression. We build upon the VQ-VAE framework and introduce several modifications. First, we replace the vanilla vector quantizer by a product quantizer. This intermediate solution between vector and scalar quantization allows for a much wider set of rate-distortion points: It implicitly defines high-quality quantizers that would otherwise require intractably large codebooks. Second, inspired by the success of Masked Image Modeling (MIM) in the context of self-supervised learning and generative image models, we propose a novel conditional entropy model which improves entropy coding by modelling the co-dependencies of the quantized latent codes. The resulting PQ-MIM model is surprisingly effective: its compression performance on par with recent hyperprior methods. It also outperforms HiFiC in terms of FID and KID metrics when optimized with perceptual losses (e.g. adversarial). Finally, since PQ-MIM is compatible with image generation frameworks, we show qualitatively that it can operate under a hybrid mode between compression and generation, with no further training or finetuning. As a result, we explore the extreme compression regime where an image is compressed into 200 bytes, i.e., less than a tweet.

研究动机与目标

• 通过统一向量量化与掩码图像建模，弥合神经图像压缩与图像生成之间的方法论鸿沟。
• 通过引入产品量化，解决标量量化在神经压缩中的局限性，实现可扩展的率失真性能。
• 通过使用掩码图像建模方法建模量化潜在码本中的空间依赖性，降低熵编码开销。
• 在无需额外训练的情况下，利用模型的生成能力，实现混合压缩-生成模式。

提出的方法

• 在 VQ-VAE 中用产品量化（PQ）替代标准向量量化，实现更广泛的率失真权衡，并支持高效的码书学习。
• 设计一个多阶段掩码图像建模（MIM）框架，其中每一阶段基于逐步增加的已观测补丁进行条件建模，以降低条件熵。
• 使用视觉 Transformer 来估计潜在码本的条件分布，该方法在编码器和解码器端对称应用，用于熵编码。
• 采用五芒星（quincunx）或基于置信度的补丁掩码策略，确定每一阶段中哪些补丁被观测，以优化率降低。
• 通过组合重建损失、感知损失和对抗性损失进行模型训练，以提升视觉质量。
• 通过复用 MIM 模型执行补丁缺失处的图像修复（inpainting），实现混合操作，证明其与图像生成的兼容性。

实验结果

研究问题

• RQ1 在 VQ-VAE 框架中，产品量化能否为神经图像压缩提供一种可扩展且有效的标量量化替代方案？
• RQ2 如何将掩码图像建模方法适配于潜在空间中，以利用补丁间的空间依赖性来改进熵编码？
• RQ3 在 FID 和 KID 指标方面，一个以感知目标训练的压缩模型能在多大程度上超越现有基线？
• RQ4 一个单一模型是否能在无需微调的情况下，有效同时支持压缩与生成两种模式？
• RQ5 在极端压缩场景下（如低于 0.6 bits per pixel），该模型的性能如何？

主要发现

• 与强基准模型（包括超 prior 方法）相比，PQ-MIM 在 PSNR 和 MS-SSIM 性能方面表现具有竞争力。
• 在结合感知损失与对抗性损失训练后，PQ-MIM 在 FID 和 KID 指标上优于 HiFiC，表明其具有更优的感知质量。
• 通过基于 MIM 的熵建模，模型显著降低了比特率，其中五芒星掩码策略在仅使用五步的情况下即实现了显著的率节省。
• PQ-MIM 展现出强大的图像修复能力，可在高达 50% 的补丁缺失情况下重建图像，同时保持结构与语义细节。
• 该模型可运行于混合压缩-生成模式，无需额外训练即可实现对缺失或损坏补丁的重建。
• 当模型容量从 XCiT-T6（350 万参数）提升至 XCiT-L6（4700 万参数）时，PSNR 性能提升约 +0.8 dB。