PaperlayerView exact revision on arXiv

Source-linked AI summary

What to Hide from Your Students: Attention-Guided Masked Image Modeling

Ioannis Kakogeorgiou, Spyros Gidaris, Bill Psomas, Yannis Avrithis, Andrei Bursuc, Konstantinos Karantzalos, Nikos Komodakis

arXiv:2203.12719v2cs.CV

TL;DR

Random masking is less challenging for images because redundant patches can preserve visual concepts, motivating informed token selection for MIM. The paper introduces AttMask, which uses a teacher's attention map to guide student masking, and reports faster learning and stronger performance across data-limited and downstream settings.

  • Problem

    Random image-token masking is not well explored and can leave visual concepts identifiable because images contain many redundant patches.

  • Method

    AttMask uses a teacher transformer encoder's self-attention map to mask salient tokens while a student reconstructs their features in distillation-based MIM.

  • Results

    AttMask improves performance over random masking across data-limited and downstream evaluations, while also accelerating learning.

  • Takeaways & Limitations

    AttMask is most beneficial in challenging settings such as limited-data evaluation and downstream use of pretrained features without fine-tuning.

Abstract

from arXiv · show

Transformers and masked language modeling are quickly being adopted and explored in computer vision as vision transformers and masked image modeling (MIM). In this work, we argue that image token masking differs from token masking in text, due to the amount and correlation of tokens in an image. In particular, to generate a challenging pretext task for MIM, we advocate a shift from random masking to informed masking. We develop and exhibit this idea in the context of distillation-based MIM, where a teacher transformer encoder generates an attention map, which we use to guide masking for the student. We thus introduce a novel masking strategy, called attention-guided masking (AttMask), and we demonstrate its effectiveness over random masking for dense distillation-based MIM as well as plain distillation-based self-supervised learning on classification tokens. We confirm that AttMask accelerates the learning process and improves the performance on a variety of downstream tasks. We provide the implementation code at https://github.com/gkakogeorgiou/attmask.

1 Introduction

The paper argues that random masking is poorly suited to images because redundant, numerous patches often leave visual concepts identifiable. It introduces AttMask, which uses transformer self-attention to mask informative tokens and create harder MIM objectives.

  • Motivation: Image-token masking is less effective when random because redundant patches can preserve visual concepts even after masking.The paper contrasts images with text, where masking words is more likely to hide semantic entities.
  • Motivation: The paper asks whether masking can make informed decisions about which image tokens to hide.
  • Approach: AttMask uses ViT self-attention to mask highly attended tokens corresponding to coherent, discriminative image regions.The intended effect is a more challenging token-reconstruction task than random masking.
  • Approach: AttMask fits distillation-based self-supervision because the teacher sees the full image while the student reconstructs from the masked image.The authors state that this incorporation incurs zero additional cost.
  • Contributions: Experiments report faster learning, better data-limited and downstream performance, and greater robustness to background changes than random masking.The reported robustness suggests reduced background dependency.

2 Related Work

Related work covers vision transformers, self-supervised learning, masked image modeling, and masking-based regularization. AttMask differs by using transformer attention to hide salient patches for dense feature reconstruction.

  • Vision Transformers: Vision transformers represent images as patch tokens and apply transformer architectures based on self-attention.The related work also notes their computational and pretraining requirements.
  • Self-supervised Learning: Self-supervised learning trains models on annotation-free pretext tasks, including patch-order prediction, rotation prediction, and contrastive objectives.
  • Masked Image Modeling: Prior MIM methods use random or block-wise random masking with discrete-token prediction, pixel regression, or autoencoding.Examples include BEiT, SimMIM, and MAE.
  • Masked Image Modeling: MST masks low-attended patches, whereas AttMask advocates masking highly attended patches in DINO- and iBOT-based distillation.
  • Regularization and Augmentation: Masking methods such as Dropout, SpatialDropout, DropBlock, and Attention Dropout regularize representations by removing activations or image regions.The paper frames AttMask as attention-guided dense reconstruction without supervision.

3 Method

AttMask derives token masks from the teacher transformer's [CLS]-attention map, then trains a student to distill teacher features on masked tokens. It primarily masks the most attended patches, with a hint variant that reveals a few salient tokens.

  • Preliminaries: The method tokenizes images into patch embeddings, prepends a [CLS] token, and processes the sequence with a vision transformer.The encoder retains positional information and outputs fixed-length token embeddings.
  • Attention Map Generation: AttMask averages attention matrices across heads and uses the [CLS] row as an attention vector over image tokens.The vector is reshaped into a two-dimensional map indicating attended image regions.
  • Preliminaries: In distillation-based MIM, a teacher generated by exponential moving average produces targets, while the student processes masked views.iBOT applies dense reconstruction loss to masked tokens alongside global [CLS]-token distillation.
  • Mask Generation: AttMask-High ranks attention values and masks the top-k tokens, where k is proportional to the token count and mask ratio.The selected indices form the high-attention mask applied to the tokenized image.
  • Mask Variants: AttMask-Low instead masks low-attended background patches and performs poorly with the considered MIM loss.
  • Mask Variants: AttMask-Hint masks highly attended tokens while randomly revealing a small subset of the most attended patches.The revealed patches provide hints about the masked object's identity, especially at high mask ratios.

4 Experiments

Experiments evaluate AttMask across distillation-based MIM, self-distillation, efficiency settings, robustness tests, and downstream tasks. AttMask-High generally outperforms random masking, accelerates training, and improves transfer performance.

  • Experimental setup: Experiments pretrain iBOT and DINO on 20% or 100% of ImageNet-1k and evaluate classification, robustness, detection, segmentation, retrieval, and video object segmentation.Evaluation uses k-NN, linear probing, fine-tuning, and feature-based protocols.
  • Masking strategies: distillation-based MIM: AttMask-High improves iBOT over default random block-wise masking by +3.0% on k-NN and +1.5% on linear probing in the 20% ImageNet setting.It outperforms the other evaluated masking strategies across the reported evaluation metrics.
  • Data and training efficiency: AttMask-High matches random block-wise masking with 42% fewer training epochs and consistently performs better when using lower percentages of ImageNet-1k.These experiments assess training and data efficiency with iBOT.
  • Masking strategies: self-distillation only: AttMask-High improves DINO k-NN by +0.5 over default DINO without an MIM reconstruction loss, showing effectiveness in self-distillation-only training.AttMask-Low is inferior, while AttMask-Hint is slightly better than AttMask-High in this setting.
  • Downstream tasks: AttMask improves transfer performance, including 0.6% APb for object detection, 0.4% mIoU for semantic segmentation, 1-3% mAP for image retrieval, and 2-6% R@1 for fine-grained classification.It also improves video object segmentation on all reported metrics without fine-tuning.

5 Conclusion

AttMask hides discriminative image cues to create harder self-supervised objectives, accelerating learning and improving downstream performance, robustness, and data-limited evaluation.

  • AttMask hides discriminative image cues from the student, producing more challenging self-supervised objectives.
  • AttMask accelerates learning and improves performance across a variety of downstream tasks.
  • AttMask increases robustness to background changes, indicating reduced background dependency.
  • Improvements are most pronounced without additional learning or fine-tuning and in limited-data settings.

A More Experiments

The appendix adds benchmarks, ablations, and visualizations to further evaluate AttMask.

  • The appendix provides additional benchmarks, ablations, and visualizations.

A.1 More Benchmarks

Additional experiments show that AttMask improves representation quality, robustness, low-shot performance, scene classification, and supervised fine-tuning outcomes across varied evaluation settings.

  • Patch Features: 9.0% k-NN accuracy separates AttMask from baseline iBOT when global average pooling uses patch features.This supports a more challenging MIM objective that improves patch representations.
  • Non-Salient Parts: AttMask yields larger gains as salient image parts are increasingly hidden, indicating better use of non-salient parts or background context.
  • Downstream Tasks using Linear Probing: 0.8% accuracy improves scene classification on Places205 with linear probing.
  • Training for More Epochs: 0.4% k-NN and 0.1% linear probing improvements remain after 300-epoch training on full ImageNet-1k.
  • Training for More Epochs: AttMask improves one-example-per-class k-NN accuracy by 1.5% over default iBOT.
  • Full fine-tuning on ImageNet-1k: 0.2% improvement follows supervised fine-tuning, increasing accuracy from 81.1% to 81.3%.

A.2 More Ablations

Ablations show that AttMask performs best with its default loss weighting and moderate mask-ratio variation, while attention visualizations compare its masking behavior with block-wise masking.

  • MIM Loss Weight: 3% k-NN accuracy is the greatest gain for AttMask at the default MIM loss weight λ = 1.Increasing the MIM loss weight produces larger k-NN gains.
  • Masking strategy and mask ratio: AttMask-High with a 10-50 mask ratio gives the best ablation results.
  • Attention visualizations: Attention maps visualize top-60%-mass patch tokens per head for default block-wise masking and AttMask.

A.3 More Visualizations

The visualizations show that AttMask focuses masking on salient image regions, whereas random strategies do not consistently hide informative content. This produces a more challenging masked image modeling task and is associated with greater foreground focus.

  • Attention Maps: With AttMask, attention heads mostly focus on salient objects or object parts rather than the background.The visualization uses models pretrained on 20% of ImageNet and compares iBOT with block-wise random masking.
  • Attention Maps: Focusing reconstruction on highly attended masked tokens appears to make the network focus more on foreground objects.
  • Masking Examples: AttMask-High and AttMask-Hint use attention to hide salient or nearly all salient image parts, unlike random masking strategies.AttMask-Hint leaves very salient patches as hints, while AttMask-High hides salient patches.

B Experimental Setup

The experimental setup combines multi-crop training, masked and classification-token losses, specified pretraining schedules, and downstream evaluations with and without finetuning. It also includes visual comparisons across masking strategies and task-specific dataset protocols.

  • Training Setup: The training objective is a weighted sum of masked-image modeling, global classification-token, and local-crop classification-token losses.The local-crop loss compares student outputs for non-masked local crops with teacher outputs for non-masked global views.
  • Training Setup: Analysis and ablation models are pretrained on 20% of ImageNet-1k for 100 epochs using ViT-S/16, AdamW, warm-up, and cosine learning-rate decay.The batch size is 240 unless otherwise stated.
  • Training Setup: The multi-crop scheme uses two 224^2 global crops and six 96^2 local crops, scaling training time by γ = 3.10.Color jittering, Gaussian blur, and solarization are also used as augmentations.
  • Evaluation: ImageNet-1k evaluation uses k-NN with k = 20 and linear probing on validation features.The k-NN protocol uses the last-layer [CLS] feature, while linear probing trains an SGD classifier for 100 epochs.
  • Downstream Evaluation: Finetuning evaluations cover CIFAR10, CIFAR100, and Oxford Flowers, with dataset-specific training durations and shared optimization settings.CIFAR10 and CIFAR100 use 500 epochs, while Oxford Flowers uses 1000 epochs.
  • Downstream Evaluation: Additional downstream evaluations cover COCO detection and instance segmentation, ADE20K semantic segmentation, retrieval datasets, Places205, and DAVIS 2017.These evaluations use task-specific protocols, including Cascade Mask R-CNN for COCO and UperNet for ADE20K.
Loading 2203.12719v2…