Elastic ViTs from Pretrained Models without Retraining

AI summary

Overview

Research area: Computer Vision — model compression and efficient inference for Vision Transformers (structured pruning, elastic inference).

Technical level: Intermediate. The paper assumes familiarity with transformer architecture, Hessian-based pruning and evolutionary strategies, though its core idea is describable in plain terms.

Scope: The paper introduces SnapViT, a label-free, retraining-free structured pruning method that turns any single pretrained Vision Transformer into an "elastic" model from which subnetworks at any sparsity level can be extracted in one shot.

What This Paper Is About

Vision foundation models such as DINOv3 are released only in a few fixed sizes (the paper lists 21M, 29M, 86M, 300M, 840M and 6.7B parameter ViTs for the DINOv3 family), so users must pick the largest model that fits their budget rather than the most suitable one. Existing remedies either require distilling into a predetermined architecture, need nested structures built in during pretraining, or require retraining for each target compute budget.

The goal is to take an already-pretrained ViT and, in a single fast pass and without labels or retraining, produce a family of smaller subnetworks spanning a continuum of compute budgets while preserving accuracy.

Key Contributions

1. An efficient single-shot structured pruning strategy for pretrained ViTs, which produces elastic models adjustable to any computational budget rather than one fixed sparsity level.
2. A novel evolutionary approximation of the Hessian's off-diagonal structure. The authors use the Exponential Natural Evolution Strategy (xNES) to model cross-block interactions that diagonal, block-diagonal and K-FAC approximations ignore, without ever forming the full Hessian.
3. A self-supervised importance scoring mechanism based on the DINO objective, which removes the need for labels or a classification head and lets the method apply to foundation models as well as supervised ones.
4. Demonstrated state-of-the-art performance under considerable pruning without retraining or labels, including compatibility with post-pruning weight correction and full fine-tuning.

Main Findings

• Speed of generation: SnapViT generates elastic models in less than five minutes on a single A100 GPU. Runtime for the full DeiT-III family is reported as 2m 35s (ViT-S/16, 22.1M parameters), 2m 55s (ViT-B/16, 86.6M), 4m 58s (ViT-L/16, 304.4M) and 11m 4s (ViT-H/14, 632.1M).
• A single run yields all sparsities: Baselines such as LLM Surgeon, SparseGPT, FPTP, LAMP, SNIP Magnitude and NViT optimize for one predetermined sparsity level, whereas SnapViT produces any target sparsity S in [0,1] from one ranking.
• DINOv1 ViT-B/16 headline result: Pruning to 40% sparsity gives a 1.58x inference speedup while keeping accuracy degradation below 5%, with 15.1% and 53.2% improvements in linear classification over SNIP Magnitude and LAMP respectively.
• Beats label-using baselines with a label-free method: At 50% sparsity on supervised AugReg ViT-B/16, the method improves by 7% over SNIP and 12.3% over FPTP; it often outperforms the LLM Surgeon, which prunes to a target sparsity in 5 shots.
• Semantic segmentation: On Pascal VOC 2012 mIoU, the method improves by 9.1% over SNIP Magnitude at 50% sparsity for AugReg and by 15.3% over NViT at 60% sparsity for DeIT.
• Large-scale pretraining makes pruning harder: DINOv3 and SigLIPv2 ViT-B/16, trained on 1.7 billion and 10 billion samples respectively, require longer optimization horizons (500 iterations versus a baseline of 50) and benefit from optimizing for six sparsity levels rather than four. Still, the method improves over the second-best method in linear classification by 21.7% and 34.3% for SigLIPv2 and DINOv2 ViT-B/16, respectively (as reported by the authors).
• Big models have a ceiling: For SigLIPv2 ViT-G/16 (1.2B parameters) and DINOv3 ViT-H+/16 (840M parameters), performance holds to 30% sparsity but degrades sharply beyond it, whereas DeIT-III ViT-H/14 remains stable up to 50% sparsity. The authors attribute this to DeIT-III's ImageNet-21k pretraining (13M samples) versus the far larger pretraining sets of the other two.
• Weight correction rescues high sparsity: A single SparseGPT-style correction step using layer-wise Hessian approximations from 1000 random ImageNet-1k training samples lets SigLIPv2 be pruned to 50% sparsity with negligible linear classification loss.
• Full fine-tuning closes the gap: A DeIT ViT-B/16 pruned to 50% sparsity and fine-tuned for 300 epochs on ImageNet-1k reaches 82.6% ImageNet-1k accuracy (versus 81.8% unpruned), with 75.4% average k-NN and 75.9% average linear accuracy. NViT scores higher on ImageNet-1k (83.3%) but is 1.7% and 3.9% lower on average k-NN and linear classification across the seven datasets.
• Pruning pattern insight: In DINO ViT-B/16, pruning favors slimming feed-forward blocks first while attention heads stay largely intact; blocks 8 to 12 are the most pruning-susceptible, while earlier blocks keep stable parameter density.
• Global interactions matter: Modeling no interactions gives 56.6% average k-NN for DINO ViT-B/16 at 50% sparsity, FFN-only interactions (12) give 60.1%, and FFN plus attention head interactions (156) give 63.5% — up to +6.9%.
• Optimization budget: For AugReg ViT-B/16 at 50% sparsity, average k-NN/linear accuracy rises from 39.3/42.2 at 50 iterations to 39.9/42.2 at 250 and 40.9/44.0 at 500, up to 1.8% in linear classification.
• Self-supervised gradients are nearly as good as labels: Using a cross-entropy loss instead of the self-supervised loss only performs marginally better for AugReg ViT-B/16.

Methodology in Plain English

The method assigns every prunable unit — row-column combinations inside feed-forward blocks, or entire attention heads — a "prunability score," then removes the lowest-scoring units until the compute budget is met. That single ranking is what makes the model elastic: any budget can be served by taking a different prefix of the ranking.

The score has two parts. The first is a local term: the squared magnitude of the self-supervised gradient for each parameter, which approximates the diagonal of the Hessian (a measure of how sensitive the loss is to that parameter) and is computed on a small set of images without needing labels. The second is a global term: a per-block scaling factor that says how strongly a block's parameters co-vary with the rest of the network, which the diagonal approximation cannot see.

Computing that global term exactly would mean an impossibly large Hessian, so the authors instead let an evolutionary algorithm (xNES) search for the scaling factors. The search maintains a Gaussian distribution over candidate scaling vectors, and for each candidate it produces a pruning mask, prunes the model, and scores how similar the pruned model's embeddings are to the original ones. Embeddings are compressed with PCA to 192 dimensions and compared by cosine similarity, averaged across several target sparsities. Crucially, the algorithm's evolving covariance matrix converges toward the inverse of the underlying curvature, so it acts as a stand-in for the off-diagonal Hessian. A single run of this loop yields one final global ranking usable at any sparsity, with a population size of 4 + 3 log(d), where d is 156 for a ViT-B/16.

Why This Matters

Research impact. The paper challenges the assumption that extraction of multiple sub-models requires architectures designed for it (nested Matryoshka structures, slimmable networks) or retraining. It also shows that an entire class of second-order information — cross-block Hessian off-diagonals, previously treated as computationally infeasible — can be approximated cheaply enough to be practical, and that self-supervised gradients can replace label-dependent pruning criteria for foundation models that have no classification head.

Real-world applications:

• Deploying vision foundation models onto edge devices, phones or embedded hardware where a fixed released model size does not match the available memory or latency budget.
• Serving a single hosted model at multiple quality tiers, matching compute to per-request cost or latency targets without maintaining separate checkpoints.
• Fine-grained adaptation for domain-specific deployments, since the pruned models are evaluated across datasets including ImageNet-1k, FGVC Aircraft, Oxford-IIT Pets, DTD Textures, EuroSAT, CIFAR 10/100 and Pascal VOC 2012.
• Reducing inference costs and energy use for large-scale vision pipelines, which the authors note could lower CO2 emissions from inference.

Industry relevance. The method's label-free, retraining-free nature matters for proprietary or non-public pretrained models where the original training data is unavailable, explicitly cited as a weakness of distillation. The sub-five-minute generation time and the reported scaling to 632.1M-parameter models make it plausible as an on-demand compression step in a model deployment pipeline, and support from Qualcomm Technologies Inc. indicates commercial interest in on-device deployment.

Future Directions

• Why large-scale pretraining resists pruning. The paper observes that DINOv3 (1.7B samples) and SigLIPv2 (10B samples) degrade beyond 30% sparsity while DeIT-III (13M samples) survives to 50%, and hypothesizes that representational knowledge is spread more evenly across parameters. Confirming and countering this effect is an open problem.
• Better behavior at extreme sparsity for very large models. Even with the method's improvements, SigLIPv2 ViT-G/16 and DINOv3 ViT-H+/16 degrade sharply past 30%, suggesting either more refined scoring or mandatory weight correction for these regimes.
• Reducing still-further the reliance on optimization horizons. Foundation models benefited from 500 xNES iterations instead of 50 and from six sparsity targets instead of four; cheaper or smarter search could remove that tuning burden.
• Extending beyond ViTs and beyond classification/segmentation. The paper evaluates classification, k-NN, linear probing and linear semantic segmentation; whether the cross-structure correlation idea transfers to other architectures or dense prediction tasks is not addressed.

Target Audience

Researchers and engineers working on model compression, efficient inference, or deploying vision foundation models under compute constraints. It is most valuable to readers already comfortable with pruning terminology and Hessian-based importance scores, since the method section builds directly on those concepts. Practitioners who need to compress a pretrained model without access to its training data, labels, or a classification head are the most direct beneficiaries.

Authors’ abstract

Vision foundation models achieve remarkable performance but are only available in a limited set of pre-determined sizes, forcing sub-optimal deployment choices under real-world constraints. We introduce SnapViT: Single-shot network approximation for pruned Vision Transformers, a new post-pretraining structured pruning method that enables elastic inference across a continuum of compute budgets. Our approach efficiently combines gradient information with cross-network structure correlations, approximated via an evolutionary algorithm, does not require labeled data, generalizes to models without a classification head, and is retraining-free. Experiments on DINO, SigLIPv2, DeIT, and AugReg models demonstrate superior performance over state-of-the-art methods across various sparsities, requiring less than five minutes on a single A100 GPU to generate elastic models that can be adjusted to any computational budget. Our key contributions include an efficient pruning strategy for pretrained Vision Transformers, a novel evolutionary approximation of Hessian off-diagonal structures, and a self-supervised importance scoring mechanism that maintains strong performance without requiring retraining or labels. Code and pruned models are available at: https://elastic.ashita.nl/