Source-linked AI summary
MST: Masked Self-Supervised Transformer for Visual Representation
Zhaowen Li, Zhiyang Chen, Fan Yang, Wei Li, Yousong Zhu, Chaoyang Zhao, Rui Deng, Liwei Wu, Rui Zhao, Ming Tang, Jinqiao Wang
TL;DR
Visual self-supervised learning has largely emphasized global, high-level representations, limiting transfer to dense prediction tasks that require local and spatial information. MST uses attention-guided token masking and a global image decoder to preserve both local context and spatial structure, achieving strong classification and dense-prediction results.
Problem
Previous visual self-supervised methods mainly learn global, high-level representations, whereas dense prediction tasks require local and pixel-level information.
Method
MST uses the teacher’s multi-head self-attention map to mask nonessential patch tokens and a global image decoder to reconstruct spatial image information.
Results
MST achieves strong performance across classification and dense prediction tasks, including 76.9% top-1 accuracy with DeiT-S after 300-epoch pre-training and 42.7% bbox mAP after 100 epochs.
Takeaways & Limitations
MST provides a versatile self-supervised visual representation that transfers to both image classification and downstream dense prediction tasks.
Takeaways & Limitations
The broader-impact discussion identifies large-scale drone surveillance and face recognition as application areas associated with privacy or societal risks.
Abstract
from arXiv · showhide
Transformer has been widely used for self-supervised pre-training in Natural Language Processing (NLP) and achieved great success. However, it has not been fully explored in visual self-supervised learning. Meanwhile, previous methods only consider the high-level feature and learning representation from a global perspective, which may fail to transfer to the downstream dense prediction tasks focusing on local features. In this paper, we present a novel Masked Self-supervised Transformer approach named MST, which can explicitly capture the local context of an image while preserving the global semantic information. Specifically, inspired by the Masked Language Modeling (MLM) in NLP, we propose a masked token strategy based on the multi-head self-attention map, which dynamically masks some tokens of local patches without damaging the crucial structure for self-supervised learning. More importantly, the masked tokens together with the remaining tokens are further recovered by a global image decoder, which preserves the spatial information of the image and is more friendly to the downstream dense prediction tasks. The experiments on multiple datasets demonstrate the effectiveness and generality of the proposed method. For instance, MST achieves Top-1 accuracy of 76.9% with DeiT-S only using 300-epoch pre-training by linear evaluation, which outperforms supervised methods with the same epoch by 0.4% and its comparable variant DINO by 1.0\%. For dense prediction tasks, MST also achieves 42.7% mAP on MS COCO object detection and 74.04% mIoU on Cityscapes segmentation only with 100-epoch pre-training.
1 Introduction
Visual self-supervised methods often learn global, image-level representations, while dense prediction requires local and pixel-level information. MST addresses this gap with attention-guided masking and image reconstruction, and reports strong results across visual tasks.
- 1 Introduction: Global-feature self-supervised methods are suboptimal for pixel-level predictions such as object detection and semantic segmentation.The cited discussion identifies limited local information as a central issue in visual representation transfer.
- 1 Introduction: MST uses a multi-head self-attention map to dynamically mask patch tokens while preserving crucial image structure.The strategy masks tokens without increasing training time, according to the cited method description.
- 1 Introduction: 76.9% top-1 accuracy is achieved with DeiT-S under the ImageNet linear evaluation protocol.The reported result uses 300-epoch pre-training and is described as state-of-the-art in the cited passage.
- 1 Introduction: A global image decoder reconstructs spatial image information, making the learned representation more suitable for downstream dense prediction tasks.The decoder recovers spatial information from the image representation.
- 1 Introduction: MST is presented as a masked self-supervised Transformer that captures local context without damaging crucial foreground structure.This contribution combines local patch masking with preservation of important image regions.
2 Related Works
Prior visual self-supervised work includes masked reconstruction and instance-discrimination approaches, while dense prediction methods apply supervision or matching at more local levels. MST combines attention-guided token masking with global image reconstruction.
- 2 Related Works: Masked reconstruction methods apply the MLM paradigm to vision by masking and reconstructing pixels or patches.iGPT masks and reconstructs pixels, whereas ViT masks and reconstructs patches.
- 2 Related Works: MST introduces attention-guided masking and an image restoration task alongside momentum design, asymmetric augmentations, and multi-crop strategies.The method combines these components within a student-teacher Transformer framework.
- 2 Related Works: Instance discrimination learns representations by distinguishing each image from other images, with MoCo, SimCLR, and BYOL modifying how representations are compared or bootstrapped.These methods use momentum encoders, same-batch negatives, or direct representation bootstrapping, respectively.
- 2 Related Works: Self-supervised dense prediction methods extend representation learning through embedding similarity, composited foreground localization, or pixel-level contrastive learning.The cited examples are Self-EMD, Insloc, and PixPro.
- 2 Related Works: Attention-guided masking preserves key image patterns more reliably than random masking by targeting nonessential regions.Figure 2 contrasts random masking, which may remove crucial features, with the proposed strategy.
3 Methods
MST combines attention-guided token masking with image restoration in a teacher–student Transformer framework, preserving crucial regions while recovering spatial information for downstream tasks.
- MST framework: MST combines momentum-based teacher–student networks, asymmetric augmentations, and multi-crop views with attention-guided masking and image restoration.The teacher and student share the same Transformer architecture, while restoration is performed on the student network.
- Masked token strategy: Random masking can remove crucial object regions, producing indistinguishable semantic features and weakening object recognition.The paper reports that crucial and nonessential tokens have the same masking probability under the random strategy.
- Masked token strategy: The attention-guided strategy uses the teacher’s final-layer, multi-head-averaged class-to-patch attention to select low-attention masking candidates.Candidate patches are selected using an attention threshold, then masked probabilistically without additional training-time cost.
- Masked token strategy: Highest-scoring patches remain visible, while selected regions are replaced with a learnable [MASK] embedding before student processing.The final masked tokens are represented as m⊙e, and the strategy preserves patches with the highest attention scores.
- Image restoration: Unlike MLM, MST reconstructs the original image with a CNN-based global decoder, encouraging pixel-level information capture and spatial-structure recovery.The restoration loss compares the input image with the decoder output, and only the student parameters are retained for pre-training.
4 Experiments
Experiments evaluate MST across transformer backbones, ImageNet classification, dense prediction, and ablations. The results show strong performance, transferability, and benefits from attention-guided masking and the decoder-compatible normalization design.
- ImageNet evaluation: 76.9% top-1 accuracy is achieved by MST with DeiT-S after 300 epochs of linear probing on ImageNet.
- ImageNet evaluation: MST outperforms MoBY by 1.8% with Swin-T under the same training epochs, supporting application across transformer architectures.Swin-T uses similar parameter counts to DeiT-S and shares its hyperparameters in this experiment.
- Semantic segmentation: 74.7% mIoU and 82.35% mAcc are achieved on Cityscapes, exceeding supervised and DINO pretrained results.The evaluation uses 300-epoch pretrained models for DINO and MST.
- Impact of different mask strategy: 73.7% top-1 accuracy with attention-guided masking exceeds 73.1% without masking, while random masking reduces performance to 63.2%.The attention-guided strategy is intended to preserve essential regions while encouraging local-context modeling.
- Impact of w/o BN: The projection-head update rule slightly improves top-1 accuracy relative to normal batch normalization, which downgrades the baseline model.The paper relates this design to structural consistency with the Conv-BN-ReLU global image decoder.
5 Conclusion
MST addresses limited local information extraction and spatial-information loss in visual self-supervised learning through attention-guided masking and global image decoding.
- MST targets the lack of local information extraction and loss of spatial information in current visual self-supervised learning.
- The attention-guided mask strategy uses the teacher model’s multi-head self-attention map to capture local relationships while preserving global semantic information.The strategy does not cause extra computation cost.
- A global image decoder recovers image spatial information under the attention-guided masking strategy, supporting dense prediction tasks.The paper describes this recovery as vital for dense prediction tasks.
- MST shows versatility and scalability across multiple downstream visual tasks.
Broader Impact
The paper identifies applications including safer autonomous driving, factory human-robot collaboration, and elderly-care robots, while noting surveillance and privacy risks.
- MST’s improved feature extraction could support pedestrian detection for improved road safety in autonomous driving.
- The method could support factory robots taking on risk-prone jobs, enabling safer human-robot collaboration and potentially saving lives.
- MST could be used in assistive robots for elderly care.
- Potential negative societal impacts include large-scale drone surveillance and privacy exposure from face recognition.
B Data augmentation
The paper specifies an augmentation pipeline, normalization choices, and evaluation details for weighted k-nearest-neighbor feature assessment.
- Data augmentation: The augmentation pipeline combines random resized cropping, horizontal flipping, color jittering, grayscale conversion, Gaussian blurring, solarization, and multi-crop.Random resized cropping and multi-crop are always applied; the other transformations are probabilistic.
- Data augmentation: MST uses SyncBN by default, with batch-normalization statistics updated from different image sets than in SimCLR.The paper notes that these implementations should lead to different results because they influence gradient variance differently.
- Feature evaluation: Weighted k-nearest-neighbor evaluation freezes the pretrained model and extracts class-embedding features for training and validation data.The paper reports that k=10 consistently gives the best accuracy across runs.
- Feature evaluation: Table 7 reports Top-1 accuracy on the ImageNet validation dataset for different k values using a 300-epoch pretrained DeiT-S model.
E Linear probing
The linear-probing setup evaluates frozen pretrained features, while the MST pseudocode combines image restoration and instance-discrimination losses with momentum teacher updates.
- Linear probing: Linear probing removes MLP heads and trains a supervised linear classifier on frozen features after self-supervised pre-training.The setup uses 100 ImageNet training epochs with SGD, batch size 1024, weight decay 0, and learning rate 0.00024.
- Algorithm: MST’s pseudocode defines backbone projection functions, a decoder, momentum coefficient, temperature coefficient, class tokens, attention maps, decoder inputs, and loss weights.
- Algorithm: For each batch, teacher and student outputs are computed, decoder reconstructions are produced, and restoration and instance-discrimination losses are combined.The combined loss is backpropagated.
- Algorithm: The teacher network is updated as an exponential moving average of the student using the momentum coefficient.
- Algorithm: The instance-discrimination loss applies temperature-scaled softmax outputs and cross-entropy from teacher predictions to student predictions.
- Linear probing: For DeiT-S, linear probing uses last-layer class tokens, while DINO’s reported comparison uses concatenated late blocks for fairness.
- Linear probing: Table 8 compares linear-probing strategies using Top-1 accuracy on ImageNet validation with a 100-epoch pretrained DeiT-S model.
F Impact of longer training
Longer training improves MST’s DeiT-S performance across linear-probing settings, consistent with earlier self-supervised learning methods.
- F Impact of longer training: Longer training improves MST performance with DeiT-S regardless of the linear-probing method.This trend is consistent with previous self-supervised learning methods.
G Implementation pseudo code
The implementation uses AdamW with specified optimization settings, followed by warmup and cosine learning-rate decay, while the referenced tables document implementation comparisons.
- G Implementation pseudo code: MST is optimized with AdamW using a 2 × 10^-3 learning rate, batch size 1024, and initial weight decay 0.04.
- G Implementation pseudo code: Table 9 documents the impact of longer training and identifies a variant using DINO linear probing.
- G Implementation pseudo code: After 10 warmup epochs, the learning rate follows cosine decay, with multi-crop augmentation and settings adapted from prior methods.
- G Implementation pseudo code: Table 10 presents differences between the proposed approach and BERT.
H Differences with BERT
The paper contrasts MST with pure MLM and random masking, while examining token placement, attention visualization, and restoration-loss weighting.
- H Differences with BERT: Pure MLM with DeiT-S reaches about 40% after 100 epochs and at best 61% after hyperparameter tuning, below DINO’s 71.6% and supervised learning’s 68.7%.
- H Differences with BERT: Random masking performs best without batch normalization at 72.6% when the masking ratio p is 0.
- H Differences with BERT: Restoration-loss coefficients from 0.2 to 0.8 produce 73.7%, 73.5%, 73.9%, and 73.6%, indicating limited sensitivity to λ2.
- H Differences with BERT: MST applies masking after linear projection, when image patches have been mapped into tokens or embeddings.
- H Differences with BERT: MST produces clearer attention-map visualizations than the supervised method.
- H Differences with BERT: Figure 3 compares the input image with supervised and MST attention maps from left to right.