Recently, Masked Pretraining (MPT) based on reconstruction pretraining tasks has risen to a promising self-supervised learning paradigm across various domains and achieves remarkable performance in multiple downstream tasks. However, the theoretical understanding of the working mechanism behind MPT is still limited. In this paper, we introduce a new theoretical framework to analyze MPT and understand the crucial role of masking in extracting meaningful representations. We establish theoretical connections between MPT and another popular self-supervised paradigm: contrastive learning. We prove that the masking technique implicitly creates positive pairs that are semantically similar and the reconstruction loss pulls them together in the feature space. Besides, as a result of the implicit alignment, we point out the dimensional collapse issue of MPT and propose a Uniformity-enhanced MPT (U-MPT) loss that can effectively address this issue and bring significant improvements in downstream tasks including linear evaluation, cross-dataset fine-tuning and out-of-distribution generalization on real-world data sets. Furthermore, we establish downstream guarantees of U-MPT and theoretically analyze the influence of masking strategies. Based on the theoretical analysis, we propose a new masking strategy which enhances the downstream performance of MPT and explains current improvements of masking strategies with our theoretical perspective.
Figures & tables
Figure 1: An illustration of the mask graph and the corresponding augmentation graph of MPT. Different colors denote different belonging classes.
Figure 2: (a) Comparison of the singular values of learned features with 1) random initialization, 2) MPT loss, and 3) our U-MPT loss. (b) The changing process of effective rank ( Roy and Vetterli, 2007 ) of the encoded features trained with different objectives (MPT and U-MPT).
CIFAR-10
ImageNet-100
ImageNet-1K
Downstream Task
loss
ViT-Tiny
ViT-Base
ViT-Base
ViT-Large
ViT-Base
ViT-Large
Linear Probing
MAE
59.6
61.7
61.2
64.4
55.4
62.2
U-MPT
68.9
70.2
67.5
72.8
58.5
65.8
Fine-tuning
MAE
89.6
90.7
86.9
87.3
82.9
83.3
U-MPT
89.4
90.8
86.8
87.3
83.0
83.2
Table 1: In-domain linear-probing accuracy (%) and fine-tuning accuracy (%) of models pretrained by MAE loss and U-MPT loss with different ViT backbones on CIFAR-10, ImageNet-100, and ImageNet-1K. The uniformity regularizer term in the U-MPT loss significantly improves the linear evaluation performance of the MAE loss without hurting the performance of fine-tuning.
Method
iNat 21
CUB
Cars
ImageNet-A
MAE
66.0
81.5
59.8
18.4
U-MPT
69.4
81.8
60.3
19.9
Table 2: Cross-dataset fine-tuning accuracies (%) of models pretrained on ImageNet using MAE and U-MPT objectives. The pretrained models are subsequently fine-tuned and evaluated on different downstream data sets.
Data Set and Corruption Type
MAE
U-MPT
ImageNet-A
0.57
0.68
ImageNet-C (Defocus Blur)
10.48
26.20
ImageNet-C (Motion Blur)
16.36
35.84
ImageNet-C (Frosted Glass Blur)
13.68
30.14
ImageNet-C (Zoom Blur)
14.24
32.48
ImageNet-C (Elastic Transform)
23.02
39.98
Table 3: Out-of-distribution generalization accuracies (%) of MAE and U-MPT on ImageNet-A and ImageNet-C. The pretrained encoder is kept frozen and evaluated on ImageNet-A and ImageNet-C using linear probing.
SimMIM
With Uniformity Loss
54.3
61.1
Table 4: Linear probing accuracy (%) of uniformity-enhanced SimMIM on ImageNet-100.
λ
0
1e-3
1e-2
1e-1
1e0
Acc
61.2
65.9
67.5
45.1
47.0
Table 5: Linear probing accuracy (%) of uniformity-enhanced MAE with different coefficients.
Figure 3: (a) Visualization of representations on random 10 classes of ImageNet-100 trained with MAE loss and our U-MPT loss. Our U-MPT loss significantly improves the class-clustering performance of the encoder. (b) The linear evaluation results during the training process. It shows that our U-MPT loss improves the downstream performance consistently along training.
Figure 4: An illustration of the toy model. Each dot represents a patch in real-world data. The green and red dots correspond to non-overlapped patches from different classes, while the yellow dots are overlapped patches.
Figure 5: An appropriate mask ratio can generate similar views from different samples in the same class.
Figure 6: (a)(b) Estimated residual eigenvalues of the augmentation graph and label error for 2000 randomly selected images from the ImageNet-100 data set, with each image subjected to 5 different random masks, under varying mask ratios. (c) The blended value of these two factors.
Setting
Data Set
Vanilla MAE
Re-weighted MAE
In-domain
ImageNet-1K
55.4
58.6
OOD
ImageNet-A
0.57
0.87
ImageNet-C (Defocus Blur)
10.48
20.51
ImageNet-C (Motion Blur)
16.36
24.28
ImageNet-C (Frosted Glass Blur)
13.68
19.03
ImageNet-C (Zoom Blur)
14.24
19.21
Table 6: Results (%) of vanilla MAE and Re-weighted MAE with γ=1.5 , where OOD means out-of-distribution. Both models are pretrained on ImageNet-1K, and the learned representations are evaluated via linear probing on ImageNet-1K for in-domain performance and on ImageNet-A and ImageNet-C for out-of-distribution performance.
Figure 7: Linear probing accuracies of ViT-Base pretrained using Re-weighted MAE on the ImageNet-100 data set for 200 epochs.
Appendix figures & tables1 asset
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Illustration of edges in an augmentation graph between intra-class images and inter-class images. The first row shows that intra-class image pairs tend to have fairly overlapped unmasked views, while the second row demonstrates that even the most similar unmasked views of inter-class image pairs are non-overlapped.