基于大型语料库训练的 Transformer 模型在自然语言处理中取得了巨大的成功,作为 Transformer 构建块,self-attention 极大地改变了计算机视觉任务。NLP 的成功不仅依赖于 Transformer 的计算效率和可扩展性,还依赖于对大量文本进行自监督学习。目前 NLP 领域存在两种主流的学习范式:基于自回归语言建模的 GPT 和基于掩码语言建模的 BERT,这两者在计算机视觉领域中也被重新设计,以便充分利用海量的网络图像。
然而,在视觉任务上设计具有相似风格的目标是具有挑战性的,因为图像作为一种高维和冗余的模态,在两个方面与文本不同:首先,文本由离散字符组成,而图像在颜色空间中呈现连续值;其次,文本中的离散 token 包含高级语义含义,而离散化的图像在像素级和 patch 级包含大量冗余 token。
因此,我们不禁会问是否有一种方法可以学习感知离散视觉 token,这将有助于图像预训练。
基于上述观察,来自中国科学技术大学、微软亚研等机构的研究者提出了学习感知 codebook( perceptual codebook ,PeCo),用于视觉 transformer 的 BERT 预训练。目前,BEiT 成功地将 BERT 预训练从 NLP 领域迁移到了视觉领域。BEiT 模型直接采用简单的离散 VAE 作为视觉 tokenizer,但没有考虑视觉 token 语义层面。相比之下,NLP 领域中的离散 token 是高度语义化的。这种差异促使研究者开始学习感知 codebook,他们发现了一个简单而有效的方法,即在 dVAE 训练期间强制执行感知相似性。
该研究证明 PeCo 生成的视觉 token 能够表现出更好的语义,帮助预训练模型在各种下游任务中实现较好的迁移性能。例如,该研究使用 ViT-B 主干在 ImageNet-1K 上实现了 84.5% 的 Top-1 准确率,在相同的预训练 epoch 下比 BEiT 高 1.3。此外,该方法还可以将 COCO val 上的目标检测和分割任务性能分别提高 +1.3 box AP 和 +1.0 mask AP,并且将 ADE20k 上的语义分割任务提高 +1.0 mIoU。
论文地址:https://arxiv.org/pdf/2111.12710v1.pdf
方法
在自然语言中,词是包含高级语义信息的离散 token。相比之下,视觉信号是连续的,具有冗余的低级信息。在本节中,该研究首先简要描述了 VQ-VAE 的离散表示学习,然后介绍如何学习感知 codebook 的过程,最后对学习感知视觉 token 进行 BERT 预训练。
学习用于可视化内容的离散 Codebook
该研究利用 VQ-VAE 将连续图像内容转换为离散 token 形式。图像表示为 x∈ R^H×W×3,VQ-VAE 用离散视觉 Codebook 来表示图像,即
其中,VQ-VAE 包含三个主要部分:编码器、量化器和解码器。编码器负责将输入图像映射到中间潜在向量 z = Enc(x);量化器根据最近邻分配原则负责将位置 (i, j) 处的向量量化为来自 Codebook 对应的码字(codewords):
然后得到如下公式:
其中 q 是量化编码器,可以将向量映射到 codebook 索引,r 是量化解码器,可以从索引重构向量。基于量化的码字为 z_q,解码器旨在重构输入图像 x。VQ-VAE 的训练目标定义为:
学习用于视觉内容的 PeCo
该研究提出,在不包含像素损失的情况下,对模型强制执行原始图像和重构图像之间的感知相似性。感知相似性不是基于像素之间的差异得到的,而是基于从预训练深度神经网络中提取的高级图像特征表示之间的差异而得到。该研究希望这种基于 feature-wise 的损失能够更好地捕捉感知差异并提供对低级变化的不变性。下图从图像重构的角度展示了模型使用不同损失的比较,结果表明图像在较低的 pixel-wise 损失下可能不会出现感知相似:
图 1. 不同损失下的图像重构比较。每个示例包含三个图像,输入(左)、使用 pixel-wise 损失重构图像(中)、使用 pixel-wise 损失和 feature-wise 损失重构图像(右)。与中间图像相比,右侧图像在感知上与输入更相似。
在形式上,假设输入图像 x 和重构图像
的感知度量可以表示为:
其中 S 表示提取特征的层数,总的目标函数为:
BERT objective 执行掩码图像建模
该研究采用 BERT objective 在离散视觉 token 上执行掩码图像建模任务,如 BEiT。对于给定的图像 x,输入 token 为不重叠的图像 patch,输出 token 是通过学习方程 (5) 获得的离散感知视觉单词。设输入为 {x_1 , x_2 , · · · , x_N },并且真值输出为
。 掩码图像建模的目标是从掩码输入中恢复相应的视觉 token,其中一部分输入 token 已被掩码掉。准确地说,令 M 为掩码索引集合,掩码输入
表示为:
其中,m 是与非掩码 token 相同维度的可学习掩码 token。掩码(masked)输入 token 被送入 L 层视觉 Transformer,最后一层的隐藏输出表示为 {h^1 , h^2 , · · · , h^N }。
实验
该研究将预训练模型应用于各种下游任务,包括 ImageNet-1K 分类、COCO 目标检测和 ADE20k 分割。
与 SOTA 模型比较
首先该研究将 PeCo 与 SOTA 研究进行比较。研究者使用 ViT-B 作为主干并在 ImageNet-1K 上进行预训练,epoch 为 300 。
图像分类任务:在 ImageNet 1K 上进行分类任务的 Top-1 准确率如表 1 所示。可以看出,与从头开始训练的模型相比,PeCo 显着提高了性能,这表明预训练的有效性。更重要的是,与之前自监督预训练模型相比,PeCo 模型实现了最佳性能。值得一提的是,与采用 800 epoch 的 BEiT 预训练相比,PeCo 仅用 300 epoch 就实现了 0.9% 的提高,并比 MAE 采用 1600 epoch 预训练性能提高 0.5%。这验证了 PeCo 确实有利于预训练。
语义分割任务:该研究将 PeCo 与 1)在 ImageNet-1K 上进行监督预训练和 2)BEiT(SOTA 性能自监督学习模型)进行比较,评估指标是 mIoU,结果如表 2 所示。由结果可得,PeCo 在预训练期间不涉及任何标签信息,却取得了比监督预训练更好的性能。此外,与自监督 BEiT 相比,PeCo 模型也获得了较好的性能,这进一步验证了 PeCo 的有效性。
目标检测与分割:如表 3 所示,在这一任务上,PeCo 获得了最好的性能:
感知 Codebook 分析
码字语义:学习的感知码字是否具有(更多)语义含义?为了回答这个问题,该研究设计实验以提供视觉和定量结果。
首先,该研究将对应于相同码字的图像 patch 进行可视化,并与两个基线进行比较:在 2.5 亿私有数据上训练而成的 DALL-E codebook;不使用感知相似性的 PeCo 模型的一个变体。结果如图 3 所示,我们可以看到该研究码字与语义高度相关,如图中所示的轮子,来自基线的码字通常与低级信息(如纹理、颜色、边缘)相关。
此外,该研究还与不使用感知相似性的变体进行了比较。如表 4 所示, 我们可以发现感知码字在线性评估和重构图像分类方面获得了更高的准确率。这表明感知 codebook 具有更多的语义意义,有利于图像重构过程。
下图为使用 BEiT 和 PeCo 在 ImageNet-1k 上重构任务的示例。对于每个样本,第一张是原始图像,第二张是对应的掩码图像,第三张是 BEiT 重构图像,最后一张是从感知 codebook(PeCo)重构的图像。PeCo 在感知 codebook 的帮助下,能够对掩码区域进行更语义化的预测。