Research map / Architectures and efficiency / Attention and transformers
Closest Neighbors are Harmful for Lightweight Masked Auto-encoders
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes NoR-MAE, which adds a neighbor-repelling decorrelation loss between visible patches and their closest neighbor patches to train lightweight ViTs from scratch with masked autoencoding, without a teacher model.
Problem
Masked autoencoder pre-training yields poor representations for lightweight vision transformers (e.g., ViT-Tiny), and existing fixes rely on costly knowledge distillation from large teachers.
Approach
Defines the closest neighbor patch (CNP) for each visible patch and penalizes off-diagonal correlations between embeddings of visible patches and their CNPs, reducing semantic aliasing between local semantic centroids; trained from scratch with negligible overhead.
Key results
- NoR-MAE with ViT-Tiny reaches 70.24% on ImageNet-1K at 200 epochs (+3.64% over MAE's 66.60%) and 78.26% on ImageNet-100 (+7.22% over MAE).
- On downstream tasks with ViT-Tiny, NoR-MAE improves over MAE by up to 5.13% accuracy, e.g., Flowers 85.80 to 88.71 and Aircraft 64.60 to 67.07.
- NoR-MAE with ViT-Base reaches 83.42% on ImageNet-1K at 200 epochs, a +0.12% gain over MAE, and the ViT-Tiny model reaches 83.96% on ImageNet-100 with 1600 epochs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add closest-neighbor repelling loss | MAE with ViT-Tiny · ImageNet-1K | supervised fine-tuning accuracy (%): 66.60 -> 70.24 (+3.64) | Helped |
| add closest-neighbor repelling loss | MAE with ViT-Tiny · ImageNet-100 | supervised fine-tuning accuracy (%): 71.04 -> 78.26 (+7.22) | Helped |
| add closest-neighbor repelling loss | MAE with ViT-Tiny · CIFAR-10 | accuracy (%): 95.01 -> 95.98 | Helped |
| add closest-neighbor repelling loss | MAE with ViT-Tiny · ADE20K | mIoU: 26.54 -> 33.54 | Helped |
| use the reported neighbor-repelling penalty settings across λ values (ablation) | NoR-MAE with ViT-Tiny · ImageNet-100 | accuracy (%): 70.24; 76.20; 78.26; 77.89; 72.17 | Mixed |
| vary patch size from 2×2 to 32×32 (ablation) | NoR-MAE with ViT-Tiny · ImageNet-100 | accuracy (%): 83.56%; 83.58%; 79.24%; 78.26%; 74.73% | Mixed |
| vary masking ratio from 50% to 85% (ablation) | NoR-MAE with ViT-Tiny · ImageNet-1K | accuracy (%): 69.27%; 70.24%; 70.16% | No clear effect |
| increase pre-training from 200 to 1,000 or 1,600 epochs (ablation) | NoR-MAE with ViT-Tiny · ImageNet-100 | accuracy (%): 78.26 -> 82.97 -> 83.96 | Helped |
| add closest-neighbor repelling loss | MAE with ViT-Tiny · CIFAR-100 | accuracy (%): 78.74 -> 80.13 | Helped |
| add closest-neighbor repelling loss | MAE with ViT-Tiny · DTD | accuracy (%): 55.07 -> 60.18 | Helped |
| add closest-neighbor repelling loss | MAE with ViT-Tiny · Flowers | accuracy (%): 85.80 -> 88.71 | Helped |
| add closest-neighbor repelling loss | MAE with ViT-Tiny · Pets | accuracy (%): 82.47 -> 85.29 | Helped |
Surprising findings
- Naively decorrelating token-level representations (instead of embedding dimensions) collapses training.
- Closest-neighbor patches are not identical to their source patches: cosine similarity ranges from 0.07 to 0.999 with mean 0.83, yet the method exploits this non-identity.
Limitations
- Gains on ViT-Base are small (+0.12% on ImageNet-1K), suggesting benefit concentrates on lightweight models.
- The hyperparameter lambda must be tuned; over-penalization hurts performance (Table 5).
- Compute hardware details are in the supplementary and not stated in the main text.
Experimental setup
Datasets: ImageNet-1K ImageNet-100 CIFAR-10 CIFAR-100 Flowers Pets DTD Food Caltech-101 Aircraft ADE20K
Benchmarks: ImageNet-1K ImageNet-100 CIFAR-10 CIFAR-100 Flowers Pets DTD Food Caltech-101 Aircraft ADE20K
Metrics: top-1 accuracy mIoU
Baselines: MAE DMAE SupMAE CrossMAE U-MAE MoCo-V3 DINO SimMIM MaskFeat PeCo BEiT
Builds on: masked autoencoder (MAE) Barlow Twins U-MAE DMAE
Closest papers by meaning
- The Dynamic Duo of Collaborative Masking and Target for Advanced Masked Autoencoder Learning (AAAI 2025 · similarity 0.89)
- Learning Mask Invariant Mutual Information for Masked Image Modeling (ICLR 2025 · similarity 0.89)
- VideoMAC: Video Masked Autoencoders Meet ConvNets (CVPR 2024 · similarity 0.89)
- Extending Video Masked Autoencoders to 128 frames (NeurIPS 2024 · similarity 0.88)
- Self-Guided Masked Autoencoder (NeurIPS 2024 · similarity 0.88)
- R-MAE: Regions Meet Masked Autoencoders (ICLR 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.