ViT和MAE模型结合初探

介绍

transfomer在NLP领域的应用已经非常广泛，但是在CV领域的应用还比较少，主要是因为CV领域的数据是二维的，而transfomer是基于序列的，因此需要将二维数据转换成序列数据，这就是ViT的由来。ViT是将图像分割成一个个patch，然后将patch展开成序列，再输入到transformer中，这样就可以将transformer应用到CV领域了。然而transformer相比于CNN缺少一个平移不变性，因此Transformer需要在模型中学习到这个归纳偏差（inductive bias），因此需要更多的数据来训练Transformer模型，来达到相同的性能。在ViT论文中，作者在使用相同的ImageNet-1k的数据集时，性能相比于当时的baseline模型还要差1~2百分点，当在更大的数据集上ImageNet-100k上训练时，ViT的性能已经能够媲美CNN的SOTA了，当在谷歌的大规模数据集JFT-300M上训练时，ViT的性能已经超过了CNN的SOTA了。因此，ViT的性能与数据集的大小是正相关的，这也是ViT的一个缺点，因为ViT的训练需要更多的数据，而且训练时间也会更长。

具有标记的数据往往相比于无标记的数据集更难以获取，并且获取的成本也要远远高于无标记数据集。因此，基于无标记数据的无监督学习成为了目前备受关注的研究方向。目前无监督学习方向主要分为两大类：对比学习和掩码学习。

• 对比学习通过设置正样本和负样本，正负样本通过模型得到各自的特征，设定loss使得特征的相似度越低越好；这使得模型能够通过关注正负样本的差别来学习到深层的表示，不同领域的对比学习设定正负样本的规则也不同，在分类任务中可以将和当前样本一致的数据视为正样本，不一致的视为负样本。
• 掩码学习通过将输入数据的一部分掩盖，然后让模型去预测掩盖的部分，这样模型就能够学习到输入数据的一些特征；掩码学习不经用于CV，也用于NLP等众多领域，如BERT、GPT等。在CV领域目前比较流行的掩码学习模型为MAE（Masked AutoEncoder）、SIMM。

无监督学习能够通过大量无标记的数据学习到数据的深层表示，而同时ViT需要大量的数据。一个很自然的想法是将ViT和MAE等无监督学习的模型结合起来，通过无监督学习得到的backbone来指导ViT的训练，从而减少ViT的训练数据，并且帮助ViT学习到更加深层的表示。

主要方法

1. 直接将MAE的训练得到的模型作为ViT的backbone，然后在ViT的训练中进行微调。从头训练一个MAE显然对于我来说是不显示的，MEA的模型可以从网上下载，由于MAE训练的模型已经取得了很好的效果，显然它已经学习到了诸如平移不变性等CNN和图像等众多的性质，因此直接在MAE得到的模型上做ViT的微调，获取可以继承MAE的性质，同时又能够学习到特定领域的特征。
2. 使用知识蒸馏的方法，将MAE的模型作为teacher模型，ViT作为student模型，通过知识蒸馏的方法来训练ViT。知识蒸馏的方法是将teacher模型的输出作为student模型的标签，通过最小化student模型的输出和teacher模型的输出的差异来训练student模型。知识蒸馏的方法可以帮助student模型学习到teacher模型的知识，从而达到和teacher模型相近的性能。和ViT一同发表的另一篇DeiT模型描述了使用知识蒸馏的方法，使用CNN作为teacher模型，transformer作为stduent模型，可以训练一个性能相比于CNN更好的模型。这说明CNN作为teacher模型，它可以传送关于一些平移不变性的归纳偏差给transformer模型，从而帮助transformer模型学习到更好的表示。因此，我们可以使用MAE作为teacher模型，ViT作为student模型，通过知识蒸馏的方法来训练ViT，从而帮助ViT学习到更好的表示。因为，我们也认为MAE已经学习到了一些关于平移不变性的归纳偏差，因此可以帮助ViT学习到更好的表示。

实验设计

由于算力的原因，我采取了一些折衷的处理方式：

• 无法从头训练一个MAE模型，因此我直接从transformer的社区获得了一个MAE模型，这个模型是在ImageNet-1k上训练的，所以为了避免我训练的数据已经在Image-1k上，我使用了人脸作为训练数据，ImageNet-1k中并不存在大量的人脸数据.
• 在知识蒸馏中，无法从头开始进行蒸馏，学习到一个新的模型，因此我只观察了一段时间的蒸馏过程，主要是观察蒸馏过程中的loss变化。

对于方法1，测试了多种不同的微调和迁移方式，例如只迁移Patch_embedding, + transformer_block.只对FC进行微调，只对Block进行微调，和全部一起训练。

对于方法2，观察相同training配置下的loss的变化情况，和ViT训练的loss情况进行对比。

为了节省显存，加快训练速度，采用了目前较为流行的fp16的压缩模型的方式，即将模型的参数转换成fp16的数据类型，这样可以减少显存的占用，加快训练速度。但是，这样也会带来一些问题，例如在使用fp16的情况下，模型的参数的梯度会变得非常小，因此需要对梯度进行放大，这样会导致梯度爆炸的问题，因此需要对梯度进行裁剪，这样会导致模型的收敛速度变慢。因此，我在训练过程中，对梯度进行了裁剪，裁剪的阈值为1.0。

实验结论

• 通过掩码学习得到的ViT模型是可以作为ViT的预训练模型，并且，只移植PatchEmbedding相比于只移植TransformerBlock，效果更好。这个原因可能是因为PatchEmbedding是是确定位置编码的，在ViT中使用的是可学习的PositionEmbedding，在ViT论文中，作者画出的Patch Embedding图像也能够索命PM层包含了模型的位置编码信息。

• 通过知识蒸馏的方法，ViT的loss下降的更快，这表明MAE确实将一些关于平移不变性的归纳偏差传递给了ViT，从而帮助ViT学习到更好的表示。为了进一步说明通过MAE进行知识蒸馏能够使得ViT更快理解平移不变性，将PM层的参数通过ViT的方法展示出来，可以证明ViT确实理解了平移不变性，PM层更类似最终的ViT的PM层。我们并没有完整的通过类似于DeiT一样训练一个ViT，知识观察了Loss曲线，并取出了训练过程中的部分层参数进行分析，所以对于最终的结果的性能还需要进一步的分析。

附录

• 本实验采用了ViT-T的模型结构；详细的结构可以参看ViT论文。

• 知识蒸馏：采用了类似DeiT的令牌来进行知识的传递，相比于传统的通过最小化模型输出的特征方式，这种方式能够在模型的每一层都进行知识传递，从而更好的帮助模型学习到更好的表示，而不仅仅是在最后一层通过梯度反向传播。

• MAE的模型结构：采用了MAE的模型结构，详细的结构可以参看MAE论文。