[Paper Review] Pruning Redundant Mappings in Transformer Models via Spectral-Normalized Identity Prior

Zi Lin, Jeremiah Zhe Liu+3 more|arXiv (Cornell University)|Oct 5, 2020

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Spectral-Normalized Identity Prior (SNIP), a structured pruning method that targets entire residual modules in Transformer models—such as attention heads, feed-forward networks, or full layers—by encouraging them to become identity mappings through function-norm thresholding. By combining spectral normalization to stabilize activation distributions, SNIP achieves state-of-the-art performance at 50% model compression, improving accuracy by 0.5–1.0% on average across GLUE tasks compared to prior methods.

ABSTRACT

Traditional (unstructured) pruning methods for a Transformer model focus on regularizing the individual weights by penalizing them toward zero. In this work, we explore spectral-normalized identity priors (SNIP), a structured pruning approach that penalizes an entire residual module in a Transformer model toward an identity mapping. Our method identifies and discards unimportant non-linear mappings in the residual connections by applying a thresholding operator on the function norm. It is applicable to any structured module, including a single attention head, an entire attention block, or a feed-forward subnetwork. Furthermore, we introduce spectral normalization to stabilize the distribution of the post-activation values of the Transformer layers, further improving the pruning effectiveness of the proposed methodology. We conduct experiments with BERT on 5 GLUE benchmark tasks to demonstrate that SNIP achieves effective pruning results while maintaining comparable performance. Specifically, we improve the performance over the state-of-the-art by 0.5 to 1.0% on average at 50% compression ratio.

Motivation & Objective

• To address the limitations of unstructured, weight-level pruning that fails to reduce inference latency despite parameter reduction.
• To explore structured pruning at the module level (e.g., attention heads, feed-forward subnetworks) to modify model architecture and improve efficiency.
• To improve pruning effectiveness by stabilizing activation distributions through spectral normalization.
• To demonstrate that structured pruning with identity-inducing priors leads to better performance-compression trade-offs than existing methods.

Proposed method

• The method applies a thresholding operator on the function norm of residual mappings to identify and prune unimportant nonlinear components.
• It introduces an identity-inducing prior that regularizes entire submodules (e.g., attention heads or FFN layers) to behave as identity functions.
• Spectral normalization is applied to the weight matrices of each module to control the largest singular value, stabilizing post-activation distributions.
• The approach is applied iteratively during fine-tuning, allowing for continuous compression curves and flexible architecture selection.
• It supports multiple pruning granularities: single attention heads, multi-head attention, feed-forward networks, or entire layers.
• The method is evaluated on BERT across 5 GLUE tasks, with performance measured at fixed compression ratios.

Experimental results

Research questions

• RQ1 Can structured pruning of entire residual modules in Transformers lead to better performance-compression trade-offs than unstructured weight pruning?
• RQ2 Does imposing an identity-inducing prior on residual mappings improve model robustness and generalization during pruning?
• RQ3 How does spectral normalization enhance the stability and effectiveness of structured pruning in deep Transformer models?
• RQ4 To what extent do different pruning granularities—single heads, multi-head attention, FFN, or full layers—affect final performance?
• RQ5 Can the remaining pruned components be task-specific, and does this flexibility improve compression efficiency?

Key findings

• SNIP achieves state-of-the-art performance at 50% model compression, improving accuracy by 0.5 to 1.0% on average across 5 GLUE benchmark tasks compared to prior methods.
• Separate pruning of attention and feed-forward subnetworks yields better performance than pruning entire layers or single heads, indicating task-specific structural redundancy.
• The method results in more sparse and regulated layer mappings during pre-training, with consistent performance improvements across diverse NLP tasks.
• Spectral normalization significantly improves the stability of activation distributions, enhancing pruning effectiveness and model generalization.
• Pruning at the module level (e.g., individual attention heads or FFN layers) leads to better inference efficiency and more flexible architecture design than weight-level pruning.
• The remaining pruned components vary across tasks, suggesting that optimal architectures are task-dependent and can be discovered through structured pruning.