预训练模型Uniform Masking | 基于金字塔Transformer的预训练模型

[189****6672]

预训练模型Uniform Masking | 基于金字塔Transformer的预训练模型

论文标题：Uniform Masking: Enabling MAE Pre-training for Pyramid-based Vision Transformers with Locality

论文地址：https://arxiv.org/pdf/2205.10063.pdf

论文代码：https://github.com/implus/UM-MAE

摘要

1. 背景介绍：MAE通过一种优雅的非对称编码器-解码器设计引领了视觉自监督领域的发展趋势，它显著优化了预训练效率和微调精度。值得注意的是，非对称结构的成功依赖于ViT的“全局”特性。

2. 现存问题： 然而，目前尚不清楚如何在MAE预训练中采用基于金字塔的高级ViT（例如PVT、SWN），因为它们通常会在局部窗口中进行注意力操作，从而难以处理部分可见视觉tokens的随机序列。

3. 解决方法： 论文提出了Uniform Masking（UM）策略，成功地为基于金字塔的具有局部性的VIT（简称“UM-MAE”）实现了MAE预训练。具体而言，UM包括一个均匀采样（US），该采样严格地从每个2×2网格中采样1个随机patch，以及一个二次掩蔽（SM），该掩蔽将已采样区域的一部分（通常为25%）再一次进行随机掩码。US在多个非重叠的本地窗口中保留了相同数量tokens，从而顺利支持流行的基于金字塔的VIT。

4. 实验结果：UM-MAE显著提高了预训练的效率（例如∼2倍，并将GPU内存减少至少∼2倍），但在下游任务中保持有竞争力（甚至更好）的微调性能。

算法

Uniform Masking（UM）是一个用于MAE的金字塔Transformer模型。

UM是一种简单的两阶段策略，它将密集图像tokens转换为稀疏图像tokens，但在空间上保持其均匀分布。

它包括**均匀采样Uniform Sampling（US）和二次掩码（SM）**两个阶段，共同保证了自监督学习表征的效率和质量。

均匀采样（US）

如图1所示，UM首先执行均匀采样（US）

均匀采样步骤如下：

1. 先将图像分割成多个2×2大小的patch
2. 在每个2×2patch中随机采样一个patch进行保留，其他进行掩码操作

均匀采样能够保证每个局部窗口中具有相同的可见tokens数量，这种性质使其能够更好的使用于基于金字塔结构的Transformer变体。

二次掩码SM

与MAE中随机掩码采样相比，US这种均匀采样的方式使得在重建过程中，能够利用这种均匀分布的地位图像特征来统计图像信息进行重建，而不是通过高维表征信息，这潜在的降低了自监督的难度和性能，这在MAE已经得到证明。

在MAE中，一种称为“网格化抽样”的类似抽样策略在很大程度上阻碍了表征学习质量。还可以在表1中观察到，均匀分布确实降低了预训练的难度：预训练损失从0.4256（MAE基线）降低到0.3858，但妨碍了下游任务学习表示的质量：语义分割任务中的微调精度下降了1个mIoU点。

为了解决以上问题，论文提出了二次掩码SM策略。

在均匀采样后的可见tokens中进行第二次掩码操作，如下图所示。

与均匀采样中的直接将掩码tokens直接丢弃不同，SM使用shared mask tokens来替代掩码tokens，这是为了基于金字塔的Transformer模型结构需要。

SM通过简单的操作就增加了自监督学习难度，保证了模型通过学习高维语义信息进行重建，而不是通过依赖相邻低维度像素图像信息来重建。

UM-MAE Pipeline with Pyramid-based ViT

上图就是整个不对称模型结构，命名为UM-MAE，主要组成部分如下所述：

编码器：

与MAE一致，每个patch大小为16×16，操作步骤如下：

1. 经过均匀采样和二次掩码获得模型输入
2. 通过金字塔Transformer模型结构，得到高维度特征输出（具有较小分辨率）
3. 使用Pytorch中的PixelShuffle来获取原始分辨率

解码器：

解码器的输入为所有tokens，包括（i）来自UM的编码tokens和（ii）掩码tokens

所有令牌都添加了位置嵌入

解码器的体系结构参考了MAE中采用的轻量级ViT。

重建目标：

通过预测均匀采样过程中的掩码patch来实现重建操作。

损失函数为均方误差（MSE），仅对掩码的patch进行计算。

实验

消融实验

可视化