Salience-Based Adaptive Masking: Revisiting Token Dynamics for Enhanced Pre-training

TL;DR

This paper tackles the sensitivity of Masked Image Modeling to masking ratios by proposing Salience-Based Adaptive Masking (SBAM), which masks tokens based on a computed token salience derived from outgoing attention weights, plus randomness for diversity. Building on SBAM, it introduces Adaptive Masking Ratio (AMR) to tailor the masking ratio per image according to salience, enabling per-sample pre-training customization. Through extensive experiments on ImageNet-1K with ViT backbones, SBAM achieves robust performance across masking ratios, faster convergence, and consistent gains across multiple baselines (e.g., MAE, BootMAE, iBoT, CMAE) and masking strategies, outperforming state-of-the-art approaches in both fine-tuning and linear probing. The work demonstrates that adaptive, salience-driven masking can significantly enhance pre-training efficiency and effectiveness in MIM, with practical impact for scalable self-supervised vision models.

Abstract

In this paper, we introduce Saliency-Based Adaptive Masking (SBAM), a novel and cost-effective approach that significantly enhances the pre-training performance of Masked Image Modeling (MIM) approaches by prioritizing token salience. Our method provides robustness against variations in masking ratios, effectively mitigating the performance instability issues common in existing methods. This relaxes the sensitivity of MIM-based pre-training to masking ratios, which in turn allows us to propose an adaptive strategy for `tailored' masking ratios for each data sample, which no existing method can provide. Toward this goal, we propose an Adaptive Masking Ratio (AMR) strategy that dynamically adjusts the proportion of masking for the unique content of each image based on token salience. We show that our method significantly improves over the state-of-the-art in mask-based pre-training on the ImageNet-1K dataset.

Figures (7)

• Figure 1: Overview of SBAM. Whereas (a) Random Masking must rely completely on chance and carefully tuned masking ratio to guarantee effective masking, (b) the proposed SBAM strategically masks tokens based on the token salience. The robustness of SBAM paved the way for the introduction of (c) AMR, which implements a tailored masking ratio for each sample in the dataset.
• Figure 2: Qualitative example of SBAM. SBAM introduces 'token salience' to prioritize and mask tokens with high significance. Hence, it is qualitatively confirmed that particularly important objects with high contribution within the image are selectively masked. Moreover, by integrating randomness with token salience, masks are probabilistically assigned to the background and less significant tokens, enriching the diversity of the token masking.
• Figure 3: Evaluation of robustness across varied masking ratios. To evaluate the robustness of SBAM, we report the comparative analysis of image classification performance on ImageNet-1K dataset deng2009imagenet against the baseline method, MAE he2022masked, using ViT-L dosovitskiy2020image as a backbone. The upper graphs display the performance of the methods at different masking ratios, while the lower graphs illustrate the Performance Improvement over Masking Ratio (PIMR) and Global PIMR. These measures indicate the extent of each model's performance enhancement as the masking ratio increases from the lowest to higher ratios. SBAM significantly outperforms MAE in every measures, demonstrating its superior effectiveness in handling various masking ratios and enhanced pre-training performances.
• Figure 4: Performance evaluation of SBAM with respect to the pre-trained epochs. We report the comparison of image classification accuracy on ImageNet-1K deng2009imagenet, pre-trained on ViT-L dosovitskiy2020image. The left graph displays fine-tuning accuracy, whereas the right graph illustrates linear probing accuracy, both over a range of pre-trained epochs. The curves illustrate that SBAM surpasses MAE he2022masked in pre-training effectiveness in every trained epoch, and also validates its quicker attainment of converged performance levels.
• Figure 5: Qualitative example of AMR. Having masking ratios that adapt throughout training is highly effective, as it allows the masking process to be tailored to each sample in the dataset, accommodating the unique composition and object sizes within each image, as shown in the above qualitative samples.