自监督学习系列(三):基于 Masked Image Modeling

前言

随着 Vision Transformer (ViT) 在 2021 年霸榜各大数据集,如何基于 ViT 构建更加合适的自监督学习范式成为了该领域的一大问题。最初,DINO 和 MoCo v3 尝试将对比学习和 ViT 相结合,取得了不错的效果。不过长期以来,由于 CV 和 NLP 领域数据和基础模型之间的差异,NLP 的 Masked Language Modeling (MLM) 掩码模式机制没能成功应用于 CV 领域,但最近 ViT 的蓬勃发展,为掩码学习机制应用于视觉自监督打开了一扇大门。

本文将介绍四篇基于掩码学习的自监督学习论文,梳理掩码学习的发展概况。

1、BEiT

论文链接:https://arxiv.org/abs/2106.08254

算法全称为 Bidirectional Encoder representation from Image Transformers (BEiT),提出了 Masked Image Modeling 自监督训练任务的概念,以此来对 ViT 进行训练。如算法概览图(下图)所示,BEiT 预训练中,每一张图片有两种视角:
一是图像块 (image patches),如每一小块图像为 16x16 像素;
二是离散的视觉标记 (discrete visual tokens)。

在预训练过程中,BEiT 先将原始图片标记化,并且对图像块进行随机掩码,并将掩码后的图片输入到编码器当中,主要的预训练目标便是基于未掩码图像块来恢复掩码图像块。

首先来看图片的表示,图像块和视觉标记。

图像块和 ViT 原文所描述的并无二致,而对于重建目标,BEiT 并没有使用原始的像素,而是通过一个 “image tokenizer” 进行离散化,遵循的是 dVAE 的思路,在 BEiT 预训练之前,先构建 “tokenizer” 和 “decoder” 进行 dVAE 的训练,并构建视觉词汇表。在 BEiT 中是直接采用 Zero-shot text-to-image generation 文章开源的代码进行训练。

对于预训练的主干网络,则是标准的 ViT,每个图像块会被线性投射为对应的 embedding 向量,同时再加上标准的可学习的绝对位置编码。而与之不同的是,BEiT 采用了 Blockwise Masking 的方式,对大约 40% 的图像块进行了掩码操作,而预训练的目标便是期望能够正确预测掩码图像块的视觉标记,从而获得可以提取图像特征向量的编码器。

在下游的分类和分割任务上,BEiT 均超过了之前的自监督算法和有监督训练模式,达到 SOTA 水准。

BEiT 可以说是将 BERT 形式的预训练方式迁移到视觉领域的开山之作,并提出 MIM 预训练的任务概念,为自监督领域做出了重要的贡献。

MAE

论文链接:https://arxiv.org/abs/2111.06377

Masked Autoencoders (MAE) 是一篇非常具有影响力的文章。MAE 相比于 BEiT,简化了整体训练逻辑,利用随机掩码处理输入的图像块,以及直接重建掩码图像块来进行训练。MAE 基于两大主要设计:一是采用了非对称结构的编码-解码器,其中编码器只计算非掩码图像块,同时采用了轻量化的解码器设计;二是遮盖大部分的图像块,如掩码概率为 75%,可以获得更加具有意义的自监督训练任务。

MAE 逻辑和其他的自编码器类似,通过编码器将原始信号映射到特定空间内的隐变量,再基于隐变量通过解码器重建原始信号,但是与传统的自编码器不同的是,MAE 采用非对称的结构和轻量级解码器。

首先看掩码部分,拆分图像块的方式和 ViT 一致,之后再随机遮盖图像的大部分,使其只留下部分可见内容(如下图),所以在训练过程中,模型不容易找到捷径解,例如插值等。

其次来看编码器和解码器,编码器即是标准的 ViT 模型,只不过只对非掩码图像块进行计算,从中提取特征,这种设计可以减少计算量和内存;而解码器则会对所有可见的图像块和掩码图像块进行计算,对于每个图像块会加上位置编码信息,以避免图像块的位置信息丢失。由于数据的输入只有掩码图像块以及编码器和解码器的非对称性,两个模块互相独立设计,所以可以大大减少训练时间。

对于重建目标,MAE 针对每个掩码图像块进行像素值预测,并计算 MSE 损失函数。

部分图像可视化如下,分别为掩码后的图像(左),重建图像(中),原始图像(右)。

在下游任务上,作者提出之前的线性评估和端到端微调具有很强的不相关性,即使在过去几年内线性评估是最受欢迎的下游评价方式。并且对基于 MAE 预训练的和 MoCo v3 所训练的 ViT-L 进行了实验对比,MAE 在线性评估中的结果要差于 MoCo v3,但是从部分微调开始,其结果都要比 MoCo v3 要更好(实验结果如下,0 代表线性评估,24 则是全量微调)。

作者提出,线性评估遗漏了一些强大但是非线性的特征,而这正是深度学习的优势。例如在 MAE 中便有更加强大的非线性的特征表示,而线性评估并不能很好的展示这一点,所以采用全量微调或者部分微调的 MAE 能取得更好的结果。

MAE 凭借简单的训练思路和 SOTA 的结果,在视觉领域是迅速走红,是一个非常漂亮的研究工作。

SimMIM

论文链接:https://arxiv.org/abs/2111.09886

在 MAE 同期,还有一篇同样影响力不俗的文章,SimMIM。这篇文章提出了一种简单掩码学习框架,相比之前 SOTA,简化了一些特殊的设计,例如 Blockwise Masking、dVAE 的 tokenizer或聚类等方法;而简单的设计,例如随机掩码、回归 RGB 像素值和采用线性预测头,也可以取得 SOTA 结果。

上图为 SimMIM 的总览图,主要可以分为四个部分:掩码部分,编码器,预测头和预测目标。

SimMIM 认为掩码部分采用随机掩码,并且适当增大图像块的分辨率即可获得很好的结果,并且文章中提供了对比试验的结果,当图像块分辨率为 32x32 时可以获得最好的效果,并将其设置为 ViT 模型的默认设置。和 NLP 领域算法及 BEiT 类似,掩码 token 为可学习的向量,并且维度和其他图像块线性映射后的维度一致。

在本文中,编码器采用了 ViT 和 Swin Transformer。

对于预测头,文章中也分别对四种架构进行了实验,包括线性层、2 层 MLP、inverse Swin-T 和 inverse Swin-B。实验结果表明使用最简单的线性层,花费最少的训练时间,却可以获得最高的准确率,所以最终便采用了最简单的线性层作为预测头,对掩码图像块进行预测。

预测目标为最简单的 L1 损失函数,即可以获得 SOTA 级别的结果。

文章中也给了一些图像重建可视化的例子,下图中展示了三种掩码方式:第一列为原图;第二列左图采用随机掩码方式,右图为重建后的结果;第三列左图采用掩盖大部分主体方式,右图为重建后的结果;第四列左图采用掩盖所有主体方式,右图为重建后的结果。

SimMIM 提供了一个非常简单的掩码学习训练框架并取得了 SOTA 结果,而且简单的框架工作更容易被人所接受和传播,是一个非常漂亮的研究工作。

MaskFeat

论文链接:https://arxiv.org/abs/2112.09133

MaskFeat 算法在整体思路上依然是重建掩码图像块的思路,只不过它的重建目标从原始像素值变成了 HOG 特征描述器。通过作者的实验,在五种不同类型的特征描述中,HOG 可使网络获得最好的结果,且训练更加高效,算法总览图如下:

MaskFeat 证明了可以直接在无标注的视频数据集上进行训练,并且具有非常优秀的迁移性能。因此,视频理解模型可以不再依靠大数据集的预训练,如 ImageNet-21K,这令视频理解类型的任务受益匪浅。

作者在文中列举了五种不同的目标特征,分别为像素值、HOG 特征、dVAE、预训练神经网络提取的特征和伪标签。基于这五种目标特征,进行了实验,图像目标特征的结果如下:

可以看到,虽然基于预训练模型的特征可以取得稍好的结果,但是其预训练所额外消耗的时间导致整个训练流程过长,所以基于 HOG 的特征可以说是当前实验结果内的最优选择。

MaskFeat 不需要依赖 ImageNet-21K 这类超大型数据集,提高了预训练的效率;另外,选取不同的目标特征进行实验也为后续的视觉自监督提供了一个新的探索方向。

猜你喜欢

转载自blog.csdn.net/weixin_43135178/article/details/126353236