TPAMI 2024 | EMAE | 具有自一致性的高效掩码自动编码器
题目:Efficient Masked Autoencoders with Self-Consistency
论文:https://arxiv.org/pdf/2302.14431
代码:(未公开)
年份:2024
创新点
本文做出了以下贡献:
- 观察到 MIM 的高随机掩码率导致两个严重问题:预训练时间延长和预训练模型的高度不一致性。
- 提出一种并行掩码策略,在保持随机性的同时实现 100% 的数据利用率,有效提高了 MIM 的预训练效率。
- 将自一致性学习引入 MIM,生成更可靠的表示。
- 大量实验证明了本文方法的有效性和更强的泛化能力。特别是,EMAE 在各种下游任务(如线性探测、微调分类、目标检测、实例分割和语义分割)中优于先前的 SOTA 方法。
摘要:
受自然语言处理任务中掩码语言建模(MLM)的启发,掩码图像建模(MIM)已被公认为计算机视觉中一种强大的自监督预训练方法。
MIM 的高随机掩码率导致两个严重问题:
1)每次迭代中图像的数据利用率不足,导致预训练时间延长;
2)预测的高度不一致导致生成结果不可靠,即相同补丁在不同掩码轮次中的预测可能不一致,从而在最终生成的结果中产生语义分歧。
为了解决这些问题,本文提出了具有自一致性的高效掩码自动编码器(EMAE),以提高预训练效率并增强 MIM 的一致性。具体而言
- 本文提出一种并行掩码策略,将图像划分为 K 个不重叠的部分,每个部分由具有相同掩码率的随机掩码生成。
- 然后在一次迭代中对所有部分并行进行 MIM 任务,模型最小化预测结果与掩码补丁之间的损失。
- 此外,本文设计了自一致性学习,以进一步保持各部分之间重叠掩码补丁预测的一致性。
总体而言,本文的方法能够更有效地利用数据并获得可靠的表示。在 ImageNet 上的实验表明,使用 NVIDIA A100 GPU,EMAE 仅用 MAE 预训练时间的 13%,就在 ViT-Large 模型上取得了最佳性能。 在多个数据集上进行预训练后,EMAE 在各种下游任务(如图像分类、目标检测和语义分割)中始终获得最先进的迁移能力。
引言
- 受掩码语言建模(MLM)成功的启发,掩码图像建模(MIM)在计算机视觉中用于预训练,并在性能上展现出显著优势。
- MIM 以一定的随机掩码率对图像的补丁进行掩码,并基于剩余的可见补丁预测被掩码补丁的原始表示。
- 由于图像和语言数据的信息密度不同,MIM 的随机掩码率远高于 MLM。
例如,MAE 采用 75% 的掩码率,而 BERT 使用 15% 的掩码率。
然而,本文发现高随机掩码率带来两个严重问题:
1)每次迭代中图像的数据利用率不足,导致预训练时间延长;
2)预训练模型的预测表现出高度不一致性,导致生成结果不可靠。
事实上,MAE 每次迭代利用整个图像的 25% 来训练模型,共训练 1600 个 epoch,而 BERT 使用文本语料库的 85% 仅训练 40 个 epoch。这种现象的原因不仅在于图像和语言数据的差异,还在于 MLM 和 MIM 方法在数据利用上的不同。
在计算机视觉的自监督学习中,MIM 的预训练效率仍然低于对比学习。例如,MoCo v3 通常对模型进行 300 个 epoch 的预训练。对比学习和 MIM 在预训练效率上的差异源于对比学习利用整个图像进行模型训练,从而实现了 100% 的数据利用率。
- 从上述描述可以得出结论,MIM 的高掩码率导致数据利用率低,从而延长了预训练时间。
- MIM 的另一个问题是高掩码率会引入不太可靠的特征。
图1. MAE[21]的不同重建结果对应不同的掩码种子。(a) 不同掩码种子采样得到的不同组合。(b) MAE的重建结果。在(b)中的三个重建结果里,只有第一个代表正常的牛,而第三个甚至重建出了一只狗。MAE的这些重建结果语义不一致。
如图 1 所示,MAE 根据不同的掩码种子生成不同的重建结果,且不同结果的语义不一致。也就是说,从原始图像中采样的不同随机可见补丁组合,对相同位置会产生不一致的预测。这种生成的不一致性是导致结果不可靠的一个因素。 同时,作为一种生成算法,如果 MAE 的生成结果表现出不一致性和不可靠性,那么它就无法获得可靠和一致的表示,从而削弱了其在各种下游任务中的迁移能力。
- 能否降低随机掩码率以提高预训练效率并改善一致性呢?答案是否定的。降低掩码率会削弱预训练模型在各种视觉下游任务中的迁移能力。
因此,在保持高掩码率不变的情况下,解决高掩码率带来的上述两个问题非常重要。在本文中提出了具有自一致性的高效掩码自动编码器(EMAE),旨在在高掩码率下提高预训练效率,并增强 MIM 的确定性和一致性。
图2. 本文的EMAE示意图。首先将整幅图像划分为 N 个补丁,然后并行策略将N个补丁划分为K个大小相同的不重叠部分 x v 1 , . . . , x v K x_{v_{1}}, ..., x_{v_{K}} xv1,...,xvK 。每个部分由 N / K N/K N/K个随机不重叠的可见图像补丁组成。然后,每个部分被输入到编码器 - 解码器架构中,并执行MIM任务,生成 x m 1 , . . . , x m K x_{m_{1}}, ..., x_{m_{K}} xm1,...,xmK 。此外,自一致性学习引导相同位置的重叠预测趋于一致。这里以 x p 1 x_{p_{1}} xp1和 x p K x_{p_{K}} xpK为例。
具体而言,如图 2 所示,本文提出一种并行掩码策略:
- 将图像划分为 K 个大小相同的不重叠部分。
- 然后,每个部分是一个包含随机不重叠可见图像补丁的集合。
- 所有部分具有相同的掩码率,并并行处理。通过这种方式,图像的数据利用率显著提高,达到了 100%。
- 同时,引入自一致性学习,鼓励模型在不同部分的相同位置输出一致的预测,从而提高了 MIM 的确定性和一致性。
- 虽然并行掩码策略增加了单次迭代的计算量,但它减少了总体训练轮数,提高了预训练效率。
本文在多种视觉预训练数据集和视觉下游任务上对本文方法进行了验证。使用 NVIDIA A100 GPU,本文的方法仅用 13% 的预训练时间,就在各种视觉下游任务中优于 MAE,并取得了最先进(SOTA)的性能。此外,在 ViT-Large 架构下,EMAE 在 COCO 目标检测和实例分割任务中,使用 ViTDet 达到了 58.1% 的边界框平均精度均值(bbox mAP)和 50.4% 的掩码平均精度均值(mask mAP)。
本文的方法
本文提出的方法流程如图 2 所示。
图2. 本文的EMAE示意图。首先将整幅图像划分为 N 个补丁,然后并行策略将N个补丁划分为K个大小相同的不重叠部分 x v 1 , . . . , x v K x_{v_{1}}, ..., x_{v_{K}} xv1,...,xvK 。每个部分由 N / K N/K N/K个随机不重叠的可见图像补丁组成。然后,每个部分被输入到编码器 - 解码器架构中,并执行MIM任务,生成 x m 1 , . . . , x m K x_{m_{1}}, ..., x_{m_{K}} xm1,...,xmK 。此外,自一致性学习引导相同位置的重叠预测趋于一致。这里以 x p 1 x_{p_{1}} xp1和 x p K x_{p_{K}} xpK为例。
EMAE 主要由两个主要组件组成。第一个组件是并行掩码策略,它实现了 100% 的数据利用率,提高了预训练效率。此外,EMAE 引入了自一致性学习,强制相同位置的不同预测之间保持一致。
关于MAE(掩码自编码器)的预备知识
图像划分
给定一个来自未标记数据集的自然图像x,将其划分为N个规则的图像补丁,记为
x
∈
R
N
×
S
x \in R^{N×S}
x∈RN×S,其中S代表补丁大小(例如在ViT中为16×16)。
代码如下
def patchify(self, imgs): #把图片划分成块(来源于MAE)
"""
imgs: (N, 3, H, W)
x: (N, L, patch_size**2 *3)
"""
p = self.patch_embed.patch_size[0] #获取Patch大小,在这里patch=16,意味着每个图像补丁是16*16=256像素
assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0 #确保图像是方形,其中imgs.shape[3]=W,imgs.shape[2]=H
# 并且确保图像可以被 p 整除(否则无法均匀分割)
h = w = imgs.shape[2] // p #计算划分过后patch的行列数h=w=imgs.shape[2](也即是H)/Patch(16)=14
x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
"""
这一步的作用是将图像从原来的imgs:(N,3,H,W)修改为imgs:(imgs.shape[0](即是N),3(3通道),h(行数=14),p(每个patch的高度=16),w(列数=14),p(每个patch的宽度=16))
例如:
imgs.shape = (N, 3, 224, 224)
p = 16
h = w = 224 / 16 = 14
x.shape = (N, 3, 14, 16, 14, 16) # 重新组织形状
"""
x = torch.einsum('nchpwq->nhwpqc', x) #使用einsum进行维度重排(目的是让 Patch 维度 h, w 移到前面,方便后续展开。)
"""
例如:
x.shape = (N, 14, 14, 16, 16, 3) # 交换维度后
"""
x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))
"""
把 h × w 个 Patch 变成一个序列(即 L=h*w)。
每个 Patch 变成一个向量:
p × p × 3(即 patch_size**2 * 3)
例如 p=16,则 16×16×3=768,每个 Patch 变成一个 768 维的向量。
最终的形状:
x.shape = (N, 196, 768) # N张图片,每张图片有196个Patch,每个Patch是768维
"""
return x
随机掩码策略
接着,用掩码
m
=
(
m
1
,
.
.
.
,
m
N
)
m = (m_1, ..., m_N)
m=(m1,...,mN)表示一个长度为N的二进制向量,其中
m
a
∈
{
0
,
1
}
m_a \in \{0, 1\}
ma∈{0,1}且
a
∈
{
1
,
2
,
.
.
.
,
N
}
a \in \{1, 2, ..., N\}
a∈{1,2,...,N},它代表对图像的掩码操作,并基于此生成x的两个互补组合:掩码补丁
x
m
x_m
xm和可见补丁
x
v
x_v
xv,如公式(1)所示:
x
m
=
x
⊙
m
∈
R
η
×
S
,
x
v
=
x
⊙
(
1
−
m
)
∈
R
κ
×
S
(1)
x_{m}=x \odot m \in R^{\eta × S}, x_{v}=x \odot(1 - m) \in R^{\kappa × S}\tag{1}
xm=x⊙m∈Rη×S,xv=x⊙(1−m)∈Rκ×S(1)
其中 η = N × p \eta = N×p η=N×p, κ = N × ( 1 − p ) \kappa = N×(1 - p) κ=N×(1−p) 。掩码m的生成遵循均匀分布,这被称为随机掩码策略。
def random_masking(self, x, mask_ratio): #随机掩码操作(来源于MAE)
"""
Perform per-sample random masking by per-sample shuffling.
Per-sample shuffling is done by argsort random noise.
x: [N, L, D], sequence
"""
N, L, D = x.shape
len_keep = int(L * (1 - mask_ratio)) #算出要保留的块的数目,保留的块的数目=总的块的数目*(1-掩码率)。即196*(1-75%)=49
noise = torch.rand(N, L, device=x.device) #生成一个 N × L 维度的随机矩阵 noise,值在 [0,1] 之间,每个样本对应 L 个随机值。
# sort noise for each sample
ids_shuffle = torch.argsort(noise, dim=1) # 对 noise 进行 升序排序,得到每个样本从小到大的索引 ids_shuffle。
# 索引含义:小的数表示靠前的位置(即这些位置会被保留),大的数表示靠后的(即这些位置会被掩码)。
#例如:
#noise = [[0.9, 0.1, 0.4, 0.7]] # 原始随机噪声
#ids_shuffle = [[1, 2, 3, 0]] # 排序后对应的索引
ids_restore = torch.argsort(ids_shuffle, dim=1) # ids_restore 是 ids_shuffle 的逆变换,即用于 恢复原来的索引顺序。
# 假设 ids_shuffle = [[1, 2, 3, 0]],那么:
# ids_restore = [[3, 0, 1, 2]]
# keep the first subset
ids_keep = ids_shuffle[:, :len_keep] #取 ids_shuffle 的前 len_keep 个索引,即保留的部分。(取ids_shuffle的前25%作为自己的ids_keep)
x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D)) #torch.gather 按 ids_keep 选择 x 的一部分,保持特征维度不变。
# generate the binary mask: 0 is keep, 1 is remove
mask = torch.ones([N, L], device=x.device) # 先生成全 1 矩阵
mask[:, :len_keep] = 0 # 前 len_keep 个位置设为 0
# 初始时,mask 为全 1(代表所有位置都被掩码)。
# 前 len_keep 个位置设置为 0(表示保留),后面的保持 1(表示被掩码)。
# unshuffle to get the binary mask
mask = torch.gather(mask, dim=1, index=ids_restore) # 通过 ids_restore 重新排列 mask,确保掩码对应到原始 x 的顺序。
return x_masked, mask, ids_restore
编码器、解码器与损失函数
MAE模型
h
=
g
∘
f
h = g \circ f
h=g∘f 采用编码器 - 解码器架构。将这些可见补丁输入到编码器
f
(
⋅
)
f(·)
f(⋅)(如ViT - Base)中,获取潜在特征。随后,解码器
g
(
⋅
)
g(·)
g(⋅)将潜在特征映射回像素空间,以重建互补组合,进而得到预测结果
x
p
x_p
xp。特别地,MAE通过最小化重建的图像补丁与掩码图像补丁
x
m
x_m
xm之间的均方误差(MSE)来优化模型,其损失函数如公式(2)所示:
L
M
A
E
(
x
)
=
L
(
x
v
,
x
m
)
=
E
∥
g
(
f
(
x
v
)
)
−
x
m
∥
2
(2)
\mathcal{L}_{MAE}(x)=\mathcal{L}\left(x_{v}, x_{m}\right)=\mathbb{E}\left\| g\left(f\left(x_{v}\right)\right)-x_{m}\right\| ^{2}\tag{2}
LMAE(x)=L(xv,xm)=E∥g(f(xv))−xm∥2(2)
- 编码器与解码器
def forward_encoder(self, x, mask_ratio): #编码器,x(输入),mask_ratio掩码率
# embed patches
x = self.patch_embed(x) #将图像输入到patch_embed操作中得到新的x
# add pos embed w/o cls token
x = x + self.pos_embed[:, 1:, :] #在当前x的基础上加入位置编码,cls_token加在第0个位置
# masking: length -> length * mask_ratio
x, mask, ids_restore = self.random_masking(x, mask_ratio)
#调用random_masking函数对输入的x进行掩码率为mask_ratio(75%)的掩码,其输出是x,mask是掩码矩阵
# append cls token
cls_token = self.cls_token + self.pos_embed[:, :1, :] #cls_token+pos_embed即加入类别嵌入,嵌在第0个位置
cls_tokens = cls_token.expand(x.shape[0], -1, -1)
"""
由于每个批次(batch)这里指的是 x.shape[0]=N 中有多个图像,所以需要扩展 cls_token 使它的形状适应批次的大小。
expand(x.shape[0], -1, -1) 将 cls_token 扩展为 [batch_size, 1, embedding_dim] 的形状,
其中 x.shape[0] 是批次大小。
"""
x = torch.cat((cls_tokens, x), dim=1)
"""
torch.cat((cls_tokens, x), dim=1) 这行代码将 cls_tokens 和 x 在 第二维度(dim=1) 进行拼接。
cls_tokens 是一个表示 分类token 的张量。它的形状是 [batch_size, 1, embedding_dim]
x 的形状是 [batch_size, num_patches, embedding_dim]
拼接后的张量 x 的形状会变成:
[batch_size, 1 + num_patches, embedding_dim]。
"""
# apply Transformer blocks
for blk in self.blocks: #对每一个block进行一次前向运算
x = blk(x)
x = self.norm(x) #对输出进行归一化得到encoder的输出
return x, mask, ids_restore
def forward_decoder(self, x, ids_restore):
# embed tokens
x = self.decoder_embed(x)
#编码器部分的输出传入decoder_embed,将encoder的输出维度1024进行线性投影降低到decoder的输入维度512
# append mask tokens to sequence
mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1)
"""
创建掩码tokens:mask_tokens 是用于恢复缺失部分的token。
首先,self.mask_token 是一个预定义的掩码token,表示被掩码的部分。
然后,repeat 操作将其扩展为需要的形状
x.shape[0] 是批次大小,ids_restore.shape[1](总的tokens数目) + 1 - x.shape[1](可见tokens数目)
是当前需要掩码的token数量(这部分根据掩码后的图像恢复)。
mask_token的维度是 [batch_size, num_masked_tokens, embedding_dim]。
"""
x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)
"""
拼接掩码tokens:将编码器输出的 x(去掉第一个分类token,即 x[:, 1:, :])与mask_tokens拼接起来,形成新的 x_
这里的 dim=1 表示在补丁维度上拼接。
拼接后,x_ 包含了所有补丁以及需要恢复的掩码token
"""
x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))
"""
恢复原始顺序:通过 torch.gather 操作将 x_ 按照 ids_restore 中的顺序恢复原始图像补丁的位置。
ids_restore.unsqueeze(-1) 是将索引的维度扩展为三维,并且 repeat(1, 1, x.shape[2]) 会将其扩展为适应 x_ 的形状。
x_ 的顺序会根据 ids_restore 恢复,确保掩码后的补丁位置重新排列。
"""
x = torch.cat([x[:, :1, :], x_], dim=1) # 把cls_token加上
# add pos embed
x = x + self.decoder_pos_embed #增加解码器位置编码
# apply Transformer blocks
for blk in self.decoder_blocks: #对decoder_block进行前向计算
x = blk(x)
x = self.decoder_norm(x) #对decoder的输出进行归一化
# predictor projection
x = self.decoder_pred(x) #把特征映射到像素空间上
# remove cls token
x = x[:, 1:, :] #除去cls_token
return x
- 损失函数
def forward_loss(self, imgs, pred, mask):
"""
imgs: [N, 3, H, W]
pred: [N, L, p*p*3]
mask: [N, L], 0 is keep, 1 is remove,
"""
target = self.patchify(imgs)
"""
图像切片:patchify 方法将原始图像 imgs 切分为多个补丁。
每个补丁是一个特征向量,target 的形状是 [N, L, p*p*3],与 pred 的形状相匹配。
这个步骤是将图像转化为补丁,以便后续的损失计算。
"""
if self.norm_pix_loss: #计算补丁的均值 mean 和方差 var。
mean = target.mean(dim=-1, keepdim=True)
var = target.var(dim=-1, keepdim=True)
target = (target - mean) / (var + 1.e-6)**.5
loss = (pred - target) ** 2
loss = loss.mean(dim=-1) # [N, L], mean loss per patch
loss = (loss * mask).sum() / mask.sum() # mean loss on removed patches
return loss
并行掩码策略
MAE的问题
如公式(1)所示,随机掩码率为
p
p
p,因此数据利用率为
1
−
p
1 - p
1−p。这意味着随机掩码率越高,数据利用率就越低。低数据利用率会因训练数据参与不足,降低预训练效率。
x
m
=
x
⊙
m
∈
R
η
×
S
,
x
v
=
x
⊙
(
1
−
m
)
∈
R
κ
×
S
(1)
x_{m}=x \odot m \in R^{\eta × S}, x_{v}=x \odot(1 - m) \in R^{\kappa × S}\tag{1}
xm=x⊙m∈Rη×S,xv=x⊙(1−m)∈Rκ×S(1)
其中 η = N × p \eta = N×p η=N×p, κ = N × ( 1 − p ) \kappa = N×(1 - p) κ=N×(1−p) 。掩码m的生成遵循均匀分布,这被称为随机掩码策略。
例如,分别以 MAE ( m a s k _ r a t i o = 0.75 , 1 − m a s k _ r a t i o = 0.25 mask\_ratio=0.75,1-mask\_ratio=0.25 mask_ratio=0.75,1−mask_ratio=0.25 ) 和 BERT ( m a s k _ r a t i o = 0.15 , 1 − m a s k _ r a t i o = 0.85 mask\_ratio=0.15,1-mask\_ratio=0.85 mask_ratio=0.15,1−mask_ratio=0.85 ) 作为 MIM 和 MLM 的代表。若两个模型都进行 M M M 个轮次(epoch)的预训练,那么MLM和MIM的数据利用率之比为 ( 17 5 ) M (\frac{17}{5})^M (517)M 。要是用对比学习取代MLM,数据利用率之比会提升到 4 M 4^M 4M,因为对比学习会利用整幅图像进行训练。
先前研究[21]表明,高掩码率(75%)能够增强预训练模型的泛化能力。所以,在高掩码率的情况下实现100%的数据利用率,对优化预训练过程、提升预训练效率来说至关重要。
直观来讲,对同一图像进行多次采样,在一定程度上可以提高数据利用率。 然而,MAE的随机掩码策略本质上是有放回采样。一旦对同一图像进行多次采样,就很容易出现采样重叠的情况。 尤其是随着
p
p
p值的增大,重叠的可能性也会增加。
在有限的采样次数下,MAE的掩码策略会引发两个问题:
- 其一,无法像对比学习或监督学习那样,保证数据的完全利用(即达到100%的数据利用率);
- 其二,同一图像中每个补丁的采样次数往往不一样。
并行掩码策略
为解决这些问题,本文提出一种并行掩码策略,在高掩码率下充分利用所有数据进行MIM预训练,并且保证在有限采样次数内,每个补丁都能得到同等的训练机会。
图2. 本文的EMAE示意图。首先将整幅图像划分为 N 个补丁,然后并行策略将N个补丁划分为K个大小相同的不重叠部分 x v 1 , . . . , x v K x_{v_{1}}, ..., x_{v_{K}} xv1,...,xvK 。每个部分由 N / K N/K N/K个随机不重叠的可见图像补丁组成。然后,每个部分被输入到编码器 - 解码器架构中,并执行MIM任务,生成 x m 1 , . . . , x m K x_{m_{1}}, ..., x_{m_{K}} xm1,...,xmK 。此外,自一致性学习引导相同位置的重叠预测趋于一致。这里,本文以 x p 1 x_{p_{1}} xp1和 x p K x_{p_{K}} xpK为例。
"""
x: the input image
K: the number of non-overlapping parts
map an image into multiple image patches
"""
x = patchify(x) #划分图像块
N, D = x.shape # length, dim
tensor = rand(N) # tensor in [0, 1]
# sort the tensor in ascending order
ids = argsort(tensor)
# acquire the position of each element
ids_tensor = argsort(ids)
# divide the whole data into K parts
for i in range(1,K+1):
# obtain the i-th visible patches
ids_i = ids[(i-1)*(N/K):i*(N/K)]
x_v_i = gather(x, dim=0, index=ids_i)
# obtain the i-th mask
m_i = ones(N)
m_i[(i-1)*(N/K):i*(N/K)] = 0
m_i = gather(m_i, dim=0, index=ids_tensor)
# obtain the i-th masked patches
x_m_i = x[m_i].reshape(-1,D)
rand: returns a tensor filled with random numbers from a uniform distribution on the
interval [0,1];
gather: gathers values along an axis specified by dim;
argsort: returns the indices that sort a tensor along a given dimension in ascending
order by value
具体实现上,先把整幅图像划分为
N
N
N个图像补丁。接着,生成长度为
N
N
N 的张量
t
t
t ,其每个值都服从
[
0
,
1
]
[0, 1]
[0,1] 区间的均匀分布。
将张量
t
t
t按值升序排序,得到排序索引
i
d
s
ids
ids ,如公式(3)所示:
其中 s ( . ) s(.) s(.)会返回按值对张量进行升序排序的索引。
i d s = s ( t ) (3) ids = s(t)\tag{3} ids=s(t)(3)
这里,把长度为 N N N的排序索引均匀分成 K K K个不重叠的部分 i d s 1 ids_1 ids1, i d s 2 ids_2 ids2,…, i d s K ids_K idsK ,如公式(4)所示:
其中 i ∈ { 1 , 2 , . . . , K } i \in \{1, 2, ..., K\} i∈{1,2,...,K}。
i d s i = i d s [ ( i − 1 ) × N K : i × N K ] (4) ids_{i}=ids[(i - 1)×\frac{N}{K}:i×\frac{N}{K}]\tag{4} idsi=ids[(i−1)×KN:i×KN](4)
这样, N N N个图像补丁就能通过索引 i d s i ids_i idsi,均匀地划分为 K K K个不重叠的部分 x v 1 x_{v_1} xv1, x v 2 x_{v_2} xv2,…, x v K x_{v_K} xvK ,如公式(5)所示:
其中 d ( ⋅ ) d(·) d(⋅)表示依据指定索引 i d s i ids_i idsi从输入 x x x中取值。
x v i = d ( x , i d s i ) (5) x_{v_{i}}=d\left(x, ids_{i}\right)\tag{5} xvi=d(x,idsi)(5)
任意一部分的掩码 m i m_i mi由公式(6)确定:
其中 m s ( ⋅ ) ms(·) ms(⋅)根据 i d s i ids_i idsi从 t t t中获取掩码。
m i = m s ( t , i d s i ) (6) m_{i}=ms\left(t, ids_{i}\right)\tag{6} mi=ms(t,idsi)(6)
任意部分 x v i x_{v_i} xvi 包含 N / K N/K N/K 个可见补丁,其对应的互补视图 x m i x_{m_i} xmi 则有 N − N / K N - N/K N−N/K 个掩码补丁,定义见公式(7)。所以,任意部分的掩码率为 ( N − N / K ) / N = ( K − 1 ) / K (N - N/K)/N = (K - 1)/K (N−N/K)/N=(K−1)/K 。算法1里有更详细的内容。
x m i = x ⊙ m i (7) x_{m_{i}}=x \odot m_{i}\tag{7} xmi=x⊙mi(7)
从上述过程可以看出,图像补丁是按照无放回采样的原则,被划分为
K
K
K个不重叠部分的。
当
K
K
K设为 4 时,每个部分的掩码率为75%(和MAE的掩码率一样)。本文的设计从理论上保证了数据的全面利用,而且在有限采样次数下,就算只迭代一次,每个补丁参与训练的次数也是相同的。有放回采样策略无法做到这点,这也是本文的设计比MAE的纯随机掩码策略数据利用率更高的原因。这种设计被称作并行掩码策略,它在理论上既能保证掩码的随机性,又能确保输入完整且严格不重叠,比MAE掩码策略的纯随机方式更可靠、更稳定。
根据图2的设计,每个包含可见补丁的部分都会作为输入,被送入编码器 - 解码器架构,并执行MIM任务,损失函数定义如公式(8)所示。
L
w
h
o
l
e
(
x
)
=
E
i
∈
[
1
,
K
]
L
(
x
v
i
,
x
m
i
)
(8)
\mathcal{L}_{whole }(x)=\underset{i \in[1, K]}{\mathbb{E}} \mathcal{L}\left(x_{v_{i}}, x_{m_{i}}\right)\tag{8}
Lwhole(x)=i∈[1,K]EL(xvi,xmi)(8)
自一致性学习
第三节介绍的并行掩码策略确保了数据的全面利用,但无法保证各部分输出结果的可靠性。根据先驱性研究[70],人类智能是一个自洽系统,有助于高效学习和纠正错误。可以合理推断,将自洽机制引入人工模型,也能提升训练效率和一致性。对于MIM预训练模型而言,其生成结果不可靠,是因为模型在处理来自同一输入但不同的不完整信息时,对相同位置会生成不一致的预测。因此,将自洽原则引入MIM,能促使预训练模型在面对同一图像的不同可见补丁输入时,对相同位置的预测保持一致。
在第三节中,每个部分包含整幅图像的
N
/
K
N/K
N/K,并生成
N
−
N
/
K
N - N/K
N−N/K个预测。各部分的预测分别为
x
p
1
,
x
p
2
,
.
.
.
,
x
p
K
x_{p_{1}}, x_{p_{2}}, ..., x_{p_{K}}
xp1,xp2,...,xpK 。显然,任意两组部分的预测之间存在一定比例的重叠,重叠比例为
(
K
−
2
)
/
(
K
−
1
)
(K - 2)/(K - 1)
(K−2)/(K−1)。 任意两组预测
x
p
i
x_{p_{i}}
xpi和
x
p
j
x_{p_{j}}
xpj(其中
i
,
j
∈
{
1
,
2
,
.
.
.
,
K
}
i, j \in \{1, 2, ..., K\}
i,j∈{1,2,...,K}且
i
≠
j
i \neq j
i=j)的重叠位置
s
i
j
s_{ij}
sij,可以通过掩码
m
i
m_{i}
mi和
m
j
m_{j}
mj来确定,其定义如公式(9)所示:
s
i
j
=
m
i
∩
m
j
(9)
s_{ij}=m_{i} \cap m_{j}\tag{9}
sij=mi∩mj(9)
因此,提出自一致性学习,以引导每个重叠位置的预测保持一致。如图2所示,自一致性学习会拉近任意两组预测
x
p
i
x_{p_{i}}
xpi和
x
p
j
x_{p_{j}}
xpj之间的重叠部分,通过最小化重叠重建结果之间的平均绝对误差来提高一致性。 自一致性损失定义如公式(10)所示,其中
s
g
[
⋅
]
sg[·]
sg[⋅]表示停止梯度。对于任意部分的每个预测,它会与其他部分的预测进行
K
−
2
K - 2
K−2次计算,以确保自一致性。
L
s
c
(
x
v
i
,
x
v
j
)
=
E
(
∥
s
g
[
x
p
i
]
−
x
p
j
∥
+
∥
x
p
i
−
s
g
[
x
p
j
]
∥
)
⊙
s
i
j
(10)
\mathcal{L}_{sc}\left(x_{v_{i}}, x_{v_{j}}\right)=\mathbb{E}\left(\left\| sg\left[x_{p_{i}}\right]-x_{p_{j}}\right\| +\left\| x_{p_{i}}-sg\left[x_{p_{j}}\right]\right\| \right) \odot s_{ij}\tag{10}
Lsc(xvi,xvj)=E(
sg[xpi]−xpj
+
xpi−sg[xpj]
)⊙sij(10)
最后,整幅图像的自一致性损失根据公式(11)计算:
L
c
o
n
s
i
s
t
e
n
c
y
(
x
)
=
E
i
∈
[
1
,
K
]
,
j
∈
[
i
+
1
,
K
]
L
s
c
(
x
v
i
,
x
v
j
)
(11)
\mathcal{L}_{consistency }(x)=\underset{i \in[1, K], j \in[i + 1, K]}{\mathbb{E}} \mathcal{L}_{sc}\left(x_{v_{i}}, x_{v_{j}}\right)\tag{11}
Lconsistency(x)=i∈[1,K],j∈[i+1,K]ELsc(xvi,xvj)(11)
通过自一致性损失引导的结果,可在图4中观察到:不同组合重建出的图像最终彼此高度相似。
目标函数
本文的EMAE由并行掩码策略和自一致性学习两部分构成。在此,仍采用MAE的像素重建损失,以保留图像的生成能力。
因此,EMAE最终的损失函数可表示为公式(12),并且为了等权重计算,将每个损失项的系数都设为1。
L
t
o
t
a
l
(
x
)
=
L
w
h
o
l
e
(
x
)
+
L
c
o
n
s
i
s
t
e
n
c
y
(
x
)
(12)
\mathcal{L}_{total }(x)=\mathcal{L}_{whole }(x)+\mathcal{L}_{consistency }(x)\tag{12}
Ltotal(x)=Lwhole(x)+Lconsistency(x)(12)
实验
数据集与评估指标
预训练数据集
在预训练阶段,选择了三个常用的视觉预训练数据集:ImageNet-1K、COCO和OpenImages ,以研究本文所提方法的通用性和迁移能力。首先,ImageNet-1K数据集是一个大规模的图像分类数据集,其训练集包含约128万张训练图像,共1000个类别。本文在ImageNet的训练集中进行了大量实验。此外,COCO数据集是一个非标志性图像数据集,其中的图像包含自然且多样的场景,具有场景中心的特点。其train2017集包含约11.8万张图像,本文在该集中进行实验。最后,OpenImages数据集同样是非标志性的,与COCO数据集相比,它包含更多的数据和类别。其训练集有大约170万张图像,本文也在该集中进行了实验。ImageNet是默认数据集。
评估数据集
为了验证预训练模型的质量,在ImageNet-1K数据集上通过线性探测和微调分类任务对EMAE进行评估,ImageNet-1K是一个常用的图像分类数据集。此外,对于计算机视觉而言,目标检测和分割是密集预测任务,因为这些任务的训练图像中每张都包含大量对象。对密集预测任务的验证可以更好地反映预训练模型的语义理解能力。因此,本文在COCO和ADE20k数据集上进行了大量实验,以验证EMAE的泛化能力和迁移能力。COCO是一个常用的目标检测和实例分割数据集,其train2017集包含约11.8万张图像,在val2017集(包含5000张图像)上进行验证。该具有挑战性的数据集包含80个类别的人工标注。此外,ADE20K也是一个具有挑战性的语义分割数据集,它包含约2万张训练图像和2000张验证图像,共150个类别。
评估指标
遵循自监督学习的常见做法,本文主要使用top-1准确率来评估预训练模型在线性探测和微调图像分类任务中的语义理解能力。同时,为了进一步验证密集视觉任务,本文采用边界框平均精度均值(box AP)和掩码平均精度均值(mask AP)来验证目标检测和实例分割任务中迁移学习的性能。最后,本文采用平均交并比(mean intersection of union)来验证语义分割任务的迁移能力,语义分割也是一个常见的密集视觉任务。
实验设置
预训练设置
训练设置与MAE相同,本文采用编码器 - 解码器结构来执行MIM任务。本文的方法适用于各种ViT骨干网络,由于计算资源的限制,大多数实验是在ViT-Base上进行的。具体来说,将224×224的图像划分为14×14个补丁,补丁大小为16×16。默认情况下,K设置为 4。批量大小设置为 4096。 同时,AdamW优化器的权重衰减、 β 1 \beta_{1} β1和 β 2 \beta_{2} β2分别设置为0.05、0.9和0.95。使用带有40个epoch热身的余弦学习率策略,基础学习率设置为 b a s e l r = 1.5 e − 4 base_lr =1.5 e^{-4} baselr=1.5e−4。不同数据集的预训练设置差异仅体现在训练轮数上:ImageNet预训练800个epoch,根据R-MAE,COCO数据集训练4000个epoch,OpenImages数据集按照长序列MAE的方法预训练800个epoch。
线性探测设置
按照常见做法,通过线性探测来评估特征表示的质量。在自监督预训练之后,移除解码器,并在冻结的特征上训练一个监督线性分类器。所有训练设置与MAE相同。
端到端微调设置
根据MAE设置
目标检测和实例分割设置
在基准测试设置中,使用相同的简单公式训练所有模型:大规模抖动(1024×1024分辨率,尺度范围[0.1, 2.0]),AdamW( β 1 \beta_{1} β1 β 2 = 0.9 , 0.999 \beta_{2}=0.9,0.999 β2=0.9,0.999),采用半周期余弦学习率衰减,线性热身0.25个epoch,以及随机深度正则化,例如ViT-Base为0.1)。所有模型训练100个epoch,批量大小设置为64。学习率为(8 e-5),权重衰减为0.1。更多细节在[37]中描述,本文严格遵循MAE的超参数。对于ViTDet的设置,输入大小为1024×1024,在训练过程中通过大规模抖动进行增强,尺度范围为[0.1, 2.0]。本文使用AdamW( β 1 \beta_{1} β1 , β 2 = 0.9 , 0.999 \beta_{2}=0.9,0.999 β2=0.9,0.999)并采用逐步学习率衰减。使用批量大小为64的线性学习率热身,共250次迭代。更多细节在[40]中描述,本文完全按照其设置进行实验。
语义分割设置
在实验中,采用UperNet框架。本文使用AdamW作为优化器。批量大小为16,层衰减率为0.65。输入分辨率为512×512。本文进行100个epoch的端到端微调。实际上,所有实验设置都与MAE一致,mmsegmentation也开源了MAE的分割代码,以帮助研究人员快速验证。
数据增强设置
根据先前工作,MAE可用作数据增强器。为了验证MAE和EMAE在数据增强方面的性能,本文使用预训练800个epoch的EMAE和预训练1600个epoch的MAE作为图像分类的数据增强器。 该实验以ResNet50为骨干网络,在ImageNet数据集上进行90个epoch的训练,并报告验证集上的 top-1 准确率。为了更好地展示MAE和EMAE的生成质量,每张图像都采用随机掩码设置,而不是注意力引导的掩码策略。其他超参数遵循[34, 84]中的训练方法。
ImageNet-1K上的图像分类
评估设置
为了进行公平比较,在图像分类实验中,本文严格遵循MAE的超参数来评估在ImageNet上预训练的模型。本文在两种监督训练设置下评估预训练编码器的性能:
1)线性探测(LP);
2)端到端微调(FT)。
在线性探测中,预训练编码器的所有参数都被冻结,仅训练最后一个分类层。
在端到端微调中,预训练编码器与分类头一起进行微调。
分类结果
【表 2】最先进的自监督学习方法对比。EMAE 在 ImageNet 训练集上进行预训练,相比之前的掩码图像建模方法,达到了最先进的性能。为进行评估,在两种监督训练设置下测试预训练模型的性能:1)线性探测;2)端到端微调。报告在 ImageNet 验证集上的 Top-1 准确率。
如表2所示,在相同的预训练轮数(300和800个epoch)下,使用ViT-Base的本文所提出的方法在线性探测中比MAE高出约6.0% - 6.7%。 300个epoch的分类结果与2400个epoch的MAE相当,而使用EMAE达到这一性能水平仅需MAE 13%的时间和约50%的计算量。 这表明本文的方法可以显著提高MIM的训练效率。额外的训练(800个epoch)进一步将线性探测结果提升到70.4%,微调结果提升到84.0%,并达到了最先进的性能。此外,值得注意的是,本文使用ViT-Large进行800个epoch训练的分类结果与使用ViT-Huge进行1600个epoch训练的MAE相当。作为一种基于MIM的方法,EMAE超过了先前基于MIM的最先进方法。由于基于对比学习的方法具有图像语义一致性假设,该假设与线性探测任务的先验一致,因此EMAE的线性探测结果略逊于基于对比学习的方法。在接下来的部分中,将通过在不同数据集和任务上进行更多实验,全面评估EMAE的迁移能力。
目标检测和实例分割
为了进一步验证在ImageNet上预训练的EMAE所学习到的视觉表示,本文在COCO目标检测和实例分割任务上进行微调。选择Mask R-CNN框架。具体来说,采用特征金字塔网络(FPNs),如[37]中所述,将特征图缩放到不同大小。本文完全遵循先前[21, 37]的策略,在COCO上进行这些实验。结果以检测的边界框平均精度均值( A P b AP^{b} APb)和分割的掩码平均精度均值( A P m AP^{m} APm) 报告在表3中。该表展示了不同自监督方法和监督训练所学习到的表示的性能。可以观察到,300个epoch的EMAE的结果优于1600个epoch的MAE。此外,800个epoch的EMAE取得了最佳结果, A P b AP^{b} APb为51.4%, A P m AP^{m} APm为45.7%,分别比1600个epoch的MAE高出1.0和0.8个百分点。
【表3】 使用 Mask R-CNN [21, 37]在 COCO 数据集上进行目标检测和实例分割的结果。各种方法的架构均采用 ViT-B。本文采用带有特征金字塔网络(FPN)[85]的 Mask R-CNN[36],并报告在 COCO val2017 数据集上的边界框平均精度均值(bounding box AP)和掩码平均精度均值(mask AP)。EMAE的性能优于先前最先进的自监督学习方法。†:这些结果来自[71]。
此外,本文还在基于ViT的最先进检测框架ViTDet(在Detectron2中实现)上进行实验,以验证MAE的迁移能力。为了进行公平比较,所有这些实验都严格采用ViTDet的训练设置。在表4中可以观察到,300个epoch的EMAE在目标检测和实例分割的各种指标上均优于1600个epoch的MAE。 同时,800个epoch的EMAE在
A
P
b
AP^{b}
APb和
A
P
m
AP^{m}
APm上分别比1600个epoch的MAE高出1.3和1.2个百分点。此外,使用级联掩码R-CNN的ViTDet的结果达到了54.8%的
A
P
b
AP^{b}
APb和47.6%的
A
P
m
AP^{m}
APm,分别比 MAE 高出 0.8 和 0.9 个百分点。
【表4】 使用ViTDet[40]在COCO数据集上微调目标检测和实例分割的结果。为进行公平比较,各种方法的架构均采用ViT-B。
本文采用ViTDet[40]作为检测框架,本文的EMAE取得了令人瞩目的性能,超越了之前最先进的自监督学习方法MAE。
此外,进一步在ViT-Large架构下使用ViTDet进行目标检测和实例分割实验。结果如表5所示。800个epoch的EMAE的结果在大型网络结构中显示出持续的改进,无论是在加载有/无级联掩码R-CNN的ViTDet中。
【表5】 使用ViTDet[40]在COCO数据集上微调目标检测和实例分割的结果。为确保公平对比,各类方法的架构均采用ViT-L。此次采用ViTDet[40]作为检测框架,本文提出的的EMAE表现卓越,性能超越了此前最先进的自监督学习方法MAE。
语义分割
本文还在另一个密集预测任务——ADE20K数据集上的语义分割任务中评估在ImageNet上预训练的EMAE。以所有语义类别的平均交并比(mIoU)作为评估指标。 具体来说,在实验中完全遵循MAE的训练设置,采用UperNet框架,并将结果报告在表6中。将本文的方法与在ImageNet-1K上的监督预训练以及最先进的自监督方法进行比较。 可以观察到,当EMAE预训练300个epoch时,其性能优于1600个epoch的MAE。此外,800个epoch的EMAE达到了最高的49.3%的mIoU,比所有其他基线方法都具有更优异的性能,进一步验证了本文提出的框架的有效性。
【表6】 使用 UperNet[81] 在 ADE20K 数据集上进行语义分割的结果。为了进行公平比较,各种方法的架构均采用ViT-B。本文报告以平均交并比(mIoU)衡量的结果 ,并且EMAE超越了先前的自监督方法,达到了最先进的性能。†表示语义分割的实验设置遵循先前的研究[71]。
此外,注意到当前基于MAE的方法,如 G e 2 Ge^2 Ge2- AE、CAE、MVP、MILAN,在下游任务中采用不同的超参数以获得更好的性能。上述实验严格遵循 MAE 的设置,结果已经证明了 EMAE 的有效性。在这里,为了进行公平比较,对预训练模型遵循 CAE 在语义分割上的设置,在 ADE20k 数据集上使用ViT-Base的结果为50.8%的mIoU,达到了最先进的性能。
数据增强器
为了展示EMAE的生成质量,在ResNet图像分类实验的数据增强过程中采用两种不同的随机掩码率:40%和75%。 最终实验结果如表7所示。在40%的随机掩码率下,MAE和EMAE都可以提高ResNet的识别准确率,这表明掩码模型作为数据增强器的可行性。然而,在75%的设置下,使用MAE生成的数据训练的ResNet性能低于基线,而使用EMAE生成的数据训练的ResNet性能仍然有所提升。这表明EMAE能够生成比MAE更可靠、更稳定的结果,确实解决了高掩码率带来的生成不一致问题。
【表7】 数据增强结果。报告ImageNet-1K的分类结果。各种方法的架构均采用ViT-L。
在COCO数据集上的预训练
为了展示EMAE的适应性,在COCO数据集上进行预训练实验。根据R-MAE在基准任务(目标检测、实例分割和语义分割)上评估预训练模型的性能。EMAE主要分别训练1000、2000和4000个epoch。为了进行公平比较,根据R-MAE采用ViTDet作为目标检测和实例分割的框架,不同方法的结果如表8所示。可以观察到,1000个epoch的EMAE的性能与8000个epoch的MAE非常接近,展示了本文提出的算法的效率。 此外,2000个epoch的EMAE的结果与8000个epoch的R-MAE相当,而4000个epoch的EMAE达到了最先进的性能。所有这些发现表明,EMAE可以获得非标志性数据的有效表示。
【表8】 COCO数据集预训练结果。本文报告COCO目标检测、实例分割以及ADE20K语义分割的结果。各种方法的架构均采用ViT-B。
在OpenImages数据集上的预训练
为了进一步展示EMAE的适应性,本文在OpenImages数据集上进行预训练实验。本文根据长序列MAE评估目标检测、实例分割和语义分割任务的性能。各种下游任务的实验设置遵循默认设置。长序列MAE是一种通过将图像分辨率从224提升到448来预训练MAE的方法。为了进行公平比较,在表9中未引入长序列MAE的结果 。与COCO数据集相比,OpenImages包含更多的数据和类别。这表明OpenImages是一个具有更复杂场景的密集对象数据集。在高掩码率的场景中,大量对象可能会被完全掩码,使得掩码重建任务更具挑战性。尽管存在这样的挑战,但根据表9,100个epoch的EMAE的结果仍然优于800个epoch的MAE。这突出了EMAE在对密集对象数据进行稳定、可靠和高效训练方面的有效性。此外,800个epoch的EMAE在COCO数据集上的性能达到 50.8 % A P b 50.8 \% AP^{b} 50.8%APb和 45.3 % A P m 45.3 \% AP^{m} 45.3%APm ,在ADE20k数据集上的 mIoU 为49.3%。在各种数据集上预训练EMAE的结果表明,它对不同场景的训练数据具有鲁棒性。
【表9】 OpenImages数据集预训练结果。报告COCO目标检测、实例分割以及ADE20K语义分割的结果。各种方法的架构均采用ViT-B。
消融实验
为了更好地研究本文提出的EMAE中不同组件的有效性,在ImageNet-1K数据集上进行消融实验。线性探测仍然是一种很好的评估方法,可以快速验证学习到的表示,MAE也根据线性探测的性能选择掩码率。因此,在消融实验中,本文采用线性探测的结果作为衡量有效性的基准。为了进行公平比较,各种方法的架构均采用ViT-Base。
整体数据利用的效果
在第三节中,讨论了所提出的并行掩码策略的优越性。如表10所示,展示了四种数据利用的设计。
- 表10(a)展示了200个epoch的MAE在每次迭代中从训练图像中进行1次纯随机掩码的结果,将其视为基线。
- 表10(b)表示200个epoch的MAE在每次迭代中根据MAE的纯随机掩码策略从训练图像中进行4次纯随机掩码的结果。
- 表10(c)展示了200个epoch的MAE在每次迭代中通过并行掩码策略从训练图像中进行4次不相交随机掩码的结果。
与表10(a)相比,表10(b)的结果表明多次采样可以在一定程度上提高数据利用率,并增强预训练模型的性能。 然而,纯随机掩码策略无法保证整体数据的利用(100%的数据利用率),并且在4次掩码过程中每个补丁的采样次数不同。值得注意的是,表10(c)的并行掩码策略在理论上可以解决这些问题,(c)的实验结果也证明了这一点:在相同的epoch数下,并行掩码策略比纯随机掩码策略性能高出0.9%。因此,本文的策略可以进一步提高数据利用率,实现100%的数据利用率,从而提高预训练效率。本文的策略的有效性在理论和实验上都得到了验证。在这里,本文设计了一个额外的实验,称为互补掩码策略,该策略对训练图像进行两次采样。第一次从训练图像中随机采样25%的数据,而第二次采样利用训练图像中剩余的75%的数据。 这种方法也允许模型利用整个数据,并且由于Transformer的计算特性,其计算复杂度甚至高于并行掩码策略。然而,互补掩码策略最终表现出比基线更差的性能。这是因为75%的部分代表了预训练过程中25%的掩码率,根据先前的工作,低掩码率会导致较差的迁移能力。 因此,采用并行掩码策略作为整体数据利用的设计。
【表10】 EMAE的消融实验:整体数据利用的效果。报告在ImageNet上进行线性评估的结果。
自一致性学习的效果
为了进一步提升一致性,本文提出了自一致性学习方法,以此鼓励模型在预训练过程中生成可靠的特征表示。在整体数据利用设计的基础上,预训练阶段引入了自一致性损失。如表11所示,
- 表11(a)展示了基于MAE的并行掩码策略(即像素重建任务)的结果。
- 表11(b)展示了基于并行掩码策略、无像素重建的纯自一致性学习结果。
该结果表明,没有像素级重建的纯自一致性模型仍能学习到出色的语义特征,其性能甚至超过了像素重建模型。表11(c)展示了结合像素重建的自一致性学习结果,其性能略优于纯自一致性学习。 由于纯自一致性学习丧失了生成和重建能力,无法扩展到图像合成任务(如参考文献[61, 64]所述),因此,仍保留像素重建损失,并采用设计(c)作为最终方案 。值得注意的是,在模型中引入自一致性学习后,从图4中可以观察到,本文提出的EMAE生成的重建图像最终彼此高度相似。综上所述,这说明了自一致性机制的有效性和优越性,它进一步提升了MIM的性能和效率。
表11 EMAE的消融实验:自一致性学习的效果。报告在ImageNet上进行线性评估的结果。
图4. 本文EMAE在不同掩码种子下的重建结果各异。通过随机种子从同一图像中采样得到可见补丁的不同组合,然后将这些组合输入到EMAE中,EMAE会生成重建图像。在此,以重建三张图像为例。这些重建结果包含相似的语义信息,并且彼此高度匹配,这证明了本文所提出的自一致性学习方法的有效性。
所提方法在不同预训练轮数下的效果
如表12所示,
- 第二行展示了MAE在ImageNet数据集上进行不同预训练轮数(例如100、200、300和800轮)的结果作为基线。
- 第三行表明,在不同预训练轮数下,将整体数据划分为4个不重叠部分。
- 表12第四行列出了自一致性学习的结果。
与基线相比,第三行的结果在相同预训练轮数下比MAE高出约4.0% - 5.2%,这表明充足的训练数据可以在不同预训练轮数下有效提升性能。这些结果也说明了充分利用整体数据的重要性,它可以有效提高数据利用率并提升预训练效率。在整体数据利用设计的基础上,表12第四行列出了自一致性学习的结果。 与整体数据利用设计相比,第四行的结果进一步提升了预训练模型的性能,即在相同预训练轮数下,第四行的结果比第三行高出约0.9% - 2.9%。因此,在相同预训练轮数下,EMAE比MAE的性能高出约6.0% - 6.7%。
【表12】 EMAE的消融实验:所提出的方法在不同预训练轮次下的效果。“+”表示在前面一行的基础上添加当前方法。报告在ImageNet上进行线性评估的结果。
更长预训练轮数的效果
为了进一步探究EMAE在更长预训练轮数下的性能,使用ViT-B在ImageNet-1K上对EMAE进行1600轮预训练。然后,将预训练模型应用于线性探测和端到端微调分类任务。 最终,其在线性探测中的结果为71.2%,微调结果为84.2%。这表明EMAE在更长的预训练轮数下仍能获得性能提升。
表13 EMAE的消融实验:更长预训练轮次的效果。报告在ImageNet上进行线性评估的结果。
K值划分的影响
根据第三节内容,K值直接决定了每个部分的掩码率,掩码率为 K − 1 K \frac{K - 1}{K} KK−1 。 从表14可以看出,在相同的预训练设置下,本文的方法在不同的高掩码率(例如0.9、 3 4 \frac{3}{4} 43 、 6 7 \frac{6}{7} 76 和 13 14 \frac{13}{14} 1413 掩码率)下始终优于MAE。同时,当K设置为4(掩码率为75%)时,本文提出的方法性能最佳,这一现象也与MAE对掩码率的观察结果相符。
【表14】 EMAE的消融实验:K值划分的影响。报告在ImageNet上进行线性评估的结果。
使用EMAE的图像生成质量
图4. 本文EMAE在不同掩码种子下的重建结果各异。通过随机种子从同一图像中采样得到可见补丁的不同组合,然后将这些组合输入到EMAE中,EMAE会生成重建图像。在此,以重建三张图像为例。这些重建结果包含相似的语义信息,并且彼此高度匹配,这证明了本文所提出的自一致性学习方法的有效性。
如图4所示,EMAE的重建结果在语义上与原始图像一致,证明了重建的可靠性。 为了进一步证实EMAE的可靠性,本文进行了实验,将MAE和EMAE分别生成的图像输入到分类器中,比较它们的分类性能。
具体来说:
- 在ImageNet验证集上使用掩码率为75%的随机掩码,分别用MAE和EMAE重建图像。
- 然后将这些图像输入到最先进的图像分类器DINO v2(搭配ViT-g)中,执行线性探测任务。
表 15 EMAE 的消融实验:生成结果的可靠性。报告在 ImageNet 上进行线性评估的结果。
分类结果如表15所示。与原始图像在DINO v2上的分类结果相比,MAE重建的图像在分类任务中表现出明显的性能下降,而EMAE重建的图像在分类任务中明显优于MAE,达到了与原始图像相当的性能。这表明EMAE缓解了不可靠和不一致的问题。
此外,进一步验证了EMAE的图像生成质量,最终结果如表16所示。在这里,结构相似性指数(SSIM)、感知损失(LPIPS)是衡量生成图像与原始图像相似性的标准指标,而弗雷歇距离(FID)是衡量生成图像的保真度和多样性的标准指标。 在这三个指标上,EMAE均表现出优于MAE的性能。然而,与专门的图像生成模型相比,仍然存在显著的性能差距。同时,MAE的扩展方法,如MAGE和RCG,已经在图像生成方面取得了最先进的性能。值得注意的是,本文的方法与这些方法是正交的。
表16 EMAE的消融实验:图像生成标准。报告结构相似性指数(SSIM)[93]、感知损失(LPIPS)[94]和弗雷歇距离(FID)[95]的结果。各种方法的架构均采用ViT-L。