Research map / Architectures and efficiency / Attention and transformers
HydraViT: Stacking Heads for a Scalable ViT
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 HydraViT, a stochastic tail-drop training method that stacks attention heads (with coupled embedding dimension) in a standard ViT so one set of weights yields many nested subnetworks for scalable inference on ImageNet-1K.
Problem
Vision Transformers need separate models of different sizes (trained, tuned and stored individually) to fit heterogeneous hardware, and the available size granularity is coarse.
Approach
Couple the embedding dimension to the number of heads (each head owns a fixed 64-dim slice) and, during training, uniformly sample a head count k and optimize only the first k heads' subnetwork, so heads become ordered by importance and any prefix is a usable model without fine-tuning.
Key results
- With 10 subnetworks (3-12 heads), HydraViT reaches 82.25% top-1 at 12 heads when trained on 9-12 heads for 300 epochs, above MatFormer's 82.04%.
- At 3 heads (192 dim), HydraViT achieves 70.56% vs SortedNet's 70.20% from DeiT-tiny init, and 10047 img/s throughput vs 5898 for SortedNet at that size.
- Training 3 subnetworks jointly gives 72.13-72.65% (3 heads), 79.45-80.08% (6 heads), 81.13-81.77% (12 heads), roughly matching separately trained DeiT-tiny/small/base.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use weighted sampling (25% 3-head, 30% 6-head, 45% 12-head) instead of uniform sampling (ablation) | HydraViT · ImageNet-1K, 12-head subnetwork | accuracy: +0.3 to nearly 0.6 p.p. | Helped |
| Use weighted sampling instead of uniform sampling (ablation) | HydraViT · ImageNet-1K, 3-head subnetwork | accuracy: a decrease of 0.5 to 1.0 p.p. | Hurt |
| Use weighted sampling instead of uniform sampling (ablation) | HydraViT · ImageNet-1K, 6-head subnetwork | accuracy: -0.2 to +0.2 p.p. | Mixed |
| Add separate classifiers for each subnetwork instead of a shared classifier (ablation) | HydraViT · ImageNet-1K, 3-, 6-, and 12-head subnetworks | accuracy: up to 0.2 percentage points | Mixed |
| Combine weighted sampling with separate classifiers (ablation) | HydraViT · ImageNet-1K, 3-, 6-, and 12-head subnetworks | accuracy: up to 0.15 p.p. higher | Helped |
| Train for 800 epochs instead of 300 (ablation) | HydraViT with 10 subnetworks · ImageNet-1K, accuracy across subnetworks | accuracy: 80.45 -> 81.93; 78.40 -> 79.84; 67.34 -> 68.78 | Helped |
| Initialize with DeiT-base instead of DeiT-tiny (ablation) | HydraViT · ImageNet-1K, 12-head subnetwork | accuracy: 80.90 -> 81.84 | Helped |
| Initialize with DeiT-base instead of DeiT-tiny (ablation) | HydraViT · ImageNet-1K, 3-head subnetwork | accuracy: 73.16 -> 70.83 | Hurt |
| Initialize with DeiT-base instead of DeiT-tiny (ablation) | HydraViT · ImageNet-1K, 6-head subnetwork | accuracy: 79.63 -> 79.62 | No clear effect |
| Train on 9–12 heads rather than the wider scalability range (ablation) | HydraViT · ImageNet-1K | accuracy: 82.25% compared to MatFormer's 82.04% | Helped |
Surprising findings
- Training from scratch, HydraViT struggles to reach good accuracy at 3 heads, which the authors attribute to sampling bias since 1- and 2-head subnetworks are excluded.
- Initializing with DeiT-base instead of DeiT-tiny lowers the average accuracy across subnetworks, even though 12-head accuracy improves.
Limitations
- Training requires more iterations to reach accuracy comparable to individually trained models; 10 losses are optimized in a shared-classifier setup.
- Evaluated only on DeiT-tiny/small/base with the same depth; not applied to larger models such as DeiT-large.
- Focus is on single-device efficiency; on-device latency is not measured directly and is only indirectly supported by other works.
Experimental setup
Datasets: ImageNet-1K
Benchmarks: ImageNet-1K ImageNet-v2 ImageNet-R ImageNet-A ImageNet-Sketch ImageNet-ReaL
Metrics: top-1 accuracy GMACs throughput RAM usage parameters
Baselines: MatFormer DynaBERT SortedNet DeiT-tiny DeiT-small DeiT-base
Builds on: DeiT MatFormer DynaBERT SortedNet Matryoshka Representation Learning
Closest papers by meaning
- RepViT: Revisiting Mobile CNN From ViT Perspective (CVPR 2024 · similarity 0.88)
- SHViT: Single-Head Vision Transformer with Memory Efficient Macro Design (CVPR 2024 · similarity 0.88)
- FasterViT: Fast Vision Transformers with Hierarchical Attention (ICLR 2024 · similarity 0.87)
- Real-time Core-Periphery Guided ViT with Smart Data Layout Selection on Mobile Devices (NeurIPS 2024 · similarity 0.86)
- GradQ-ViT: Robust and Efficient Gradient Quantization for Vision Transformers (AAAI 2025 · similarity 0.86)
- You Only Need Less Attention at Each Stage in Vision Transformers (CVPR 2024 · 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.