How Mask Matters: Towards Theoretical Understandings of Masked Autoencoders
Abstract
Masked Autoencoders (MAE) based on a reconstruction task have risen to be a promising paradigm for self-supervised learning (SSL) and achieve state-ofthe-art performance across different benchmark datasets.
However, despite its impressive empirical success, there is still limited theoretical understanding of it.
In this paper, we propose a theoretical understanding of how masking matters for MAE to learn meaningful features.
We establish a close connection between
MAE and contrastive learning, which shows that MAE implicit aligns the maskinduced positive pairs.
Built upon this connection, we develop the first downstream guarantees for MAE methods, and analyze the effect of mask ratio.
Besides, as a result of the implicit alignment, we also point out the dimensional collapse issue of MAE, and propose a Uniformity-enhanced MAE (U-MAE) loss that can effectively address this issue and bring significant improvements on real-world datasets, including CIFAR-10, ImageNet-100, and ImageNet-1K. Code is available at https://github.com/zhangq327/U-MAE .
Introduction
Recently, self-supervised learning (SSL) has been proposed as a promising paradigm for learning data representations without access to labels.
Beside the popular contrastive learning methods [4, 14, 12, 27], recently there has been a renaissance of reconstruction-based autoencoders for SSL, for example, MAE [15], BEiT [1], iBOT [32] and SimMIM [29] which demonstrate state-of-theart performance on various downstream tasks.
Taking MAE as an example, it learns to reconstruct the masked patches from the unmasked context with an encoder-decoder architecture, from which we can see that there are two main components in MAE: masking and autoencoders.
While for autoencoders, its canonical one, even equipped with expressive neural networks, is still less competitive than modern
SSL methods [4], which indicates that autoencoders may not be the key factor for the success of
MAE.
Then for masking, MAE needs a very large mask ratio, e.g., 75% patches of the input image are masked out.
Such a large mask ratio will lead to the loss of most image contents that are hardly recoverable from the rest.
Further considering that MAE deliberately chooses a weak decoder with a shadow architecture, we believe that reconstruction might not be the ultimate goal of MAE in the learning process.
Then, some natural questions are raised here:
What is the role of masking in MAE?
How does it affect the
downstream
performance?
To answer these questions, firstly, we build a close connection between MAE and contrastive learning.
In particular, we find that masking could produce implicit positive pairs while MAE’s reconstruction loss is lower bounded by an alignment loss on these positive pairs.
Further for feature uniformity, we notice that although MAE will not fully collapse, it still suffers from dimensional collapse where features lie in a low-dimensional space and become highly alike.
Inspired by the uniformity loss in contrastive learning, we propose a Uniformity-enhanced MAE ( U-MAE ) to promote the feature diversity.
Empirically, the proposed U-MAE attains significant improvements over MAE on the linear probing task across different benchmark datasets (CIFAR-10, ImageNet-100, and ImageNet-1K) and different ViT backbones, and it effectively eliminates the feature collapse issue of MAE.
Secondly, the connection between MAE and contrastive learning enables us to establish the first theoretical guarantee for downstream classification among MAE methods.
Our guarantees suggest that a large mask ratio is indeed necessary for bridging semantically similar samples together.
We summarize our contributions as follows:
• We propose a new theoretical understanding of MAE by establishing a formal connection between MAE and contrastive learning.
In particular, we show that a small reconstruction loss implies better alignment of mask-induced positive pairs.
• Built upon this connection, we establish the first theoretical guarantee on the downstream performance among MAE methods, which helps to understand the role of high mask ratio.
• We point out the dimensional collapse issue of MAE, and propose U-MAE that enhances feature diversity with a uniformity loss.
Empirically, U-MAE improves the linear probing accuracy of MAE by a large margin across different real-world datasets and backbones.
Related work
Self-Supervised Learning.
Canonical deep learning relies on labeled data ( e.g., ImageNet) to train deep models.
Instead, self-supervised learning (SSL) aims to learn meaningful representations from fully unlabeled data.
A particular example is contrastive learning (CL) that learns to align augmented samples in the feature space, which achieves remarkable success and largely closes the performance gap between supervised and self-supervised learning [4, 14, 12, 27].
Recently, Masked
Image Modeling (MIM), steming from the Masked Language Modeling (MLM) paradigm widely adopted in NLP ( e.g., BERT [9]), also shows promising results in visual representation learning, to name a few, BEiT [1], iBOT [32], SimMIM [29], and MAE [15].
Among these MIM methods, the former three (similar to BERT) have close connections to contrastive learning (their objective can be directly formulated as contrastive loss [20]), while MAE has several key differences to BERT-like methods.
As the name suggests, MAE is more like an autoencoder rather than a token predictor: it adopts a pixel-level reconstruction loss, omits the masked tokens in the encoder input, and utilizes a fully non-linear encoder-decoder architecture.
In this work, we mainly provide a theoretical understanding of the working mechanism of MAE, in particular, how it learns generalizable features.
Meanwhile, our analysis can also be applied to other MIM methods by different specifications of the encoder-decoder and their inputs.
Understanding MAE.
Despite the impressive success of MAE, the theoretical understanding of it is largely underexplored.
Among several concurrent attempts, Cao et al.
[2] mostly focus on the attention mechanism of MAE through an integral kernel perspective.
A similar theoretical analysis [23] also shows that autoencoders can preserve useful semantics in the pretraining data.
However, we notice that these MAE analyses mostly focus on the autoencoder architecture, while they largely ignore the role of masking.
As shown by He et al.
[15], the patchwise masking strategy is a key component that distinguishes MAE from standard autoencoders, and different mask ratios ρ have a large impact on the downstream performance of pretrained features of MAE.
Therefore, our work aims to explain how the masking-based reconstruction task learns meaningful representations.
Downstream Generalization of SSL.
Motivated by the empirical success of SSL, many researchers try to theoretically understand how SSL works.
Arora et al.
[25] establish downstream guarantees of contrastive learning representations.
Wang et al.
[28] revise their bounds by resolving the class collision issues, and develop a new understanding of contrastive learning from the perspective of augmentation overlap.
Haochen et al.
[13] propose an augmentation graph framework and characterize the downstream performance of the eigen-decomposition solution.
While these prior works focus on the contrastive learning method, there are little theoretical understanding on the downstream performance on MAE.
In this work, we establish the first theoretical guarantee on the downstream performance among MAE methods and discuss the effect of mask ratios theoretically and empirically.
Masked Autoencoders Perform Implicit Contrastive Learning
In this section, we establish a formal connection between MAE and contrastive learning by showing that MAE implicitly aligns positive input pairs induced from the masking mechanism.
Specifically, in Section 3.2, we show how a small MAE loss implies a small alignment loss.
Motivated by this connection, in Section 3.3, we further study its feature collapse issue and point out that MAE suffers from dimensional feature collapse.
To address this issue, in Section 3.4, we propose a
Uniformity-enhanced (U-MAE) loss to further promote feature diversity.
3.1 Mathematical Formulation for MAE
We begin by introducing a mathematical formulation of MAE [15].
Given a natural image ¯ x from an unlabeled dataset D u , we first reshape it into n patches, denoted as ¯ x ∈ R n × s where s denotes the patch size ( e.g., 16 × 16 in ViT [10]).
Then, we draw a random binary mask m ∈ { 0 , 1 } n (drawingwith probability ρ , i.e., the mask ratio), and obtain two complementary masked views of ¯ x :
[Equation]
where n 1 = n (1 − ρ ) , n 2 = nρ are integers satisfying n = n 1 + n 2 .
We denote this random masking process as drawing x 1 , x 2 from the joint distribution M ( x 1 , x 2 | ¯ x ) , whose marginal distributions are
M 1 ( x 1 | ¯ x ) and M 2 ( x 2 | ¯ x ) , respectively.
The MAE model h = g ◦ f is an encoder-decoder architecture, where an encoder f maps inputs x 1 to a latent feature z 1 = f ( x 1 ) , and a decoder g maps the latent feature z 1 back to the pixel space to reconstruct the complementary target view x 2 .
Specifically, MAE adopts a simple mean square error (MSE) loss:
[Equation]
where the decoder output ˆ x 2 = h ( x 1 )
=
g ( f ( x 1 )) is assumed to be l 2 -normalized following the original paper of MAE saying that normalized target x 2 yields better performance [15].
Although our analysis is based on MAE, it is quite general and can be naturally transferred to other Masked Image
Modeling (MIM) frameworks, such as BEiT [1], iBOT [32], and SimMIM [29].
This is because their differences mainly lie in the implementation details, which is discussed in detail in Appendix C.
The Bipartite Mask Graph G M of MAE.
We notice that MAE essentially learns to pair the two complementary views x 1 , x 2 via the reconstruction task, which can be modeled by a mask graph G M .
Denote the set of all unmasked views as X 1 = { x 1 } and the set of all masked views as X 2 = { x 2 } , where the two sets are assumed to be finite 3 (can be exponentially large), i.e.,
[Equation]
The mask graph
G M over the joint set X =
X 1 ∪ X 2 is defined here:
• Node: each view x ∈ X .
• Edge: the edge weight w x 1 ,x 2 between any x 1 , x 2 ∈ X is defined as their joint probabil- ity M ( x 1 , x 2 ) =
E ¯ x M ( x 1 , x 2 | ¯ x ) .
In other
words,
there is
an
edge between two views ( i.e., w x 1 ,x 2 > 0 )
if and only if they are complementary views generated by masking.
Considering the masking process in Eq.
1, there only exist edges between the two sets X 1 , X
and there is no edge within each set itself.
Thus, mask graph G M is a bipartite graph.
Its adjacency matrix can be simply defined as A M ∈ R N 2 × N 1 where ( A M ) x 2 ,x 1 = w x 1 ,x 2 for x 1 ∈ X 1 , x 2 ∈ X 2 .
As long as ρ ̸ = 0 .
5 , N 1 ̸ = N 2 and A M is not necessarily a square matrix.
We can define the normalized adjacency matrix as ¯ A M = D − 1 /
AD − 1 / 2, where D 1 , D 2 are the diagonal degree matrices with elements d x 1 = (cid:80) x 2 w x 1 ,x 2 and d x 2 = (cid:80) x 1 w x 1 ,x 2 , respectively.
3.2 MAE Implicitly Aligns Positive Input Pairs
As the name suggests, MAE is basically composed of two key designs: mask and autoencoders.
For simplicity, we begin by assuming that the encoder-decoder architecture of MAE is capable of accomplishing the vanilla autoencoder task, i.e., reconstructing the original input.
Assumption 3.1.
For any non-degenerate decoder g , we assume the existence of a pseudo-inverse encoder f g ∈ F such that the resulting pseudo autoencoder
h g =
g ◦ f g satisfies E x ∥ h g ( x ) − x ∥ 2 ≤ ε , where x represents either unmasked data x 1 or masked data x 2 .
[Table] [Table]
[Table] [Table]
[Table] [Table]
</
[Table]
    
 6LQJXODUYDOXHVLQGH[     /RJ  R I V L QJX O D U  Y D O XH V 5DQGRP,QLW0$(80$(
[Figure]

     (SRFK      ( II H F W L Y H  U DQ N 0$(80$(
[Figure]
This assumption is likely to hold in practice, since deep networks already demonstrate excellent performance on the autoencoding task [19, 17].
Besides, some works show that the Transformer network adopted in MAE has universal approximation ability [30].
However, we note that vanilla autoencoder task cannot learn as meaningful representations as MAE.
For example, without applying any masks, the linear probing accuracy of MAE drops dramatically from 61.2% to 17.4% on
ImageNet-100.
This indicates that the autoencoding ability (studied in [2, 23]) is not enough to explain the efficacy of MAE, which motivates us to further explore the masking mechanism in MAE.
MAE Performs Asymmetric Input-output Alignment on the Mask Graph.
First, we show that the MAE loss can be lower bounded by an asymmetric alignment loss between the two complementary views x 1 , x 2 using two-branch autoencoders h (real) and h g (pseudo), respectively.
Theorem 3.2.
Under Assumption 3.1, the MAE loss can be lower bounded by
[Equation]
and
[Equation]
where H denotes the output matrix of h on X 1 whose x 1 -th row is
H x 1 = (cid:112) d x 1 h ( x 1 ) , and H g denotes the output matrix of h g on X 2 whose x
As a result, a small MAE loss implies a small alignment loss (as a lower bound), which indicates that
MAE will implicitly align the masked and unmasked views with an encoder-decoder architecture.
However, it is a little bit confusing why aligning the two complementary views helps learn meaningful features.
For a closer look, we find that via masking, MAE actually produces implicit connections among different input samples in the form of 2-hop connectivity .
Consider a pair of 2-hop input neighbors x 1 , x +1 ∈ X 1 that share a common complementary target view x 2 ∈ X 2 (more likely to happen under a larger mask ratio ρ ), as illustrated in Figure 1(a).
By enforcing x 1 , x +1 to reconstruct the same output x 2 , MAE will implicitly map their features together.
In this way, the 2-hop input neighbors serve as positive pairs that will be implicitly aligned as in contrastive learning.
Motivated by this observation, we seek to establish a formal connection between MAE and contrastive learning below.
In particular, we will show that like the data augmentations in contrastive learning, the masking mechanism of MAE also implicitly introduces affinities between input samples and create (implicit) positive pairs .
For a formal exposure, we define an augmentation graph G A for modeling the relationship between all input samples in X 1 , which is different from the mask graph
G M for modeling the input-output relationship between X 1 and X 2 .
The Augmentation Graph of MAE.
We construct an augmentation graph G A for modeling the maskinduced affinity between all unmasked views in X 1 4 .
Specifically, for any two views x 1 , x +1 ∈ X 1 , we define their edge weight in the adjacency matrix A as the probability of having the same target view, i.e., A ( x 1 , x ′ 1 ) =
E x 2 M ( x 1 | x 2 ) M ( x ′ 1 | x 2 ) 5 .
Note that the augmentation graph has the same degree matrix D 1 as the mask graph.
Based on this formulation, we define a symmetric alignment loss w.r.t.
the positive pairs ( x 1 , x +1 ) ∼ A ( x 1 , x +1 ) drawn according to the augmentation graph:
[Equation]
MAE Performs Symmetric Input Alignment on the Augmentation Graph.
Built upon the augmentation graph constructed above, we theoretically verify the intuition on 2-hop connectivity by establishing the relationship between the asymmetric input-output alignment loss on the mask graph and the symmetric input alignment loss on the augmentation graph in the following theorem.
Theorem 3.3.
The asymmetric alignment loss on the mask graph (Eq.
(4)) can be lower bounded by the symmetric alignment loss on the augmentation graph (Eq.
[Equation]
Proof Sketch.
We provide a proof sketch of this inequality as it is the key to our analysis.
We first symmetrize the asymmetric alignment loss L asym ( h ) with an arithmetic inequality, and then establish its equivalence to the symmetric alignment loss L align ( h ) .
The derivation highlights the intrinsic connection between the mask graph ( A M ) and the augmentation graph ( A ).
[Equation]
( reformulated to the mask graph (Theorem 3.2 ))
[Equation]
(transformed to the augmentation graph)
[Equation]
(following the definition in Eq.
Combining Theorem 3.2 and Theorem 3.3, we arrive at the main theorem showing that MAE’s reconstruction loss can be bounded by the symmetric alignment loss of the positive input pairs ( x 1 , x +1 ) drawn according to the augmentation graph.
Theorem 3.4.
Under Assumption 3.1, MAE’s reconstruction loss (Eq.
(2)) can be lower bounded by the alignment loss between positive pairs ( x 1 , x +1 ) ∼ A ( x 1 , x +1 ) ,
[Equation]
In this way, we establish a close relationship between the two mainstream SSL paradigms (MAE and contrastive learning) by showing that a small MAE loss will imply a small alignment loss of positive input pairs as in contrastive learning.
Leveraging this connection, we could establish guarantees for the downstream generalization of MAE (discussed in Section 4).
Comparison to Contrastive Learning.
Comparing the alignment loss (Eq. (5)) to that of contrastive learning [22, 4, 13], we notice that the main difference lies in that contrastive learning aligns features in the latent space of the encoder f , while MAE aligns features in the output space with an encoder-decoder architecture h
= g ◦ f .
Nevertheless, we also note that most variants of contrastive learning apply a nonlinear projection head g upon the encoder f before calculating the alignment loss [4, 5, 14, 12, 6], and MAE also adopts a shadow decoder g .
Therefore, we may also regard the role of MAE’s decoder as the projection head in contrastive learning.
Theoretically, if we further assume the bi-Lipschitzness of the decoder, we can show that the MAE loss is further lower bounded by an alignment loss defined in the feature space.
Corollary 3.5.
Under Assumption 3.1 and the assumption that the decoder is L -bi-Lipschitz, i.e., ∀ ( x 1 , x 2 ) , 1 /L ∥ x 1 − x 2 ∥ 2 ≤ ∥ g ( x 1 ) − g ( x 2 ) ∥ 2 ≤ L ∥ x 1 − x 2 ∥ 2 .
Then, the MAE reconstruction loss can be lower bounded by the alignment loss w.r.t.
the encoder outputs:
[Equation]
3.3 The Feature Collapse Issue in MAE
We have revealed that the MAE loss is closely related to an alignment loss.
However, it is well known that in contrastive learning, simply aligning the positive pairs will result in the full feature collapse, because the alignment loss can also be minimized when the encoder produces a constant feature for all inputs.
There, various techniques are proposed to address the issue by incorporating an additional loss to encourage feature uniformity or feature decorrelation [4, 14, 31], or by asymmetric structural designs [12, 6, 3].
MAE does not implement any of these techniques, but its latent features still do not fully collapse.
One would wonder how MAE attains this property.
In the following theorem, we show that minimizing the MAE loss can provably get rid of the full feature collapse.
Theorem 3.6.
When the encoder fully collapses, i.e., ∀ x ∈ X 1 , f ( x ) = c , the MAE loss has a large lower bound:
[Equation]
where Var( x 2 ) denotes the variance of masked targets computed on the training dataset.
MAE Can Avoid Full Feature Collapse.
As the training data contain diverse images, the variance
Var( x 2 ) will be relatively large.
Therefore, unlike the alignment loss in contrastive learning, the
MAE loss cannot be minimized (to a small value) by a collapsed encoder.
The main reason is that alignment loss operates in a fully flexible latent space that allows a collapsed encoder to minimize the loss, while in MAE, the reconstruction loss adopts a parameter-invariant and sample-dependent target x 2 .
In this case, a fully collapsed encoder can no longer minimize the reconstruction loss w.r.t.
x 2 , which, in turn, means that MAE will not fully collapse.
MAE Still Suffers from Dimensional Collapse.
Although MAE can avoid full feature collapse, it could still suffer from dimensional feature collapse where the learned features lie in a low dimensional subspace [16, 18], which also limits its representation power.
We empirically examine this issue on
ImageNet-100.
Figure 1(b) shows that after learning, MAE’s features indeed become more collapsed as there are fewer large singular values (indicating non-collapse dimensions).
Quantitatively, Figure 1(c) shows that the features of MAE gradually collapse during the training process, as the effective rank [24] becomes smaller and smaller.
This shows that MAE suffers from an increasing degree of dimensional feature collapse.
Next, we introduce an explicit regularization to address this issue.
3.4 U-MAE: Enhancing Feature Diversity with an Explicit Uniformity Regularization
To further enhance the feature diversity of MAE, inspired by the uniformity loss in contrastive learning [22, 4, 13], we propose the following Uniformity-enhanced MAE (U-MAE) loss with an explicit regularization on feature uniformity through a coefficient λ > 0 ,
[Equation]
where L
[Equation]
and x − 1 denotes an independently drawn unmasked view from X 1 .
Intuitively, the spectral uniformity loss L unif ( f ) [13] encourages a small feature similarity between random unmasked views, which could effectively promote the feature diversity of all samples.
As a preview of the results, we can observe from Figures 1(b) & 1(c) that U-MAE effectively addresses the dimensional feature collapse issue and improves the effective feature dimensionality by a large margin.
Theoretically, combined with Corollary 3.5, we can show that the U-MAE loss is lower-bounded by the spectral contrastive loss with a specific choice of λ .
Theorem 3.7.
Denote the spectral contrastive loss (SCL) from Haochen et al.
[13] as
[Equation] [Equation]
Under Assumption 3.1 and the assumption that the decoder is L -bi-Lipschitz, when λ = 1 / (4 L ) , the
U-MAE loss can be lower bounded by the SCL loss:
[Equation]
As a result, minimizing the U-MAE loss will implicitly minimize the spectral contrastive loss among input views.
Leveraging this inherent connection, in the next section, we further explain the important role of masking on the downstream generalization of MAE (as well as U-MAE).
Downstream Generalization of Masked Autoencoders
In Section 3, we establish a close relationship between MAE and contrastive learning.
Therefore, leveraging existing understandings of contrastive learning, this connection could provide new insights to understand how the masking strategy affects MAE’s downstream generalization.
In Section 4.1, we first provide theoretical bounds on downstream performance based on the masking-induced augmentation graph introduced in Section 3.2.
Built upon these theoretical insights, in Section 4.2, we further investigate the effect of mask ratios and explain the optimal mask ratio of MAE.
4.1 Theoretical Guarantees on Downstream Classification
In this part, leveraging the lower bound in Theorem 3.7 as an intermediate bridge, we can theoretically characterize the downstream performance of U-MAE on the linear classification task [25, 13].
On the downstream task, we mainly care about the prediction accuracy of natural images ¯ x w.r.t.
their labels y (¯ x ) .
For simplicity, we adopt the mean classifier c f ( x ) =
arg max W f f ( x ) , where W f ∈ R c × k is the weight of the linear classification head, and for y ∈
[ c ] , the y -th row W y =
E x 1 |
y f ( x 1 ) ⊤ contains the mean representation of the class y .
Arora et al.
[25] empirically show that the mean classifier can obtain comparable performance to learnable linear heads.
Besides, we define the Laplacian matrix of the augmentation graph G A as ¯ L =
I − ¯ A , and let 0 ≤ λ 1 ≤ · · · ≤ λ N 1 denote its eigenvalues.
Theorem 4.1.
Denote the masking-induced label error as α = E ¯ x,x 1 1 [ y ( x 1 ) ̸
= y (¯ x )] , the downstream error can be upper bounded by the U-MAE pretraining loss:
[Equation]
where c 1 , . . .
, c 6 are constants.
This theorem provides the first formal guarantee on the downstream error of U-MAE.
In particular, we provably show that a small U-MAE loss would imply a small downstream classification error, which explains the good downstream generalization ability of MAE [15].
Meanwhile, we also notice that the contrastive loss L SCL ( f ) serves as an intermediate bridge between the downstream error and the U-MAE pretraining loss: a small U-MAE loss also implies a small SCL loss, which further guarantees a small downstream classification error.
This close connection further enables us to characterize the learned features using the spectrum of the masking-induced augmentation graph G A .
Theorem 4.2.
Denote the optimal encoder learned by the SCL loss as f ∗ ( · ) and its predictor as c ∗ ( x ) =
C f ∗ f ∗ ( x ) .
Then, the downstream error of c ∗ ( x ) has the following upper bounds:
[Equation]
Theorem 4.2 shows that the downstream error can be minimized with a small vanilla autoencoding error ε , a small label error α , and a large k + 1 -th eigenvalue that indicates the connectivity of the augmentation graph (according to the spectral graph theory [7]).
In other words, MAE with large capacity will be likely to have a smaller downstream error with a smaller ε , and in the meantime the masking ratio ρ should be properly chosen such that the label error α is small and λ k +1 is large.
Indeed, He et al. show that the mask ratio has a decisive influence on the downstream performance of MAE.
Therefore, in the next part, we discuss how the choice of mask ratio would affect the downstream generalization of MAE.
4.2 The Effect of Mask Ratio on Downstream Generalization
A surprising fact in MAE is that the optimal mask ratio is pretty high, e.g., ρ = 0 .
75 , where most image contents are lost.
Therefore, we are wondering why such a high mask ratio is necessary and how it helps (instead of damaging) the learning of data representations.
In this part, we provide both theoretical and empirical insights on this problem based on our proposed theory.
Theoretical Insights.
A higher mask ratio will increase the connectivity of the augmentation graph, which, according to Theorem 4.2, will further lead to a smaller downstream error.
Specifically, according to the spectral graph theory [7], a large eigenvalue λ k +1 in Theorem 4.2 stems from a strong connectivity of the augmentation graph G A .
Thus, MAE requires a large masking ratio to bridge semantically similar samples together.
As illustrated in Figure 2(a), we find that when the mask
Mask Ratio 0.
Car
[Figure] [Figure]
    
0DVN5DWLR   
 
L 2  ' L V W DQ F H LQWHUFODVVLQWUDFODVV
[Figure]
samples     
0DVN5DWLR    
 

L 2  ' L V W DQ F H   L Q W U D  F O D VV  L Q W H U F O D VV
[Figure]
ratio increases, many dissimilar patterns are masked out and the left patches can have an increasing level of similarity, for example, the tires of the two cars.
Meanwhile, Figure 2(a) also shows that a too large mask ratio, e.g., ρ = 0 .
95 , could make it hard to tell out the belonging class of the left patches.
This suggests that a too large mask ratio will cause a high label error α .
Therefore, Theorem 4. suggests that there exists a tradeoff in the mask ratio ρ , and we should choose a properly large one to balance the effect on λ k +1 and α .
Below, we further examine this understanding on real-world data.
Empirical Investigation.
To verify our perspective, we conduct experiments on CIFAR-10 with different mask ratios.
We evaluate the distance between two images by calculating the average ldistance between every pair of patches.
As shown in Figure 2(b), we find that the average distance of intra-class images decreases with the increasing of mask ratio, which means a higher mask ratio leads to generating more edges between intra-class samples with better connectivity.
Meanwhile, we also notice that a high mask ratio will also lead to a smaller inter-class distance, which suggests an increasing label error α .
Therefore, a high mask ratio will simultaneously bring an increase of λ k + (connectivity) and α (labeling error).
The Sweet Spot for Mask Ratio.
The discussion above reveals that mask ratio has an effect on both intra-class and inter-class connectivity, and a good mask ratio should be selected to balance the two sides such that intra-class distance is relatively high while inter-class distance is relatively low.
Motivated by this, we compute the relative distance (intra-class over inter-class).
As shown in Figure 2(c), the relative distance decreases first with the increasing mask ratio, showing that the intra-class distance decreases faster than the inter-class distance under small mask ratio.
When ρ > 0 . 7 , the relative distance becomes larger again, indicating that the difference between intra-class and inter-class edges disappears under too large mask ratio.
The sweet spot lies in ρ = 0 .
7 , which is pretty close to the optimal mask ratio of MAE ( ρ = 0 .
This shows that our theoretical analysis of the effect of mask ratio aligns surprisingly well with the practice of MAE.
Experiments
In this section, we first present the main empirical results of our proposed U-MAE loss on different real-world datasets with different backbones.
Then we conduct a series of experiments to understand how well the U-MAE loss works.
5.1 Evaluation on Benchmark Datasets
To evaluate the effectiveness of the proposed U-MAE loss, extensive experiments are conducted on
CIFAR-10 [21], ImageNet-100 [8], and ImageNet-1K [8].
Setup.
We mainly follow the basic setup of MAE [15]: for the encoder, we adopt different variants of
ViT [10], i.e., ViT-Tiny, ViT-Base, and ViT-Large.
For the decoder, we use a flexible one following [15].
The mask ratio is set to 0.75.
For U-MAE, the coefficient of the uniformity term is set to 0.001.
On CIFAR-10, we pretrain the model for 2000 epochs with batch size 4096 and weight decay 0.05.
On ImageNet-100 and ImageNet-1K, we pretrain the model for 200 epochs with batch size 1024 and
The uniformity regularizer term in the U-MAE loss significantly improves the linear evaluation performance of the MAE loss without hurting the performance of fine-tuning accuracy.
[Table]
[Table]
[Table]
weight decay 0.05. We conduct both linear evaluation and non-linear fine-tuning on the pretrained encoder.
For linear evaluation, we train a linear classifier on the frozen pretrained encoder.
As for non-linear fine-tuning, we train both the pretrained encoder and the linear classifier with the cross entropy loss.
Effectiveness of the Proposed U-MAE Loss.
In Table 1, we compare the performance of original
MAE loss and U-MAE loss on different benchmarks.
We find that, on linear evaluation results, our proposed U-MAE loss increases 8.9% on CIFAR-10, 7.2 % on ImageNet-100, and 3.4% on
ImageNet-1K with two different backbones.
On fine-tuning results, our proposed U-MAE loss will not hurt the performance of fine-tuning results of MAE.
The experimental results verify the effectiveness of the proposed U-MAE loss, which achieves better performance than the original MAE loss across different datasets and different backbones.
Moreover, to quickly evaluate the training process of the encoder, we set an online linear classifier to monitor the linear accuracy whose results can be found in Appendix B.5.
Extention to Other MIM Frameworks.
We also verify our proposed method in another MIM framework SimMIM [29].
For SimMIM, we use ViT-Base as the encoder and use the linear decoder as in [29].
We use the recommended mask ratio 0.6.
Similar to U-MAE, Uniformity-enhanced
SimMIM loss (U-SimMIM) adds a uniformity regularization term to the original reconstruction loss of SimMIM, where the coefficient of the uniformity term is set to 0.0001.
For ImageNet-100, we pretrain the model for 200 epochs with batch size 128 and weight decay 0.05.
Linear evaluation is conducted in Table 2.
We can see that the linear accuracy increases 6.8% with the uniformity regularizer, which further verifies the general property of our approach.
[Table]
[Table]
5.2 Empirical Understandings
Visualization of Representations.
To intuitively understand the improvement of our U-MAE loss on clustering intra-class samples, we use t-SNE [26] to visualize the representations trained with
MAE loss and U-MAE loss on ten random class of ImageNet-100 (detailed classes are introduced in Appendix B.4).
We find that with our uniformity regularizer term, the samples are much betterclustered corresponding to their ground-truth labels.
To be specific, the red class (“hens”) and the gray class (“indigo birds”) are separated from others, this is because most of other classes are the animals living in the oceans while these two classes are more like the birds living on the land or the sky.
Thus, these two classes are easier to be distinguished, especially with our uniformity regularizer.
Different Coefficients of the Regularizer Term.
The most important hyper-parameter of our proposed U-MAE loss is the coefficient of the uniformity regularizer term.
In Table 3, we present the results of linear evaluation on ImageNet-100 trained with the U-MAE loss with different coefficients of the regularizer term.
We can see that the downstream performance increases when the coefficient
[Figure]
Accuracy
[Figure]
[Figure]
Linear Accuracy : 67. [Figure]
    (SRFK     7 H V W $ FF X U D F\     2ULJLQDO0$(/RVV 80$(/RVV
[Figure]
Training process.
To further compare the performance between the original MAE loss and U-MAE loss, we plot the linear evaluation accuracy on ImageNet-100 during the training process in Figure 3(b).
We can observe that our proposed U-MAE loss improves the performance of MAE with all different training epochs, which verifies the effectiveness of our proposed U-MAE loss.
Conclusion
In this paper, we proposed a new theoretical understanding of MAE by formally connecting MAE and contrastive learning.
In particular, we show that a small MAE loss implies a good alignment between the mask-induced positive samples.
Based on this connection, we further analyze the dimensional collapse issue of MAE and propose a new variant, named Uniformity-enhanced MAE (U-MAE), through adding an explicit feature uniformity regularization.
Theoretically, this connection enables us to establish theoretical guarantees on the downstream classification error, which also provides theoretical insights on the choice of large mask ratio in MAE.
Empirically, the proposed U-MAE successfully mitigate the dimensional collapse issue of MAE, and achieves consistent improvement on the linear probing task on CIFAR-10, ImgaeNet-100, and ImageNet-1K.
Acknowledgment
Yisen Wang is partially supported by the NSF China (No. 62006153), Project 2020BD006 supported by PKU-Baidu Fund, and
Huawei Technologies Inc.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[No] (c) Did you discuss any potential negative societal impacts of your work?
[No] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] (b) Did you include complete proofs of all theoretical results?
[Yes] 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No] Release upon acceptance.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Section 5.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[No] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[No] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether
and how consent was obtained from people whose data you’re using/curating?
[No] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[No] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Proofs A.1 Proof of Theorem 3.
Proof.
With Assumption 3.1, we have
[Equation]
We formulate the features as two matrices H , H g .
We denote H ( x 1 ) =
(cid:112) d x 1 h ( x 1 ) as the x 1 -th column of the matrix H and H g ( x 2 ) =
(cid:112) d x 2 h ( x 2 ) as the x 2 -th row of the matrix H g .
As defined before, ( A M ) x 2 ,x 1 is the joint distribution of x 1 and x 2 , i.e., ( A M ) x 2 ,x 1 = w x 1 ,x 2 .
We denote the normalized form of A M as ¯ A M , i.e., ( ¯ A M ) x 2 ,x 1 = w
[Equation]
1 ,x 2 √
[Equation]
x 1 √ d x 2 .
Then we can reformulate the reconstruction loss,
[Equation]
As the output of decoder is normalized, i.e., ∥ H g ∥ 2 = 1 .
we obtain
L MAE
[Equation]
Then we element-wise compute ¯ A ⊤ M ¯ A M and HH ⊤ , we have
[Equation]
As the trace is the sum of the diagonal value of the matrix, we consider x 1 -th diagonal value of ( ¯ A ⊤ M ¯ A M HH ⊤ ) , i.e.,
[Equation]
With that, we can element-wise expand Eq.
[Equation]
where ˆ p ( x 1 , x +1 )
= (cid:80) x
[Equation]
x
[Equation]
+
[Equation]
x.
Similarly, we define a matrix H g , where ( H g ) x
2 = (cid:112) d x 2 h g ( x 2 ) .
We let ( ¯ A M ) x 2 ,x 1 = w x 1 ,x 2 √ d x 1 √ d x
and we obtain
[Equation]
where ¯ p
[Equation]
A.2 Proof of Corollary 3.
With Theorem 3.
[Equation]
we know that
[Equation]
As the decoder is L -bi-Lipschitz
[Equation]
we obtain
[Equation]
So,
[Equation]
normalized)
[Equation]
A.3 Proof of Theorem 3.
Proof.
When the encoder fully collapses, the encoder f maps all the features to the same point c , i.e.,
[Equation]
Then,
[Equation]
We then select a g(c) to make Equation (22) minimal.
According to KKT conditions, it has a closed-form solution q ⋆ , satisfying
[Equation]
i
[Equation]
e., Q ⋆ = E
x 2 x 2 .
Then
[Equation]
A.4 Proof of Theorem 3.
With Corollary 3.
[Equation]
we have
[Equation]
[Equation]
Then we set λ = 14 L and we obtain
[Equation]
[Equation] A.5 Proof of Theorem 4.
Proof.
We compose the marginal distribution of x 1 as a matrix D and D x 1 = d x 1 is the x 1 -th row of
D .
And we denote U as the matrix composed of encoder features, i.e., U x 1 = (cid:112) d x 1 f ( x 1 ) .
Recall that A x 1 ,x +
=
A ( x 1 , x +1 ) and ¯ A is the normalized form of A , i.e., A x 1 ,x
+1 =
A ( x 1 ,x +1 ) (cid:113)
d x 1 · d x + .
Then we reformulate the downstream error,
[Equation]
where C x 1 ,j = (cid:113) ( d i ) 1 y x 1 = j .
Then we consider the relationship between the downstream error and the augmentation graph, we element-wise consider the matrix ( D 1 / 2 Y − ¯ AC ) ,
[Equation]
[Equation]
So when j = y x 1 ,
[Equation]
When j ̸
[Equation]
[Equation]
We define β x 1 = (cid:80)
[Equation]
[Equation] ( x 1 , x +
[Equation]
1 y x +1 ̸ =
y x 1 , and we have
[Equation]
With that
[Equation]
we obtain ∥ D 1 / 2 Y − ¯ AC ∥
= (cid:80) x 2 β 2 x 1 d x 1 .
As we assume that E ¯ x ∼P d ( A ( x 1 | ¯ x ) 1 y x 1 ̸ =¯ y ) ≤ α , so (cid:88)
[Equation]
≤
[Equation]
Then we have
[Equation]
( definition of β x 1 )
[Equation]
( Equation (30) )
Then we obtain
E x,y ∥
[Equation]
=
[Equation]
=
[Equation]
[Equation]
) ≤ 4(
[Equation]
[Equation] A
[Equation]
UU [Equation]
[Equation]
≤ 4(
[Equation]
[Equation] UU
[Equation]
[Equation] [Equation]
[Equation]
) = 4 L SCL ( f ) + const + ∥ ( U T C
[Equation] W f )
[Equation]
+
[Equation]
α ( Lemma B.8 in [13] ) ≤ 16 L · L U-MAE ( f ) + 8 E y [ E x | y f ( x ) − b y ] 2 + 8 α + 16 ε + const.
( Theorem 3.7 ) ≤ 16 L · L U-MAE ( f ) + 8 α + 16 ε + const.
( W f is a mean classifier)
In the next step, we analyze the prediction error.
We denote ¯ y as the ground-truth label of original data ¯ x .
We first define a ensembled linear predictor c ′ f .
For an original sample, the predictor ensembles the results of all different views and choose the label predicted most.
With the definition, ¯ y ̸ = c ′ f
(¯ x ) only happens when more than half of the views predict wrong labels.
So
[Equation]
9 in [
[Equation]
[Equation]
(definition of α )
[Equation]
[Equation] const
[Equation]
128 L · L U [Equation]
const
[Equation]
A.6 Proof of Theorem 4.
Proof.
With Theorem 4.1, we have
[Equation]
We set L m ( U ) =
∥
[Equation]
¯ A − UU T ) ∥ 2 .
When U ⋆ is the minimizer of L m ( U ) , according to the analysis in [11], we obtain
[Equation]
[Equation]
where λ k +1 · · · λ N 1 are the n − k largest eigenvalues of matrix ¯ A .
So we obtain
[Equation] [Equation]
Similarly,
[Equation]
Claim B.9 in [13]) [Equation]
B Additional Experiment B.1 Experiment Details of Computing Effective Rank
We conduct experiments on ImageNet-100 and use ViT-Base as the backbone.
We compare two kinds of pretrained encoders: 1) the encoder trained with original MAE loss, 2) the encoder trained with U-MAE loss ( λ = 0 . 0001 ).
Then we store the normalized encoded features.
We construct
[Table]
a feature matrix A with n (n=100) random samples and compute its singular values σ 1 , . . .
, σ n .
Then, we compute the effective rank [24] of the feature matrix as follows.
We first compute the distribution of sigular values, i.e., p k
= σ k ∥ σ ∥ 1 .
Then we can obtain the effective rank with that:
[Equation]
B.2 Experiment Details of Section 4.
We conduct the verification experiment of Section 4.2 on ImageNet-100.
We set patch size toand use random masking.
When computing the distance, we compute the max l 2 distance between the patches of two images and denote it as the distance between two images.
Then we compute the average distance of different images.
When computing the distance of intra-class samples, we only compute the distance between the samples in the same class.
While for inter-class distance, we only compute the distance between the samples of different classes.
To reduce the amount of calculation, we random select 10 classes of ImageNet-100 for this experiment.
B.3 Additional Experiment Details of Section 5.
For CIFAR-10, we train ViT-Tiny on 1 × NVIDIA V100 with 1600 epochs, which needs abouthours, and ViT-Base on 1 × NVIDIA V100 with 1600 epochs, which needs about 25 hours.
For
ImageNet-100, we train ViT-Base on 4 × NVIDIA V100 with 200 epochs, which needs abouthours, and ViT-Large on 4 × NVIDIA V100 with 200 epochs, which needs about 42 hours.
B.4 Details of Visualization
Results in Figure 3(a)
We use t-SNE to to visualize the representations learned with MAE and U-MAE loss on randomclasses of ImageNet.
The ten classes are 0) "cock", 1) "hen", 2) "tiger shark, Galeocerdo cuvieri", 3) "tench, Tinca tinca", 4) "goldfish, Carassius auratus", 5) "hammerhead, hammerhead shark", 6) "electric ray, crampfish, numbfish, torpedo", 7) "stingray", 8) "great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias", 9) "indigo bunting, indigo finch, indigo bird,
Passerina cyanea",
B.5 Online Linear Evaluation Experiments
Besides the offline linear evaluation results in Section 5, we also present the results obtained by an online linear classifier.
Specifically, we train a linear head along the MAE training process and detach its gradients.
From Table 4, we find that our promoting loss increases 14.71% for linear evaluation results on CIFAR-10 with two different backbones and increases 7.35 % on ImageNet-100 with two different backbones.
And we could see that our U-MAE also significantly outperforms MAE on large-scale datasets by improving 6.77% Top-1 accuracy on ImageNet-1K.
B.6 Additional Experiments on SimMIM
For SimMIM, we use ViT-Base as the encoder and use the linear decoder as proposed in SimMIM [29].
We use the recommended mask ratio, 60%.
Similar to MAE, we propose a Uniformity-promoting
SimMIM loss (U-SimMIM) which adds a uniformity regularizer term to the original reconstruction loss of SimMIM.
For the uniformity term of our proposed loss, we set the coefficient of the uniformity term to 0.01.
For ImageNet-100, we pretrain the model for 200 epochs with batch size 128 and weight decay 0.05.
We conduct linear evaluation on the unsupervised pretrained encoder.
From Table
[Table]
5, we could see that our U-SimMIM also significantly outperforms SimMIM on large-scale datasets by improving 8.46% Top-1 accuracy on ImageNet-1K.
C Extension to Other MIM Methods
The basic paradigm of current MIM frameworks is to reconstruct the masked patches from unmasked ones, while their differences mainly exist in the implementation details [1, 32, 29, 15], e.g., 1) the input of encoder: including masked patches (SimMIM, iBOT, BEiT) or not (MAE); 2) the decoder: one-layer (SimMIM, iBOT) or multi-layer (MAE, BEiT); and 3) the target: RGB (SimMIM, MAE) or tokenized discrete tokens (iBOT, BEiT).
In fact, our theoretical framework is quite general, and can be easily extended to these variants:
• Input of Encoder.
If masked patches are adopted, the input of encoder could be modeled as ˜ x 1 = ( x 1 , p x 1 + x 2 ) and ˜ x 2 = ( x 2 , p x 1 + x 2 ) , where p represents the position embeddings.
In this case, both of them have the access to the entire position encoding p x 1 + x 2 .
• Decoder.
One-layer decoder is a special case of our multi-layer formulation.
Moreover, the one-layer decoder obtains the invertibility and our analysis can be simplified.
• Target.
The tokenized target only corresponds to a specific format of x 2 , which does not affect our analysis of their alignment property.