# Differentiable Subset Pruning of Transformer Heads

arXiv: https://arxiv.org/abs/2108.04657 · HTML: https://mlsanity.com/paper/2108.04657/

## Key facts

- Cited by 47 papers in the map (data through 14 Oct 2026).
- Builds on 47 of its 64 references that are in the map.
- First posted on arXiv 10 Aug 2021.

## Builds on (most cited first, 30 of 47)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [A Broad-Coverage Challenge Corpus for Sentence Understanding through Inference](https://mlsanity.com/paper/1704.05426/) (2017)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Learning both Weights and Connections for Efficient Neural Networks](https://mlsanity.com/paper/1506.02626/) (2015)
- [fairseq: A Fast, Extensible Toolkit for Sequence Modeling](https://mlsanity.com/paper/1904.01038/) (2019)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Pruning Filters for Efficient ConvNets](https://mlsanity.com/paper/1608.08710/) (2016)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Are Sixteen Heads Really Better than One?](https://mlsanity.com/paper/1905.10650/) (2019)
- [Channel Pruning for Accelerating Very Deep Neural Networks](https://mlsanity.com/paper/1707.06168/) (2017)
- [Learning Efficient Convolutional Networks through Network Slimming](https://mlsanity.com/paper/1708.06519/) (2017)
- [To prune, or not to prune: exploring the efficacy of pruning for model compression](https://mlsanity.com/paper/1710.01878/) (2017)
- [Learning Sparse Neural Networks through L0 Regularization](https://mlsanity.com/paper/1712.01312/) (2017)
- [Learning Structured Sparsity in Deep Neural Networks](https://mlsanity.com/paper/1608.03665/) (2016)
- [Variational Dropout and the Local Reparameterization Trick](https://mlsanity.com/paper/1506.02557/) (2015)
- [ThiNet: A Filter Level Pruning Method for Deep Neural Network Compression](https://mlsanity.com/paper/1707.06342/) (2017)
- [Reducing Transformer Depth on Demand with Structured Dropout](https://mlsanity.com/paper/1909.11556/) (2019)
- [Movement Pruning: Adaptive Sparsity by Fine-Tuning](https://mlsanity.com/paper/2005.07683/) (2020)
- [What BERT is not: Lessons from a new suite of psycholinguistic diagnostics for language models](https://mlsanity.com/paper/1907.13528/) (2019)
- [EIE: Efficient Inference Engine on Compressed Deep Neural Network](https://mlsanity.com/paper/1602.01528/) (2016)
- [Variational Dropout Sparsifies Deep Neural Networks](https://mlsanity.com/paper/1701.05369/) (2017)

## Built on by (most cited first, 30 of 47)

- [Continual Pre-training of Language Models](https://mlsanity.com/paper/2302.03241/) (2023)
- [Structured Pruning Learns Compact and Accurate Models](https://mlsanity.com/paper/2204.00408/) (2022)
- [Transformers are Multi-State RNNs](https://mlsanity.com/paper/2401.06104/) (2024)
- [X-Pruner: eXplainable Pruning for Vision Transformers](https://mlsanity.com/paper/2303.04935/) (2023)
- [Rethinking the Role of Scale for In-Context Learning: An Interpretability-based Case Study at 66 Billion Scale](https://mlsanity.com/paper/2212.09095/) (2022)
- [Model Compression and Efficient Inference for Large Language Models: A Survey](https://mlsanity.com/paper/2402.09748/) (2024)
- [Discrete Prompt Compression with Reinforcement Learning](https://mlsanity.com/paper/2308.08758/) (2023)
- [Transformers in Speech Processing: A Survey](https://mlsanity.com/paper/2303.11607/) (2023)
- [Activation Scaling for Steering and Interpreting Language Models](https://mlsanity.com/paper/2410.04962/) (2024)
- [Neural Architecture Search for Parameter-Efficient Fine-tuning of Large Pre-trained Language Models](https://mlsanity.com/paper/2305.16597/) (2023)
- [Faster Causal Attention Over Large Sequences Through Sparse Flash Attention](https://mlsanity.com/paper/2306.01160/) (2023)
- [Pruning Self-attentions into Convolutional Layers in Single Path](https://mlsanity.com/paper/2111.11802/) (2021)
- [Adapting a Language Model While Preserving its General Knowledge](https://mlsanity.com/paper/2301.08986/) (2023)
- [Controllable Context Sensitivity and the Knob Behind It](https://mlsanity.com/paper/2411.07404/) (2024)
- [How Much Does Attention Actually Attend? Questioning the Importance of Attention in Pretrained Transformers](https://mlsanity.com/paper/2211.03495/) (2022)
- [Mokey: Enabling Narrow Fixed-Point Inference for Out-of-the-Box Floating-Point Transformer Models](https://mlsanity.com/paper/2203.12758/) (2022)
- [Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers](https://mlsanity.com/paper/2505.13737/) (2025)
- [Interpreting and Exploiting Functional Specialization in Multi-Head Attention under Multi-task Learning](https://mlsanity.com/paper/2310.10318/) (2023)
- [Tailoring Instructions to Student's Learning Levels Boosts Knowledge Distillation](https://mlsanity.com/paper/2305.09651/) (2023)
- [Joint Localization and Activation Editing for Low-Resource Fine-Tuning](https://mlsanity.com/paper/2502.01179/) (2025)
- [Identifying and Evaluating Inactive Heads in Pretrained LLMs](https://mlsanity.com/paper/2504.03889/) (2025)
- [Probing via Prompting](https://mlsanity.com/paper/2207.01736/) (2022)
- [Sub-network Discovery and Soft-masking for Continual Learning of Mixed Tasks](https://mlsanity.com/paper/2310.09436/) (2023)
- [Just CHOP: Embarrassingly Simple LLM Compression](https://mlsanity.com/paper/2305.14864/) (2023)
- [PreQuant: A Task-agnostic Quantization Approach for Pre-trained Language Models](https://mlsanity.com/paper/2306.00014/) (2023)
- [Sparsity-Accelerated Training for Large Language Models](https://mlsanity.com/paper/2406.01392/) (2024)
- [Auto-Compressing Subset Pruning for Semantic Image Segmentation](https://mlsanity.com/paper/2201.11103/) (2022)
- [Token Sparsification for Faster Medical Image Segmentation](https://mlsanity.com/paper/2303.06522/) (2023)
- [Modular Transformers: Compressing Transformers into Modularized Layers for Flexible Efficient Inference](https://mlsanity.com/paper/2306.02379/) (2023)
- [From PEFT to DEFT: Parameter Efficient Finetuning for Reducing Activation Density in Transformers](https://mlsanity.com/paper/2402.01911/) (2024)

## Abstract (arXiv)

Multi-head attention, a collection of several attention mechanisms that independently attend to different parts of the input, is the key ingredient in the Transformer. Recent work has shown, however, that a large proportion of the heads in a Transformer's multi-head attention mechanism can be safely pruned away without significantly harming the performance of the model; such pruning leads to models that are noticeably smaller and faster in practice. Our work introduces a new head pruning technique that we term differentiable subset pruning. Intuitively, our method learns per-head importance variables and then enforces a user-specified hard constraint on the number of unpruned heads. The importance variables are learned via stochastic gradient descent. We conduct experiments on natural language inference and machine translation; we show that differentiable subset pruning performs comparably or better than previous works while offering precise control of the sparsity level.