Research map / Architectures and efficiency / Attention and transformers
The Need for Speed: Pruning Transformers with One Recipe
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 OPTIN, a one-shot, retraining-free transformer pruning framework that scores weights by how masking them perturbs downstream-layer features and logits (a "trajectory" metric), then selects masks via a mask search, evaluated on GLUE, ImageNet, CIFAR-10, Cityscapes, and CNNs.
Problem
Compressing pre-trained transformers to meet FLOP budgets across modalities typically requires expensive re-training or architecture-specific designs, limiting practical one-shot compression.
Approach
Measures parameter importance as the cumulative change in downstream-layer features (manifold distillation loss) plus a KL-divergence on logits when each weight is masked in a gradient-free forward pass, then uses a polynomial-time mask search under a FLOP constraint; also extends to token reduction schedules via ToMe.
Key results
- On GLUE with BERT-Base at ~40% FLOPs reduction, OPTIN shows an average accuracy drop of at most 2% without re-training and outperforms PTF's mask search at the same compression rate.
- On ImageNet-1K, OPTIN is reported to be 0.5% better than SAViT at about 5% lower FLOPs, and its best DeiT-Small result beats CP-ViT by 0.4% at roughly 10% higher FLOPs reduction.
- On Cityscapes with Mask2Former (Swin-Ti), roughly 24% encoder FLOPs reduction and 47% parameter reduction with competitive mIoU.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use subsequent-layer trajectory [i+1,N] rather than measuring only at [i] (ablation) | OPTIN on BERT-Base and DeiT-Ti · MNLI and ImageNet-1K | accuracy: 78.89 -> 81.90; 68.91 -> 71.25 | Helped |
| Use sum rather than mean to aggregate layer distillation errors (ablation) | OPTIN · MNLI and ImageNet-1K | accuracy: 80.75 -> 81.90; 71.15 -> 71.25 | Helped |
| Add KL loss to manifold distillation loss (ablation) | OPTIN · MNLI and ImageNet-1K | accuracy: 80.91 -> 81.90; 68.85 -> 71.25 | Helped |
| Set lambda to 0.01 rather than the default 0.1 (ablation) | OPTIN · MNLI | accuracy: 81.90 -> 82.12 | Helped |
| Set lambda to 0.01 rather than the default 0.1 (ablation) | OPTIN · ImageNet-1K | accuracy: 70.99 -> 71.25 | Helped |
| Use temperature 4 rather than temperature 1 (ablation) | OPTIN · ImageNet-1K | accuracy: 70.54 -> 71.25 | Helped |
| Use temperature 4 rather than temperature 8 (ablation) | OPTIN · ImageNet-1K | accuracy: 71.00 -> 71.25 | Helped |
| Use temperature 8 rather than temperature 4 (ablation) | OPTIN · MNLI | accuracy: 81.90 -> 82.14 | Helped |
| Use LayerNorm features rather than IM-Dense features (ablation) | OPTIN · MNLI and ImageNet-1K | accuracy: 81.83 -> 81.92; 70.90 -> 71.27 | Helped |
| Change the importance-score batch size from 32 to 128 (ablation) | OPTIN · MNLI | accuracy: 81.90 -> 82.20 | Helped |
| Change the importance-score batch size from 32 to 16 (ablation) | OPTIN · ImageNet-1K | accuracy: 71.25 -> 70.53 | Hurt |
| Use the OPTIN schedule instead of the default schedule for random pruning (ablation) | OPTIN token reduction with Random Prune · ImageNet-1K | accuracy drop (%): ↓ 8.47 -> ↓ 7.11 | Helped |
Surprising findings
- Accumulating manifold distillation error over deeper layers and using the dense output layer for embeddings works best, and the paper finds sum aggregation beats mean.
Limitations
- Some comparisons still rely on re-training-based baselines or on setups with additional components, making direct comparison inconsistent.
- Token reduction component requires the ToMe merging scheme; the framework's pruning search space is task-dependent and hand-designed per task.
- Results on CNNs require fine-tuning following HRank's procedure, so the method is not fully retraining-free there.
Experimental setup
Datasets: GLUE (MNLI, QNLI, MRPC, STS-B, QQP) ImageNet-1K CIFAR-10 Cityscapes
Benchmarks: GLUE ImageNet-1K CIFAR-10 Cityscapes
Metrics: accuracy mIoU FLOPs reduction throughput (im/s) latency ratio
Baselines: PTF CoFi DynaBERT CP-ViT SAViT EViT DynamicViT ToMe HRank VTP PoWER HVT
Builds on: PTF (Post-Training-Framework) mask search ToMe token merging manifold distillation KL-divergence logit distillation
Closest papers by meaning
- Resource-Efficient Transformer Pruning for Finetuning of Large Models (CVPR 2024 · similarity 0.89)
- SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization (ICML 2024 · similarity 0.87)
- Transformer Fusion with Optimal Transport (ICLR 2024 · similarity 0.87)
- Zero-TPrune: Zero-Shot Token Pruning through Leveraging of the Attention Graph in Pre-Trained Transformers (CVPR 2024 · similarity 0.87)
- Viewing Transformers Through the Lens of Long Convolutions Layers (ICML 2024 · similarity 0.86)
- HybridNorm: Towards Stable and Efficient Transformer Training via Hybrid Normalization (NeurIPS 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.