[论文解读] Masked autoencoders are effective solution to transformer data-hungry

Jiawei Mao, Honggu Zhou+2 more|arXiv (Cornell University)|Dec 12, 2022

Advanced Neural Network Applications被引用 8

一句话总结

本文提出小样本掩码自编码器（SDMAE），一种通过简化MAE解码器、引入位置预测任务和类别标记感知对比学习任务来提升视觉Transformer（ViTs）在小样本场景下性能的方法。SDMAE在Tiny-ImageNet、CIFAR-100和医学影像基准等小样本数据集上实现了最先进性能，通过减少过拟合并提升局部不变性和表征学习能力，且无需修改ViT架构。

ABSTRACT

Vision Transformers (ViTs) outperforms convolutional neural networks (CNNs) in several vision tasks with its global modeling capabilities. However, ViT lacks the inductive bias inherent to convolution making it require a large amount of data for training. This results in ViT not performing as well as CNNs on small datasets like medicine and science. We experimentally found that masked autoencoders (MAE) can make the transformer focus more on the image itself, thus alleviating the data-hungry issue of ViT to some extent. Yet the current MAE model is too complex resulting in over-fitting problems on small datasets. This leads to a gap between MAEs trained on small datasets and advanced CNNs models still. Therefore, we investigated how to reduce the decoder complexity in MAE and found a more suitable architectural configuration for it with small datasets. Besides, we additionally designed a location prediction task and a contrastive learning task to introduce localization and invariance characteristics for MAE. Our contrastive learning task not only enables the model to learn high-level visual information but also allows the training of MAE's class token. This is something that most MAE improvement efforts do not consider. Extensive experiments have shown that our method shows state-of-the-art performance on standard small datasets as well as medical datasets with few samples compared to the current popular masked image modeling (MIM) and vision transformers for small datasets.The code and models are available at https://github.com/Talented-Q/SDMAE.

研究动机与目标

• 解决视觉Transformer（ViTs）在小样本数据集（如医学和科学图像）上对数据量的高需求问题。
• 探究掩码自编码器（MAE）是否能在不修改ViT架构的前提下有效降低ViT对数据的依赖性。
• 识别MAE在小样本数据集上训练时防止过拟合的最优、轻量级解码器配置。
• 通过自监督任务引入定位性和不变性等归纳偏置到MAE中。
• 通过对比学习任务实现在MAE预训练过程中对ViT类别标记的有效预训练。

提出的方法

• 提出小样本掩码自编码器（SDMAE），一种专为标准ViT模型在小样本数据上训练而设计的改进型MAE框架。
• 通过广泛的消融实验，识别出一种最小解码器配置（深度=1，d_model=128），可有效避免在小样本数据集上的过拟合。
• 引入一个轻量级的位置预测头，可预测被掩码标记的位置，参数量可忽略不计且无需人工标注。
• 设计一种对比学习任务，以增强特征不变性，并实现对ViT中类别标记的专用预训练。
• 在统一的预训练目标中结合掩码图像建模、位置预测和对比学习。
• 在所有实验中采用一致的预训练和微调协议，以确保公平比较。

实验结果

研究问题

• RQ1 MAE是否能在不修改ViT架构的前提下有效适配小样本数据集？
• RQ2 在小样本数据集上训练时，MAE的最优解码器配置是什么，以防止过拟合？
• RQ3 在小样本设置下，轻量级位置预测任务在恢复ViT的局部归纳偏置方面有多有效？
• RQ4 对比学习任务是否能改善表征学习并实现基于MAE的ViT中类别标记的预训练？
• RQ5 将掩码图像建模与额外的自监督任务相结合，是否能在小样本和医学数据集上实现最先进性能？

主要发现

• SDMAE在Tiny-ImageNet上达到72.24%的top-1准确率，优于标准MAE以及现有的ViT和CNN基线模型，在小样本数据集上表现更优。
• 消融实验表明，同时结合位置预测和对比学习可获得最高性能，当移除位置预测任务时，准确率下降1.3%。
• 最优解码器配置（深度=1，d_model=128）显著减少了过拟合，并提升了小样本数据集上的泛化能力。
• SDMAE在医学影像任务（如APTO 2019）上实现了最先进性能，优于先进的CNN模型和其他MIM方法。
• 医学数据上的注意力可视化结果表明，SDMAE能够聚焦于临床上相关的区域，体现了其实际应用价值。
• 对比学习任务实现了对类别标记的有效预训练，而该任务在标准MAE方法中常被忽略。