谷歌、MIT提出统一框架MAGE：表征学习超MAE，无监督图像生成超越 Latent Diffusion

识别和生成是人工智能领域中的两大核心任务，如果能将二者合并到一个统一的系统中，这两个任务应该能实现互补。事实上，在自然语言处理中，像 BERT [1] 这样的模型不仅能够生成高质量的文本，还能够提取文本中的特征。

然而，在计算机视觉领域，目前的图像生成模型和识别模型大多是分开进行训练，没有充分利用这两个任务的协同作用。这主要是由于图像生成和图像识别的模型通常具有本质上的结构差异：图像生成的输入是低维度的特征或噪声，而输出是高维度的原始图像；与之相反，图像识别的输入是高维度的原始图像，而输出是低维度的特征。

最近，来自 MIT 和 Google Research 的研究人员提出了一种基于图像语义符掩码的表征学习方法，首次在一个统一的框架中实现了图像生成和表征学习，并在多个数据集上取得了 SOTA 表现。研究论文已被 CVPR 2023 接收，相关代码与预训练模型已开源。

• 论文地址：https://arxiv.org/abs/2211.09117
• 代码地址：https://github.com/LTH14/mage

在 CVPR 2022 上，MAE [2] 提出了一种基于图像掩码（MIM）的表征学习方法，并在多个子任务上取得了非常好的效果。在高达 75% 的掩码率下，MAE 可以重构出与原图语义十分贴合的图像，并借此让网络能够自监督地学习图像中的特征。然而，如图 1 所示， MAE 重建的图像虽然具有与原始图像相似的语义信息，但会出现严重的模糊与失真问题。类似的问题也出现在所有基于 MIM 的表征学习方法中。同时，目前的生成模型，不管是扩散模型还是 GAN，都缺乏提取高质量图像特征的能力。

图 1：MAE 与 MAGE 重构对比

方法概述

针对上述问题，本文作者提出了 MAGE（Masked Generative Encoder），首次实现了统一的图像生成和特征提取模型。与MIM直接作用于图像的掩码方法不同，MAGE 提出了基于图像语义符的 masked image token modeling 方法。如图所示，MAGE 首先使用 VQGAN [3] 编码器将原始图像转换为离散的语义符。之后，MAGE 对其进行随机掩码，并使用基于 transformer 的 encoder-decoder 结构对掩码进行重构，重构后的语义符可以通过 VQGAN 解码器生成原始图像。通过在训练中使用不同的掩码率，MAGE 可以同时进行生成模型（接近 100% 掩码率）和表征学习（50%-80% 掩码率）的训练。如图 1 所示，MAGE 重建出的图像不仅具有与原始图像一致的语义信息，还能够同时保证生成图像的多样性与真实性。

图 2：MAGE 结构图

实验结果

MAGE 在多个图像生成与图像识别任务上都达到或超过了 SOTA。

在 ImageNet 的无监督图像生成任务中，MAGE 的 FID 从之前的 > 20 降至 7.04，甚至达到了有监督图像生成的水准（有监督 Latent Diffusion 在 ImageNet 上的 FID 为 3.60）：

图3：MAGE 无监督图像生成样例

MAGE 还能够进行各类图像编辑工作，包括 image inpainting、outpainting、uncropping：

图 4：MAGE 图像编辑样例

在表征学习方面，MAGE 在 ImageNet linear probing、少样本学习、迁移学习等任务中，相较于目前的 MIM 方法有了大幅提升，并且可以达到或超过目前最优的自监督学习方法的水平。

结语

本文旨在将图像生成与表征学习统一起来。为此，本文作者提出了 MAGE，一种基于图像语义符掩码的自监督学习框架。该框架简洁、高效，并首次在图像生成和表征学习上都达到或超越了 SOTA 的表现。感兴趣的读者可以查看论文原文，以了解更多研究细节。