[论文解读] Pre-training of Lightweight Vision Transformers on Small Datasets with Minimally Scaled Images

Jen Hong Tan|arXiv (Cornell University)|Feb 6, 2024

Advanced Neural Network Applications被引用 4

一句话总结

本论文表明，通过使用最小图像缩放的掩码自编码器（MAE）预训练，轻量级视觉Transformer（ViTs）在CIFAR-10和CIFAR-100等小数据集上可超越ResNets。通过采用带有可学习位置嵌入和75%掩码率的MAE，作者在参数量低于370万个、计算成本极低的情况下，实现了SOTA精度——CIFAR-10达到96.41%，CIFAR-100达到96.01%。

ABSTRACT

Can a lightweight Vision Transformer (ViT) match or exceed the performance of Convolutional Neural Networks (CNNs) like ResNet on small datasets with small image resolutions? This report demonstrates that a pure ViT can indeed achieve superior performance through pre-training, using a masked auto-encoder technique with minimal image scaling. Our experiments on the CIFAR-10 and CIFAR-100 datasets involved ViT models with fewer than 3.65 million parameters and a multiply-accumulate (MAC) count below 0.27G, qualifying them as 'lightweight' models. Unlike previous approaches, our method attains state-of-the-art performance among similar lightweight transformer-based architectures without significantly scaling up images from CIFAR-10 and CIFAR-100. This achievement underscores the efficiency of our model, not only in handling small datasets but also in effectively processing images close to their original scale.

研究动机与目标

• 探究轻量级视觉Transformer（ViTs）在仅使用最小图像缩放的情况下，是否能在小数据集上与ResNet等CNN模型相媲美或超越。
• 通过自监督预训练缓解ViTs与CNN在低数据场景下的性能差距。
• 证明ViTs无需架构修改（如卷积归纳偏置或图像上采样）即可实现SOTA性能。
• 评估掩码自编码器（MAE）预训练在小尺寸图像（32×32）上、仅使用最小增强时的有效性。
• 建立高效、小参数量ViT模型的基准，使其在小数据集上与更大或更复杂架构具有竞争力。

提出的方法

• 采用掩码自编码器（MAE）框架，使用75%的块掩码率，仅将25%的块输入编码器。
• 在编码器和解码器中均使用可学习位置嵌入，而非正弦位置嵌入，以提升空间感知能力和适应性。
• 设计独立于编码器的解码器，采用可学习位置嵌入和灵活架构，以增强重建质量。
• 计算组合损失函数，包含掩码块重建损失（使用MSE）和未掩码块的折扣损失，以稳定训练过程。
• 使用标准交叉熵损失和标准优化协议，在下游分类任务上微调预训练的ViT模型。
• 采用最小图像缩放进行训练——仅将分辨率从32×32提升至36×36，以保持原始数据保真度并减少计算开销。

Figure 1: Modified MAE: Implementation using separate learnable positional embeddings for Encoder and Decoder. The Encoder and the Decoder in this figure only consist of transformer layers, a layer of layerNorm and a linear projection layer.

实验结果

研究问题

• RQ1 在不进行图像上采样的情况下，轻量级视觉Transformer能否在CIFAR-10和CIFAR-100等小数据集上实现优于ResNet的性能？
• RQ2 在低数据场景下，仅使用最小数据增强和小图像分辨率的MAE预训练，是否能使ViTs克服其缺乏归纳偏置的问题？
• RQ3 在编码器和解码器中均使用可学习位置嵌入，对重建质量和下游分类性能有何影响？
• RQ4 在某一数据集上预训练并在另一数据集上微调的ViTs是否能有效泛化，还是域内预训练更为关键？
• RQ5 组合损失函数（掩码+未掩码块损失）对训练稳定性和模型泛化能力有何影响？

主要发现

• 所提出的MAE预训练ViT模型Mae-ViT-C10和Mae-ViT-C100在CIFAR-10和CIFAR-100上分别实现了96.41%和96.01%的top-1准确率，超越了当前SOTA轻量级ViT模型。
• 参数量（372万个）和MACs（0.26G）与MAE模型相近的ViT-Lite-7/4模型，仅达到93.57%和73.94%的准确率，证明了MAE预训练策略的优越性。
• 即使未引入额外的卷积归纳偏置，该模型仍优于其他ViT变体，如CVT-7/4（CIFAR-100上94.01%）和CCT-7/3×2（CIFAR-10上95.04%）。
• 在某一数据集上预训练并在另一数据集上微调，虽收敛更快，但最终性能略有下降，表明域内预训练更具有效性。
• 重建质量随训练逐步提升，至第3020个周期时图像细节已明显改善，证实了有效特征学习。
• 对未掩码块使用折扣损失提升了训练稳定性，并在早期训练阶段显著促进了泛化能力。

Figure 2: The architecture of our MAE. ’B’ stands for batch size, ’N’ for number of embeddings