Masked Autoencoders are Robust Data Augmentors

TL;DR

Masked-Reconstruct Augmentation (MRA) introduces a model-based data augmentation strategy that leverages a lightweight masked autoencoder to generate nonlinear, semantics-preserving views of images. By applying attention-driven masking to retain foreground patches and reconstructing masked regions, MRA provides robust augmentation that improves supervised, semi-supervised, and few-shot classification, while also offering greater occlusion robustness. The approach is label-free during pretraining and demonstrates consistent gains across ImageNet, fine-grained, long-tail, STL-10, and miniImageNet benchmarks, with ablations guiding design choices such as masking ratio and model size. Overall, MRA offers a practical, effective augmentation paradigm that enhances generalization by constraining generated variants within high-level semantic regions.

Abstract

Deep neural networks are capable of learning powerful representations to tackle complex vision tasks but expose undesirable properties like the over-fitting issue. To this end, regularization techniques like image augmentation are necessary for deep neural networks to generalize well. Nevertheless, most prevalent image augmentation recipes confine themselves to off-the-shelf linear transformations like scale, flip, and colorjitter. Due to their hand-crafted property, these augmentations are insufficient to generate truly hard augmented examples. In this paper, we propose a novel perspective of augmentation to regularize the training process. Inspired by the recent success of applying masked image modeling to self-supervised learning, we adopt the self-supervised masked autoencoder to generate the distorted view of the input images. We show that utilizing such model-based nonlinear transformation as data augmentation can improve high-level recognition tasks. We term the proposed method as \textbf{M}ask-\textbf{R}econstruct \textbf{A}ugmentation (MRA). The extensive experiments on various image classification benchmarks verify the effectiveness of the proposed augmentation. Specifically, MRA consistently enhances the performance on supervised, semi-supervised as well as few-shot classification.

Figures (4)

• Figure 1: An overview of Mask-Reconstruct Augmentation (MRA). We first mask out the image patches via the attention-based masking strategy. Using class token as query, we calculate attention weights with the keys of each image patch. We remain patches of top-k greatest attention visible. Then, the pretrained masked autoencoder (MAE) he2021masked completes the original image relying on the visible image patches. The reconstructed image can be viewed as robust augmentation for a number of classification tasks, such as supervised, semi-supervised and few-shot classification.
• Figure 2: Visualization of different masking strategies. The first row visualizes the masked images under different masking strategy. The second row shows the reconstructed images by pretrained MAE-Mini. The gradient-weighted class activation mapping (Grad-CAM) is shown in the last row.
• Figure 3: Top-1 accuracy on ImageNet with different masking ratio when pretraining MAE-Mini.
• Figure 4: Top-1 error rate on ImageNet validation set with boundary occlusion.