Downstream Task Guided Masking Learning in Masked Autoencoders Using Multi-Level Optimization
Authors: Han Guo
Ramtin Hosseini
Ruiyi Zhang
Sai Ashish Somayajula
Ranak Roy Chowdhury
Rajesh K. Gupta
Pengtao Xie
Venue: Accepted by TMLR ()

Abstract: Masked Autoencoder (MAE) is a notable method for self-supervised pretraining in visual representation learning. It operates by randomly masking image patches and reconstructing these masked patches using the unmasked ones. A key limitation of MAE lies in its disregard for the varying informativeness of different patches, as it uniformly selects patches to mask. To overcome this, some approaches propose masking based on patch informativeness. However, these methods often do not consider the specific requirements of downstream tasks, potentially leading to suboptimal representations for these tasks. In response, we introduce the Multi-level Optimized Mask Autoencoder (MLO-MAE), a novel framework that leverages end-to-end feedback from downstream tasks to learn an optimal masking strategy during pretraining. Our experimental findings highlight MLO-MAE's significant advancements in visual representation learning. Compared to existing methods, it demonstrates remarkable improvements across diverse datasets and tasks, showcasing its adaptability and efficiency. Our code is available at https://github.com/Alexiland/MLO-MAE.

Published in Transactions on Machine Learning Research (03/2025)
Downstream Task Guided Masking Learning in Masked Au- toencoders Using Multi-Level Optimization
Han Guo UC San Diego
Brcusouucdd..mm edu
Ramtin Hosscini UC San Diego
Chossein&hucsde edu
Ruiyi Zhang UC San Diego
rua0488ucsd edu
Sai Ashish Somayajula
ssomayaj@ucsd. edu
UC San Diego
Ranak Roy Chowdhury UC San Diego
eechoddhrussd edu
Rajesh K. Gupta UC San Diego
rgupta edu
Pengtao Xie
p plzie Bucsd. edu
UC San Diego
Reviewed OpenReview: https: // openreview. net/ forum? id= cFimonkD55
Abstract
Masked Autoencoder (MAE) is a notable method for self-supervised pretraining in visual representation learning. It operates by randomly masking image patches and reconstructing these masked patches using the unmasked ones. key limitation of MAE lies in its disregard for the varying informativeness of different patches, as uniformly selects patches to mask. To overcome this, some approaches propose masking based patch informativeness. How- ever, these methods often do not consider the specific requirements of downstream tasks, potentially lcading to suboptimal representations for these tasks. In response, we intro- duce the Multi-level Optimized Mask Autoencoder (MLO-MAE), novel framework that leverages end-to-end feedback from downstream tasks to learn an optimal masking strategy during pretraining. Our experimental findings highlight MLO-MAE's significant advance- ments in visual representation learning. Compared to existing methods, it demonstrates remarkable improvements across diverse datasets and tasks, showcasing its adaptability and efficiency. Our code is available at https://gitee com www./ooo/... wwwp//AexxlandMMMMMM.E.....
1 Introduction
In the rapidly evolving field self-supervised learning (Balestriero al., 2023: Gui al., 2023), particularly in visual representation learning, Masked Autoencoder (MAE) (He al., 2022) has emerged a prominent approach, which draws inspiration from the succossful masked language models like BERT (Devlin et al.,
*Corresponding author

---

Published in Transactions on Machine Learning Research (03/2025)
2018) and RoBERTa (Liu et al., 2019). Similar to how BERT learns textual representations by predicting randomly masked tokens, MAF is designed to learn visual representations by masking random patches of an image and then reconstructing them using the remaining unmasked ones.
Although MAE has shown empirical success, it applies a uniform random approach to mask patches, over- looking the varying distribution of information across different image regions (Chen et al., 2023; Kong & Zhang, 2023; Wang et al., 2023; Liu et al., 2023). It assumes equal informativeness across all parts of an image, an assumption that does not always hold true. Some image areas may hold more critical information than others, factor not considered in MAE's current design. Such oversight might hinder the model's capa- bility and efficiency in learning representations. This lack of distinction between more and less informative regions in the MAE could lead to disproportionate allocations of computational resources. Consequently, the model might spend excessive effort on less significant areas while inadequately processing and capturing the nuances in regions that contain more valuable information.
To mitigate this limitation, various strategies have been suggested for masking patches contingent on their informativeness. Key approaches include masking regions with high attention scores to prioritize areas of interest (Li et al., 2021; Kakogeorgiou et al., 2022); employing semantic segmentation to identify and mask regions rich in information (Li et al., 2022); automatically learning masking module (Madan et al., 2024); and learning a differentiable mask generator via adversarial training (Chen et al., 2023). These approaches aim to refine the masking process by prioritizing patches based on the level of information they contain, rather than treating all patches uniformly.
Although these methods are promising, they mask patches without incorporating feedback from downstream tasks. Their process involves two separate stages: initially employing a spocific masking strategy to pretrain an image encoder, then using this encoder to perform downstream tasks (via fineluning (He et al., 2022) or linear probing (He et al., 2022)), with the hope that the encoder pretrained using this strategy will be effective for these tasks. During this process, the design of the masking strategy is not influenced by the requirements of the downstream tasks. As a result, the representations developed through this strategy may not be well-aligned with the needs of these tasks, which could limit their effectiveness.
To bridge this gap, we propose a dounstream task guided masking strategy learning framework based on multi-level optimization (MLO) (Vicente & Calamai, 1994). Our approach utilizes feedback from downstream tasks to autonomously learn the optimal masking strategy. It pretrains an image encoder and applies the pretrained encoder to perform a downstream task in an end-to-end manner, allowing the downstream task's performance to directly influence the masking process during pretraining. Our method Icarns a masking network to mask patches. It processes an input image to identify specific patches for masking. Our method consists of three interconnected stages. In the first stage, a preliminary version of the masking network masks certain patches, followed by the pretraining of an image encoder tasked with reconstructing these masked patches. In the second stage, we utilize this encoder to construct downstream model, which is subsequently trained using the training dataset specific to a downstream task. The final stage involves evaluating the downstream model using held-out validation datasct. The effectiveness of the masking network is indirectly measured by the downstream model's validation performance. An inferior masking network might fail to correctly identify the optimal patches for masking, leading to ineffective pretraining of the image encoder. When applied the downstream task, the encoder's inadequate representation Icarning capabilities will lead lo suboptimal validation performance of the downstream model. To prevent this, we continuously refine the masking network, ensuring it maximizes downstream validation performance. Fach of these stages is formulated one level of optimization problem in our MLO framework. The three levels of optimization problems are mutually dependent on each other and solved jointly. This enables the three stages to be conducted end-to-end, where the downstream validation performance closely guides the learning of the masking network.
The major contributions of this work include:
We propose a multi-level optimization based end-to-end framework to learn an optimal masking strategy in Masked Autoencoder by leveraging feedback from downstream tasks.
2

---

Published in Transactions on Machine Learning Research (03/2025)
Our approach outperforms range of leading-edge methods in learning representations, as evidenced across various datasets such as CIFAR-10, CIFAR-100, and ImageNet-1K.
Our method showcases remarkable transfer learning abilities, in fine- fine-grained classification, seman- tic segmentation, and object detection tasks, demonstrated datasets including CUB-200-2011, Stanford Cars, iNaturalist 2019, ADE20K, and MS-COCO.
2
Related works
2.1 Masked autoencoders
Following the success of masked language models in the field of natural language processing (Devlin et al., 2018), various masked image models have been proposed (Chen et al., 2020; Bao et al., 2022). Among them, Masked Autoencoder (MAE) has become promising methodology for generic visual pretraining (He et al., 2022). MAE is a denoising autoencoder that randomly masks the input image and tries to reconstruct the missing pixels. It uses a high masking ratio (75% in MAE compared to 15% in BERT) and light decoder architecture that forces the encoder to learn meaningful visual representations. Zhang et al. (2022) propose a theoretical framework to understand the role of masking MAE, and introduce Uniformity- enhanced MAE (U-MAE) to address the dimensional collapse issue. Despite MAE's effectiveness, recent works underscore the importance of replacing the random patch masking method in MAE with more so- phisticated masking strategies (Kakogeorgiou et al., 2022; Shi al., 2022). For instance, MST (Li et al., 2021) utilizes attention maps to guide the masking process, selectively obscuring less attended regions to maintain important information. SemMAE (Li et al., 2022) combines a StyleGAN- decoder with the MAE decoder and leverages attention maps from the StyleGAN decoder to provide semantic cues for patch masking. Furthermore, some recent methods propose to use a learnable masking module to generate masking strategies and optimize the masking module in the pretraining process. For example, AutoMAE (Chen al., 2023) links differentiable mask generator with MAE using Gumbel-Softmax (Jang et al., 2016), following similar two-stage setup as in SemMAE. CL-MAE (Madan et al., 2024) leverages curriculum learning to enhance MAE by progressively increasing the complexity of the masks generated from learnable masking module. Compared to these existing methods, the key distinction of our approach lies in the utilization of feedback from downstream tasks to inform the development of masking strategies, a mechanism absent in the current methodologies.
2.2 Bi-level and multi-level optimization
Recently, Bi-level Optimization (BLO) and Multi-level Optimization (MLO) techniques have been widely applied for meta-learning (Feurer et al., 2015; Finn et al., 2017), neural architecture search (Cai et al., 2019; Xie et al., 2019; Xu et al., 2020; Hosscini et al., 2021) and hyperparameter tuning (Feurer et al., 2015; Baydin et al., 2017). BLO, a formulation that consists of two levels of nested optimization problems, has been broadly applied in numerous machine learning applications (Liu et al., 2018; Liang et al., 2019). BLO basod methods have enabled automatic and efficient lcarning of upper-level parameters, such as meta parameters and neural architectures, thereby reducing the need for extensive hyperparameter luning through manual efforts. Following the success of BI.O, MT.0 which has more than two levels of nested optimization problems (Hosseini & Xie, 2022; Hosseini et al., 2023; Sheth et al., 2021; Garg et al., 2021) has been used to solve machine learning tasks with more complicated dependencies. These works develop multi-stage pipelines, with each stage corresponding to one level of optimization problem (OP). Different stages are executed end-to-end by solving all levels of interdependent OPs jointly. Despite its effectiveness, MLO based methods increase memory and computation costs due to their growing number optimization levels. To tackle this challenge, Choe al. (2022) develop software that integrates multiple approximation algorithms to efficiently compute the hypergradients within BLO and MLO problems.
3 Methods
3.1 Overview
We introduce the Multi-level Optimized MAE (MLO-MAE), an end-to-end visual representation learning method which leverages the guidance from downstream task to learn an optimal masking strategy auto- matically. Conceptually, MLO-MAF, reimagines the visual representation learning by coupling pretraining
3

---

Published in Transactions on Machine Learning Research (03/2025)
Masking network
Image encoder
Decodle
Classification
Pretrained
encoder
head
image encoder
Stage I: Pretrain image encoder by reconstructing masks
Training Image
Stage Train classification head
Updated masking network
Trained classification
Classification head
Groundtruth
Validation
Update
Masking network
encodor
Predicted label
loss
image
Stage Validate classifier and update masking network
MLO-MAE Pretraining
Frozeh network
Pretrained encader
Classification head
Tiger
Image encoder
Segmentation head
encoder
Image classification
Semantic segmentation
Apply pretrained encoder for downstream tasks
Figure An overview of MLO-MAE, which consists of three stages performed end-to-end. Modules with learnable parameters indicated in orange, and those with frozen parameters in blue.
with downstream task performance. MLO-MAE introduces a learnable masking strategy that prioritizes patches based on their relevance to a downstream task. This enables the model to focus its representation learning on features that are most critical for achieving better downstream performance.
As illustrated in Figure 1, the architecture of MLO-MAE consists of three key components: a Vision Trans- former (VIT) (Dosovitskiy et al., 2020)-based image encoder E, classification head C, and masking net- work T. The masking network processcs input images to identify patches for masking. The image encoder then extracts representation by reconstructing the masked patches using information from the unmasked ones, serving as the backbone for visual representation learning. The classification head C' predicts class label based on the image representation extracted by the encoder. The pretraining process is carriod out on an unlabeled dataset D.,, similar other self-supervised methods. However, the learning of the masking network guided by a downstream image classification task using a labeled dataset D, which is further divided into training subset Dr and validation subset Deal
MLO-MAE operates in three interconnencted stages. In the Stage I, the masking network T generates preliminary mask for input images from Du, and the image encoder E is pretrained on these masked images by minimizing reconstruction loss. In the Stage 11, the pretrained encoder extracts representations for images in the training subset D These representations, along with their associated labels, are used to train the classification head by minimizing classification loss. Finally, Stage III, the pretrained encoder extracts representations for the validation subset Doal, which are passod the classification head for label predictions. Validation loss, computed by comparing the predicted labels with the actual labels, serves as feedback mechanism for evaluating the effectiveness of the masking network T. By focusing minimizing this validation loss, the masking network T is iteratively refined.
To seamlessly integrate these stages, MLO-MAE employs a multi-level optimization (MLO) framework, where each stage corresponds to one level optimization. Optimal parameters obtained at each lower level serve as inputs for the loss functions at the subsequent upper levels. Conversely, non-optimal parameters from the upper levels are utilized to define the loss functions at lower levels. These nested optimization problems are solved jointly, enabling the three stages to operate in end-to-end manner. This multi-level
4

---

Published in Transactions on Machine Learning Research (03/2025)
approach ensures that the pretraining process dynamically aligns with the downstream task requirements, allowing MI.O-MAF to learn masking strategy that optimally enhances visual representation learning.
3.2 Multi-level optimization framework
The framework of our proposed MLO-MAE is structured into three interconnected stages. These three stages are integrated within multi-level optimization framework.
Stage I: pretrain image encoder. Given an imput image X E Du divided into N non-overlapping patches of equal size, denoted as (P:) the masking network T takes X as input and generates probability (P,,X;77 for each patch P3 which represents the likelihood that P; should be masked. Given a masking ratio r, a hyperparameter dictating the proportion of patches to be masked, we first rank all patches in descending order based on their masking probabilities. We then select the top N x r patches denoted as Tr),, those with the highest probabilities, and mask them. The remaining patches, denoted as X --M(X: r), are unmasked. Then feed the unmasked patches into an autoencoder (He et al., 2022), which consists of the image encoder E and a decoder D, to reconstruct the masked patches M();T,r). In detail, the unmasked patches are first processed by the image encoder E, which is responsible for extracting their representations. Then, these representations are input into the decoder D. The decoder's role is to accurately predict the pixel values of the masked patches. To evaluate the performance of this reconstruction, we employ a reconstruction loss, Crees defined as the squared differences between the predicted and ground truth pixel values of the masked patches. Importantly, the reconstruction loss for the j-th masked patch M r) is weighted according to its masking probability (X; r), X; This probability reflects the likelihood of patch being masked and thus, guides the autoencoder to prioritize the reconstruction of patches deemed more likely to be masked.
In this stage, we provisionally hold the masking network T constant, and focus on training the image encoder E and decoder D, by solving the following optimization problom:
Nxr
DT) argmin Σ E C
(1)
EDD XED
The notation E (T) indicates that the optimal solution E is function of T, as E* is determined by the loss function which in turn depends on T.
Stage II train classification head. Utilizing the pretrained image encoder ET(T) from Stage I, we develop an image classification model for a downstream task. This model comprises the encoder E (T) and the classification head C. For any given input image, is first processed by the encoder to generate a representation. This representation then input into the classification head to determine the class label. In this stage, we keep the encoder parameters fixed and focus on training the classification head. This is achieved by minimizing cross-entropy classification loss Lete the training dataset Drr:
C* (*(E((T) arggmin (D', E (T),)).
(2)
C
Stage III: update masking network. In Stage III, we assess the classification model developed in Stage II on the validation set Deal This model integrates the image encoder, (T),, which was pretrained in Stage I, and the classification head, C* (T),, trained in Stage II. The validation loss serves as an indirect mcasure of the efficacy of the masking network T. Our objective is to enhance the performance of T by minimizing this validation loss:
min (D E (E))))))
(3)
Multi-level optimization. Integrating the three optimization problems together, we have the following multi-level optimization problem:
min ("D) T
s.t. (E)( = argmin (e,(
Nxr Σ (MM,( MM (4)
C
E' argmin Σ
E
XEP m11
5

---

Published in Transactions on Machine Learning Research (03/2025)
Table 1: Top-1 accuracy (%) on the test sets of CIFAR-10, CIFAR-100, and ImageNet, in fine-tuning experiments. The baseline methods SemMAE and AutoMAE are not included in the comparison on CIFAR- 10 and CIFAR-100, due the absence of reported results for these datasets in their original publications and the unavailability of their implementation code for conducting evaluations on these datasets.
(No Pretraining)
(Random Masking)
(Teernable Masking)
ViT
MAE
U-MAE
SemMAE
AutoMAE
MLO-MAE (Ours)
CIFAR-100
56.4
64.0
64.6
79.4
CIFAR-10
82.3
93.7
94.3
96.2
ImageNet-1K
77.9
83.6
83.0
83.3
83.3
84.8
In this formulation, the three levels of optimization problems are mutually dependent. The first. level's output, E (T), defines the loss function in the second level. Both the outputs of the first and second levels are fed into the loss function of the third level. Simultaneously, the third level's optimization variable, T, influences the loss functions in the first two levels. By concurrently solving these optimization problems across all three levels, we enable an integrated, end-to-end execution of the three stages.
Optimization algorithm. Inspired by Liu et al. (2018), we develop an efficient hypergradient-base method to solve the problem in E4)( First, we approximate the optimal solutions E (T) and D* by executing several iterations (termed as unrolling steps) of gradient descent updates the loss function at the first level. The approximation of (T) is then plugged into the second-level loss, and (T)) is similarly approximated using multiple steps of gradient descent updates of C against this approximate loss. The approximations of (T) and C* are then applied to the third-level loss, enabling the gradient descent update of T. This iterative process of updating contimues until convergence is achieved. Details of this optimization algorithm are deferred to Appendix A and Algorithm 1.
and D against
4 Experiments
4.1 Experimental settings
Model setup and hyperparameters. In our method, the masking network is structured with multiple layers. Initially, there is linear layer, where the input size is determined by the product of the number of patches (196 for ImageNet and 256 for CIFAR) and the embedding dimension (we used the patch embedding method in VIT, with a dimension of 768), and it has a hidden size of 512. This is followed by a ReLU layer. Next, there is another linear layer, which takes an input size of 512 and produces an output size equivalent to the number of patches. Finally, a sigmoid activation function is applied to the output to generate probabilities in the range of 0 to 1. Implementation details are described in Appendix B.2. Following MAE (He et al., 2022), an asymmetric VÝT (Dosovitskiy et al., 2020) encoder-decoder architecture was used for mask roconstruction. Recognizing the constraints of computational resources, we primarily employed the ViT-B (Dosovitskiy et al., 2020) as the image encoder, ensuring balance between efficiency and performance. The classification head consists of a single linear layer. It is intentionally made simple to focus on evaluating the offectiveness of the Iearned representations. The patch size was set to 2 for CIFAR-10 and CIFAR-100, and 16 for ImageNet. For all experiments, unless otherwise specified, we used the default mask ratio of 75% as suggested MAE (He et al., 2022).
The number of unrolling steps in the algorithm for solving the MLO problem was set to 2. We employed the AdamW optimizer (Loshchilov & Hutter, 2017) with 3 values of 0.9 and 0.95 for optimizing all parameters. The learning rates were set specifically for different components: 1e - 4 for the image encoder, and 4e - 5 for both the classification head and the masking network. We used a batch size of 256. For training, we set the epoch number to 50 for the ImageNet dataset and to 200 for the CIFAR datasets. All experiments were conducted on Nvidia A100 GPUs. Further information on our experimental settings can be found in Appendix B.
Basclines. We conducted comparisons with several basclines, including: 1) vanilla MAE (He et al., 2022) and U-MAE (Zhang et al., 2022), which employ uniform random masking of images; 2) SemMAE (Li et al., 2022) and AutoMAE. (Chen al., 2023), which mask patches according to their informativeness; and 3)
6

---

Published in Transactions on Machine Learning Research (03/2025)
Table 2: Test accuracy (%) in linear probing experiments.
MAE
U-MAE
SeMMAE
AutoMAE
MLO-MAE (Ours)
CIFAR-100
46.6
50.4
63.8
CIFAR-10
73.5
77.1
84.3
ImageNet-1K
68.0
58,5
65.0
68.8
70.2
Table 3: Accuracy (%) on fine-grained image classification datasets. All methods use ViT-B as the backbone with patch size of 16.
Table 4: ADE20K.
Semantic segmentation results on
Method
mloU
Method
iNaturalist
CUB
Cars
Supervised Pretraining
45.3
MAE
79.5
83.3
92.7
MAE
48.1
SemMAE
79.6
82.1
92.4
SemMAE
46.3
AutoMAE
79.9
83.7
93.1
AutoMAE
46.4
MLO-MAE (Ours)
80.1
84.0
93.4
MLO-MAE (Ours)
49.8
the Vision Transformer (ViT) (Dosovitskiy et al., 2020) without pretraining by MAE methods (i.e., directly trained for classification from scratch). ViT-B was used the image encoder in these methods. Following their original papers, the number of pre-training epochs for MAE, U-MAE, SemMAE, and AutoMAE on ImageNet are 1600, 200, 800, and 800 respectively. The patch size in all methods is 16 for ImageNet.
Evaluation protocols. In the literature on self-supervised learning, including MAE methods, there are two standard approaches for evaluating pretrained image encoders in downstream classification tasks (He et al., 2022). The first one is fine-tuning, which fine-tunes the pretrained encoder (Logether with training randomly initialized classification head) by minimizing classification loss on the downstream training data. The second approach is linear probing, which keeps the pretrained encoder fixed and only trains the classification head. Our experiments used both protocols. For each datasct D, pretraining was conducted on unlabeled images in D; fine-tuning and linear probing were conducted on as well, utilizing both images and their associated labels.
It is important to note that our method does not unfairly utilize more labeled data than the baselines. The labeled data used in Stage II and IIT of our framework is identical to that used in the fine-tuning phrase of the basclincs.
4.2 Main results
Fine-tuning results. Table 1 shows the results. On the CIFAR-100 dataset, MLO-MAE demonstrates superior performance, achieving a test accuracy of 79.4%, substantially outperforming MAE's 64% and U-MAE's 64.6%. This trend of outperformance is also evident on the CIFAR-10 dataset, where MLO- MAE surpasses both MAE and U-MAE by 2.5% and 1.9% (absolute percentage) respectively. Moreover, on the ImageNet dataset, MLO-MAE performs better than all baseline methods. Specifically, MLO-MAE outperforms AutoMAE and SemMAE by 1.5% (absolute) improvements in top-1 accuracy.
These outcomes underscore MLO-MAE's strong capability in learning effective visual representations across datasets of varying scales, from the large-scale ImageNet to the smaller-sized CIFAR datascts. The supe- riority of MLO-MAE over MAE and U-MAE stems from its advanced masking strategy that selectively targets informative patches, a significant enhancement over the indiscriminate, random masking approach of the two baselines. Furthermore, ML.O-MAE, surpasses AutoMAE, and SemMAF by integrating feedback from downstream classification tasks into its masking process. This dynamic adaptation contrasts with the static masking strategies of AutoMAE and SemMAE, which do not account for the specific requirements of downstream tasks, limiting their effectiveness.
Linear probing results. Table 2 shows the linear probing results. Our method MLO-MAE demonstrates superior performance compared to basclines across various datasets. Specifically, on ImageNet, MLO-MAE achieves accuracy of 70.2%, substantially surpassing MAF's accuracy of 55.4% and U-MAE's 58.5%.
7

---

Published in Transactions on Machine Learning Research (03/2025)
Table 5: Object detection result of MLO-MAE and baselines on MS-COCO detection task.
Method
AutoMAE
MAE
MLO-MAE
Apbox (%)
50.5
50.3
51.1
Similarly, on the CIFAR-100 datasct, MLO-MAE continues to outperform, attaining an accuracy of 63.8%, significantly higher than the 46.6% accuracy of MAE and 50.4% accuracy of U-MAE.
The superiority of MLO-MAE compared to baseline methods stems from its unique approach integrating the pretraining of the image encoder and linear probing in seamless, end-to-end workflow. Specifically, MLO-MAE conducts pret.rining al the first level and linear probing the second level within unified framework. This integration allows the linear probing performance, evaluated at the third level, to directly inform and enhance the pretraining process. Consequently, this leads to a pretrained encoder that is more effectively tailored for the downstream linear probing task. In contrast, baseline methods handle pretraining and linear probing as distinct, separate stages, where the performance of linear probing does not impact or contribute to the pretraining phase.
4.3 Transfer learning
4.3.1
Fine-grained image classification
In MLO-MAE, the masking network is trained using a specific downstream classification dataset, raising concerns about potential overfitting and limited generalizability to other datasets. To address this, we performed transfer learning experiments on labeled fine-grained classification datasets including iNaturalist 2019, CUB-200-2011, and Stanford Cars. Classification accuracy results, as presented in Table 3, show that MLO-MAE surpasses all baselines across these datasets. This indicates that the masking network, as learned by MLO-MAE for a particular downstream classification datasct, is capable of generalizing its effectivencss to additional datasels, rather than being overly tailored to that specific downstream classification dalaset.
4.3.2 Semantic segmentation and object detection
We also explored the transferability of the masking network, initially learned through a downstream classi- fication task, to other tasks including semantic segmentation and object detection. Given a ViT-B model pretrained on ImageNet using MLO-MAE or a bascline and subsequently fine-tuned on ImageNet, to transfer it for semantic segmentation, we integrated as a backbone model into the UPerNet. (Xiao el al., 2018) semantic segmentation framework. It was then further fine-tuned on the challenging ADE20K dataset (Zhou et al., 2017) containing 25K images spanning 150 semantic categories. The fine-tuning was conducted by the AdamW optimizer for over 160,000 iterations, with batch size of 8 and learning rale of 0.0001. In Table 4, MLOOMAA. showcases significant improvement over the baselines. Specifically, MLOOMMA, attains enhancements in mean Intersection over Union (mloU) by margins of 3.7% and 3.4% when compared to these baselines. This performance highlights the capability of MLO-MAE in executing dense prediction tasks.
Following MAE, we also adapt ML.O-MAE pretrained VIT-B model for the use of an FPN backbone in Mask R-CNN As shown in Table 5, MLO-MAE pretrained backbone model performs better than all baselines (51.1 comparing to 50.5 and 50.3, Apbox). We did not inchude SemMAE and U-MAE as they did not report on MS-COCO detection. Due to space limits, we defer the results of object detection on PASCAL VOC 2007 to Appendix D.1.
4.4
Continued pretraining
We further investigated the effectiveness of MLO-MAE in continued pretraining setting. Starting with ViT-B model pretrained by MAE on ImageNet, we applicd MLO-MAE pretraining to 2 datascts, PDDB (Barbedo et al., 2018) and PAD-UFES (Pacheco et. al., 2020). Table 6 compares the test accu- racy (%) across three settings: (1) no continued pretraining, where the model is directly fine-tuned using labels; (2) continued pretraining on target dataset using MAE, followed by fine-tuning: (3) MLO-MAE pre- training solely on target dataset, followed by fine-tuning; and (4) continued pretraining on target dataset using MT.O-MAF, followed by fine-tuning. Our results show that continued pretraining with MLO-MME.
8

---

Published in Transactions on Machine Learning Research (03/2025)
Table 6: Continued pretraining on PDDB and PAD-UFES. All settings are initialized with weights of ViT-B
pretrtained by MAE on ImageNet. Runtime measured in GPU hours on A100.
PDDB
PAD-UFES
Method
Acc(%)
Runtime
Acc(%)
Runtime
No continued pretraining
88.6
2132
75.0
2132
MAE continued pretraining
89.3
2132 + 39
75.4
2132 + 3
MLO- MLO-MAE pretraining (no IN pretrain)
90.2
2132
+ 37
75.9
2132 + 3
MLO-MAE continued pretraining
92.7
2132
37
77.6
2132 + 3
83.8
80.7 83 88.77
84.3
Ln. probing (%)
84
Lin. probing aco (%)
80
77.3
78.3
74.5
12.1
70
82
81.5
60
96.5
20
30
40
60
70
80
90
MLO-MAE
BLO-MAE
Masking ratio
Figure 2: Comparison of Figure 3: Impact of different masking ratios linear probing perfor- BLO-MAE and MLO-MAE CIFAR-10.
mance CIFAR-10.
significantly outperforms both MAE-based pretraining and no protraining in both datascts. This highlights the practical advantage of MLO-MAE: by leveraging MAE pretraining once, subsequent tasks can benefit from fast, efficient continued pretraining with MLO-MAF, delivering substantial performance gains without requiring the computationally expensive process of large-scale pretraining from scratch for each downstream task.
4.5 Ablation studies
Reduction to two levels. To investigate the importance of maintaining three levels in the MLO-MAE framework, we simplified it to two levels, by combining the first and second levels, leading to the following bi-level optimization (BLO) problem (referred to as BLO-MAE):
min Leto (T),, T
ET),, D"," arggiin (D"; E.D.C (x(()))r
Σ
XED.
where y is a tradeoff parameter. Here, the image encoder is trained using a multi-task learning strategy. which involves minimizing the weighted sum of the pretraining loss and the downstream classification loss. Figure 2 presents the results. The BLO-MAE method leads to notable decrease in accuracy by 2.8% compared to MLO-MAE. This highlights the significance of employing three-stage process
two-stage In the BLO-MAE approach, the lower level addresses multi-task learning challenge by optimizing weighted sum of losses from two distinct tasks. This scenario often lcads to task competition, where minimizing the loss for one lask inadvertenty increases the loss for the other. Balancing these competing losses requires meticulous adjustment of the tradeoff parameter 7, a process that both challenging and timo-consuming. In contrast, our MLO-MAE method tackles this challenge through sequential process integrated into an end-to-end framework. Initially, the method involves pretraining the encoder. Following this, the pretrained encoder transitioned to the next stage, where the classification head is trained. The

---

Published in Transactions on Machine Learning Research (03/2025)
Table 7: Linear probing accuracy (%) on CIFAR- 10 under different unrolling steps.
Table 8: Linear probing accuracy (%) on CIFAR- 10 under different patch sizes.
Unrolling steps
2
5
Patch size
2
4
8
Accuracy (%)
83.8
84.3
84.6
Accuracy (%)
84.3
82.1
81.9
Table 9: Mean intersection over Union (MIoU) in using semantic segmentation as the downstream task on ADE20K
Method (mloU)
MAE
MLO-MAE (Ours)
ADE20K
28.8
29.1
pretraining task in MLOOMMF, aids the classification task by providing an effective image encoder, instead of competing with the classification task.
Unrolling steps. In this study, we explored how the number of unrolling steps in the Optimization Algorithm, as detailed in Section 3.2, affects the final performance. The study was performed on CTFAR- 10, with linear probing as the evaluation protocol. Table 7 shows linear probing accuracy under different unrolling steps. The results reveal that an increase in the unrolling steps leads to a gradual improvement in accuracy. This enhancement can be altributed to the fact that higher number of unrolling steps allows for more frequent updates to the image encoder weights (with more iterations in Stage 1), prior to any updates being made to the classification head and masking network. Consequently, this yields more refined gradient estimation for the parameters of the classification head and masking network, as these estimations are based on the image encoder that has undergone more extensive training. However, it is important to note that increasing the number of unrolling steps also brings in a notable computational overhead. In this context, our observations indicate that while increasing the unrolling step from one to two leads to a 0.5% boost in CIFAR-10 linear probing performance, the gain diminishes to just 0.3% when the unrolling steps are further raised from two to five. This suggests diminishing return in performance improvement relative to the increased computational demand.
Masking ratios. We studied how the masking ratio during pretraining affects downstream task perfor- mance, using the CIFAR-10 dataset and a linear probing protocol. The results, illustrated in Figure 3, reveal that intermediate masking ratios deliver optimal performance. When the masking ratio is too low, the reconstruction task becomes overly simple, failing to push the image encoder to develop robust represen- tations. Conversely, an excessively high masking ratio makes the task overly challenging, which also impedes the learning of effective representations. Notably, our method demonstrates resilience to changes in masking ratios ranging from 60% to 80%. Within this interval, there is minimal fluctuation in the results of linear probing. Additional experiment on dynamic mask ratio can be found in Appendix D.4.
Patch sizes. In this study, we investigated the impact of different image patch sizes on the performance of our method when applied to the CIFAR-10 dataset, utilizing the linear probing evaluation protocol. Our experiments focused on three patch sizes: 2 x 2, 4 x 4, and 8 x 8. The results, presented in Table 8, show that the 2 x 2 patch size achieves a linear probing accuracy of 84.3%, which is 2.4% (absolute) higher than that obtained with the larger 8 8 patch size. These findings suggest that MLO-MAE is more effective when employing a larger number of smaller patches, particularly for small images, such those in the CIFAR-10 dataset with size of 32 x 32. The reason is that smaller patches allow for more precise and detailed candidate masks. This leads to better feature representation learning. Moreover, smaller patches provide a finer grid over the image, allowing the model to capture more detailed and subtle features. This is particularly beneficial for small images, where each pixel can carry significant. information.
Semantic sogmentation as downstream feedback. In the basc MLO-MAE framework, image classi- fication was used as the downstream task to guide the masking network. In this section, we evaluate the generalizability of MLO-MAE by adapting semantic segmentation as the downstream task in Stages II and 111. Specifically, we adopted the ADE20K dataset for both pretraining and semantic segmentation evaluation. The results, summarized in Table 9, demonstrate that when semantic segmentation is used as the down- stream feedback, MT.O-MAF. consistentlyy outperforms the baseline. This highlights that the affectivenese of MLO-MAE is not limited to a specific downstream task, underscoring its versatile and robust design.
10

---

Published in Transactions on Machine Learning Research (03/2025)
Table 10: Pretraining time (GPU hours measured on A100).
MAE
SemMAE
AutoMAE
MLO-MAE (Ours)
Ima 2132 hrs
1154 hrs
1344 hrs
1083 hrs
4.6 Computational costs
Although MT.O-MAE introduces additional computational overhead due to its multi-level optimization, this is balanced by its lower epoch requirement for achieving convergence. In contrast to the 800 epochs needed for standard MAE and SemMAE, MLO-MAE converges in just 50 epochs. This siguificant reduction in the number of epochs effectively offsets the increased computational demands. The total GPU hours (on Nvidia A100 GPU) for MLO-MAE on ImageNet, as shown in Table 10, amount to 1083. This number less than those baselines while our method achieves better test accuracy than baselines as shown Tables and
5 Conclusion
In this paper, we proposed MLO-MAE, a method that automatically learns an optimal masking strategy in Masked Autocncoder (MAE) by leveraging feedback from downstrcam tasks. Unlike the vanilla MAE which applies uniform patch masking irrespective of their informativeness, MLO-MAE adaptively concen- trates more informative image regions. Different from other MAE methods that do not utilize feedback from downstream tasks for masking-strategy optimization, MLO-MAE uniquely capitalizes such feed- back to refine its masking approach. Our experiments across various datasets demonstrate that MLO-MAE outperforms MAE baselines by learning downstream task guided masking strategies. Furthermore, the rep- resentations generated by MLO-MAE exhibit high transfcrability to a range of downstream tasks, including fine-grained classification, semantic segmentation, and object detection, highlighting our method's versatility and effectiveness.
6 Broader impact
Our paper presents the Multi-level Optimization for Masked Autoencoders (MTO--AAE) framework, enhanc- ing self-supervised learning in visual data processing, MLO-MAE can have a broad impact across various fields like medical imaging, autonomous vehicles, and content moderation. In the healthcare sector, it could improve disease detection and diagnosis. Similarly, in the realm of autonomous driving, it may enhance object recognition for safer vehicle automation. However, we also recognize the ethical considerations and potential risks associated with the application of our work. The increased capability of image processing models can lead to concerns around privacy, surveillance, and the potential misuse of technology in unau- thorized or harmful ways. The deployment of such technologies must be guided by ethical principles and regulatory frameworks to protect individual privacy and prevent misusc.
7 Limitations
Our multi-level optimization framework introduces two main limitations. First, the hypergradient computa- tion incurs significant computational overhead. We mitigate this by using gradient approximation, where we control the approximation level through the number of unrolling steps. This trade-off balances computational cost with optimization accuracy. Second, our approach increases GPU memory usage due to the inclusion of classification head and the need to load three separate training datasets sinultaneously. Although our current implementation does not explicitly address this, techniques like gradient accumulation could reduce memory demands. Future work can focus on developing more efficient multi-level optimization algorithms to systematically overcome these limitations.
Acknowledgments
This research was supported by NSF TIS2405974 and NSF IIS2339216.
11

---

Published in Transactions on Machine Learning Research (03/2025)
References
Randall Balestriero, Mark Tbrahim, Vlad Sobal, Ari Morcos, Shashank Shekhar, Tom Goldstein, Florian Bordes, Adrien Bardes, Gregoire Mialon, Yuandong Tian, et al. A cookbook of self-supervised learning. arXiv preprinl Xrx::22304. 12210, 2023.
Fan Bao, Guoqiang Wu, Chongxuan Li, Jun Zhu, and Bo Zhang. Stability and generalization of bilevel programming in hyperparameter optimization. Advances in neural informalion processing systems, 34: 4529 4541, 2021.
Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. BETT: BERT pre-training of image transformers. Tn International Conference on Learning Representations, 2022. URL https:://ppenrevior wtt.://oiiw.ottoooonnniiiti p-BhZS25904.
Jayme Garcia Arnal Barbedo, Luciano Vieira Koenigkan, Bernardo Almeida Halfeld-Vicira, Rodrigo Veras Costa, Katia Lima Nechet, Claudia Vieira Godoy, Murillo Lobo Junior, Flavia Rodrigues Alves Patricio, Viviane Talamini, Luiz Gonzaga Chitarra, et al. Annotated plant pathology databases for image-based detection and recognition of discascs. IEEE Latin America Transactions, 16(6):1749-1777, 2018.
Atilim Gunes Baydin, Robert Cornish, David Martinez Rubio, Mark Schmidt, and Frank Wood. Online learning rate adaptation with hypergradient descent. arXiv preprinl arvi: vrXXi:1703.000 703.7 04782, 2017.
Han Cai, Ligeng Zhu, and Song Han. Proxylcssnas: Direct noural architecture search on target task and hardware. In ICLR, 2019.
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.
2021.
Haijian Chen, Wendong Zhang, Yunbo Wang, and Xiaokang Yang. Improving masked autoencoders by learning where to mask. arXiv preprint arXiv:2303.06580 2023.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu,, Heewoo Jun, David Luan, and Nya Sutskever. Genera- tive pretraining from pixels. In Hal Dauné III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedinggs of Machine Learning Research, pp. 1691-1703. PMLR, 13- 18 Jul 2020. URL https: httpc:///ooceeeeeeeee mlr. ppess/vv19/ chenn0s html.
Sang Keun Choe, Willie Neiswanger, Pengtao Xie, and Eric Xing. Betty: An automatic differentiation library for multilevel optimization. arXiv preprint arXi::2207.088/1 2022.
Sang Keun Choe, Sanket Vaibhav Mehta, Hwijeen Ahn, Willie Neiswanger, Pengtao Xie, Fmma Strubell, and Eric Xing. Making scalable meta learning practical. arXiv preprint arXiv:2310.0661, 2023.
MMSSggentation Contributors. MMSegmentation:: Openmmlab semantic segmentation toolbox and bench- mark. https:///github.commm http://tttt...ccooooooaaaaaaaaaa
2020.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hicrarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255. Icee, 2009.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. orXiv preprint 1810: 04805, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Un- terthiner, Mostafa. Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth Transformers for image recognition scale. arXiv preprint arXic:2010. 11929, 2020.
16x16 words:
M. Everingham, L. Van Gool, C. K. Williams, J. Winn, and A. Zisserman. The PASCAL Visual Object Classes Challenge 2007 (VOC2007) Results. network.org/challenggg..
http://www.pptt
hwtp:///woowoo//oooooo workshoppindexx.hmm html.
12

---

Published in Transactions on Machine Learning Research (03/2025)
Matthias Feurer, Jost Tobias Springenberg, and Frank Hutter. Initializing bayesian hyperparameter opti- mization via meta-learning. In Proceedings of the Twenty-Ninth AAAT Conference on Artificial Intelli- gence, AAAF15, pp. 1128 1135. AAAI Press, 2015. ISBN 0262511290.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pp. 1126 1135. PMLR, 2017.
Bhanu Garg, Li Lyna Zhang, Pradyumna Sridhara, Ramtin Hosseini, Eric Xing, and Pengtao Xie. Learning from mistakes framework for neural architecture search. In AAAI Conference on Artificial Intelligence, 2021.
Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep leedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249- 256. JMLR Workshop and Conference Proceedings, 2010.
Riccardo Grazzi, Luca Franceschi, Massimiliano Pontil, and Saverio Salzo. On the iteration complexity of hypergradient computation. In International Conference on Machine Learning, pp. 3748-3758. PMLR, 2020.
Jie Gui, Tuo Chen, Qiong Cao, Zhenan Sun, Hao Luo, and Dacheng Tao. A survey of self-supervised learning from multiple perspectives: Algorithms,, theory, applications and future trends. arXio preprinl 20X1:23001.05712, 2023.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 16000 16009, 2022.
Ramtin Hosseini and Pengtao Xie. Saliency-awand aware neural architecture search. Advances in Newral Information Processing Systems, 35:14743 14757, 2022.
Ramtin Hosseini, Xingyi Yang, and Pengtao Xie. Dsrna: Differentiable search of robust neural architectures. In Proceedings of the IEEE/CVI Conference on Compuler Vision and Pallern Recognilion, pp. 6196 -200, 2021.
Ramtin Hosseini, Li Zhang, Bhanu Garg, and Pengtao Xie. Fair and accurate decision making through group-aware learning. In International Conference on Machine Learning, pp. 13254 13269. PMLR, 2023.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:161.00114,, 2016.
Kaiyi Ji, Junjie Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. International conference on machine learning, pp. 4882 4892. PMLR, 2021.
loannis Kakogeorgiou, Spyros Gidaris, Bill Psomas, Yannis Avrithis Andrei Bursuc, Konstantinos Karantza- los, and Nikos Komodakis. What to hide from your students: Attention-guided masked image modeling. In European Conference on Computer Vision, pp. 300-31 318. Springer, 2022.
Xiangwen Kong and Xiangyu Zhang. Understanding masked image modeling via learning occlusion invariant feature. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 6241-6251, 2023.
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained catego- rization. In Proceedings of the IEEE international conference computer vision workshops, pp. 554-561, 2013.
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-10 (canadian institute for advanced rescarch). URL http://wwe /www cs. wwwottorooooooomo httatt///a//.kkii/iifaaa.....
Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. Cifar-100 (canadian institute for advanced research). URL http://wwee cs. www.onto.eeu..kriii/../ html.
13

---

Published in Transactions on Machine Learning Research (03/2025)
Gang Li, Heliang Zheng, Daqing Liu, Chaoyue Wang, Bing Su, and Changwen Zheng. Semmae: Semantic- guided masking for learning masked autoencoders. Advances in Newral Information Processing Systems. 35:14290 14302, 2022.
Zhaowen Li, Zhiyang Chen, Fan Yang. Wei Li, Yousong Zhu, Chaoyang Zhao, Rui Dong, Liwei Wu, Rui Zhao, Ming Tang, et al. Mst: Masked self-supervised transformer for visual representation. Advances in Newral Information Processing Systems, 14-13155-3317 2021.
Hanwen Liang, Shifeng Zhang, Jiacheng Sun, Xingqiu He, Weiran Huang, Kechen Zhuang, and Zhenguo Li. Darts+: Improved differentiable architecture search with early stopping. arXiv preprint orXiv:1909.06000 2019.
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 EisionEECVV 2014: 18th European Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13, pp. 750-755. Springer, 2014.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv:1806.0900 2018.
Yinhan Liu, Myle Out, 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 xix: 1 11692, 2019.
Zhengqi Liu, Jie Gui, and Hao Luo. Good helper is around you: Attention-driven masked image modeling, In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 1799-1807, 2023.
lya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ar preprint arviv 111.11111100011 2017.
Neelu Madan, Nicolae-Cätälinn Ristea, Kamal Nasrollahi, Thomas B Moeslund, and Radu Tudor Ionescu. Cl-mae: Curriculum-learee masked autoencoders. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision, pp. 2492-2502, 2024.
AGC Pacheco, GR Lima, AS Salomao, В Krohling, IP Biral, GG de Angelo, FCR Alves Jr, JGM Esgario, AC Simora, PBC Castro, et al. Pad-ufes-20: skin lesion datasct composed of patient data and clinical images collected from smartphones. data brief 32: 106221, 2020.
Sebastian Prillo and Julian Eisenschlos. Softsort: A continuous relaxation for the argsort operator. International Conference on Machine Learning, pp. 7793-7802 PMLR, 2020.
Parth Sheth, Yueyu Jiang, and Pengtao Xie. Learning by teaching, with application to neural architecture search. arXiv preprint arXic:2103.0 20000009,,, 2021.
Yuge Shi, N Siddharth, Philip Torr, and Adam R Kosiorek. Adversarial masking for sule-vie learning. International Conference on Machine Learning, pp. 20026 20040. PMLR, 2022.
Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pietro Perona, and Serge Belongie. The inaturalist species classification and detection dataset. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 8769 8778, 2018.
Luis N. Vicente and Paul H. Calamai. Bilevel and multilevel programming: bibliography review, 1994.
Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The caltech-ncsd birds- 200-2011 datasct. 2011.
Haochen Wang, Kaiyou Song, Junsong Fan, Yuxi Wang, Jin Xie, and Zhaoxiang Zhang. Hard patches mining for masked image modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pallern Recognition, pp. 10375- 10385, 2023.
14

---

Published in Transactions on Machine Learning Research (03/2025)
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-434, 2018.
Sirui Xie, Hehui Zheng, Chunxiao Liu, and Liang Lin. SNAS: stochastic neural architecture search. In ICLR, 2019.
Yuhui Xu, Lingxi Xie, Xiaopeng Zhang, Xin Chen, Guo-Jun Qi, Qi Tian, and Hongkai Xiong. PC-DARTS: partial channel connections for memory-effiiint architecture search. In JCLR, 2020.
Qi Zhang, Yifei Wang, and Yisen Wang. How mask matters: Towards theoretical understandings of masked autoencoders. Advances in Neural Information Processing Systems, 35:27127 27139, 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 computer vision and pattern recognition. pp. 633-644 2017.
15

---

Published in Transactions on Machine Learning Research (03/2025)
A Optimization Algorithm
We develop an efficient optimization algorithm to solve the MLO-MAE problem demonstrated in Figure 1. Notations are given in Table 11.
Table 11: Descriptions of notations used in optimization algorithm in Appendix A.
Notation
Description
E
Backbone encoder that takes input patches and generates representation embedding
D
Backbone decoder that takes in representation embedding and generates the reconstructed image
C
Classification head
T
Masking network
D
Image classification dataset
Dir
Train set split based on
D
Validation set split. hased D
Da
Unlabeled dataset based on D
X
An arbitrary image from
M(.)
Masked image
Erec())
MAE image reconstruction loss
(.)(.)
Cross entropy image classification loss
a()
Sigmoid activation function that produces masking probability given our masking network
M
Learning rate for updating the encoder
nC
Learning rate for updating the classification head
ar
Learning rate for updating the the masking network
Masking ratio
A.1 Implicit Differentiation for Gradient Computation
In the MLO-MAE framework, we adopt implicit differentiation as key tool for computing gradients in scenarios characterized by complex, nested optimization structures. This approach is particularly effective when dealing with variables implicitly interconnected. To solve the MLO-MAF, optimization problem, we utilize robust algorithm first introduced in (Choc et al., 2022). This algorithm is underpinned by solid theoretical framework, and its convergence properties have been extensively examined in recent scholarly contributions. Each stage the optimization process necessitates the identification of an optimal solution, denoted with an asterisk (*) and positioned on the left-hand side of the equation. Computing this exact optimal solution typically resource- insouree-intensive. To manage this efficiently, we apply the strategy proposed in (Liu et al., 2018), which involves approximating the optimal solution via one-step gradient descent update. This approximation is then integrated into the next level of the optimization process. In our analysis, the d. symbol a. is used to represent partial derivatives, while d- d. significs ordinary derivatives. The term V fXX,Y) denotes the sccond-order partial derivative of f(X,Y) with respect to Y and X, formalized as F(x.Y. The axay
initial phase of our methodology involves approximating E (T) as follows:
Nxt Σ XED.
(' TE VE Σ
(6)
where WE denotes the learning rate. Subscquently, E' substituted into (E,) Ct") to yield an approximated objective function. Similarly, C* (T)( is approximated using single-step gradient descent with respect. to this approximated objective:
16

---

Published in Transactions on Machine Learning Research (03/2025)
(( C VcLa, ('()))
(7)
Finally, (T) and (T)) are incorporated into Lato (T)) Deal) to obtain an approximated version of the objective function. The parameter is then updated using gradient descent:
T T---mp 1T VrLas(D"m (DD,
(8)
By applying the chain rule to this approximation, we obtain:
dCodl
aC δ
dc 8C' dE' OC' dF' or
(9)
OT
O OT
where:
OE'
Nxr
0T
-
V
E,T
Σ XED
Σ
(10)
and
OC"
= nJc Cr
Ltr
(11)
Here, D
Leds ( C(E((T)))) and L. = (D"
A.2 Finite Difference Approximation for Gradient Estimation
To reduce the complexity of solving MLO-MAE, we utilize the Finite Difference Approximation (FDA) method, particularly in cstimating gradients where analytical differentiation is challenging. Specifically, directly computing Jacobian vector multiplication with MLO problems computationally expensive, which can be efficiently approximated by FDA methods (Choe al., 2022).
FDA approximates the gradient of function by computing the change in the function value for small perturbation in the input. For example, for function f(x), the gradient approximation is given by:
f(x + 6x) f()
V f((x) sJ
(12)
dx
In our experimental framework, extended the formula above to compute hypergradients in MLO problems. Specifically, incorporated the bi-directional finite difference approximation (FDA) method for gradient istitinion within complex, nested optimization contexts. This technique is particularly pertinent in scenarios where traditional analytical gradient computation is either impractical or excessively resource-intensive.
The bi-directional FDA extends the conventional finite difference approach by introducing perturbations both positive and negative directions relative the current parameter values. This methodology provides a more nuanced and accurate gradient estimation compared to one-sided finite difference methods.
Each parameter in current optimization problem undergoes initial positive perturbation, determined by predefined small step size c. Post this perturbation, we compute the loss and the corresponding gradients with respect to the preceding optimization problem's parameters. subsequent negative perturbation, amounting to double the initial epsilon, shifts the parameters below their original values for a re-evaluation of the loss and re-calculation of gradients.
In our Multi-level Optimization (MLO) framework, the bi-directional FDA offers substantial benefits. It enables efficient gradient estimation in situations where conventional backpropagation may fail, especially in
17

---

Published in Transactions on Machine Learning Research (03/2025)
handling the complex, implicit dependencies between parameters. This approach is instrumental in enhancing
our optimization techniques within intricate MT.0 settings.
Our implementation of the bi-directional FDA is finely tuned to strike a balance between computational efficiency and the accuracy of gradient estimation. The selection of the epsilon value is critical, aiming minimize numerical instability while maintaining sensitivity to changes in parameters. This balance is essential for ensuring the robustness of the gradient estimation process in the stochastic realm of machine learning models.
A.3 Integration of Methods
Implicit differentiation and finite difference approxination are integrated to balance theoretical accuracy with computational feasibility. This combination enhances the robustness and efficiency of our optimization process in the MLOO-AA. framework. In the application of these optimization methods, several key consid- erations are taken into account. Firstly, the choice AT in the finite difference approximation a critical factor, as it directly influences the accuracy and stability of the gradient estimation. An appropriate value for AT' ensures balance between precision and numerical stability. Secondly, we address the computational complexity problem inherent in implicit differentiation. This aspect is particularly relevant. in deep network architectures, where computational resources be a limiting factor. To mitigate this, optimize the use of implicit differentiation to balance computational demands with the need for accurate gradient computa- tion. Tastly, maintaining numerical stability is paramount in both methods. Techniques such as gradient normalization and careful arithmetic handling are employed to ensure that the computations remain stable, especially scenarios where small numerical errors can significantly impact the overall results. This compre- hensive optimization algorithm pivotal in enabling the efficient training and validation of the MLO-MAE framework.
Algorithm MLO-MAE Optimization Algorithm
1: Input: Training dataset Drr validation dataset Deat
2: Output: Optimized parameters (T),,, (T),), (F T
3: procedure MAE (RETRAINING(P.s
4: Initialize encoder E and decoder for each image X, E Du do Patchify image into X; Compute masked image using (MM,(X;I,))
Update
E,
D
by
minimizing
Σ Σ x) XED ==1
end for
10: return (T),), (T)
11: end procedure
12: procedure CLASSIPICATION HEAD TRAINING (F),), D.,)
13: Freeze (T), initialize classifier C
14: for each image E do
15: Update by minimizing Leta (Dr;
16: end for
17: return (T))
18: end procedure
19: procedure VALIDATION OPTIMIZATION( (ET), D
20:
Freeze ETTT), (E)
21: Optimize T by minimizing Lass (Dr vat: E* (*
22: return T
23: end procedure
18

---

Published in Transactions on Machine Learning Research (03/2025)
A.4 Differentiability of MLO-MAE Framework
In the initial phase of our Multi-level Optimization Masked Autoencoder (MLO-MAE) framework, the focus is on the reconstruction loss and its differentiability relative to the parameters of the learnable masking network. This network is pivotal, as it determines the masking patterns for the input data, directly impact- ing the autoencoder's reconstruction loss. The key to effective gradient-based optimization lies in ensuring the differentiability of this reconstruction loss with respect to the masking network's parameters. Due to the inherently discrete nature of mask selection, integrating the masking network directly into the MAE's reconstruction loss initially leads to non-differentiability issues. To circumvent this, our approach employs a sigmoid activation function, denoted as o)); to generale soft. masks. These soft masks assign a continuous value between 0 and to each image patch, indicating the likelihood of that patch being masked. This likeli- hood is learned from both masked and unmasked patches. In this first phase, we tackle non-diffcrentiabiii by utilizing the MAE reconstruction loss logether with the masking probability. This is achieved as illus- trated in Equation (1), which allows the proportional contribution of each patch to influence the overall loss, facilitating differentiation with respect to the network parameters through the application of a(.).
To further address the non-diiferentiaablityy issue in our MLO-MAE framework incurred by mask selec- tion, the SoftSort (Prillo Fisenschlos, 2020) technique a differentiable approximation of sorting opera- tions replacing the conventional, non-differentiable "argsort' operation, can be integrated. SoftSort enables the learning of contimuous relaxation of sorting, vital for gradient-based optimization and for crafting nuanced, performancc--enhancing masks. This represents a significant leap in refining the effectiveness of the reconstruction process. However, given the computational demands of SoftSort, we propose a simpli- fied approach that approximates the base Jacobian with an identily matrix similar to (Choe et al., 2023), thereby simplifying the gradient during backpropagation and enhancing the efficiency by 'jumping over' the non-differentiable argsort operation. By treating non-differentiable operations as identity functions during backpropagation, this method allows gradients from the reconstruction loss to flow as if the masking op- erations were inherently differentiable. This strategy significantly speeds up training by enabling the use of standard gradient- based optimization techniques. In our MLO-MAE, the masking network employs soft masks, which, through a sigmoid activation function o(),, assign each image patch a continuous value from O to 1. This assignment reflects the probability of masking, informed by both masked and unmasked patches. By approximating the gradients for dinddffferentabble functions as identity functions during backpropaga- tion, our method enables the differentiation of the MAE reconstruction loss, Lree; relative to the masking network parameters T. This differentiation is facilitated by incorporating the output probabilities of (.)) into the reconstruction loss Lree; detailed in Equation (1)-
A.5 Convergence Property of the MLO-MAE framework
Previous research has extensively analyzed the convergence properties of optimization algorithm for bi-level optimization (BLO) problems. For instance, Ji et al. (2021) establish sharper non-asymptotic convergence rates for implicit differentiation (AID) and iterative differentiation (TTD) methods, addressing both deter- ministic and stochastic bilevel problems with improved computational complexity. Grazzi et al. (2020) ana- lyze the iteration complexity of hypergradient computation, providing explicit non-asymptotic convergence bounds and demonstrating the efficiency of AID with conjugate gradient under contraction assumptions. Bao et al. (2021) focus on the stability and convergence of bilevel optimization by deriving expectation bounds for the generalization gap, emphasizing how regularization stratogics improve convergence. Bao et al. (2021) ad- dress bilevel problems with nonconvex lower-level structures, introducing the IAPTT-GM framework, which ensures convergence even in the absence of lewer-level convexity. Zhang et al. (2021) propose convergence- guaranteed approach for bilevel optimization using stochastic implicit gradients, ensuring stationary point convergence while reducing computational overhead.
While the analyses in the aforementioned works are primarily on bi-level optimization problems, the conver- gence proof can be readily extended to the multi-level cases. In addition our emipirical results demonstrate that MLO-MAE framework consistently converges and outperforms the baseline methods as shown in Sec- tion 4. The training loss curves in Fig 4 also corroborate with our empirical convergence finding.
19

---

Published in Transactions on Machine Learning Research (03/2025)
Edita
(a) Masking Loss
(b) Finetune Loss
(c) Pretrain Loss
Figure 4: Loss curves for masking, fine-tuning, and pretraining.
B Detailed Experimental Settings
B.1
Datasets
CIFAR-10 (Krizhevsky et al., a): CIFAR-10 is a fundamental dataset for image classification, com- prising 60,000 32x32 color images across 10 classes, with 6,000 images per class. It is widely used in machine learning research because of its manageable size and diversity of images. CIFAR-10 tests the ability of our MLO-MAE framework to capture essential features in small-scale images and generalize across a variety of everyday objects.
CIFAR-100 (Krizhevsky et al., b): CIFAR-100 is similar to CIFAR-10 in image size and total number of images but significantly more challenging due to its 100 classes, each containing 600 images. The increased number of classes in CIFAR-100 allows us to evaluate the capability of MLO-MAE in more granular classification context, providing insights into how well the model differentiates between larger number of categories with fewer examples per category.
ImageNet-1K (Deng et al., 2009): ImageNet, subset of the larger ImageNet databaso, is one of the most influential datasets in the field of image classification. It contains 1.3M training images and 50K validation images categorized into 1,000 classes. The dataset's extensive size and diversity present a rigorous test for any machine learning model. Our use of ImageNet is aimed at assessing the scalability and robustness of the MLO-MAE framework in handling complex, largo-scale classification tasks.
CUB-200-2011 (Wah et al., 2011): The CUB-200 11-B-20-211 200-2011 dataset is a specialized collection designed for fino-grained bird species classification, developed by the California Institute of Technology. It contains 11,788 images, representing 200 bird species, with focus on North American birds. Each image in the dataset is accompanied by detailed annotations, including bounding baxes, part locations, and attribute labels, making it ideal for detailed image analysis tasks. The dataset is widely used in computer vision research, particularly in tasks that require distinguishing between visually similar sub-categorics.
Stanford Cars (Krause et al., 2013): The Stanford Cars dataset is a large collection of car images created by rescarchers at Stanford University, it contains 16,185 images of 196 classes of cars. The data is split into 8,144 training images and 8,041 testing images, where each class has been split roughly in 50-50 split. Classes are typically at the level of Make, Model, Year, ex. 2012 Tesla Model S or 2012 BMW M3 coupe.
iNaturalist 2019 (Van Horn et al., 2018): The NNturraist 2019 datasot is a comprchensive biodiversity collection used for machine learning and research, created by the iNaturalist project, a collaboration bet the California Academy of Sciences and the National Geographic Society. It features over 859,000 high-quality images of 1,010 species, cach with detailed metadata including species name, observation location, date, and time. This extensive dataset is crucial for species identification and, distribution modeling, and is utilized in the annual iNaturalist Challenge to enhance automated species recognition technologies.
20

---

Published in Transactions on Machine Learning Research (03/2025)
ADE20K (Zhou et al., 2017): ADE20K is comprehensive dataset for semantic segmentation, contain- ing more than 20,000 images annotated for variety of scenes and objects, making it one of the most diverse datasets available for this task. Each image in ADE20K is annotated with pixel-level segmentation masks, encompassing a wide range of objects and scene categorics. The complexity and richncss of ADE20K make it an ideal choice for testing the efficacy of the MTO-MAE framework understanding and segmenting complex visual scenes. The dataset challenges the model to not only recognize a diverse array of objects but also understand their spatial relationships and boundaries within various contexts.
MS-COCO (Lin et al., 2014): MS-COCO is a comprehensive largo-scale dataset used for tasks such as object detection, segmentation, keypoint detection, and image captioning, containing a total of 328,000 images. Initially released in 2014, the dataset was split into 83,000 training images, 41,000 validation images, and 41,000 test images. In 2015, an expanded test set was introduced, adding 40,000 new images to the existing test set for total of 81,000 test images. Responding to feedback from the research community, the 2017 version adjusted the training/ validation split to 118,000 training images and 5,000 validation images, while maintaining the same images and annotations as previous versions. The 2017 test set consists of 41,000 images, a subset of the 2015 test set, and the release also includes a new unannotated dataset of 123,000 images.
wraining/validation
PASCAL VOC 2007 (Everingham et al.): The PASCAL VOC 2007 dataset is a prominent bench- mark in the field of object detection and image segmentation. It consists of 9,963 images annotated with 24,640 objects across 20 distinct classes, including everyday items like cars, cats, and chairs. The dataset provides comprehensive bounding box annotations for each object, facilitating the training and evaluation of object detection models. Additionally, it includes segmentation annotations, although they are more limited compared to later versions. The images in PASCAL VOC 2007 vary in size and aspect ratio, reflecting diverse real-world conditions. This dataset is divided into training, validation, and test. sets, with the test set publicly available, making it an invaluable resource for benchmarking and comparing the performance of different object detection algorithms. Its widespread adoption in the research community has contributed to significant advancements in computer vision, serving as a foundational dataset for evaluating the effectiveness of new detection methods and models.
B.2 Masking Network
To facilitate the mask generation process, we design a lightweight masking network with two linear layers, one ReLU layer in between, and one sigmoid activation. The first linear layer can be expressed in Py ToTch code as nn. Linear(numm patches x emb dim, 512), where num patches is the number of patches from the original image (c.g. ImageNet sample of size 224 224 with patch size of 16 16 will generate 14 x 14 = 196 patches) and emb dim is the embedding dimension from VrT-B (768 in our case). We use a hidden size of 512 for the output dimension of the first and the second lincar layer. The ReLU layer is expressed nn neLU))(() The second linear layer can be expressed as mm.iinear(512, num patches) where the output is the masking probability for each image patch with corresponding order. We pass the resulting tensor from the second linear layer through the sigmoid activation to generate values between 0 and 1.
B.3 ImageNet
We adopt the default ViT-B model that has been employed in the original MAE. In total, we train 50 epochs for all three stages. For ImageNet experiments, we use an image size of 224 x 224.
Pretraining Detailed three-stage MLO-MAE pretraining setting is in Table 12. We follow MAE settings on data augmentations by only using RandomRessizdCChp and RandomHorizontalFiiip. For MLO-MAE in Table 1 experiment, we train MLO-MAE using xavier uniform (Glorot & Bengio, 2010). We follow the linear Ir scaling as used in the MAE. We randomly split the training set of ImageNet by a ratio of 80/20 to be the new training set and the new validation set. We use the same new training sot Stage and Stage II for training, while use the new validation set in Stage III for training the masking network. We use the original Net validation set to report validation accuracy in Stage IT.
21

---

Published in Transactions on Machine Learning Research (03/2025)
Table 12: MLO-MAE ImageNet Pretraining Settings.
MLO-MAE Stage
Config
value
model optimizer
ViT-B AdamW
base learning rate
le 4 0.05
weight decay
image size
224
Stage I
patch size
16 16
optimizer momentum
B1, 82 = 0.95
batch size learning rate scheduler unrolling steps mask ratio
256
cosine anneal
0.75
learning rate weight decay
4e - 5
0.05
image size
224
Stage II
optimizer momentum
B1, B2 = 0.9, 0.95
batch size
256
learning rate scheduler
cosine anneal 1
unrolling steps
learning rate
4e 5
Ln.Linear(num patches x emb dim, 512)
nn.ReLu
masking network
nn.Linear(512,, num patches) torch.sigmoid() 0.05
Stage III
weight decay
image size
224
optimizer momentum batch size learning rate scheduler unrolling steps
81, 32 = 0.95 256
0.95
cosine anneal 1
Due to the inherent design of our three-level optimization framework, we do not conduct separate linear probing and directly report the Stage II testing accuracy (test dataset not seen in MLO-MAE training). Therefore, we report the training setting as in Stage shown in Table 12.
Linear Probing
We directly followed the MAE fine-tuning experiments and did not make additional changes. Detailed settings can be found in Table 9 MAE (He et al., 2022).
Fine-tuning
B.4 CIFAR
We adopt the same ViT-B architecture as in B.3 but change the input image size and patch size to be 32 and 2 respectively. In total, we train MLO-MAE 50 epochs for all three stages.
Pretraining We pretrain our model from scratch (i.c. no pretrained initialization from other datasots) using the MLO-MAE method on two CIFAR datasets. Table 13 shows the detailed training setting for CIFAR-10 and CIFAR-100 experiments. We adopted similar setting from the ImageNet experiment, with
22

---

Published in Transactions on Machine Learning Research (03/2025)
Table 13: MLO-MAE 10F00-10/100 Pretraining Settings.
MLO-MAE Stage
Config
value
model optimizer
ViT-B AdamW
base learning rate
le 3
weight decay
0.05
image size
32
Stage I
patch size
x
optimizer momentum
B1, 92 = 0.9, 0.95
batch size learning rate scheduler unrolling steps mask ratio
128
cosine anneal 2 0.75
learning rate weight decay
1e - 3
0.05
image size
32
Stage II
optimizer momentum
B1, B2 = 0.9, 0.95 128
batch size
learning rate scheduler
cosine anneal 1
unrolling steps
learning rate
le - 3
Ln.Linear(num patches x emb dim, 512)
nn.ReLu
masking network
nn.Linear(512,, num patches) torch.sigmoid() 0.05
Stage III
weight decay
image size
32
optimizer momentum batch size learning rate scheduler unrolling steps
81, 32 = 0.95 128
cosine anneal 1
minor modifications on learning rate, data augmentation, image size, and patch size. We use conventional RandomCrop and RandomHorizontalFiip on both CIFAR-10 and CIFAR-100. We pretrain with MLO-MAE for 200 epochs.
Linear Probing Similar to section B.3, we report the training setting as in Stage shown in Table 13.
Finc-tuning For CIFAR fine-tuning, we use Ir=le 4, weight_dccay=55e docay=5e 5, Adtimizer=AdamW, and epoch=100. We use default data augmentation on CIFAR, including RandomCrop, Resize, and PanddomHorizontalFiip. We maintain the image size to be 32. This experiment is performed top of MLO-MAE CIFAR-10/ pretrained weights respectively.
batch size=64, batch_size==4
B.5 Classification on Fine- graiggr Datasets
We perform image classification using MLO-MAE pretrained VIT-B on CUB-200-2011 (Wah et al., 2011). Stanford Cars (Krause et al., 2013), and iNaturalist 2019 (Van Horn et al., 2018) fine-grained datasets. We
23

---

Published in Transactions on Machine Learning Research (03/2025)
Class labels
Rock crab
Brambling
Cardoon
Eel
Barracouta
Original image
MAE
masking
SemMAE maeking
AutoMAE masking
MLO-MAE
masking
Figure 5: Visualization of the masking patterns of ML.O-MAF. and baselines on randomly sampled ImageNet images.
follow the setting from MAE (He et al., 2022) with minor adjustments on learning rate and epochs. These experiments are performed on top of MLO-MAE ImageNet-1K pretrained weights respectively.
B.6 Semantic segmentation ADE20K
We the semantic segmentation code implementation of MAE by MMSegmentation (Contributors, 2020). Given VIT-B model pretrained on ImageNet using MLO-MAE or baseline and subsequently fine-tuned ImageNet, to transfer it for semantic segmentation, integrated it as backbone model into the UPer- Net (Xiao et al., 2018) semantic segmentation framework. It was then further fine-tuned on the challenging ADE20K dataset (Zhou et al., 2017) containing 25K images spanning 150 semantic categories. The fine- tuning was conducted by the AdamW optimizer for over 160,000 iterations, with a batch size of and a learning rate of 0.0001.
C Visualization
C.1 Masking Pattern Visualization
We randomly sampled five images from ImageNet and visualized the masked patches learned by our method MT.O-MAF. We also included visualizations for baseline methods, including MAF, SemMAE, and AutoMAE.
24

---

Published in Transactions on Machine Learning Research (03/2025)
Original
Arnolafoon
MLO-MAE
Original
Arnotations
MLO-MAE
Original
Arrolation
MAE
MLO-MAE
Figure 6: Visualizing 2D activations of MAF and ML.O-MAE pretrained VTT-B weights on 12 examples of ADE20K validation set using GradCAM. For each set of images, from left to right, are original image, ground truth annotations, MAE activation map, and MLO-MAE activation map. For both MAE and MLO-MAE models, features are extracted from the norm1 layer from the last VIT block. Red colors high activation region.
with a masking ratio of 10%. As shown Figure 5, the majority of the masked patches learned by MT.O- MAE are on foreground objects directly relevant to the class labels of these images. In contrast, MAE, SemMAE, and AutoMAE place the majority of masked patches on background regions irrelevant to image class labels. These results indicate that MLO-MAE encourages the encoder network to focus on learning effective representations for objects rather than background regions. By focusing on correctly reconstructing the masked patches in object regions, the encoder can effectively capture the intrinsic properties of the objects.
MLO-MMA achieves this ability by leveraging the downstream classification task to guide the pretraining of the encoder and the learning of the masking strategy. Minimizing the validation loss of the downstream classification task in Stage ITT of ML.O-MAE encourages the pretrained encoder to learn discriminative distinguish between different classes. To learn these discriminative representations,
representations that MLO-MAE emphasizes masking and reconstructing patches in object regions, as these objects are directly related to the class labels. In contrast, SemMAE and AutoMAE use the atlention maps produced by StyleGAN and adversarial learning to mask patches. These attention maps are created without leveraging guidance from the class labels of the downstream classification task, resulting in SemMAE and AutoMAE being less effective at masking class-label-relewaa objects compared to MLO-MAE.
It is worth noting that MLO-MAE emphasizes masking objects directly related to the image class label rather than any objects. For instance, in the fourth image, which contains two types of objects eel and starfish, MLO-MAE places more masked patches on the eel because the class label of the image is eel. Although starfish are prominent objects in this image, ML.O-MAE masks fewer patches on it since the image's class label is not starfish. Again, this targeted masking strategy is learned with guidance from the downstream image classification lask, aimed at enhancing classification accuracy.
C.2 GradCAM visualization
To delve into the representation learning prowess of MLO-MAE, also present visual analysis of activa- tion maps generated by both the pretrained MAE and MLO-MAE ViT-B models. Specifically, we examine 12 examples from the ADE20K validation set, as depicted in Figure 6. These activation maps are derived from features extracted from the norm1 layer within the final VIT block of the backbone architecture. Our comparative analysis reveals that MLO-MAE consistently produces activation maps of notably higher se- mantic coherence compared MAE. This enhancement attributed to the tailored guidance provided by task-specific masking during the pretraining phase. Notably, ML.O-MAE, demonstrates proficiency in pin-
25

---

Published in Transactions on Machine Learning Research (03/2025)
MAE
DINO
MLC-MAE
Classi
Figure 7: t-SNE visualizations of representations learned by MAE, DINO, and MLO-MAE for CIFAR-10 test images.
pointing highly informative regions, demonstrating the efficacy of the MLO-based pretraining methodology refining visual representation learning.
C.3 t-SNE Visualization
Table 14: Ratio of average intra-class similarity to inter-class similarity for representations extracted by encoders pretrained using MLO-MAE, MAE, and DINO on the test. images of CIFAR-10, CIFAR-100, and ImageNet.
Method
Ratio on CIFAR-10
Ratio on CIFAR-100
Ratio on ImageNet
MAE
1.25
1.19
1.08
DINO
1.31
1.27
1.11
MLO-MAE 1.59
1.52
1.43
We utilized the encoders learned by MLO-MAF, MAE, and fully supervised method, DINO (Caron et al., 2021), to extract representations for the test images of CIFAR-10. These representations were then visualized using t-SNE, as shown in Figure 7. The visualization indicates that in the MLO-MAE representation space, different classes are better separated, with images from the same class grouped together. In contrast, the MAE and DINO representations show mixing of different classes. Furthermore, we mcasured the ratio between intra-class similarity and inter-class similarity. For intra-class similarity, similarity between the representations of each pair of images within the same class and averaged these values, Likewise, for inter-class similarity, computed the average cosine similarity for pairs of images from different classes. The results are in Table 14. MLO-MAE achieves the highest ratios across all datasets, demonstrating its learned representations can better distinguish between different. classes. This can be attributed to our method's use of validation loss from the downstream classification task to guide pretraining, resulting in more discriminative representations.
calculated the cosine
Additional Experiments
D.1 Object Detection
We evaluate the transfer ability MLO-MAE VIT-B model to object detection task on PASCAL VOC 2007 dataset.
Following similar setup as in Scction 4.3, we directly train the pretrained model for the object detection task using the following procedure. The dataset downloaded and organized into training, validation, and test. sets. Images preprocessed to fixed size (e.g., 512x512 pixels) and augmented with
26

---

Published in Transactions on Machine Learning Research (03/2025)
Table 15: Object detection result of MLO-MAE on PASCAL VOC 2007 detection task.
Method
VIT-B
MAE
MLO-MAE
mAP (%)
79.1
82.8
83.5
techniques such as random cropping, flipping, and normalization to enhance model robustness. The ViT- Base backbone pretrained using MAE and MLO-MAE on large corpus of unlabeled images learn rich visual representations. An object detection head is then attached to the backbone, consisting of fully con- nected layers for predicting bounding boxes and class labels. The training process employs combination of classification loss (cross-entropy) and bounding box regression loss (smooth L1), optimized using the AdamW optimizer with & learning rate of le-4 and weight decay of le-4. We use a batch size of 16 and train the model for 50-100 epochs, incorporating early stopping based on validation performance. During training, the backbone initialized with .re--pretd weights, and forward passes are performed extract features and make predictions. The total loss is computed and backpropagation is used to update the model weights. Periodic validation monitors performance and guides hyperparameter adjustments. Evaluation metrics in- clude mean Average Precision (mAP), precision, and recall at different loU thresholds, alongside inference speed. For final evaluation, the trained model predicts bounding boxes and class labels on the PASCAL VOC 2007 Lest set, and performance is benchmarked using the mAP metric. Results may be submitted to the PASCAL VOC evaluation server for standardized comparison with other models, demonstrating the efficacy of the MAE-prctrained ViT-Base backbone in object detection tasks. Table 15 shows the result. MLO-MAE surpasses supervised ViT-B and MAE with 4.4% and 0.7%, respectively.
D.2 Robustness
Table 16: Robustness evaluation on ImageNet variants. We use IN-1K finctuned ViT-B and directly reported from Table 1 without further training. Results are top-1 accuracy.
Dataset
MAE
MT..-MAE
ImageNct-A
35.9
46.2
ImageNet-B
18.4
46.3
ImageNet-C
51.7
55.5
ImageNet-R
48.3
55.6
ImageNet-S
34.5
41.8
We evaluate the robustness of our VIT-B models on different variants of ImageNet validation sets. Following MAE's setup, we use the fine-tuned model from Table without further training and only run inferences on the ImageNet robustness variants. Table 16 shows our ML.O-MAE surpasses MAE in all variants with large margin.
D.3 Computational Efficiency
Table 17: Test accuracy on CIFAR-100 with different update frequencies.
Update frequency
Per epoch runtime (GPU hrs)
Test accuracy on CIFAR-100 (%)
Every iteration Every 5 iterations
0.6
79.4
0.4
79.1
In Soction 4.6 of the main paper, we compared the computational cost of our method to that of bascline methods, including MAE, SemMAE, and AutoMAF. The overall runtime of method is similar to that
27

---

Published in Transactions on Machine Learning Research (03/2025)
Table 18: Ablation on curriculum masking ratio from 0.5 to 0.9.
Masking ratio
Test Accuracy on CIFAR 100
0.75
79.4
0.5-0.9, linear increase
76.5
Table 19: Test accuracy CIFAR-100 with full finc-tuning in stage II.
Method
Test Accuracy on CIFAR 100
Runtime (days on GPUs)
Linear probing in stage II
79.4
0.7
Full fine-tuning in stage II
79.5
1.1
of the baselines. Our method converges in fewer epochs but has a higher per-epoch runtime compared to the baselines. To reduce the per-epoch training time, we can decrease the update frequency of the masking network; instead of updating every iteration (mini-batch), we update every few iterations (e.g., every five iterations), while still updating the encoder and linear head at each iteration. Calculating the hypergradient of the masking network requires computing Jacobian matrices and performing their multiplication with vectors, which is more computationally intensive than calculating the gradient of the encoder and linear head. By reducing the update frequency of the masking network, we can significantly lower the overall computational costs. We experimented with this approach on CIFAR-100, parallelized across 8 GPUs as shown in Table 17. As can be seen, the per epoch training time significantly reduced. Meanwhile, we empirically found that reducing the update frequency of the masking network does not significantly impact classification accuracy. This likely because once an intermodiate masking strategy lcarned, it can be used for a while to pretrain the encoder without needing frequent updates.
D.4 Curriculum Mask Ratio
Experiments performed in section 4 use fixed masking ratio, following the baseline setup. To further explore the impact of different masking behavior (fixed and dynamic masking ratio) on MLO-MAE, we experimented with a curriculum masking ratio setting. We dynamically increased the masking ratio of our method from 0.5 to 0.9 as training progressed, with a linear schedule. The results on CIFAR-100, shown in the Table 18, indicate that using a dynamic ratio does not outperform a fixed ratio of 0.75.
D.5 Full Fine-tune Stage
min ( (FF) T
s.t. (T) = argmin Las F,C
(13)
Nxr
E' argmin Σ Σ E.D
NEP.JLL
We experimented with full fine-tuning of the encoder during Stage II, parallelized across 8 GPUs. In this setup, full fine-tuning of the pretrained encoder involves training an encoder to have a small L2 distance from. Equation 13 is the formulation for performing full fine-tuning of the pretrained encoder during Stage II in MLOO-ME, Table shows the results on CIFAR-100. As can be seen, full fine-tuning the encoder during stage II does not significantly outperform linear probing but incurs much higher computational costs. Therefore, it is preferable to fix the encoder during Stage IT.
28

---

Published in Transactions on Machine Learning Research (03/2025)
E Further Discussion
E.1 Utilizing Unlabeled Data
The proposed MLO-MAE framework relics on downstream feedback to guide its masking network. This dependency naturally raises the question of how to effectively utilize unlabeled data. Here, we would like to consider two scenarios: 1) using MLO-MAE the continued pretraining stage, and 2) direct transfer learning from the MLO-MAE pretrained model. While MLO-MAE's current formulation leverages downstream task labels to guide masking during pretraining, the framework can still operate in label-scarce settings through the following mechanisms:
MLO-MAE in the continued pretraining stage MLO-MAE can be initialized with vanilla MAE pretrained on unlabeled data. Subsequent MT.O-MAF, continued pretraining on the labeled data can be used to further improve the performance on downstream tasks (shown in Section 4.4), demonstrating MLO-MAE's flexibiliiyy under diverse scenarios. Note that in this scenario, the continued pretraining is happening in the finetuning stage, therefore we can assume to have access to the downstream labeled data.
Direct Transfer Learning from MLO-MAE When domain-specific labels are unavailable, MLO-MAE leverage public labeled datasets (e.g., ImageNet) to learn generalizable masking strategies. First, initialize MLO-MAE with a standard MAE pretrained unlabeled domain-specific data. We then pretrain it on ImageNet using its labels to guide task-aware masking, and finally transfer the model to downstream tasks without requiring their labels during pretraining. Our experiments in Section 4.3.1 show that MLO- MAE pretrained on one labeled dataset (o.g., ImageNet) transfers effectively to unseen tasks /atks/datssoss without requiring their labels during pretraining. This confirms that the learned masking strategy generalizes beyond the specific downstream task used for guidance.
E.2 Fairness in Comparing with Baselines
While the masking network T supervised by the labels, we would like to emphasize that the comparison with bascline methods presented in Section remains fair with the following clarifications:
Equivalent Label Utilization All baseline methods (MAE, SemMAE, etc.) require labeled data for fine-tuning lincar probing post-pretraining. In MLOOMME, the same labeled data is used only during pretraining (Stage II and Stage III) to guide masking, not for direct encoder updates. Thus, no additional labels are introdueel-compaaiioon remain fair (Section 4.2).
Empirical Validation Transfer learning results (Tables 3, 4, 5) demonstrate that MLO-MAE's gains per- sist even when pretrained on one task (e.g., ImageNet classification) and evaluated on others (e.g., ADE20K segmentation). This confirms that performance improvements stem from better representation learning.
29


Reviews:
[{"review_id": "mpougsYiSl", "reviewer": "Reviewer_Q4kg", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Advantages:\nThis paper includes labels from down-stream tasks into the training stage of the MAE encoder. This is well motivated since the random masking strategy is non-optimal to fit the distribution of down-stream data. This paper proposes to learn a masking generator network to generate mask according to specific data. The authors also formulate the proposed learning framework as a multi-level optimization problem, where each stage corresponds to one level of optimization. The three levels are mutually dependent and solved jointly.  This nested optimization structure enables the model to learn a masking strategy that is directly influenced by downstream task performance, rather than relying on a fixed or heuristic masking approach.\n\n\nWeakness:\n(1)\tDependency on Downstream Tasks. On the one hand, the masking network learned by MLO-MAE is heavily dependent on the specific downstream task used during pretraining. What if all the available data are unlabeled during the pretraining stage? Will this lack-of-label situation make the MLO-MAE algorithm non-usable? On the other hand, although the encoder network is not directly affected by the gradient generated by the image labels, the training of the masking network T is supervised by the labels. This indirect supervision from the labels implicitly boost the representation learning of the MAE encoder, making the comparisons with other baseline methods somewhat unfair.\n(2)\tHigh Computational Cost. The multi-level optimization introduces great computational complexity due to the nested optimization problems. Although the authors propose methods to mitigate this, the overall computational cost remains high compared to other MAE methods, which might limits the scalability of MLO-MAE to very large datasets or backbone models.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Multi-level Optimized Mask Autoencoder (MLO-MAE), a novel framework for improving the representation ability of MAE model by optimizing the masking strategy using the feedback from downstream tasks, such using image labels from the classification task. The author designs a multi-level optimization (MLO) framework to jointly learn multiple mutually-independent modules via an end-to-end three stage training pipeline.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See the weaknesses.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "8HkU5Wabsv", "reviewer": "Reviewer_n4Bm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Multi-level optimization formulation of the problem is novel and experimental results look promising - in particular the use of the performance on a downstream task as a feedback is an elegant idea\n- End-to-end pipeline is an interesting a very generalizable approach\n- In depth discussion of experimental results allows to get a good understanding of the behavior of the method in different settings and its properties (especially sections 3.5, 3.6, 3.7)\n\nWeaknesses:\n- One of the main issues with the paper is the lack of clarity: end-to-end pipeline is not explained well and it takes some time for the reader to understand the method\n- Likewise, the paper can be better organized, e.g. section 3 is very broad - from method overview to experimental results discussion and property analyses.\n- Lack of theoretical analysis of the optimization problem - would have liked an analysis of the optimality of the formulation, e.g. bounds or at least a discussion on the fact that it can be seen as an approximation approach to jointly solving the 3 stages at once\n- No discussion on the potential limitations of the method", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In the context of training Masked Autoencoders, the paper introduces a novel method based on Multi level optimization to enhance the effectiveness of the training. Unlike previous approaches based on independent strategies or heuristics, the presented method the MLO-MAE method tries to focus the most informative regions - using the feedback from a downstream task to evaluate this informativeness. The training process is in 3 optimization stages:\n1. Pretraining an encoder using a tentative version of the mask\n2. Training a classifier head using the representation learned by the encoder\n3. Update the mask from step 1 with the performance of the classifier on a validation set\n\nAfter presenting the method in details, the paper showcases its performance in a variety of settings, and discusses its benefits in terms of better masking and representation learning.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses sections.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "uaLu8Jd3qv", "reviewer": "Reviewer_L73A", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- This paper is well-written, with clear expression, well-presented formulations and well-organized figures/tables.\n- The motivation to incorporate feedback from downstream tasks/datasets is well-justified.\n- The experimental design is thoughtfully executed, providing evidence to support the advantages of the proposed MLO-MAE.\n\nWeakness:\n- Some experimental details are insufficiently explained, as outlined in the \"Requested Changes\" section below.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents MLO-MAE, an end-to-end method for learning an optimal masking strategy for Masked Autoencoder (MAE) that incorporates feedback from downstream tasks/datasets. The training process consists of three stages, during which the masking network is optimized using Multi-level Optimization (MLO). While MLO introduces higher computational costs, the authors claim that MLO-MAE converges more quickly than baseline methods. Although the core ideas of the optimization algorithm are well-presented, I am not an expert in optimization and cannot fully guarantee its correctness.\n\nExperimental results demonstrate that incorporating feedback from downstream datasets can enhance performance by producing more effective masking networks. Additionally, the experiments highlight MLO-MAE's ability to generalize to other datasets within the same task, transfer between different tasks, and perform in a continued pretraining setting. The ablation studies provide valuable insights into the effectiveness of three-level optimization, the diminishing returns of increasing unrolling steps, the selection of appropriate masking ratios, and the superior performance achieved with smaller patch sizes.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Overall, I find this paper well-written and easy to follow. The following items are suggestions/questions that would simply strengthen the work:\n\n- In Section 3.5.2, the results demonstrate the capability of transferring from classification to semantic segmentation. I wonder how the method would perform if trained solely using the feedback from semantic segmentation. \n- In Section 3.6, it would be helpful to include results for MLO-MAE's performance without pretraining on MAE.\n- In Section 3.8, could the authors provide an explanation for why MLO-MAE converges in fewer epochs compared to baseline methods?", "broader_impact_concerns": "No other broader impact concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces Multi-level Optimized Mask Autoencoder (MLO-MAE), a novel framework for improving the representation ability of MAE model by optimizing the masking strategy using the feedback from downstream tasks. The method is well-motivated since the random masking strategy is obviously not an optimal solution to fit the downstream data. To achieve this, the authors proposed an end-to-end pipeline and the experimental results are promising.

All reviewers support the novelty of the work. During the discussion period, the authors well addressed the issues related to training convergence and computational cost. I am willing to follow the majority as all reviewers lean toward acceptance.
Decision: Accept (poster)