Downstream Task Guided Masking Learning in Masked Autoencoders Using Multi-Level Optimization

Q7: Confusion on optimal masking strategy

We apologize for the confusion. MLO-MAE determines the masking regions by directly leveraging the classification loss, incorporating the optimization of the masking network in Stage III. This approach eliminates the need for human-designed heuristics, relying instead on a task-driven strategy. The term "optimal strategy" refers to the masking strategy used in Stage I, which is optimized using the downstream classification loss defined in Stage III. For example, if the masking network in MLO-MAE determines that the masking region for a specific image should cover all foreground objects, this implies that reconstructing the entire masked foreground is beneficial for downstream classification performance. Such a strategy remains desirable.

Q8: Elaborate on the implementation of Stage 3

We thank the reviewer for this comment. In Stage III, we optimize the trainable masking network by minimizing the validation loss, leveraging the optimized classification head and encoder backbone obtained from Stage I and II. The validation loss is computed as the cross-entropy classification loss on the validation set, using the frozen vision encoder and classification head. It is important to note that this validation set is created from the ImageNet training set, while the original ImageNet validation set is reserved as the test set to evaluate the overall performance of our method. In the MLO-MAE framework, the vision encoder, classification head, and masking network are implicitly interconnected. As described in Equation (8), we utilize the chain rule to propagate gradients, with the vision encoder and classification head treated as approximate optimal solutions from Stages I and II, respectively. Finally, the masking network is updated using one-step gradient descent.

We appreciate your constructive feedback and provide our response to your questions as follows.

Q1: Deviation from the motivation

We thank the reviewer for pointing this out. Our transfer learning experiments in Section 4.4.2 demonstrate that MLO-MAE, even with Stage II focused on a classification task, is still capable of learning effective and generalizable representations. These representations lead to superior performance on semantic segmentation and object detection tasks, highlighting the versatility of our approach. Additional object detection experiments can also be found in Appendix D.1 Our ultimate goal is to enhance the quality of representations learned by the backbone network, ensuring they generalize well across a wide range of downstream tasks.

Q2: Unfair comparison

We would like to clarify that our method does not unfairly use more labeled data than the baseline approaches. The labeled data used in Stages II and III of our framework is identical to that employed during the fine-tuning phase of the baselines, ensuring a fair comparison. In MLO-MAE, we directly leverage the downstream objective to guide and optimize the mask-reconstruction process during the MAE pretraining stage. It is also important to note that labels are utilized during downstream evaluation for all methods, maintaining fairness in the comparison. Moreover, the effectiveness of MLO-MAE is further demonstrated through transfer learning experiments, where neither the data nor the labels from the downstream tasks are used during the pretraining of either MLO-MAE or the baselines. The superior performance of MLO-MAE in these settings empirically highlights its strong generalizability.

Q3: Clarification on MLO-MAE’s masking strategy

We apologize for any confusion caused. MLO-MAE determines the masking regions by directly leveraging the classification loss, incorporating the optimization of the masking network in Stage III. It integrates three interconnected optimization objectives: the MAE pretraining objective, the optimization of the classification head on the classification task, and the optimization of the masking network on the same task. These three levels of optimization are mutually dependent, as formalized in Equation (4). To address this complex problem, we employ an efficient hypergradient-based method to solve Equation (4), enabling gradient descent updates for the masking network. This formulation allows the masking network, which is optimized on the cross-entropy loss for class label prediction in Stage III, to generate downstream-informed masking patterns in Stage I. Consequently, the resulting masking patches are directly aligned with the class labels, ensuring their relevance to the downstream tasks.

Q4: Ablation on ImageNet

We thank the reviewer for this suggestion. Our primary classification results in Table 1 and Table 2 show the effectiveness of MLO-MAE on both ImageNet and CIFAR dataset. Due to the resource limit, we conducted the ablation study on the CIFAR experiment.

Q5: Unclear of the update of masking network and unrolling steps

We apologize for the confusion. In Stage III, we optimize the masking network by minimizing the validation loss, leveraging the optimized classification head and encoder backbone obtained from Stages I and II. This process is carried out using gradient-based methods. The detailed computation of implicit gradients is outlined in Appendix A.1, and the pseudo-code for the integrated approach is provided in Algorithm 1.

The unrolling steps serve as hyperparameters to approximate the optimal weights in Equations (1) and (2). In Equation (1), $E^*(T)$ represents the optimal weights of the encoder. However, finding the exact optimal solution for $E^*$ is computationally infeasible. To address this, we approximate $E^*$ using a fixed number of gradient descent steps, referred to as unroll steps. In Equation (2), $C^*$ is approximated using the same approach. In most of our experiments, we set the unroll step to 2 for $E^*$ amd 1 for $C^*$ , as this configuration empirically provides a good tradeoff between training time and performance. Additional details on the impact of different unroll step values can be found in our ablation studies in Section 4.6. We hope this clarifies our approach and the rationale behind our chosen hyperparameters.

Q6: Choosing classification as Stage II

Image classification is a fundamental computer vision task. Our transfer learning experiments on fine-grained classification, semantic segmentation, and object detection show the effectiveness of the classification guided masking network in generating the mask patterns that benefit other downstream tasks as well.

We appreciate your constructive feedback and provide our response to your questions as follows.

Q1: Comparison with ColorMAE

We thank the reviewer for this helpful suggestion. MLO-MAE achieves superior performance than ColorMAE in ImageNet1K classification, semantic segmentation on ADE20K, and object detection on COCO as shown in the table below.

Q2: Transferability and applicability of MLO-MAE

We appreciate the reviewer for this insightful comment. As described in Section 4.1, we use the original ImageNet1K validation set as the test set and split the training set into a new training set and a new validation set. The masking network is then optimized based on the performance on the newly created validation set, following Equation (4). Our transfer learning experiments in Section 4.4, conducted on fine-grained classification tasks (iNaturalist 19, CUB, and Stanford Cars), semantic segmentation (ADE20K), and object detection (MS-COCO), demonstrate the transferability of the visual representations learned by MLO-MAE across a diverse range of downstream tasks, beyond classification alone. Additionally, we tested the MLO-MAE approach in the medical domain using a continued pretraining setup, as detailed in Section 4.5. Empirical results on two medical datasets, PDDB and PAD-UFES, further highlight the effectiveness of MLO-MAE as a continued pretraining technique for domain-specific datasets.

We appreciate your constructive feedback and provide our response to your questions as follows.

Q1: Rationale behind MLO-MAE

We appreciate the reviewer for this insightful suggestion.Thank you for the thoughtful suggestion, and we deeply appreciate the opportunity to elaborate on the key distinctions between MLO-MAE and existing adaptive mask generation methods, such as SemMAE and AutoMAE. These differences can be summarized as follows:

1. Learnable Masking Network with Downstream Guidance: MLO-MAE employs a learnable masking network that is directly optimized based on downstream performance, specifically guided by the cross-entropy loss, rather than relying on heuristic or standalone approaches.

2. Hierarchical Optimization Framework: MLO-MAE introduces a hierarchical design of mutually dependent optimization problems, solved using implicit differentiation, to address the complexity of this framework effectively.

Unlike other methods, MLO-MAE simultaneously solves three optimization problems: optimizing the weights of the vision encoder (Stage I), the classification head (Stage II), and the masking network (Stage III). This structure ensures that the optimization processes are interdependent—for example, optimizing the vision backbone depends on the masking network, and optimizing the masking network requires the vision backbone and classification head to be optimized. This mutual dependency allows MLO-MAE to integrate downstream guidance directly into the mask selection during the mask-reconstruction process (Stage I). In contrast, SemMAE employs a standalone semantic encoding module to generate masking patterns, while AutoMAE relies on an external ViT to extract attention maps indicative of semantic regions. These methods depend heavily on external semantic modules, making them potentially sensitive to different data distribution. MLO-MAE, by design, avoids this limitation, ensuring greater robustness and adaptability across different datasets.

Q2: Potential ache and memory issue

The cache and memory challenges can be effectively addressed using gradient checkpointing, a feature supported by most standard deep learning libraries. While MLO-MAE introduces additional computational complexity due to the calculation of the best-response Jacobian, this overhead is offset by the reduced number of training epochs required to achieve comparable or superior performance on ImageNet benchmarks, as demonstrated in Table 9.

Q3: Analysis of the generated masks

We thank the reviewer for pointing this out. We have provided a detailed analysis of the masking patterns learned by MLO-MAE in Appendix C.1, alongside a comparison with AutoMAE and SemMAE. Unlike MAE, SemMAE, and AutoMAE, which predominantly mask background regions that are irrelevant to image class labels, MLO-MAE strategically guides the encoder network to focus on learning effective representations for objects. These findings highlight the advantages of MLO-MAE in prioritizing meaningful features during pretraining.

Q3: Unfair evaluation

We would like to clarify that our method does not unfairly use more labeled data than the baseline approaches. The labeled data used in Stages II and III of our framework is identical to that employed during the fine-tuning phase of the baselines, ensuring a fair comparison. In MLO-MAE, we directly leverage the downstream objective to guide and optimize the mask-reconstruction process during the MAE pretraining stage. It is also important to note that labels are utilized during downstream evaluation for all methods, maintaining fairness in the comparison. Moreover, the effectiveness of MLO-MAE is further demonstrated through transfer learning experiments, where neither the data nor the labels from the downstream tasks are used during the pretraining of either MLO-MAE or the baselines. The superior performance of MLO-MAE in these settings empirically highlights its strong generalizability.

Q4: Insights on smaller number of epoch of MLO-MAE

We appreciate the reviewer for this keen observation. The training process for MLO-MAE involves jointly solving three interconnected optimization problems using implicit differentiation, which introduces additional computational complexity compared to the standard MAE optimization. In practice, we approximate E* and C* by unrolling a limited number of gradient steps, which further increases the computational complexity per iteration. It is important to note, however, that the number of training epochs for MLO-MAE is not directly comparable to that of MAE, as MAE does not involve implicit differentiation or gradient unrolling.

Q5: Difficulty handling discrepancies in patch size between pre-training and fine-tuning

We thank the reviewer for this helpful comment. We primarily adhere to the downstream evaluation protocol used by MAE, in which the classification head from MLO-MAE is not utilized for downstream tasks. Instead, the classification head is independently initialized during the fine-tuning phase. Additionally, to maintain consistency with ImageNet-pretrained encoders, we could resize CIFAR images to a resolution of 224x224, and therefore effectively apply the ImageNet pretrained MLO-MAE on CIFAR dataset.

Q6: Data split used by MLO-MAE

We thank the reviewer for pointing this out. The new data split was performed using an 8:2 random split within each class of the ImageNet training set. While we did not specify a random seed, the large size of the dataset ensures that the performance of our method remains consistent when following the procedures described in Section 4.1 and Appendix B.3. To ensure a fair comparison with MAE on datasets where predefined train, validation, and test sets exist, MLO-MAE uses the original validation set in Stage III. Meanwhile, MAE is trained on the combined train and validation sets. Both methods are then evaluated on the test set to maintain consistency and fairness in the comparison.

Q6: Different masking patterns comparing to baselines and minimal changes to make learnable masks beneficial

We thank the reviewer for this insightful comment. It is important to note that the MAE ImageNet result reported in Table 1 is based on a model pretrained for 1600 epochs, whereas an MAE pretrained for 800 epochs achieves an accuracy of 83.2% during fine-tuning [1]. This suggests that the learnable masking strategies employed by SemMAE and AutoMAE are still somewhat effective for improving downstream performance. However, the semantic masking approach used by SemMAE and the adversarially trained masking strategy of AutoMAE yield only marginal improvements (~0.1%) in downstream ImageNet fine-tuning, raising questions about the efficacy of directly learning a masking strategy guided by downstream tasks. SemMAE employs pretrained attention maps to adaptively mask both within and across semantic parts, while AutoMAE utilizes an adversarially trained generator and a joint optimization technique to update the mask generator. Notably, both methods rely on a standalone pretrained vision encoder (iBoT for SemMAE and a pretrained ViT for AutoMAE) to extract semantic information. In contrast, MLO-MAE integrates the learning of masking patterns directly into a unified optimization framework using multi-level optimization, as detailed in Appendix A. This approach eliminates the need for a separate pretrained vision encoder for semantic extraction.

In Section 4.6, we explored a simplified variant of MLO-MAE, referred to as BLO-MAE, where the optimization problem is reformulated as shown in Equation (5). While BLO-MAE achieves 2.8% lower accuracy compared to MLO-MAE, it still outperforms the MAE baseline (81.5% vs. 73.5%). This demonstrates that even minimal modifications, such as the BLO-MAE setup, which combines Stage I and Stage II, can effectively implement a downstream-guided masking strategy. [1] SemMAE: Semantic-Guided Masking for Learning Masked Autoencoders, NeurIPS 2022

Q7: Effectiveness and the convergence of BLO-MAE

We thank the reviewer for this astute observation. While BLO-MAE has lower CIFAR accuracy than MLO-MAE, it still achieves better performance than MAE baseline (73.5%, according to Table 2) and shows convergence across all stages, indicating that BLO-MAE, compared to the MAE, has a positive effect on the performance. It is not surprising that MLO-MAE achieves better performance than BLO-MAE because the lower level of BLO-MAE approach involves optimizing a weighted sum of losses from two distinct tasks ($C^*$ and $E^*$). This scenario often leads to task competition and balancing the competing losses requires meticulous adjustment of the tradeoff parameter, which is challenging.

Q8: Additional experiment on different size of ViT model

We thank the reviewer for this helpful suggestion. We are currently running MLO-MAE on a different model size. However, due to resource constraint and time limitation during the rebuttal period, we are unable to finish the experiment. We will include the result into the paper once it finishes.

We appreciate your constructive feedback and provide our response to your questions as follows.

Q1: Insights of design behind learnable masking in MLO-MAE

Our work is inspired by the idea of directly incorporating downstream information to guide masking decisions through multi-level optimization. In Appendix C.1, we provide a comprehensive analysis of the masking patterns learned by MLO-MAE, comparing them to those produced by MAE, SemMAE, and AutoMAE. Unlike these approaches, which predominantly mask background regions irrelevant to image class labels, MLO-MAE strategically focuses masking on areas that encourage the encoder network to learn effective representations for objects rather than background regions. This highlights the ability of MLO-MAE to prioritize meaningful features in the learning process.

Q2: Motivation and insight explaining the improvement of MLO-MAE

We thank the reviewer for this helpful comment. Our MLO-MAE explores the effectiveness of introducing the downstream insight in guiding the pretraining task by integrating a learnable masking network through a hypergradient-based multi-level optimization (MLO) framework. Through our transfer learning and continued pretraining experiment, MLO-MAE has shown its effectiveness and generalizability to diverse domains and various vision tasks. Due to the page limit, we have included detailed discussion of our implicit differentiation setup for our MLO gradient computation and of the differentiability of the MLO-MAE framework in the Appendix.

Q3: MLO-MAE utilizes image class labels as a multi-objective training method

We appreciate the reviewer’s keen observation on the mechanism of our MLO-MAE. However, we would like to point out that MLO-MAE is not a multi-objective training method (MOO). MOO aims to optimize multiple objectives simultaneously and disregards the inherent hierarchical dependencies that lie within the definition of the optimization problem. The convention of PT-FT pipeline explicitly defines a hierarchical relationship between pretraining and fine-tuning objectives, since the fine-tuning is directly depending on the pretraining. MLO-MAE leverages this explicit hierarchical dependency and introduces a learnable masking network on top of the existing dependency to form a multi-level optimization objective. The design of MLO has multiple benefits over MOO. Firstly, MLO explicitly models the dependency between levels, ensuring that the solution at one level supports and enhances the objectives of the next level. Secondly, MLO avoids explicitly balancing competing objectives within a single level.

Our transfer learning experiments on additional dense prediction tasks, as well as continued pretraining setting, demonstrate that MLO-MAE achieves superior performance compared to baseline models. This highlights the model's ability to learn more robust representations and exhibit greater generalizability. It worth to mention that, under the transfer learning setting—where neither the dataset nor the labels used for evaluation are present during MLO-MAE's pretraining or the baseline's pretraining stages—the comparison remains fair. This ensures that MLO-MAE's enhanced performance does not stem from leveraging additional information unavailable to the baseline. In addition, we explored MLO-MAE as a continued pretraining paradigm against baseline methods on domain specific dataset in Section 4.5, and our empirical results show that MLO-MAE is effective in continued pretraining domain.

Q4: Presentation of formulas in Section 3.2

We apologize for any confusion. Due to the page limit, we have included more detailed explanations of formulas from Section 3.2 in Appendix A.1 and A.3.

Q5: Improvement in downstream tasks over ImageNet is small

We thank the reviewer for this observation. Table 3 contains transfer learning experiments on iNat 19, CUB, and cars. Similar marginal improvements are also reported in SemMAE, which also tends to have larger improvements over ImageNet than the fine-grained datasets (1.7% improvements on ImageNet, and 0.3%, 0.2%, 0.6 on iNat 19, cub, and cars, respectively) under transfer learning setting. The experimental result of MLO-MAE on ADE20K semantic segmentation (Table 4) and MS-COCO object detection (Table 5) showed non-trivial performance boost, indicating the improvement is not limited to the pretraining dataset.

We appreciate your constructive feedback and provide our response to your questions as follows.

Q1: Unfair comparison with baseline methods

We thank the reviewer for this helpful comment. It is important to emphasize that our method does not utilize more labeled data than the baseline approaches. The labeled data used in Stages II and III of our framework is identical to that employed during the fine-tuning phase of the baselines, ensuring a fair comparison. In MLO-MAE, we directly leverage the downstream objective to guide and optimize the mask-reconstruction process within the MAE pretraining procedure. Additionally, during downstream evaluation, labels are used consistently across all methods, maintaining the fairness of the comparison. The effectiveness of MLO-MAE is further highlighted in transfer learning experiments, where neither the data nor the labels from downstream tasks are incorporated during the pretraining of MLO-MAE or the baselines.

In addition, we conducted transfer learning experiments on a fine-grained dataset using ImageNet fine-tuned checkpoints to facilitate a fair comparison. Results are shown in the table below. In this setting, both the baseline method and the MLO-MAE vision backbone have access to the same ImageNet training data and labels, ensuring a fair evaluation. Notably, MLO-MAE consistently outperformed the baseline methods across all three datasets, demonstrating its robustness and effectiveness.

Table 1: Performance of ImageNet fine-tuned MAE, SemMAE, and MLO-MAE across fine-grained classification.

Q2: Confusions on transfer learning experiments

We apologize for the confusion. To clarify, the experimental settings in Section 4.4 did not involve a separate fine-tuning process. To ensure a fair comparison and incorporate ImageNet label information into the baseline experiments, we also conducted transfer learning using ImageNet fine-tuned baseline checkpoints. The results, provided in the table above, demonstrate that starting from ImageNet fine-tuned checkpoints, MLO-MAE consistently achieves comparable or superior performance on fine-grained classification tasks.

Q3: Masking patterns of MLO-MAE

We appreciate the reviewer for this helpful comment. We have provided a detailed analysis of the masking patterns learned by MLO-MAE in Appendix C.1, alongside a comparison with the patterns produced by MAE, SemMAE, and AutoMAE. Unlike MAE, SemMAE, and AutoMAE, which predominantly mask background regions that are irrelevant to image class labels, MLO-MAE strategically guides the encoder network to focus on learning effective representations for objects. These results highlight the advantages of our approach in prioritizing meaningful features during pretraining.