🤖 AI

Rethinking Depth Pruning for Vision Transformers: A Heterogeneity-Aware Perspective

This paper introduces HetDPT, a heterogeneity-aware depth pruning method that overcomes the accuracy recovery challenges of existing approaches by addressing inter-layer heterogeneity, thereby achieving significant speedups with minimal accuracy loss on Vision Transformers across multiple benchmarks.

Original authors: Zhenfeng Su, Kang Zhao, Han Bao, Tao Yuan, Zhongzhe Hu, Xianzhi Yu, Wenxuan Wang

Published 2026-07-07

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhenfeng Su, Kang Zhao, Han Bao, Tao Yuan, Zhongzhe Hu, Xianzhi Yu, Wenxuan Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a Vision Transformer (ViT) as a highly sophisticated, multi-story factory designed to recognize images. Each floor of this factory represents a "layer" where the image gets processed. Some floors are dedicated to Attention (scanning the whole picture to see how different parts relate), and others are dedicated to Activation (applying a specific filter or "non-linearity" to sharpen the details).

For years, engineers trying to make these factories run faster on smaller devices have mostly focused on Width Pruning. This is like hiring fewer workers on each floor. You get a smaller factory, but it still has the same number of floors, so the product still has to travel through every single level.

The paper argues that a better way to speed things up is Depth Pruning: simply removing entire floors. If you remove 50% of the floors, the product gets to the end twice as fast. However, previous attempts to do this failed miserably. The factory would collapse, and the quality of the output (the image recognition) would tank.

The authors of this paper, HetDPT, discovered why the factory was collapsing and built a new method to fix it. Here is the breakdown in simple terms:

The Problem: The "Mismatch" and the "Different Personalities"

The paper identifies two main reasons why simply deleting floors doesn't work:

The Dimension Mismatch (The Broken Conveyor Belt):
Imagine the factory has a conveyor belt system. The "Attention" floors and "Activation" floors are connected in a specific way. If you just randomly rip out an Activation floor, the conveyor belt from the previous floor arrives at a machine that doesn't exist anymore, or the machine that does exist expects a package of a different size. The paper calls this a "dimension mismatch." It's like trying to plug a square peg into a round hole because you removed the adapter in between.
- The Fix: The authors realized they could remove the "Activation" floors entirely and simply merge the two "Linear" (machine) floors that were on either side of it. This creates a new, seamless machine that fits the conveyor belt perfectly, keeping the factory running smoothly even with fewer floors.
The Heterogeneity (The Different Personalities):
This is the paper's biggest insight. The authors realized that Attention floors and Activation floors are not the same; they have different "personalities."
- Attention Floors are like the Senior Managers. If you fire a few managers, the factory runs a bit slower at first, but it takes a long time to train new ones to get back to full speed. They are hard to replace quickly.
- Activation Floors are like the Junior Interns. If you fire them, the factory crashes immediately (accuracy drops to near zero). However, if you give them a quick 10-minute training session (fine-tuning), they bounce back to full performance almost instantly.
- The Mistake of Old Methods: Previous methods treated all floors the same. They looked at the "gradients" (a measure of how much the floor contributes) and fired the ones with the lowest numbers. Because the Senior Managers (Attention) naturally have different "voltage" than the Interns (Activation), the old methods kept firing the wrong people, leading to a collapse.

The Solution: HetDPT (The Smart Factory Manager)

The authors created a two-step process called HetDPT (Heterogeneity-Aware Depth Pruning) to manage this factory:

Step 1: The Budget Allocation (The Strategy)
Before firing anyone, the manager uses a "Model Accuracy Predictor" (a smart calculator). Instead of asking "Which specific floor is the worst?", it asks, "If we fire 3 Senior Managers and 5 Interns, will the factory still work?"

It tests different combinations of firing Managers vs. Interns.
It finds the perfect balance (e.g., "We can fire 10 floors total, but we must fire 6 Interns and only 4 Managers to keep the quality high").
This avoids the mistake of comparing Managers and Interns directly, which is like comparing apples and oranges.

Step 2: The Execution (The Cleanup)
Once the budget is set (e.g., "Fire 6 Interns"), the method goes into each group separately.

It looks at all the Interns and fires the specific ones that are least needed.
It looks at all the Managers and fires the specific ones that are least needed.
Crucially: It merges the Linear machines around the fired Interns so the conveyor belt fits perfectly (solving the mismatch).
Finally, it gives the remaining staff a quick "refresher course" (fine-tuning) to get them back to 100% performance.

The Results: Faster Without Losing Quality

The paper tested this on several standard image datasets (like ImageNet, which is like a giant photo album of 1 million pictures).

Speed: They achieved 1.58x faster processing for a standard model (DeiT-B) while keeping the accuracy exactly the same as the original.
Extreme Compression: When they combined this method with other techniques (width pruning), they created a super-lightweight model that was 5.19x faster than the original, with almost no loss in accuracy.
Versatility: It worked not just on standard models, but also on more complex ones (Swin Transformers) and even massive models with billions of parameters (DINO-V2-Giant).

Summary Analogy

Think of a long assembly line making cars.

Old Way: You try to speed it up by making the workers on the line work faster (Width Pruning), or you randomly fire workers from different stations, causing the car chassis to fall off the line because the next station isn't ready (Failed Depth Pruning).
HetDPT Way: You realize that some stations (Interns) can be removed entirely if you weld the two machines on either side together. You also realize that firing the Station Managers (Attention) hurts the line for a long time, while firing the Interns hurts it briefly but they recover fast. So, you carefully calculate exactly how many of each to fire, weld the machines together, and give the remaining crew a quick pep talk. The result? The car comes out the end of the line twice as fast, and it's still a perfect car.

The paper concludes that by respecting the different "personalities" of the layers and fixing the mechanical mismatch, we can make these powerful AI models run significantly faster on everyday devices without losing their intelligence.

Technical Summary: Rethinking Depth Pruning for Vision Transformers: A Heterogeneity-Aware Perspective

1. Problem Statement

Vision Transformers (ViTs) have achieved state-of-the-art performance across various domains but suffer from high computational costs and large parameter counts, leading to extended inference latency. While structured pruning is an effective compression strategy, existing methods face significant limitations:

Width Pruning: Reduces channels or attention heads within layers but often yields limited speedups on mainstream hardware.
Depth Pruning: Removes entire layers, offering significantly higher speedup potential (as shown in Figure 2a). However, it is notoriously difficult to recover accuracy, especially with aggressive layer removal.
The Core Bottleneck: The paper argues that the failure of existing depth pruning methods stems from two overlooked issues:
1. Neglect of Heterogeneity: ViTs contain distinct layer types (e.g., Attention layers vs. Activation function layers within Feed-Forward Networks). Existing methods often treat these layers uniformly or rely on gradient-based metrics that fail due to the significant disparity in gradient scales between these types.
2. Dimension Mismatch: Simultaneously pruning Attention layers and the Linear layers in the Feed-Forward Network (FFN) creates structural incompatibilities. Since FFNs utilize an expansion-contraction mechanism (e.g., $D \to 4D \to D$ ), removing a Linear layer breaks the tensor shape continuity required by subsequent layers (illustrated in Figure 3).

2. Methodology: HetDPT

The authors propose HetDPT (Heterogeneity-aware DePTh pruning), a two-stage framework designed to address heterogeneity and avoid dimension mismatch.

Key Insights

The method is grounded in two critical observations regarding layer heterogeneity:

Gradient Disparity: Attention layers and Activation layers exhibit vastly different gradient magnitudes during backpropagation. Using standard gradient-based importance metrics leads to biased pruning decisions, consistently favoring the pruning of one type over the other.
Recovery Asymmetry:
- Pruning Attention layers causes moderate initial accuracy drops but requires extensive retraining for recovery.
- Pruning Activation layers causes severe initial degradation (often $\ge 90\%$ accuracy loss) but enables rapid recovery during fine-tuning.
- Implication: Short-sighted metrics that evaluate immediate post-pruning accuracy are ineffective.

The Two-Stage Pipeline

Stage 1: Identification of Redundant Layers
This stage determines how many and which specific layers to prune without comparing heterogeneous layer types directly.

Step 1: Pruning Budget Allocation (Model Accuracy Predictor - MAP): Instead of comparing layers directly, the method uses a lightweight data collection process (iterative prune $\to$ fast fine-tune $\to$ evaluate) to build a dataset. A polynomial function approximates the Model Accuracy Predictor (MAP), which predicts the final accuracy of a pruned model given the pruning ratios of Attention ( $\tilde{m}_a$ ) and Activation ( $\tilde{m}_g$ ) layers. This allows the system to find the optimal allocation of the pruning budget between the two layer types.
Step 2: Specific Redundant Layer Removal: Once the budget is allocated, the method identifies specific layers to remove within each homogeneous group (e.g., only among Attention layers). It uses learnable importance parameters and gradient propagation to rank layers, ensuring no cross-type comparison occurs.

Stage 2: Pruned Model Optimization

Fine-tuning: The pruned model is fine-tuned, optionally using self-distillation guided by the original unpruned ViT to restore accuracy.
Merging: To resolve dimension mismatch and accelerate inference, adjacent Linear layers that were separated by a pruned Activation layer are merged. For example, if an Activation layer is removed, $L_{l,1}$ and $L_{l,2}$ are merged into a single linear layer $L_l$ . This reduces model depth while maintaining rigorous dimension alignment.

HetDPT+: The authors further integrate HetDPT with established width pruning methods to create HetDPT+, a pipeline for extreme ViT compression that addresses both depth and width dimensions.

3. Key Contributions

Pruning Activation Functions: The paper proposes the first strategy to remove Activation function layers (GELU) in ViTs. By removing these non-linearities, adjacent Linear layers can be naturally merged, solving the dimension mismatch problem inherent in heterogeneous depth pruning.
Heterogeneity Analysis and Solution: The authors reveal and address gradient disparity and recovery asymmetry. They introduce a two-stage framework (HetDPT) equipped with a Model Accuracy Predictor (MAP) to manage these heterogeneities, supported by theoretical analysis guaranteeing the existence of an optimal configuration and the consistency of the predictor under noisy fast-finetuning.
State-of-the-Art Records:
- Depth Pruning: HetDPT achieves a 1.58 $\times$ speedup for DeiT-B and 1.39 $\times$ for DeiT-S with near-lossless accuracy, outperforming existing depth pruning works.
- Extreme Compression: HetDPT+ sets a new benchmark for extreme ViT compression. For the Isomorphic-Pruning-2.6G configuration, it enhances the acceleration ratio from 4.24 $\times$ to 5.19 $\times$ while maintaining near-lossless accuracy.

4. Experimental Results

The method was evaluated on ImageNet-1K, CIFAR-100, COCO, and ADE20K.

Accuracy-Speedup Trade-off: On ImageNet-1K, HetDPT consistently outperforms isolated pruning of attention or activation layers. For DeiT-B, it achieves 81.8% Top-1 accuracy (matching the baseline) with a 1.58 $\times$ speedup.
Comparison with SOTA: HetDPT+ surpasses existing extreme compression methods (e.g., Isomorphic Pruning, NViT). At ~81% accuracy, HetDPT+ achieves 5.19 $\times$ speedup compared to 4.24 $\times$ for Isomorphic Pruning. In ultra-extreme settings (1.1G MACs), it achieves 9.19 $\times$ speedup.
Transferability: The method demonstrates robust transfer learning on CIFAR-100 and maintains performance on dense prediction tasks (COCO object detection and ADE20K semantic segmentation) with minimal accuracy degradation.
Scalability: The approach scales to billion-parameter models, demonstrated on DINO-V2-Giant, where it prunes 15 layers with only a 0.06% accuracy drop.
Orthogonality: HetDPT is orthogonal to token pruning methods (e.g., GTP-ViT, ToMe), providing additional speedups when combined.

5. Significance and Claims

The paper claims that the primary contribution is enabling more efficient deployment of Vision Transformers by addressing the specific structural and dynamic challenges of depth pruning.

Democratization: By substantially reducing computational requirements (up to 9.19 $\times$ speedup), the method makes state-of-the-art vision models accessible to researchers and practitioners with limited resources.
Sustainability: Reduced energy consumption in compressed models contributes to more environmentally sustainable AI development.
Edge Deployment: The efficiency gains facilitate real-time deployment on edge devices for applications requiring low latency, such as healthcare diagnostics and autonomous systems.

The authors note a limitation: the Model Accuracy Predictor (MAP) is currently refitted for each architecture family, and extending the method to language models or other modalities remains future work. However, the work establishes a new theoretical and practical foundation for heterogeneous depth pruning, moving beyond the "one-size-fits-all" approaches of previous literature.