Research map / Architectures and efficiency / Attention and transformers
Accelerating Augmentation Invariance Pretraining
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 an architecture-preserving acceleration framework for augmentation-invariance SSL of ViTs that combines randomized token dropout and flexible patch scaling on a cost-adjusted gradient-error-optimal schedule, achieving 4x (MoCo-v3), 3.3x (SimCLR), and 2.5x (DINO) speedups on ImageNet.
Problem
Contrastive and distillation-based self-supervised pretraining of Vision Transformers is computationally expensive, with long schedules needed for competitive performance. Existing accelerations often change the architecture or the pretraining objective.
Approach
Accelerate SSL by compressing the ViT input sequence (random token dropout plus flexible patch scaling) without changing architecture or loss, choosing the compression mix via a cost-adjusted MSE (bias-variance) analysis of gradient estimates, and adjusting the schedule over training so early training uses aggressive compression and later training uses less.
Key results
- On ImageNet-1k, dynamic acceleration for MoCo-v3 reaches 75.9% LP vs 76.0% baseline with 1542M vs 6168M budget (about 4x).
- SimCLR accelerated reaches 3.3x speedup and DINO 2.5x speedup on ImageNet-1k with comparable accuracy.
- On ImageNet-100, the dynamic schedule at 150M budget reaches 74.7 NN / 84.1 LP, exceeding the non-accelerated 520M baseline's 70.0 NN, with 84.4 LP for baseline.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use dynamic compression scheduling instead of no acceleration | MoCo-v3 on ImageNet-100 · ImageNet-100 | NN accuracy: 58.1 -> 74.7 | Helped |
| Use dynamic compression scheduling instead of no acceleration | MoCo-v3 on ImageNet-1k · ImageNet-1k | LP accuracy: 75.9 vs 76.0 | No clear effect |
| Apply dynamic compression scheduling | SimCLR on ImageNet-1k · ImageNet-1k | speedup: 3.3x | Helped |
| Apply dynamic compression scheduling | DINO on ImageNet-1k · ImageNet-1k | speedup: 2.5x | Helped |
| Keep the key sequence uncompressed while applying token dropout to the query (ablation) | MoCo-v3 · ImageNet-100, 100M budget | NN accuracy: 65.4 -> 69.6 | Helped |
| Increase token dropout to query length 20 rather than query length 50 (ablation) | MoCo-v3 · ImageNet-100, 100M budget | NN accuracy: 69.6 -> 63.9 | Hurt |
| Use symmetric patch scaling to patch size 30 rather than no compression (ablation) | MoCo-v3 · ImageNet-100, 100M budget | NN accuracy: 58.1 -> 68.1 | Helped |
| Increase patch size beyond 30 (ablation) | MoCo-v3 · ImageNet-100, 100M budget | NN accuracy: 68.1 -> 62.0 | Hurt |
| Apply constant token dropout throughout a larger training budget rather than dynamic scheduling (ablation) | MoCo-v3 · ImageNet-100 | NN accuracy: 69.6 at 100M -> 65.3 at 200M | Hurt |
| Use dynamic scheduling rather than constant token dropout (ablation) | MoCo-v3 · ImageNet-100, 104M budget | NN accuracy: 69.6 -> 71.2 | Helped |
| Use dynamic scheduling rather than constant patch scaling (ablation) | MoCo-v3 · ImageNet-100, 104M budget | NN accuracy: 68.1 -> 71.2 | Helped |
| Apply token dropout to both sequences in SimCLR (ablation) | SimCLR · ImageNet-1k | NN accuracy: 50.7 vs 50.2 | No clear effect |
Surprising findings
- Constant acceleration overfits at higher budgets and fails to reach peak baseline performance, while a decreasing-compression schedule recovers it.
- Asymmetric token dropout (uncompressed keys) works better than symmetric, whereas symmetric patch scaling works better.
Limitations
- Linear complexity assumption ignores quadratic attention cost, which only matters beyond ~400 tokens.
- Schedule is derived from intermediate checkpoints offline rather than on-the-fly.
- Batch size reduced to 1024 on ImageNet-1k due to hardware limits.
Experimental setup
Datasets: ImageNet-100 ImageNet-1k
Benchmarks: ImageNet-100 ImageNet-1k
Metrics: linear probe accuracy nearest neighbor accuracy finetuning accuracy cost-adjusted MSE
Baselines: MoCo-v3 SimCLR DINO ViT pre-training without acceleration
Builds on: MoCo-v3 SimCLR DINO FlexiViT
Closest papers by meaning
- Structured Initialization for Vision Transformers (NeurIPS 2025 · similarity 0.87)
- GradQ-ViT: Robust and Efficient Gradient Quantization for Vision Transformers (AAAI 2025 · similarity 0.86)
- Elastic ViTs from Pretrained Models without Retraining (NeurIPS 2025 · similarity 0.86)
- Configuring Data Augmentations to Reduce Variance Shift in Positional Embedding of Vision Transformers (AAAI 2025 · similarity 0.86)
- Vision Transformers with Self-Distilled Registers (NeurIPS 2025 · similarity 0.86)
- RePaViT: Scalable Vision Transformer Acceleration via Structural Reparameterization on Feedforward Network Layers (ICML 2025 · similarity 0.86)
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.