PaperlayerView exact revision on arXiv

Source-linked AI summary

Exploring Target Representations for Masked Autoencoders

Xingbin Liu, Jinghao Zhou, Tao Kong, Xianming Lin, Rongrong Ji

arXiv:2209.03917v3cs.CV

TL;DR

Masked autoencoders reconstruct masked inputs using target representations, but it is unclear whether those targets must be carefully chosen. The paper proposes multi-stage masked distillation with bootstrapped, randomly initialized teachers; dBOT reports favorable results across classification, detection, and segmentation.

  • Problem

    The paper investigates whether masked autoencoders require careful selection of target representations for learning good visual representations.

    Source: §1 · p. 1

  • Method

    dBOT repeatedly distills masked inputs across stages, bootstrapping each teacher from the preceding student and allowing random initialization.

    Source: §1 · p. 2, §4 · p. 4

  • Results

    Across classification, object detection, and semantic segmentation, dBOT achieves favorable performance, including 84.5% ImageNet-1K top-1 fine-tuning accuracy with ViT-B/16 while surpassing MAE.

    Source: §1 · p. 2, §8 · p. 9

  • Takeaways & Limitations

    With multi-stage distillation, careful target design is not necessary for learning good visual representations, and random teachers avoid an extra pretraining stage.

    Source: §8 · pp. 8–9

    The CLIP-teacher comparison uses 400M image-text pairs, a drastically different data setup from multi-stage distillation on ImageNet-1K alone.

    Source: §C.3 · p. 13

Abstract

from arXiv · show

Masked autoencoders have become popular training paradigms for self-supervised visual representation learning. These models randomly mask a portion of the input and reconstruct the masked portion according to the target representations. In this paper, we first show that a careful choice of the target representation is unnecessary for learning good representations, since different targets tend to derive similarly behaved models. Driven by this observation, we propose a multi-stage masked distillation pipeline and use a randomly initialized model as the teacher, enabling us to effectively train high-capacity models without any efforts to carefully design target representations. Interestingly, we further explore using teachers of larger capacity, obtaining distilled students with remarkable transferring ability. On different tasks of classification, transfer learning, object detection, and semantic segmentation, the proposed method to perform masked knowledge distillation with bootstrapped teachers (dBOT) outperforms previous self-supervised methods by nontrivial margins. We hope our findings, as well as the proposed method, could motivate people to rethink the roles of target representations in pre-training masked autoencoders.The code and pre-trained models are publicly available at https://github.com/liuxingbin/dbot.

1. Introduction

The paper asks whether masked autoencoders need carefully chosen reconstruction targets and finds that multi-stage distillation makes teacher choice less consequential. It introduces dBOT, which bootstraps teachers from students and can start from random initialization.

1. Introduction

  • MIM masks part of an input and reconstructs the masked portion using a transformed target representation.The target transformation may be generated by a parameterized network or a traditional image-feature method.Source: §1 · p. 1
  • The study compares teachers from contrastive learning, masked autoencoding, supervised learning, and autoregressive generation, all pretrained on ImageNet-1K.It reports that students distilled from these different teachers behave similarly after several stages.Source: §1 · pp. 1–2
  • After several distillation stages, students from different teachers show rapidly shrinking performance variance and similar weight and output-feature properties.The paper concludes that careful target design is unnecessary for learning good visual representations with multi-stage distillation.Source: §1 · p. 2
  • dBOT repeatedly distills across stages, assigning the student’s weights to the teacher and reinitializing the student at each stage boundary.The pipeline permits pretraining from a randomly initialized teacher.Source: §1 · p. 2
  • dBOT significantly surpasses MAE on ImageNet-1K, reaching 84.5%, 86.6%, and 88.0% top-1 fine-tuning accuracy with ViT-B/16, ViT-L/16, and ViT-H/14, respectively.It also reports 52.7 and 56.0 APbox on COCO and 49.5 and 54.5 mIoU on ADE20K with ViT-B/16 and ViT-L/16, respectively.Source: §1 · p. 2

2. Related work

Related work spans masked visual reconstruction with targets ranging from learned representations to pixels and hand-crafted features, alongside knowledge-distillation methods. The paper identifies a lack of system-level study on target choice in masked knowledge distillation.

2.1. Self-Supervised Visual Learning

  • Masked visual learning uses DALL-E tokens, online-teacher targets, pixels, or HOG features; MAE’s asymmetric encoder-decoder improves training efficiency.data2vec reconstructs representations from multiple neural-network layers and reports generalization across vision, speech, and language.Source: §2.1 · p. 2

2.2. Knowledge Distillation

  • Prior masked knowledge-distillation examples include BEiT with dVAE, MaskFeat with HOG, and MVP with CLIP, while conventional feature distillation showed negligible gains on MIM-trained models such as MAE.Related work also describes feature distillation matching contrastive models with MIM-trained ones.Source: §2.2 · p. 2
  • The authors identify no prior system-level study of how to choose target representations or teacher networks for masked knowledge distillation.They propose a randomly initialized teacher bootstrapped across stages as an alternative practice.Source: §2.2 · p. 2

3. Does hφ(·) Matter in MKD?

The section tests whether teacher choice affects student transfer performance, weights, and outputs, and what reduces differences between students. It finds that multi-stage distillation narrows teacher-dependent performance gaps and that random teachers can work comparably to pretrained teachers.

3. Does hφ(·) Matter in MKD?

  • The study asks whether teacher choice changes transfer performance or model weights and outputs, and what matters more if teacher choice does not.It uses a standard masked-autoencoder framework for its system-level study.Source: §3 · p. 2

3.1. Preliminary Study

  • Students distilled from different pretrained teachers have closer downstream performances than the teachers, with performance variance dropping from 2.24 to 0.37 after the first stage.The four teacher types are DeiT, DINO, DALL-E, and MAE; after stage one, students consistently outperform their teachers.Source: §3.1 · p. 3, Table 1 · p. 3

3.2. Distillation with Multiple Stages

  • With multi-stage distillation, teacher-dependent performance gaps shrink; ImageNet classification variance falls from 2.24 to 0.37, 0.07, and 0.04 across stages.The paper says similar properties of model weights and outputs across layers support its conclusion that teacher choice has little influence on downstream performance.Source: §3.2 · p. 3, Table 1 · p. 3
  • Random-teacher students after multiple stages perform comparably to students from pretrained teachers, with gains of 6.1%, 20.4, and 21.3 performance points over random initialization on classification, detection, and segmentation.Object detection and semantic segmentation empirically require one more stage to saturate than classification.Source: §3.2 · pp. 3–4
  • dBOT updates the teacher at stage breakpoints by copying student parameters and reinitializing the student.Figure 1 contrasts this schedule with a continuously frozen offline teacher and an online teacher updated by exponential moving average.Source: §3.2 · p. 3, Figure 1 · p. 4

4. MKD with Bootstrapped Teachers

dBOT performs masked knowledge distillation in stages, keeping each teacher fixed within a stage and bootstrapping the next teacher from the trained student. At breakpoints it resets the student, differing from both offline fixed-teacher and continuously updated online-teacher approaches.

4. MKD with Bootstrapped Teachers

  • At each breakpoint, dBOT copies the trained student’s parameters into the teacher and reinitializes the student; this repeats until downstream performance saturates.The teacher remains fixed within each stage to provide a stable visual representation.Source: §4 · p. 4
  • Pretrained-teacher approaches require an extra pretraining stage and may involve a different computation pipeline or additional data.The teacher is kept fixed in this offline approach.Source: §4 · p. 4
  • Unlike EMA-based online teachers, dBOT sets m = 0 at breakpoints and m = 1 otherwise.Here m denotes the momentum for updating teacher parameters from student parameters.Source: §4 · p. 4
  • EMA bootstrapping can improve teacher representation quality but is described as unstable and sensitive to hyperparameters.MAE’s pixel-identity target functions similarly to a fixed random teacher, but does not provide teacher-student synergy.Source: §4 · p. 4

5. Experiments

The experiments evaluate dBOT across ImageNet classification, semi-supervised learning, downstream prediction tasks, and ablations. Results show strong performance across model sizes and tasks, while ablations assess stage scheduling and design choices.

5.1. Pre-Training

  • Classification uses two 800-epoch distillation stages by default, while object detection and semantic segmentation use three stages.The pre-training setup uses batch size 4096 and scales the initial learning rate as 1.5e-4 × batch size / 256.Source: §5.1 · pp. 4–5, Table 2 · p. 5

5.2. ImageNet Results

  • 84.5%, 86.6%, and 87.4% top-1 accuracy with ViT-B, ViT-L, and ViT-H, respectively, exceed MAE by 0.9%, 0.7%, and 0.5%; at 448 image size, ViT-H reaches 88.0%.These are ImageNet-1K end-to-end fine-tuning results; Table 2 reports the comparison.Source: §5.2 · p. 5, Table 2 · p. 5
  • With 1% and 10% of ImageNet-1K labels, dBOT outperforms MAE by 1.7 and 1.4 points, respectively, indicating higher label efficiency.Table 3 compares semi-supervised results at the two label fractions.Source: §5.2 · p. 5, Table 3 · p. 5

5.3. Downstream Tasks

  • With ViT-B on COCO, dBOT reaches 52.7 APbox and 45.7 APmask, exceeding supervised pre-training by 2.9 and 2.5 points; gains with ViT-L are 4.8 and 3.6 points.Table 4 reports object detection and instance segmentation results using Cascade Mask R-CNN.Source: §5.3 · p. 5, Table 4 · p. 5
  • On ADE20K, dBOT improves mIoU from 47.4 to 49.5 with ViT-B and from 49.9 to 54.5 with ViT-L, gains of 2.1 and 4.6 points over supervised pre-training.Table 5 reports semantic segmentation results using UperNet.Source: §5.3 · p. 6, Table 5 · p. 6
  • On smaller transfer datasets, dBOT performs comparably to or better than prior best methods, with gains over supervised pre-training of 4.7% on iNaturalist18 and 3.3% on iNaturalist19.Table 6 reports ViT-B transfer accuracy across datasets and gives average results.Source: §5.3 · p. 6, Table 6 · p. 6

5.4. Ablation Study

  • For classification, two-stage distillation performs best at 84.5%; three stages reduce accuracy by 0.1%, and evenly splitting two stages achieves the best performance.With 800 epochs in stage one, 800 epochs suffice for stage two; extending stage two to 1,200 epochs brings no gain.Source: §5.4 · p. 6
  • Ablations favor the vanilla momentum update, Layer Normalization on target patch tokens, re-initializing the student between stages, and a 75% mask ratio.The 75% optimum matches MAE; target patch tokens without Layer Normalization are less suitable for guiding student learning.Source: §5.4 · p. 6

6. Property Analysis

The property analyses find that multi-stage dBOT students from different teachers share similar attention and representation patterns. In unsupervised object localization, students using different teachers achieve similar results and consistently outperform their teachers, while dBOT is slower than MAE but faster than data2vec and BEiT.

Averaged attention distance

  • Distilled students exhibit similar attention patterns across teachers and more local attention than their teachers.Higher average attention distance indicates more global attention.Source: §6 · p. 6, Figure 2 · p. 7

Unsupervised object detection

  • dBOT students using different teachers achieve very similar CorLoc results and consistently outperform their teachers.The evaluation uses SVD-based feature decomposition and CorLoc on POC-VOC 2012 trainval sets.Source: Table 8 · p. 7, §6 · pp. 7–8

Singular value decomposition

  • Students distilled from different teachers exhibit similar singular-value patterns across layers.The analysis reports the percentage of top-k singular values for k=1 to 5; higher percentages indicate less-correlated outputs and greater spatial-representation redundancy.Source: §6 · pp. 6–7, Figure 3 · p. 8

Training efficiency

  • With an asymmetric encoder-decoder, dBOT trains more slowly than MAE but much faster than data2vec and BEiT.The advantage over data2vec and BEiT becomes more pronounced for larger models; training times are reported per epoch on the same setting with 32 NVIDIA A100-80G GPUs.Source: Table 9 · p. 8, §6 · p. 8

7. Distill from Bigger Teachers

Using larger teachers yields no substantial classification gains, but distilling from ViT-H improves dense-prediction results. The authors also explore data-richer teachers such as CLIP, reporting new state-of-the-art self-supervised results.

Setup

  • The larger teachers are first distilled for two stages, then used to distill new students for one stage.ViT-H inputs are resized to 196×196 to match the output length of ViT-B/L.Source: §7 · p. 8

Results

  • +0.8 APbox and +1.3 mIoU are reported for ViT-B students distilled from ViT-H, with +0.1 APbox and +0.7 mIoU for ViT-L students.The authors report no substantial classification gains from larger teachers.Source: Table 10 · p. 8, §7 · p. 8

Exploratory teachers

  • Exploratory masked distillation with data-richer teachers such as CLIP sets new state-of-the-art self-supervised results.The paper refers readers to Appendix C for details.Source: §7 · p. 8

8. Conclusion

The conclusion argues that careful target-representation design is unnecessary when masked distillation is performed in multiple stages. It presents randomly initialized teachers as a practical alternative and reports favorable performance across classification, object detection, and semantic segmentation.

Conclusion

  • With multi-stage distillation using bootstrapped teachers, careful target-representation selection is not necessary for learning good visual representations.The conclusion frames this as a finding about masked autoencoders.Source: §8 · pp. 8–9
  • The method uses randomly initialized teachers instead of pre-trained models, without an extra pre-training stage.The conclusion describes this as a simple practice.Source: §8 · p. 9
  • dBOT achieves favorable performance on image classification, object detection, and semantic segmentation.Source: §8 · p. 9

A. Implementation Details

The appendix specifies pre-training and downstream-task implementation choices for classification, object detection and instance segmentation, and semantic segmentation.

A.1. Pre-Training

  • Pre-training uses Xavier Uniform initialization for the Vision Transformer and asymmetric stochastic drop-path rates for students and teachers.For larger-teacher distillation, larger teachers are trained for two stages and then used to distill students for one stage.Source: §A.1 · p. 11

A.2. Classification

  • Classification follows a common self-supervised ViT fine-tuning recipe, also used when distilling from bigger teachers.The recipe is shown in Table A2.Source: §A.2 · p. 11

A.3. Object Detection and Instance Segmentation

  • Object detection and instance segmentation use ViT with Cascade Mask R-CNN on COCO, adding relative position embeddings for variable image sizes.The authors use the same setup and sweep learning rate and stochastic drop-path rate across methods.Source: §A.3 · p. 11

A.4. Semantic Segmentation

  • Semantic segmentation uses ViT with UperNet on ADE20K and adds relative position embeddings for variable image sizes.The authors use the same setup and sweep learning rate and layer-wise decay across methods.Source: §A.4 · p. 11

B. Additional Experiments

Additional experiments find that random-network patch-token targets perform comparably to pixel targets, while dBOT outperforms other methods on COCO detection and instance segmentation. Linear probing is comparable to MAE, and the paper also explores distillation from data-richer teachers.

B.1. Pixels vs. Random Mapping of Pixels

  • After 400-epoch ViT-B pre-training, randomly mapped patch-token targets achieve results comparable to pixel targets.The targets include tokens from the last transformer block and from a linear projection without a transformer block.Source: §B.1 · pp. 11–12

B.2. Object Detection with Mask R-CNN

  • dBOT outperforms other methods by a large margin on COCO object detection and instance segmentation with Mask R-CNN.The results cover ViT-B and ViT-H.Source: §B.2 · p. 12, Table B5 · p. 12

B.3. Linear Probing

  • dBOT achieves linear-probing performance comparable to MAE under the ViT-B evaluation setup.Source: §B.3 · p. 12

C. Distill from Data-Richer Teachers

C.1. Pre-Training

For data-richer teachers, the authors find that training recipes and distillation stages affect results, and report strong downstream performance from CLIP-distilled students. They note that the comparison with multi-stage distillation remains unresolved because CLIP’s in-house training data was unavailable.

C.1. Pre-Training

  • The authors report that a vanilla architecture without an asymmetric encoder-decoder is optimal, and hypothesize that a decoder may restrain encoder learning when targets contain high-level semantics.They contrast semantic targets such as DINO and CLIP with low-level targets such as pixels and random pixel mappings.Source: §C.1 · pp. 12–13
  • Momentum used for bootstrapping models helps distillation with random teachers but hurts performance with stronger pre-trained teachers.Source: §C.1 · p. 13
  • For stronger teachers, multi-stage MKD’s gains over one-stage MKD decrease, and with CLIP the two-stage result is 0.9% lower than the one-stage result.The comparison is 2×800 versus 1×1600 epochs.Source: §C.1 · p. 13

C.2. Downstream Tasks

  • With CLIP-B as teacher, ViT-B achieves 85.7% top-1 accuracy on IN1K and 52.9 mIoU on ADE20K; with CLIP-L, ViT-H at resolution 448 achieves 89.1% top-1 accuracy on IN1K.The CLIP-B student surpasses previous methods, and the CLIP-L result sets a new image-recognition state of the art.Source: §C.2 · p. 13, Table C8 · p. 13

C.3. Conflict with Main Conclusion

  • The CLIP comparison is not directly resolved: CLIP uses 400M image-text pairs, and the authors lack access to its in-house data for multi-stage distillation.They leave testing whether the results match on that data to future work.Source: §C.3 · p. 13

Related papers

Loading 2209.03917v3…