Structured Neuron-level Pruning to Preserve Attention Scores in Vision Transformers

Core Concepts

The authors propose a novel graph-aware neuron-level pruning method, Structured Neuron-level Pruning (SNP), to effectively compress and accelerate Vision Transformer (ViT) models while preserving their attention scores.

Abstract

The paper introduces a novel graph-aware neuron-level pruning method called Structured Neuron-level Pruning (SNP) to accelerate and compress Vision Transformer (ViT) models. Key highlights: SNP prunes neurons with less informative attention scores and eliminates redundancy among heads in the Multi-head Self-Attention (MSA) module of ViTs. SNP preserves the overall attention scores by pruning graphically connected query and key layers having the least informative attention scores, while pruning value layers to eliminate inter-head redundancy. SNP achieves significant acceleration while maintaining the original performance on several ViT models. The compressed DeiT-Small outperforms the smaller DeiT-Tiny model in both accuracy and latency. SNP can be combined with conventional head or block pruning approaches to further compress and accelerate ViT models. Extensive experiments and ablation studies demonstrate the effectiveness and robustness of the proposed SNP method.

Stats

The DeiT-Small with SNP runs 3.1× faster than the original model and achieves 21.94% faster and 1.12% higher performance than the DeiT-Tiny. SNP with head pruning could compress the DeiT-Base by 80% of the parameters and computational costs and achieve 3.85× faster inference speed on RTX3090 and 4.93× on Jetson Nano.

Quotes

"SNP prunes neurons with less informative attention scores and eliminates redundancy among heads." "SNP preserves the overall attention scores by pruning graphically connected query and key layers having the least informative attention scores, while pruning value layers to eliminate inter-head redundancy." "SNP achieves significant acceleration while maintaining the original performance on several ViT models."

Key Insights Distilled From

SNP: Structured Neuron-level Pruning to Preserve Attention Scores

by Kyunghwan Sh... at arxiv.org 04-19-2024

https://arxiv.org/pdf/2404.11630.pdf

Deeper Inquiries

How can the proposed SNP method be extended to other types of Transformer-based models beyond Vision Transformers

The proposed SNP method can be extended to other types of Transformer-based models beyond Vision Transformers by adapting the pruning criteria to suit the specific architecture and requirements of the model. For instance, in models like BERT or GPT, where attention mechanisms play a crucial role, SNP can be applied to prune neurons in the self-attention layers while preserving the overall attention scores. By customizing the importance scoring and pruning criteria based on the structure of the model, SNP can effectively compress and accelerate a wide range of Transformer architectures.

What are the potential challenges in integrating SNP with the training process of large Transformer models to further improve efficiency

Integrating SNP with the training process of large Transformer models to further improve efficiency may pose several challenges. One challenge is the computational complexity of calculating importance scores for a large number of neurons in the model, especially during the training phase. This could lead to increased training time and resource consumption. Additionally, ensuring that the pruning process does not negatively impact the model's performance during training and fine-tuning is crucial. Balancing the trade-off between model size reduction and performance retention requires careful optimization and validation. Moreover, adapting SNP to handle dynamic input sizes and varying attention patterns in different vision tasks may require additional adjustments to the pruning criteria and methodology to accommodate the diverse requirements of different tasks.

How can the SNP method be adapted to handle dynamic input sizes or varying attention patterns in different vision tasks

Adapting the SNP method to handle dynamic input sizes or varying attention patterns in different vision tasks can be achieved by incorporating adaptive pruning strategies. One approach could involve dynamically adjusting the pruning criteria based on the input data characteristics or task requirements. For dynamic input sizes, the importance scores could be recalculated or updated based on the input dimensions, allowing for flexible pruning of neurons to adapt to different input sizes. Similarly, for varying attention patterns in different tasks, the pruning criteria could be modified to prioritize preserving specific attention patterns or relationships that are crucial for the task at hand. By incorporating adaptive mechanisms into the SNP method, it can be tailored to effectively handle the variability in input sizes and attention patterns across different vision tasks.