Learning Mask Invariant Mutual Information for Masked Image Modeling
Authors: Tao Huang
Yanxiang Ma
Shan You
Chang Xu
Venue: ICLR 2025 Poster (2025)
Keywords: Masked image modeling
Self-supervised learning
Visual pretraining
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: This paper proposes MI-MAE, a masked image modeling method that learns mask invariant mutual information based on information bottleneck theory.

Abstract: Masked autoencoders (MAEs) represent a prominent self-supervised learning paradigm in computer vision. Despite their empirical success, the underlying mechanisms of MAEs remain insufficiently understood. Recent studies have attempted to elucidate the functioning of MAEs through contrastive learning and feature representation analysis, yet these approaches often provide only implicit insights. In this paper, we propose a new perspective for understanding MAEs by leveraging the information bottleneck principle in information theory. Our theoretical analyses reveal that optimizing the latent features to balance relevant and irrelevant information is key to improving MAE performance. Building upon our proofs, we introduce MI-MAE, a novel method that optimizes MAEs through mutual information maximization and minimization. By enhancing latent features to retain maximal relevant information between them and the output, and minimizing irrelevant information between them and the input, our approach achieves better performance. Extensive experiments on standard benchmarks show that MI-MAE significantly outperforms MAE models in tasks such as image classification, object detection, and semantic segmentation. Our findings validate the theoretical framework and highlight the practical advantages of applying the information bottleneck principle to MAEs, offering deeper insights for developing more powerful self-supervised learning models.

Published as a conference paper at ICLR 2025
LEARNING MASK INVARIANT MUTUAL INFORMATION FOR MASKED IMAGE MODELING
Tao Huang! Yanxiang Ma
Shan You
Chang Xu-t
1 Schoool of Computer Science, Faculty of Engineering, The University of Sydney
TSenseTime Research
ABSTRACT
Masked autoencoders (MAEs) represent prominent self-supervised learning paradigm in computer vision. Despite their empirical success, the underlying mechanisms of MAEs remain insufficiently understood. Recent studies have at- tempted to elucidate the functioning of MAES through contrastive learning and feature representation analysis, yet these approaches often provide only implicit insights. In this paper, we propose a new perspective for understanding MAEs by leveraging the information bottleneck principle in information theory. Our theo- retical analyses reveal that optimizing the latent features to balance relevant and irrelevant information is key improving MAE performance. Building upon our proofs, we introduce MI-MAE, novel method that optimizes MAEs through mu- tual information maximization and minimization. By enhancing latent features retain maximal relevant information between them and the output, and minimizing irrelevant information between them and the input, our approach achieves better performance. Extensive experiments standard benchmarks show that MI-MAE significantly outperforms MAE models in tasks such as image classification, ob- ject detection, and semantic segmentation. Our findings validate the theoretical framework and highlight the practical advantages of applying the information bot- tleneck principle MAEs, offering deeper insights for developing more powerful self-supervised learning models.
INTRODUCTION
Masked autoencoders (MAEs) (He et al. 2022 Xie et al. 2022 Bao et al. 2022) have emerged as a powerful self-supervised learning paradigm, particularly in the realm of computer vision. Inspired by the success of masked language models like BERT (Devlin et al., 2019) in natural language processing. MAEs leverage a similar masking and reconstruction strategy learn meaningful visual representations. The fundamental concept involves masking portion of the input image and training a model to predict the missing parts, thereby enabling the model to capture the underlying structure and semantics of the visual data. This approach has proven effective in numerous applications (Li al. 20226, Kirillov et al. 2023 Tong et al., [022] Fang et al., (022b)) showcasing the potential of MAES to learn robust and generalizable features from unlabeled data.
Despite their success, the understanding of how MAEs function and why they perform well remains an open question. Recent research has sought to demystify the inner workings of MAEs, providing valuable insights into their operation. Several studies have approached this task from various per- spectives, including contrastive learning (Zhang al. [022 Kong & Zhang 2023 Huang et al.] 2023) and feature representation analysis (Xie et al., 2023 Pan et al. 2023). Specifically, (Kong & Zhang 2023) proposed that MAEs inherently learn oceluston-invarian features by treating masked patches as a form of data augmentation. This approach, also suggested by (Zhang et al.) 2022 Yue al. 2023), aligns MAEs with contrastive learning frameworks, where the models learn align Teatures between different masked views of the same image. Other studies such as (Xie et al. 2023
Pan al. 2023) analysed the latent feature representations learned by MAES to understand how these models capture and organize visual information. However, these efforts often provide only
"quual contributions. 'Corresponding author.

---

Published as a conference paper at ICLR 2025
implicit insights and do not fully address the need for comprehensive and systematic understand- ing of the learning objectives and framework of MAES.
In this paper, we propose new perspective for understanding MAEs by leveraging information theory, specifically the information bottleneck (IB) principle (Tishby & Zaslavsky 2013). This perspective provides systematic and comprehensive framework for optimizing MAES, offering theoretical insights that can guide the development of more effective models. The IB principle posits that any deep neural network can be understood as a system that balances the trade-off' between retaining relevant information and compressing irrelevant information. By applying this principle to MAEs, we aim to provide more robust understanding of their mechanisms and to identify key areas for improvement.
Based on our findings, we introduce a novel masked image modeling method, dubbed MI-MAE, to learn Mask Invariant Mutual Information for MAEs through the lens of the information bottleneck theory. Our method systematically optimizes the latent features produced by the encoder, ensuring that they contain maximal relevant information and minimal irrelevant information on the informa- tion bottleneck of MAE. Concretely, we introduce two aspects mutual information based losses on the latent feature: (1) Mutual information maximization. To optimize the autoencoder in the latent space, we derive a loss to maximize the mutual information between the latent features of multiple orthogonal masks (2) Mutual information minimization. We optimize an upper bound of mutual information between the input and latent space to minimize the irrelevant information in the latent features and thus maximize the capacity of relevant information. This comprehensive optimization strategy helps in achieving better feature representations and improved performance.
We conduct a series of evaluations on standard benchmarks, showing that our method performs significant improvements over MAE in various tasks, including image classification, object detec- tion, and semantic segmentation. For example, our 400-epoch model achieves 83.9% accuracy on ImageNet- IK, surpassing the 1600-epoch MAE by 0.5%. The experimental results validate our the- oretical findings and highlight the practical benefits of applying the information bottleneck principle to masked autoencoders. Additionally, by providing a new perspective and a rigorous analytical framework, our work paves the way for future research in this area, offering insights that can drive the development of even more powerful self-supervised learning models.
2 RELATED WORKS
Contrastive learning. Contrastive learning (Chen et al.
2020 He et al. 2020 Chen & He 2021
Grill et al. 2020; Caron et al. 2021) stands out as the leading self-supervised representation learn- ing approach in computer vision, achieving invariance by comparing different augmentations of the same image. A notable example is SimCLR (Chen et al. 2020), which enhances semantic repre- sentations by increasing the similarity between various views of the same image in the latent space. MoCo v3 (Chen et al. 2021) applies contrastive learning techniques to pre-train vision transformers. DINO (Caron et al. 2021) delves into novel properties of celf-superviee vision transformers.
Masked image modeling. Masked image modeling (MIM) has gained significant traction in the field of computer vision as an effective sulf-supe learning paradigm. Recently, with the widespread use of vision transformers (ViTs) (Dosovitskiy et al. 2001; Liu et al.) (021). ), a series of notable methods such as BEiT (Bao et al., 2022), MAE (He et al. 2022). and SimMIM (Xie et al.) 2022) have been proposed to pre-train ViT's following the BERT-style masked modeling paradigm used in natural language processing (NLP) (Devlin et al., 2019 Liu et al.] 2019). Many follow-up
works extend masked pre-training by exploring data augmentations (Chen et al. 2023 Fang et al. 2023ah, mask strategies (Li et al. 2022a; Wang et al. 2023 2024). and hierarchical structures (Xic et al. 2022; Huang et al. 2022; Woo et al. 2023). Additionally, there is growing interest in un- derstanding MAE and its connection with contrastive learning (Zhang et al., 2022] Xie et al. [202]]
Huang et al., 2023 Kong & Zhang. 2023 Pan et al. 2023). In this paper, we further investigate MAH from an information bottleneck perspective.
Information bottleneck. Under information theory, any closed system can be quantified by the mutual information between bottleneck and output variables (Tishby et al., 2000). A DNN with a
Here, "orthogonal" means the inner productions between each mask are 0, which means we are completely dividing the image into visible parts of multiple masks.
2

---

Published as a conference paper at ICLR 2025
given imput can be considered closed system that introduces no other information. During the forward propagation, the complexity of the intermediate variables usually decreases in a general prediction model, as does the amount of information they contain. It is possible to measure the goodness of each layer and even the whole prediction network by the mutual information that can be used between the intermediate variables the outputs and the network's prediction target (Tishby & Zaslavsky 2015).
3 PRELIMINARIES
Masked Autoencoders (MAEs) (He et al., 2022) are a type of self-supervised model designed to reconstruct masked patches in images. These autoencoders consist of two main components: an encoder, which encodes the image into latent features, and decoder, which predicts the masked patches. During training, each input image first embedded into a feature representation X, which is divided into multiple patches. A random mask is then generated to select a subset of these patches as visible patches. The visible patches, represented as X x- m)) are concatenated with a learnable class token and fed into a Vision Transformer (ViT) encoder to obtain the latent feature z.
Subsequently, the latent feature z is concatenated with a set of learnable mask tokens representing the masked patches and passed into a decoder to predict the original unmasked patches X. Finally, a linear projector is used to generate the reconstructed image of (X) The loss function for training MAEs is based on the reconstruction errors between the original masked pixels and the predicted masked pixels:
Lrer = I|o(X m) o(X mm)]l..
where o(X - m) denotes the original pixel values of the image associated with the masked patches X m.
Approaches in understanding MAEs. Recent studies have provided several insights into the func- tioning of Masked Autoencoders (MAEs), with works mainly understanding MAEs from the per- spective of contrastive learning. For instance, (Kong & Zhang 2023) proposed that MAES inherently learn pellusion-invarian features by treating masked patches as form of data augmentation. This perspective aligns MAES with contrastive learning frameworks, suggesting that MAES implicitly align features between different masked views of the same image. Similarly, U-MAE (Zhang et al.) 2022) established theoretical connection between MAEs and contrastive learning, showing that the positive pairs, thereby
reconstruction loss of MAEs aligns well with the aligmment of enhancing feature uniformity and diversity. However, these papers only provide implicit analyses of MAES by introducing additional views of the image to justify that the MAES implicitly align with contrastive learning, and the introduced methods only perform par with the original MAE For comprehensive understanding of MAEs, further analyses of the learning objective and autoencoder framework are needed.
Taking the above approaches into account, we conclude that optimizing the latent features produced by the encoder is crucial for improving MAEs. This motivates us to perform an in-depth analysis of MAEs and the latent features using information theory. In this paper, we provide more comprehen- sive and theoretically sound analysis following the information bottleneck principle and show that the key to improving MAES is maximizing relevant information while compressing irrelevant infor- mation in the latent space. Through our analysis with information theory, we find that the contrastive learning latent space can help minimize the IB distortion
4 METHOD
In Section 4.1 we introduce the information theorem to explain the workings of MAEs. The in- formation bottleneck principle, as introduced by (Tishby & Zaslavsky 2013). is employed and cus- tomized for MAES to elucidate the overall training objective via the information bottleneck frame- work. In Section 4.2 we break down the overall objective and propose two new loss functions for MAEs, based on the assumptions considering the information bottleneck within MAES.
3

---

Published as a conference paper at ICLR 2025
774
Lye
Encader
Decoder
m
m
Encoder
Reversed
Masked token Latent feature Gradient
X-(L-m)
CZ
Assumption
/ Figure I: Pipeline of MI-MAE for each mask We introduces two losses (max_mi) and 1 on the latency to maximize the relevant information and minimize the irrelevant information respec- tively, and Lree denotes the original MAE loss. The top sequence in the figure denotes forward propagation and the bottom denotes back propagation. m denotes mask and X denotes the original map. ? the inverse function of decoder, is the output of the reduced target map the MAE on y, and Z is defined as latent feature small neighbourhood of Ç, and their bias &₂ is decided by Ey. in backpropagation represents gradients, while Vn the gradient in layer it of the encoder.
(min.mi)
4.1 MAE WITH INFORMATION BOTTLENECK PRINCIPLE
In information theory, deep neural networks suffer from information distortion the information complexity of intermediate variables decreases (Tishby et al. 2000 Tishby & Zaslavsky 2015).
Definition 1. Based on the definition of notation in Section 3 the information distortion in MAEs is defined as
Dj 1(X (1 (1- m); X m -m),
(2)
where is the prediction and I(-:-) denotes the mutual information between two variables.
The information distortion describes the portion of the mutual information between the masked image and the unmasked image that is not captured by the recovered image. For a given preset, training the neural network reduces the information distortion. The internal variable that captures all the mutual information between the masked image and the recovered image is called the effective description, with the one having the least information complexity referred to as the simplest effective description (Tishby & Zaslavsky 2015). For any given MAE and training data, the information
distortion is determined by the simplest effective distortion, denoted by X - (1) The simplest effective distortion is considered the information bottleneck (IB) in MAEs. Thus, information
distortion is limited by the IB D [(X (1 m)) X m (1 - -))) According (Tisbyy (et al. 2000, Tishby Zaslavsky 2013), the goal the MAE can be re-interpreted as minimizing Lagrangian term that includes D IBB. formulated as:
l[p(E(=)) T(X (1: X m)) BDD
In this Lagrangian term, the first sub-term represents the complexity of the simplest effective de- scription of the samples, and the second sub-term represents the information distortion of the given network. It is challenging to precisely find X (1- m) by training MAE given data distribu- tion (Tishby & Zaslavsky 2015). The MAE can only find sub-optimal effective deseription in the neighborhood of X (1-m) m).
Theorem 2. Denote X (1 m) + as a biased simplest effective description found through train- ing, where r is the bias. Let the predicted latent feature for the MAE be E. The latent feature is
4

---

Published as a conference paper at ICLR 2025
the information bottleneck for the MAE, and thus 2 = (1
1-11mmm + r. The mutual information
[(X -(1 (1- m);X. - m)) can be upper bounded by a generalization bound as:
K T
T(X (1- X.xxmii m)) & T(2; X. m) O
T(ē; X ./l)),
(4)
where K₂ = |X] denotes the complexity of X -(1 -11-1) (1) og is the size of X (1 m). and I is the empirical estimate of the mutual information from the given training set.
From the upper bound, it can be seen that mitigating the bias on information bottleneck helps achieving better MAE performance. Unfortunately, optimizing the latent feature encounters the problem of difficulty in finding the optimal latent feature. In the following, we will analyse how to learn using a sub-optimal latent feature.
The proof of Theorem is in Appendix A.I
4.2 MULTIPLE OBJECTIVES WITH INFORMATION THEOREM
Define an optimal simplest effective description as Ç X (1-11mmm) mation between the latent feature and Ç will help reduce I X mm/r). Unfortunately, it is chal- lenging to find a precise Ç in the latent space for all samples. One approach is to use sub-optimal latent feature.
(1-m)) m). Maximizing the mutual infor-
Assumption 3. The MAE loss has been minimized on the given training set, i.e. Free & Ej. Denote as the ground truth of the MAE, and Y as the prediction. An upper bound of the information distortion at the output layer be found
H(Y cY,
where HY ) is the information in Y, Et is a small constant, and ey is determined by El.
Proof of the validity of Assumption 3 is shown in the Appendix. Under Assumption 3 we can find set of sub-optimal latent features for X (1 m). defined as Z, as the target for the samples whose latent feature similar Z. Furthermore, 1 [] & <z. This idea is similar to contrastive learning. For a set of samples with similar relevant information in the ground truth, the relevant information in the information bottleneck also similar the same MAE. To make such a set, consider generating multiple masks for a certain image before training and keeping the masks invariant, Consider the following case: for a given image X, generate mutually orthogonal masks pre-selected set M {m1, {m1,72,--.. ..... ... my). }. For each generated mask m,, the input is denoted X; = = 1x-( m.). In particular, represent Xo as the part of that not included in any input, i.e.,
j=0
Xo = X Σ X,, Then the ground truth can be expressed as X m,
Σ Xj. Let each mask nitht
m. correspond to optimal latent feature S, whose prediction is 2, Note that the prediction of the latent feature the biased simplest effective description of X,, and the optimal latent feature the simplest effective description of X;.
Corollary 4. With Assumption standing, there exists certain mask my whose predicted latent feature & = Z. The mutual information between the prediction and the optimal latent feature for the other masks is
T(2x; =k) < 4 T(2x) - T(E;; Xolzx)
Σ [(Ex
I( Xy1=0).
(6)
1=1
where 4, = I(24; +
Σ I(2;;X)). For a given 24 4 is a fixed value. I(2x) 2x) can be
maximized only when the three following conditions are satisfied: (1). I(2 2.) is maximized:
(2). T(2;; Xn) is minimized; (3). for any j that satisfies j E Zn [I, and j $ (i,k),
1=1
Σ I(2x: Xz1æs) is maximized, where Z denotes the set of integers.
5

---

Additionally, we find that Published as a conference paper at ICLR 2025
Algorithm 1 Self-supervised pre-training with MI-MAE. Our changes to MAE are marked with *
Input: Encoder IE, decoder variational distribution approximation network V with parameters 0, training dataset Drr, number of masks per image N. for iteration in total iterraioon do
2: X f Dr::
# Sammle batch of images from training set 3: * Generate orthogonal masks M = {m,, 112; mxx for each image;
4: Encode the masked images, 2, < B(X (1 m²)), V1 S i S N;
5: Decode the latents, y (
6:
Predict variational distribution go (2(x)) + 22 ()) o(X;
7:
Compute Lrec. Lmax_mi, and Lmin.mi with X, M, qo X);
8: Optimize encoder E' with V V(A1 [ree + X2 mmmmammi X mmin.mi) (Eq. 11
9: Optimize decoder D with CreCreci
10:
Optimize approximation network V with V Capproxxi
11: end for
Output: Trained encoder E and decoder TD.
From the first condition of the mutual information maximization in Corollary we can adopt In- foNCE, a widely-used loss to maximize the lower bound of mutual information (Oord al. 2018). i.e.,
mmax.m) log Lik
exp(sim((,, $x)/T)
NB Eo
(esti) exp(sim &c)/T)²
where sim(u, v) w//l |||l||||| denotes the cosine similarity between two feature vectors, 1[c o i] is an indicator function that evaluates to 1 if and only if c + i, B denotes the batch size, NB is the total number of masked images with masks per image, and t is a temperature factor. We set 0.07 in all experiments. Therefore, the final MI maximization loss among all the image pairs is formulated as
N
Lima.mm
1 N2
(max.mi) [ese 1,k
k=
Considering the first term in Eq. 3 we also need to minimize the mutual information between the la- tent feature and the masked image. Unlike (/ ((max_mi) the mutual information about 2y and X, cannot be represented by the cosine similarity, since they are not in the same feature space. Therefore, for the minimization of MI, we use the Mutual Information Neural Estimator (MINE) (Belghazi et al.) 2018) to represent the mutual information in KL divergence, This ( X,) representation requires prior probability p(č X)). Since the prior probability is intractable, we follow (Kingma & Welling 2013 Cheng et al.] (2020) and use an approximation neural network to estimate the variational distri- bution of p(2j Xj). where the loss function for minimizing it is the negative log-likelihood between 24 and X;, i.e.
1
L Capprox
log go (xy));),
where 0 is the parameters in the approximation network. With the estimated posterior probability, use the upper bound of MI presented in CLUB (Cheng al., 2020 to minimize [(E;; Xz). for all j. we aim to minimize:
N
Lmin
1
mmii..ml) with mmin. log go (ε)1X,)
log go (2xx
(10)
N
N
k=
Detailed derivations for the upper bound can be found in Appendix A.4 by optimizing Eq. I0 the third condition in Corollary is also satisfied.
In addition, Assumption needs the MAE loss to be limited a small value. Thus, we should also add the original MAE loss as part of the training loss. Considering Assumption with both parts of the Lagrangian term minimized, our final loss becomes
LMIMAE >1 + X2mmax.mi AgEmiin,i:
(11)

---

Published as a conference paper at ICLR 2025
Table 1: Results on ImageNet classification task.
The backbone for SimMIM-based methods is (Dosovitskiy et al., 2021). *: The 800-epoch based on running the official code of MAE.
Swin-B (Liu et al. (022)) while others are ViT-B
MAE results are reported by MFF (Liu al. 2023
ImageNet ACC (%)
Method
Epochs
FT
LIN
F
Supervised
81.8
DINO Caron el al.
2021
800
82.8
78.2
MoCo v3 (Chen al 2021
300
83.2
76.7
63.4
BEiT [Bao 2022
800
83.2
C-MAE Kong & Zhang 202.3
400
83.2
SemMAL (Li al. 2022a
800
83.3
65.0
MFF (Liu ét al. 2023)
800
83.6
67.0
48.0
MAE' (He et al.
2022
800
83.3
65.6
45.4
MI-MAE
200
83,9 (+0.6)
67.9 (+2.3)
48,2 (+2.8)
MAE (He al.
2022)
1600
83.6
68.0
51.1
MI-MAE
400
84.1 (+0.5)
69.3 (+1.3)
52.3 (+1.2)
PiXMIM (Liu al.
2024
800
83.5
67.2
47.9
SimMIM [XIC 2022
800
83.8
56.7
MI-SimMIM
400
84.1 (+0.3)
59.1 (+2.4)
49.1 (+1.2)
where '1 A2 and λ are hyper-parameters for balancing the loss terms. Considering all the losses, the pipeline our MI-MAE is shown in Fig. [ To be more specific, for the major MAE, we add Lminami and mmax.mi to the latent features after the forward propagation of the encoder. Be- fore Lmin mi is calculated, the approximation network is deployed to get the posterior probabilities {03))(x;) and go(2x x)). Note that the gradient of Linin. mi and L mi will only influence the encoder in back propagation. After that, the decoder will recover the latent feature to the image space. Lree will then be used influence the whole MAE in back propagation. The training process of our model is illustrated in Algorithm T Specifically, the size of Xo in our method is set to 0 to satisfy the second condition in Corollary
EXPERIMENTS
5.1 EXPERIMENT SETUP
To sufficiently validate the efficacy of our method, conducted series of experiments on image classification, object detection, and semantic segmentation tasks.
Image classification. Our method is developed based the official code of MAE (He et al. 2022). We strictly adhere to the original pre-training and fine-tuning settings on ImageNet-IK, (Rus- sakovsky et al. 2015).
Pre-training, Since our method samples four masks for every image. for fair comparisons, we reduce our training epochs to on-rurrter of the epochs used in the compared models. Specifically, we pre-train the models using an AdamW optimizer (Loshchilov & Hutter 2019) with 3, = 0.9, 3₂ = 0.95, and a weight decay of 0.05. The total batch size is 1024 (equivalent 4096 as we augment each image with masks inside the model). We use a cosine decay learning rate schedule with a 10-epoch warmup and base learning rate of 1.5 x 10-4 For the hyper-parameters introduced by our MI-MAE, we set λ = X2 = and A = 10. For Assumption 3 él is set to 0.5. This means that we only use Lree before Eree is less than 0.5, and the entire loss is used after Lree meets the assumption. We set N = 4, which means four orthogonal masks are generated in each iteration for each image with a masking ratio of 0.75. We also conduct experiments on SimMIM (Xie et al. 2022) architecture with masking ratio of 0.5, and the N is set to 2 accordingly.
Fine-tuning. The base learning rate for fine-tuning is set 1 x 10-3 We warm up the learning rate for five epochs and train the models for a total of 100 epochs with an overall batch size of 1024. Stronger augmentations and regularization such as RandAug-
7

---

Published as a conference paper at ICLR 2025
Table 2: Results COCO instance segmentation and ADE20K semantic segmentation. The back-
bone of all methods is ViT-B. The results of MoCo v3 and BEiT are from MAE (He et al. [2022). The COCO results of MAE are from ViTDet (Li 2022b). and our method uses the same archi- tecture and training strategy.
COCO
ADE20K mloU
Method
Pre-train data
Ap
Apmask
Supervised
INIK w/ labels
47.9
42.9
47.4
MoCo v3
INIK
47.9
42.7
47.3
BEiT
I
49.8
44.4
47.1
MAE
INIK
51.2
45.5
MI-MAE
INIK
52.0
46.1
49.3
Cubuk al. 2020), label smoothing (Szegedy al. 2016), and mixup (Zhang al., are adopted.
2018
Linear probing. For linear probing, we use the pre-trained and fixed feature of the class token to learn a linear predictor. We use the LARS optimizer (You et al., (2077) with base learning rate of 0.1 and a batch size of 16384. The weight decay is set to O. We train the linear probing for 90 epochs with 10-epoch warmup. All our experiments use NVIDIA V100 GPUs
Object detection. We transfer the pre- pre-trrined ViT models to COCO Lin et al. 2014) dataset. We adopt Mask R-CNN framework (He et al. (017), which predicts detections and instance segmen- tations simultaneously. We follow the model setup and training strategy used in ViTDet (Li al. 2022b)
Semantic segmentation. We conduct semantic segmentation experiments the ADE20K (Zhou et al. 2017) dataset, using the same settings as in MAE (He et al. (022). Specifically, we fine-tune UperNet (Xiao al. 2018) for 160k iterations with batch size of 16.
5.2 MAJOR RESULTS
Image classification. We conduct pre-training the ImageNet dataset and report the fine-tuning, linear probing, and low-shot (1 % samples) fine-tuning accuracies in Tab. [ After pre-training ViT-B for 200 epochs (equivalent to 800 epochs of MAE), our method achieves significant improvements of 0.6, 2.3, and 2.8 percentage points on fine-tuning, linear probing, and low-shot fine-tuning, re- spectively. Compared the optimal 1600-epoch MAE, our method still surpasses it by 0.5, 1.3, and 1.2 percentage points, obtaining an outstanding final fine-tuning accuracy of 84. %. We also imple- ment our method on SimMIM (Xie et al., 2022) framework, another typical masked image modeling method with hierarchical Swin model (Liu et al., 2021). The results show that, our MI-SimMIM achieves 84.1 1 % accuracy, outperforming previous methods such as PixMIM and SimMIM.
To evaluate the transfer learning performance of our method, we apply pre-trained 400-epoch model to downstream tasks on the COCO and ADE20K datasets.
Object detection and instance segmentation. Tab. 2 reports the bounding box AP and mask AP performance on COCO detection. Compared to MAE, our method achieves significant improve- ments of 0.8 and 0.6 in Aphux and Apmusk, respectively. This validates our method's superiority in dense prediction tasks.
Semantic segmentation. We also report the performance the ADE20K segmentation task in Tab. [ Our method achieves a notable improvement 1.2 mloU compared to MAE, demonstrating its superiority in discriminating semantic pixels.
5.3 ABLATION STUDY
To investigate the contributions of each innovation in our method and aid in determining the design choices, we conduct ablation experiments on our method. All experiments are performed with ViT- on ImageNet-1KK We pre-train the models for 50 epochs, while for MAE baseline, train
8

---

Published as a conference paper at ICLR 2025
Table 3: Ablation experiments with VrT-B ImageNet-IK. All the models pre-trained for 50 epochs and fine-tuned for 100 epochs. We run the original MAE for 200 epochs for comparison. Default settings are marked in gray
(a) Losses.
(b) Loss weights.
(c) Mask generation.
Combination
Lee 1,
mmi f(miammm) Lapprox
ACC
A2
ACC
Type ACC
82.2
82.6
Independent 82.6
82.5
82.8
Orthogonal 82.8
(c)
Collapse
1
82.7
(d)
v
82.4
0.1
10
82.5
(e)
82.4
0.5
10
82.7
V
82.8
10 10 82.6
the model for 200 epochs to match the number of mask samples in our methods. We compare the 100-epoch fine-tuning accuracies of the models.
Ablation on the proposed losses. reported in Tab. 3 (a), we pre-train the models with different combinations of losses proposed in our method and obtain the following findings:
(I) Compared to the original MAE (a), maximizing the mutual information between the latent fea- tures of image in (b) results in 0.3 percentage point improvement.
(2) Using 4,e (max.mi) only in (c) results in collapse of training, iee. the autoencoder cannot reconstruct the image and has low linear probing accuracy, as the effective maximization of [[2]|X Xx x -1 m)) requires small reconstruction loss as per Assumption
(3) Further adding m in (d) reduces the accuracy of (b) by 0.1 percentage points, since the approximation network needs to be trained by Lupprox to predict the correct conditional relation of XX m)).
(4) Optimizing the complete losses of mutual information minimization boosts the accuracy by 0.2 percentage points in (e) and by 0.3 percentage points our complete method (f).
Loss weights. We conduct experiments to tune the loss weights X2 and A for (max, mi) and ! (min.mm)
respectively. For the loss weight A of the original reconstruction loss, we keep it at 1. As summa- rized in Tab. (b), the optimal result occurs when A2 and
Orthogonal masks vs. independent masks. In our method, we generate four orthogonal masks for image, with each randomly covering 75% of the pixels, and the remaining 25% of the visible pixels from each mask do not overlap (i.e., when combined, these masked images reveal the entire original image). We compare the results of the orthogonal masks and the independently generated masks. As summarized in Tab. (c), the pre-training with independently generated masks drops the accuracy by 0.2%, suggesting that complete mutual information learning all the image patches helps MI-MAE to learn better representations. This phenomenon makes sense since in the indepen- dent case, there is no guarantee that Xo is Thus, the second condition in Corollary worsens compared to the orthogonal case.
Masking ratios. We design experiments to ex- plore the influences of masking ratios. Un- like the original MAE, which generates one mask for each image, MI-MAE generates 4 masks for each image with a masking ratio of 0.75. For other masking ratios, we design two masking strategies: (1) Complete masking: the number of orthogonal masks determined by min(1/(1 - ratio), 2). For the number N lower than 2, we set to to utilize our losses. This ensures every iteration processes all the patches of each image. (2) Fixed masks: all ratios use the same 4 masks. The training epochs are ad- justed to match the same number of mask sam- ples.
83.0
Complete masking
82.8
82.6
82.5-
82.4
82.2
ACC 08
82.0
85.0
8..
81.5
81.0#
0.1
0.2
0.3
0.4
0.5
0.6
0.7
0.8
0.9
1.0
Masking ratio (9)
Figure 2: Ablation of masking ratios.

---

Published as a conference paper at ICLR 2025
As summarized in Fig. 2 we observe optimal performance with a masking ratio of 0.75. For other ratios, complete masking shows superiority at a ratio of 0.9, while for smaller ratios, the fixed masks strategy gains advancements. The difference between these two strategies is that complete masking generates 10 masks at a 0.9 ratio, while only 2 masks at ratios 0.1 ~ 0.5. This indi- cates that more positive samples in MI maximization loss 4,6 ((maa))))) leads to better performance. We also investigate the selected 0.75 ratio with 6 and 8 masks, finding minor differences in the results (82.9%, 82.7%). This is reasonable as we find the f of all the numbers are close and the third condition of Corollary 4 is well satisfied for each N. This means that the third condition of Corollary only relies on / (Ex; Xz)2z), which is already considered in miin.mi) For simplicity, we keep our complete masking strategy.
We also find that the reconstruction loss does not change much, while /// (fminmmm) remains at low level. This means ej in Assumption 3 can always be satisfied as observed Appendix A.2
Pre-train with ViT-S. To validate our efficacy on other model size, we perform pre-training on a smaller ViT-S model. As summarized in Tab. 4 with ViT-S, our MI-MAE obtains 79.8% accuracy on ImageNet, which still obvi- ously outperforms MAE by 1.4%. This further demonstrate the effectiveness of our method.
Table 4: ImageNet results on different ViTs.
Model
Method
FT
LIN
MAE
83.3
65.6
VIT-B
MI-MAE
83.9
67.9
MAE
78.4
50.5
VIT-S
MI-MAE
79.8
53.1
5.4 ARCHITECTURE OF APPROXIMATION NETWORK
Following CLUB (Cheng et al. 2020), we design an approximation network to estimate the pos-
terior distribution of p(2, X2). Similar to VAE works (Kingma & Welling 2013 Pu et al. 2016).
the approximation network has two branches to predict the mean j((x) and variance a xX,) of
the Gaussian distribution, respectively. Then, the approximation is determined by go 2 (X))
((x)), (X,)).
We use simple architectures in these two branches to predict the mean and variance. For the pi X,) branch, taking X, as input, we adopt a multi-layer perceptron (MLP) consisting of two fully- connected layers and an intermediate GELU activation function to encode the feature, then predict (X,)) by another fully-connected layer and a LeakyReLU activation. Similarly, the (X,) branch has similar architecture but with the last activation LeakyReLU replaced by ReLU. We train the approximation network simultaneously with the autoencoder as in Algorithm
Note on the computation cost. The approximation network is light-weight and is optimized by reusing the input X and latent feature: 2 obtained in autoencoder training. Hence, we did not observe noticeable increment in the training time.
CONCLUSION
In this paper, we introduced a new perspective for understanding and improving masked autoen- coders (MAEs) by leveraging the information bottleneck (IB) theory. Building on these insights, we proposed MI-MAE, a novel method that enhances MAEs through mutual information maximization and minimization losses on the latent features. We conducted extensive theoretical and empirical analyses of our method, and experiments on tasks such as image classification, object detection, and semantic segmentation demonstrated its effectiveness. Our findings validate the theoretical framework and highlight the practical advantages of applying the information bottleneck principle MAES, providing deeper insights for developing more powerful self-supervised learning models. Future research could build on our findings to further explore and enhance the capabilities of MAES, potentially leading new advancements in self sef-ssupervieed learning and computer vision.
ACKNOWLEDGEMENTS
This work was supported in part by the Australian Research Council under Projects DP240101848 and FT230100549.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Hangho Bao, Li Dong, Songhao Piao, and Furu Wei. BEit: BERT pre-training of image trans-
formers. In International Conference Learning Representations, 2022. URL https:
//openzeview. bnbrenreview.ne///eee www.nneeet/eineeeeeee www.:mai1----2pp555000- Bh2sz5904
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International conference on machine learning, pp. 531-540. PMLR, 2018.
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 9650-9660, 2021.
Kai Chen, Zhili Liu, Lanqing Hong, Hang Xu, Zhenguo Li, and Din-Yan Yeung. Mixed autoencoder for self-supervised visual representation learning. In Proceedings of the IÉEE/CVF conference on computer vision and pattern recognition, pp.
2023.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMLR, 2020.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 15750-15758, 2021.
Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. In Proceedings of the IEEE/CVE international conference on computer vision, pp. 9640-9649, 2021.
Pengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu, Zhe Gan, and Lawrence Carin. Club: A contrastive log-ratio upper bound of mutual information. In International conference on machine learning, pp. 1779-1788. PMLR, 2020.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops, pp. 702-703, 2020.
Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Com- putational Linguistics. doi: 10. 10.18853/1//N19-11 1109--1422 1423.
URL https: https::/aclanthhlll... http://acaannohoyy...ooo orq
N19- 1423
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https: //openzeview. htt.:/ee//eooeoooo wwwp//ert//r/eeeeeee d=Y CNTTT
Yuxin Fang, Li Dong, Hangbo Bao, Xinggang Wang. and Furu Wei. Corrupted image modeling for self-supervised visual pre-training. The Eleventh International Conference on Learning Representations, 2023a. URL htRt::t::: httpt:pp:///peeeeiiiiiee htpp//ppevvveeeeeeeeeeeeeeeeeeee
Yuxin Fang, Wen Wang, Binhui Xie, Quan Sun. Ledell Wu, Xinggang Wang, Tiejun Huang, Xinlong Wang, and Yue Cao. Eva: Exploring the limits of masked visual representation learning at scale. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 19358-12 19369, 2023b.
dean-Bastion Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Beraaro Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 33:21271-21288 2020.

---

Published as a conference paper at ICLR 2025
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In Proceedings the IEEE international conference on computer vision, pp. 2961-2969, 2017.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9729-9738 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders are scalable vision learners. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common cor- ruptions and perturbations. Proceedings of the International Conference on Learning Represen- tations, 2019.
Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. CVPR, 2021.
Lang Huang, Shan You, Mingkai Zheng, Fei Wang, Chen Qian, and Toshihiko Yamasaki. Green hierarchical vision transformer for masked image modeling. Advances Neural Information Processing Systems, 35:19997-20010, 2022.
Zhicheng Huang, Xiaojie Jin, Chengze Lu, Qibin Hou, Ming-Ming Cheng, Dongmei Fu, Xiaohui Shen, and Jiashi Feng. Contrastive masked autoencoders are stronger vision learners. IEEE Transactions on Puttern Analysis and Machine Intelligence, 2023.
Diederik P Kingma and Max Welling. arXiv:1312.00 14, 2013.
Auto-encoding variational bayes.
arXiv preprint
Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al. Segment anything. In Proceed- ings of the IEEE/CVF International Conference on Computer Vision, pp. 4015-4026, 2023.
Xiangwen Kong and Xiangyu Zhang. Understanding masked image modeling via learning occlusion invariant feature. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6241-6251, 2023.
Gang Li, Heliang Zheng. Daqing Liu, Chaoyue Wang, Bing Su, and Changwen Zheng. Semmae: Semantic- -uuiddd masking for learning masked autoencoders. Advances in Neural Information Processing Systems, 35:14290-14302, 2022a.
Yanghao Li, Hanzi Mao, Ross Girshick, and Kaiming He. Exploring plain vision transformer back- bones for object detection. In European Conference on Computer Vision, pp. 280-296. Springer, 2022b.
Tsung- Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft coco: Common objects in context. In Computer Vision-ECCV 2014: 13th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 740-755. Springer, 2014.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1107, 1692, 2019.
Yuan Liu, Songyang Zhang, Jiacheng Chen, Zhaohui Yu, Kai Chen, and Dahua Lin. Improving pixel-based mim by reducing wasted modeling capability. In Proceedings of the IEEE/CVE In- lernational Conference on Computer Vision, pp. 5361-5372, 2023.
Yuan Liu, Songyang Zhang, Jiacheng Chen, Kai Chen, and Dahua Lin. PixMIM: Rethinking pixel reconstruction in masked image modeling. Transactions on Machine Learning Research, 2024. ISSN 8835-8856. URL TRI.httss. petttps://openeveee /ttp///pwvvvii/eeeeee net forum? (etForrmmiddoooe fd=prrmooiiggagx00rggr
12

---

Published as a conference paper at ICLR 2025
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer. Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision, pp. 10012-10022, 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Represenlations, 2019. URL https:://ppereviiee net.://eweeeeeeeeeeeeeeeeeeeeeee mer/forum? ?id=
Bkg6RicqY7
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprini arXiv:1807.03744 2018.
Jiachun Pan, Pan Zhou, and Shuicheng YAN, Towards understanding why mask reconstruction pretraining helps in downstream tasks. In The Eleventh International Conference on Learning Representations, 2023. URL https: http///pwwww..eee///ee
/ net/ wtt.te///rouunnaaaPPP-///...
Liam Paninski. Estimation of entropy and mutual information. Neural computation, 15(6): 1253, 2003.
Yunchen Pu, Zhe Gan, Ricardo Henao, Xin Yuan, Chunyuan Li, Andrew Stevens, and Lawrence Carin. Variational autoencoder for deep learning of images, labels and captions. Advances in neural information processing systems, 29, 2016.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International, journal of computer vision, 11::211-252, 2015.
Christian Szegedy, Vincent Vanhoucke, Sergey loffe, Jon Shlens, and Zbigniew Wojna. Rethink- ing the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2818 2-266, 2016.
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 IEEE Information Theory Workshop (ITW), pp. 1-5, 2015. doi: 10.110//T..2015.713331.1
Naftali Tishby, Fernando Pereira, and William Bialck. The information bottleneck method, 2000.
Zhan Tong, Yibing Song, Jue Wang, and Limin Wang. Videomae: Masked autoencoders are data- efficient learners for self-supervised video pre-training. Advances in neural information process- ing systems, 15:1007--10093, 2022.
Haochen Wang, Kaiyou Song, Junsong Fan, Yuxi Wang, Jin Xie, and Zhaoxiang Zhang. Hard patches mining for masked image modeling. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pp. 10375-10385, 2023.
Haochen Wang, Junsong Fan, Yuxi Wang. Kaiyou Song, Tong Wang, and ZHAO-XIANG ZHANG. Droppos: Pre-training vision transformers by reconstructing dropped positions. Advances in Neu- ral Information Processing Systems, 36, 2024.
Sanghyun Woo, Shoubhik Debnath, Ronghang Hu, Xinlei Chen, Zhuang Liu, In So Kweon, and
Saining Xie. Convnext v2: Co-designing and scaling convnets with masked autoencoders. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16133- 16142, 2023.
Tete Xiao, Yingcheng Liu, Bolei Zhou, Yuning Jiang, and Jian Sun. Unified perceptual parsing for scene understanding. In Proceedings of the European conference on computer vision (ECCV), pp- 418 418-434,72 2018
Zhenda Xie, Zheng Zhang, Yue Cao. Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: simple framework for masked image modeling. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 9653-9663, 2022.
Zhenda Xie, Zigang Geng, Jingcheng Hu, Zheng Zhang, Han Hu, and Yue Cao. Revealing the dark secrets of masked image modeling. In Proceedings of the IEEE/CVF conference on computer 2023.
vision and pattern recognition, pp.
13

---

Published as a conference paper at ICLR 2025
Yang You, Igor Gitman, and Boris Ginsburg. Lange batch training of convolutional networks. arXiv preprint arXiv:1708. 100.00388888888 2017.
Xiaoyu Yue, Lei Bai, Meng Wei, Jiangmiao Pang. Xihui Liu, Luping Zhou, and Wanli Ouyang. Understanding masked autoencoders from local contrastive perspective. arXiv preprint arXiv:2310.01994, 2023.
Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond em- pirical risk minimization. In International Conference on Learning Representations, 2018. URL
https::/opeerevial ftt..///oooommmmarorroror....
Qi Zhang, Yifei Wang, and Yisen Wang. How mask matters: Towards theoretical understandings of masked autoencoders. Advances in Neural Information Processing Systems, 35:27127-271
2022.
Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 633-641, 2017.
A APPENDIX
A.1 PROOF OF THEOREM 2
A.1.1 PROOF OF HOw EQUATION 4 INFLUENCE THE LAGRANGIAN TERM
From Eq. to minimize the Lagrangian term, I(X -1- X (1 m)) should be minimized.
Since 1(X (1 m); X (1- (1-m)) = 1(X m,; X (1-1-m)) (H m) - ((X -(1 (1- m1] (9.XXmm) Xm)))
As 1(X m (1 m)) is determined only by the data, we minimize [(X -(1 (1-m); m)) which is the left-hand side of Eq.
A.1.2 PROOF OF BIAS IN THE ESTIMATION OF MUTUAL INFORMATION
According to the IB principle Tishby & Zaslavsky 2015), where I(X; S f(X:))++++++) 0 Kla
vt we derive the mutual information bound for decoder MAE, In this case, the decoder takes the latent feature X (1 m) as the input and X as the output, generalization bound for the mutual information between the simplest effective description and the output as,
T(X (1 11-mi m1=m);X- X.m) m) & f(x (1-m::11m: X m) + O Y
(12)
n
However, it is hard to find the precise simplest effective description X (1 m) due to the follow-
ing reasons: (1) The computation of X (1 m) is based on empirical data, which is influenced by sample size and distribution. This introduces biases and approximations into the optimization pro-
cess. (2) The prediction of (1 - m) is constrained by the model capacity of the encoder-decoder structure, which limits its ability to fully capture the optimal representation. As result, in lim- ited data distribution and model capacity, we can only find empirical estimation of biased simplest effective description as [ X m), where å X
m) + and is the bias term.
We now provide formal proof to the existence of bias The mutual information between the optimal effective description X (1 m) and the observed unmasked data m is defined as:
(1-m)) ((X - (1-1;; m) = H(X (1 - m)) - (X (1 - m) X - m))
(13)
where H is the entropy function. For the predicted latent feature ξ, the empirical estimation of the mutual information
ī(2; X m) = Ar()) f= .m),
(14)
where A and (-I-)) are empirical estimates of entropy and conditional entropy, respectively.
14

---

Published as a conference paper at ICLR 2025
We compare the true and empirical mutual information
AMI == E m)] T(X m); X m)
(15)
= = ( (1 m)]] [[ m)] H(X .(1-.11111111 m)|X . m)].
As result, the bias arises from two components:
Bias in the entropy term: [[] H( (1- m)).
2. Bias in the conditional entropy term: [[ m)] H (1 m) m).
From the paper entropy estimation (e.g., Paninski
2003
Belghazi al.
2018), we have
B[H( - H(X (1 m)) = 0 K
(16)
where K2 is the complexity of the latent space representation ê, and is the number of training samples. Similarly, for the conditional entropy term, we have
K
( m)] [X m) m) =0
(17)
Combining the above results, the total bias the estimation of mutual information can be bounded which can prove that there indeed exists bias.
O
K
A.1.3 PROOF OF GENERALIZATION BOUND VIA BIASED BOTTLENECK
Unlike general discriminative networks, the target effective dimension of MAE is usually considered higher than the imput effective dimension. Thus the bottleneck should be in the middle. The relevant information contained in the intermediate variables of the decoder is derived from the latent feature extracted by the encoder, despite the increment the complexity. Intuitively, the latent feature the information bottleneck in MAE. To further analyse the changes in effective information, we analyse the encoder and decoder of MAE separately.
Denote by the latent feature of the latent space, whose simplest effective description is denoted by For the decoder, 2 is the input and (X m) is the ground truth. Denote the empirical estimation of mutual information by f(-), and according to section 4.1 the upper bound on the input-ground truth mutual information is,
K Tz=
T(2; X m) & i(2; X m) + 0
(18)
where ny is the size of the output, and K2 is the complexity of the simplest effective descrip- tion. Eq. 18 the term O K=|X|) describes the IB distortion, and goes worse with K. Since vnz
I(X- m): m) [(2 r: m) I(2: X m) I(2: m/r). We can expand the
left-hand side of the formula 18 get upper bound of -(X (1 m); X m)
Kg Y
I(X - -1 (1- m); x m) < f(;; m) + 0 VTz
((2; X mlr),
(19)
It can be seen that this upper bound is better the size of decreases. The decoder needs a to capture as much information as possible, while the size z has to stay on the smallest possible scale. Thus, for decoder, the optimal case is that the simplest effective description of the hidden layer is the latent feature itself. Using the internal variables in encoder to take the place of latent feature, the decoder will suffer worse generalization bound by a worse O K [X]
Vn.
15

---

Published as a conference paper at ICLR 2025
A.1.4 PROOF OF BIASED GENERALIZATION BOUND
This means that for a given masked image, there exists an unknown but optimal latent feature, denoted by Ç. Assume that there exists ideal mapping process from ground truth X optimal solution of the hidden layer as 7(-), the decoder, which takes latent feature as the input and approximates X m, can interpreted as the inverse process of 7(-) (see the bottom part of Fig.
the
such that on well optimized MAE, we can get an approximated latent Z ~ 7(X m) from the decoder and m, which satisfies
I(;:X m) I[2; m) E
(20)
where 6; is small constant.
Meanwhile, using the IB theory, taking the X (1 - m) as the input of the encoder, whose sim- plest effective description is denoted by X. and the predicted latent feature of the encoder, the generalization bound of the encoder
T(x & ((6() + O Kakl n
(21)
where Th denotes the size of the sample X (1 m). and K₂ denotes the length of the simplest effective description. Thus, considering both Eq. 20 and Eq. 21 and using the internal variables in decoder to take the place of latent feature, the encoder will also suffer worse generalization bound by worse O KaKly V
In conclusion, the latent feature is the simplest effective description for the
whole MAE.
A.2 PROOF OF ASSUMPTION
0.8
Reconstruction loss
0.74
0.6
C = 0.5
0.5
0.406
0.4
0
125
250
375
Epochs
Figure 3: The curve of reconstruction loss Lree during the 400-epoch training of MI-MAE. We set to 0.5.
Fig. shows that Lree satisfies Assumption 3 This means the L2 distance between Y and Ý can be limited. The mutual information between Y and Ý is negatively relevant to the To distance between them. Assume that when Lree Ex, I(Y;Y)
When Lree & er, I(Y;Y) > When Y is
given, H(Yl)) be calculated from T(Y:Y) as,
H(Ylý) H(Y)-1(Y;))
(22)
where H(Y) is fixed. Since T(Y 2 for ey H(Y) have H(Yl eey
16

---

Published as a conference paper at ICLR 2025
A.3 PROOF OF COROLLARY
For normal MAE consider my E M, the mutual information between the prediction Y, and the ground truth Y, can be denoted as [Y;; Yi). We have,
= =T(Y; Xo
Σ X²)
j=l;jıi
(23)
<I((x +
Σ
1=1:f
With Assumption 3 we can easily find another sample whose input X (1 m&) and output is Ys, where my E M is orthogonal to m,, THe mutual information between Y3 and Y2 is,
(YYx; in (Y%: Xi),
(24)
where Lg Xo)
The upper bound described in Eq. is approximate the upper bound described in Eq. The objectives of the other samples in M differ from those of i by less than 7(Y;: X;) T(Y;: X
which is also a small constant. Thus, for networks with the same initial state, the information about the intermediate variables in the network inference process is very close when sampling different masks M. This shows the possibility using latent feature a neighbourhood as objective.
write the similar mutual information on latent space. For and k that satisfies & E Z [ and j * i, the mutual information ( 24) is,
Encouraged by Theorem 2
I(2xizk)
=I(Ak
Σ
Xj) + [x)
(o,k)
(25)
<I(Ex;Xo) +
[(2x Xy) ((bs) Es 18xx X.).
(i.k)
Taking the known 2; into Eq the upper bound can then be rewritten as,
I(2k5 zk)
<[[(z,; Xo) I(2k: Xol==)] +1
Σ (z); -
Σ
1(2;;, 1((
(i.()
(26)
[[(2k;2k) < + T(2x; 2/) T(2;; X 2x)
[T X,(z)) - T(Ex: X,1=;)]-
Rearranging the above equation gives the equation in the corollary. l is decided only by [(2;; 22). I(å,, Xol=:)] is decided only by Xo when z; given. So small ((x)) can help get bigger T(âx; =x). For any that satisfies j E & n [1, N and j & ki I(2;; X, X,]=s)], when 2; is given, T(E;; and T(4;; Xyl=)) fixed, while 7 [(g; X,1æ;) T(zx; X2) o T(zx: zg). Thus, the three conditions about maximizing [(2x; zk) given in Corollary holds.
A.4 DERIVATIONS OF THE UPPER BOUND OF MI
To minimize the mutual information between predicted latency X, and its corresponding input X,, leverage upper bound of MI defined in CLUB (Cheng et al. 2020) is (we will show the proof
17

---

Published as a conference paper at ICLR 2025
of this upper bound later): 1() S
X,)1 Decs X2)]
(27)
However, the conditional relation (X,] is intractable. We can instead use variational distribu- tion ge (8 (X3) parameterized by 0 to approximate it. Consequently, the upper bound 2,) mutual information in Eq. (7 becomes:
2,) of
I(X,, 23) :- Bp(x,,,)) qa (2) [X;]] p )p(, go (23 [x]]_
(28)
Nevertheless, [(X). 22) in Equation 28 no longer guarantees an upper bound of (Xx; 2,) due the variational approximation. Fortunately, we can prove that f(X,, &2) can be a reliable upper bound estimator when the difference between p(X,, 2g) and qo (Xj, 2g) is small.
We first compare the difference between them as
A = [X,, 2,) - (()
(29)
With (X,) being the entropy of variable Xj, and using the Mutual Information Neural Estimator (MINE) [Belghazi et al. 2018). we can rewrite mutual information J(X;, 2j) as
((x;;2)) == H(X)) (X))2))
(30)
= -xx( ((1,
p(2)).
Therefore, with (Xx, 2,) (x(2) (Xx,) being the variational joint distribution induced by t0 (e)))X;),, Eq. 29 can be reformulated by Eq. 30 and Eq. 28
A -> 2j) 2g) 2n)
-Ep(x,,2, [lo (log ((, log p(2;)) E (x log 20 (X,) + (xx)))((s) llog go (2)1x,,)
ll ((,) -log (0(2)|X,) Ept Ex llg (xxx)) flog go ( (2)1 -KL(p(X;;,))le (X,, - KL ( 11 2,))
(31)
The above equation shows that, (1) when KT,(p(X, 25))
22)), we can directly get 1() < [(x 2z), and [(xj, 24)) already an upper bound of MI. (2) Otherwise,
(L 1 (X,
(L)((X))p()))llee 2y), by learning good variational approximation go (X,, 2,) that closes to p(X,, 2y), have minimized KL(p(X,, 2, 1y))la (x), Cg. then !ř(X,. 2,) I(X, 2x,,,,) Cu.
f(X,, 22) can become an MI estimator whose absolute error is bounded by the approximation performance KT(p(X, X2, &, l (X,, 2,)).
Derivation Lapprxx in Eq. 9 We show that KL(p(X,, 2,, (go, 22)) can be minimized by minimizing the negative log- log-likeling lihoooood of (o Xý). because of the following equation:
min KL & (X,, 2z)
min Ep(x,,,,)
Enxn, (log(p X (X))) log(ge (2 (2)lX;)p(Xx))
(32)
min E,xx, (og X,)) Fx (lg go (2: X,)]-
Eq. 32 equals to maximizing the second term max, Eprx, .,) [log go (0()) X2)) ) as the first term has no relation to 0, and hence the learning object Lapprox of 0 is
1
Lapprox
log qo (2gl)))).
(33)
Derivation of { miin.mi)
Eq. 10
The MI upper bound İ ((X;; &5) has an unbiased estimation as
i i(X;, 21) log go (2g1()))
log (0x IX;),
(34)
N k=1
which reflects our fmin..mm) in Eq.
10
18

---

Published as a conference paper at ICLR 2025
A.5 ROBUSTNESS EVALUATION
Compared to MAE, our method with explicit mutual information optimization based on information bottleneck, would have better robustness, the IB helps guide latent features to suppress noise while retaining semantic information. To validate this, we test our fine-tuned model on two popular ImageNet variants for robustness evaluation, ImageNet-A (Hendrycks al., 2021) and ImageNet- C (Hendrycks & Dietterich 2019). As summarized in Tab. 5 we report the top-1 accuracy on ImageNet- IK and ImageNet-A, and mean corruption error (mCE, lower is better) on ImageNet-C. The results show that, our MI-MAE significantly improves the performance of MAE, indicating better robustness.
Table 5: Robustness evaluation results.
Method
ImageNet-IK ACC ImageNet-A ACC
ImageNet-C mCE
MAE
83.3
51.7 35,9
MI-MAE
83.9
37.4
49.5
19


Reviews:
[{"review_id": "hUGZC1X0dP", "reviewer": "Reviewer_w8mj", "summary": "This work proposes a theoretical analysis of masked autoencoders (MAEs) based on information bottleneck theory. Based on the analysis, it indicates that MAE requires to balance the influences of relevant information and irrelevant information when optimizing the latent space. Therefore, the authors propose two loss functions to satisfy the information bottleneck constraints (One is to retain maximize relevant information between them and the output, and the other is to minimize irrelevant information between them and the input). Experiments show that the method outperforms MAE on image classification, object detection, and semantic segmentation", "questions": "1.\tCan you further illustrate how two new losses (l_max_mi, l_min_mi) work to constrain the hidden space that can separate relevant and irrelevant variables?\n2.\tWhat is the difference of the work compared with VAE, GAN, or diffusion model that can constrain latent space as well?\n3.\tThe performance improvement of MI-MAE is not significant enough. I guess the assumptions of this work are too idealized, which is not a better method for application. Besides, although the theoretical analysis has revealed the relationship of MAE and information bottleneck theory, the solution is similar to other methods with latent constraints. Can you provide more experimental evidence to show the effectiveness of MI-MAE and the difference of the model with other methods with latent constraints, such as case studies, comparisons with more baselines.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThis work proposes a theoretical analysis on the performance of MAE. It has proven that MAE under information bottleneck theory can achieve better performance theoretically.\n2.\tA novel but simple architecture is proposed to apply information bottleneck theory to MAE, which can improve its performance.\n3.\tExperiments are conducted on diverse tasks, which is convinced to prove the claim.", "weaknesses": "1.\tThe illustration of the model in Section 4.2 is not very clear. I am not sure about why the architecture can achieve the separation of the relevant and irrelevant parts of the latent space. \n2.\tThere are several studies that introduce the isolation of the latent space with VAE, GAN, or diffusion model. Therefore, I am not sure about the novelty of the proposed MI-MAE in visual tasks.\n3.\tExperiments only cover several general visual tasks and the results seem not to be significantly better than MAE and other baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BdgGbefsao", "reviewer": "Reviewer_MLQd", "summary": "This paper studied Masked Autoencoders via information bottleneck: MAE can be interpreted as obtaining the simplest effective distortion to capture all information between masked and recovered images. Then it proposes to add two losses to improve MAE: maximize the mutual information between the latents of different masked views of the same image, and also minimize the mutual information between the latents and the input. The authors show a 0.5% improvement on ImageNet-1K with 400-epoch training compared to 1600-epoch MAE, and outperforms MAE on transfer tasks such as instance segmentation and semantic segmentation.", "questions": "Where is the proof of Eq. (12)? There is no clear description of the IB distortion term in Eq. (12); how did the authors derive such an error bound, by what theorems? There is the same issue for Eq. (14). Detailed proofs by stating the theorems used will be helpful. \n\nMINE is known for high variance; have the authors considered other alternative estimators?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality: this is the first paper to study MAE under information bottleneck. It re-interprets MAE as minimizing a Lagrangian term that includes two terms: the simplest effective description of input and distortion of the network. It then argues that the MAE can only find a suboptimal solution. Finally it introduced two terms, based on explicit terms (MI maximization and minimization) to enforce IB principle, to improve MAE.\n\nSignificance: the knowledge of interpreting MAE from an IB viewpoint is useful for the community.", "weaknesses": "Clarity: \n\nThe overall motivation is clear, but many small explanations are missing.\n\n(a) First, it is unclear from the discussions after Eq. (3), why MAE can only find a sub-optimal effective description. The reviewer would appreciate more explanations, such as specific constraints or limitations, preferably with formal proofs, that prevent MAE from finding the optimal solution\n\n(b) Sorry if the reviewer has missed it, but it is not clear, after Eq. (4), why mitigating the bias $r$ would help MAE. What is the exact effect of $r$ on the upper bound? Why is improving that upper bound helpful for the LHS of Eq. (4)? And how does the LHS of Eq. (4) directly impact MAE, as the LHS in Eq. (4) is neither the $D_{IB}$ term nor the first MI term in RHS of Eq. (3)? Any mathematical derivation showing these would be very helpful.\n\n(c) Following up on the last question, there is no clear theoretical link to how the proposed losses can *directly* improve Eq. (3). It is briefly explained in Lines 226 - 227, that “maximizing the mutual information between the latent feature and $\\zeta$ will help reduce $I(\\hat{z}; X \\cdot m | r)$”, but why? Are there any possible direct proofs? Even though this is true, how does reducing $I(\\hat{z}; X \\cdot m | r)$ directly affect the IB formulation in Eq. (3)? Any concrete steps showing these would be very helpful.\n\nQuality: \n\nThe empirical improvement is unfortunately not substantial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LPEuUF9j2k", "reviewer": "Reviewer_PnEF", "summary": "This paper proposed MI-MAE, extending the MAE framework by maximizing relevant and minimizing irrelevant information in latent representations. MI-MAE introduces mutual information-based losses in the encoder's latent space to enhance feature representation. The method optimizes two main loss types, maximizing mutual information across orthogonal masks to retain relevant information and minimizing mutual information between the input and latent space to filter out unnecessary data. Additionally, MI-MAE's setup includes generating multiple orthogonal masks per image, which are reconstructed to validate the relevant mutual information content across different patches. Experiments reveal that MI-MAE outperforms standard MAE configurations across some benchmarks, including ImageNet and COCO.", "questions": "Please refer to the Weakness section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This work provides some new perspective in analyzing MAE with MI backed motivations, and resulting improvements demonstrated the practicability of applying mutual information maximization and minimization within latent representations and between inputs.\n2. The paper demonstrates MI-MAE’s efficacy across a variety of vision tasks, including image classification, object detection, and semantic segmentation. In reported results, MI-MAE shows better efficiency in terms of number of training epoch comparing to MAE with comparable accuracy. \n3. This paper provides detailed ablations on effects of different components, such as mask generation strategies, loss functions, and loss weight configurations.", "weaknesses": "1. Increased complexity in training. Additional loss terms $l^{max_mi}$ and $l^{min_mi}$ requires weighting parameters introduced (i.e. $\\lambda_1,\\lambda_2, \\lambda_3$) which are empirically determined, this makes the optimization more complicated than vanilla MAE.\n2. This method uses an approximation network to estimate variational distributions for mutual information minimization. It also introduces another layer of approximation, which may not capture the true complexity of the mutual information in the latent space accurately. This could lead to sub-optimal representation learning if the approximation fails.\n3. This paper assumes that the model can effectively minimize information distortion in intermediate layers as data progresses through the encoder-decoder structure, which might lead to the over-compression relevant information.\n4. Lack of robustness analysis. An important aspect of using information bottleneck is that it can increase the robustness of the pre-trained model. It is of interest to test out how this method performs in ImageNet-A/C validation set.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8O57AzzEgN", "reviewer": "Reviewer_3dfm", "summary": "This paper proposes to interpret Masked Autoencoders (MAE) using information bottleneck principle. It first conduct a theoretical analyses to show that balance the relevant and irrelevant information in the latent features is the key to improve MAE performance. Then it introduces an improved MI-MAE by maximizing the relevant information between input and latent features and minimize the relevant information between output and latent features. This is achieved by introducing two new losses functions beside the original reconstruction loss. Experiments on image classification, object detection and segmentation reveal the effectiveness of the proposed MI-MAE", "questions": "- As demonstrated in Table 1, the MI-MAE generally performs better in linear probing (LIN) or fine-tuning with 1% (FT1%) of the data compared to full fine-tuning (FT). Could the authors provide more details on why this occurs from the perspective of information bottleneck?\n\n- In the caption of Figure 1 Line 177-178, why the notions of two losses l(max_mi) and l(min_mi) doesn't follow \\mathcal{L}_{\\text{rec}} as the same format? I would recommend  standardizing the notation for consistency, if there isn't a specific reason for the difference.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ This paper provides a new perspective in understanding the MAE using information bottleneck in information theory, which distinguishes itself with other methods. Moreover, this paper provide detailed proof on how to understand MAE and how to improve MAE with the idea of information bottleneck.\n\n+ The paper introduces two types of mutual information based losses on the latent space. This is derived and supported by the theoretical  proof. \n\n+ Validation on various experiments show the effectiveness of MI-MAE. Specifically, MI-MAE can achieve better results compared with MAE even with 4X fewer pretraining epochs. The method can also be generalized to other mask image modeling method such as SimMIM.", "weaknesses": "- The paper mentions that during each training iteration, it has 4 masks for each image, which can be regraded as data augmentation during training. For a fair comparison with MAE 400 epoch, it might be worth trying to run MAE with 4 augmented masks but without the proposed information losses to truly ablate the effectiveness of the information losses. To be more specific, we could compare (1) Standard MAE (already have). (2). MAE with 4 masks per image but no information losses (3). Full MI-MAE with 4 masks and information losses(already have).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)