# Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers

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

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Builds on 57 of its 83 references that are in the map.
- First posted on arXiv 21 Oct 2025.

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

- [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)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Aggregated Residual Transformations for Deep Neural Networks](https://mlsanity.com/paper/1611.05431/) (2016)
- [A Baseline for Detecting Misclassified and Out-of-Distribution Examples in Neural Networks](https://mlsanity.com/paper/1610.02136/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [LAION-400M: Open Dataset of CLIP-Filtered 400 Million Image-Text Pairs](https://mlsanity.com/paper/2111.02114/) (2021)
- [Conditional Prompt Learning for Vision-Language Models](https://mlsanity.com/paper/2203.05557/) (2022)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [Evidential Deep Learning to Quantify Classification Uncertainty](https://mlsanity.com/paper/1806.01768/) (2018)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Do Vision Transformers See Like Convolutional Neural Networks?](https://mlsanity.com/paper/2108.08810/) (2021)
- [ImageNet-21K Pretraining for the Masses](https://mlsanity.com/paper/2104.10972/) (2021)
- [LiT: Zero-Shot Transfer with Locked-image text Tuning](https://mlsanity.com/paper/2111.07991/) (2021)
- [Importance Estimation for Neural Network Pruning](https://mlsanity.com/paper/1906.10771/) (2019)
- [ViM: Out-Of-Distribution with Virtual-logit Matching](https://mlsanity.com/paper/2203.10807/) (2022)
- [BatchEnsemble: An Alternative Approach to Efficient Ensemble and Lifelong Learning](https://mlsanity.com/paper/2002.06715/) (2020)
- [Open-Set Recognition: a Good Closed-Set Classifier is All You Need?](https://mlsanity.com/paper/2110.06207/) (2021)
- [Stochastic Gradient Hamiltonian Monte Carlo](https://mlsanity.com/paper/1402.4102/) (2014)
- [Why ReLU networks yield high-confidence predictions far away from the training data and how to mitigate the problem](https://mlsanity.com/paper/1812.05720/) (2018)

## Built on by (most cited first, 5 of 5)

- [UAT-LITE: Inference-Time Uncertainty-Aware Attention for Pretrained Transformers](https://mlsanity.com/paper/2602.02952/) (2026)
- [Improving Semantic Uncertainty Quantification in LVLMs with Semantic Gaussian Processes](https://mlsanity.com/paper/2512.14177/) (2025)
- [Regularized Meta-Learning for Improved Generalization](https://mlsanity.com/paper/2602.12469/) (2026)
- [Leveraging Visual Signals for Robust Token-Level Uncertainty in Vision-Language Generation](https://mlsanity.com/paper/2605.27136/) (2026)
- [Controllable Diversity in Normalization-Based Implicit Ensembles via Softmax-Temperature Modulation](https://mlsanity.com/paper/2607.23860/) (2026)

## Abstract (arXiv)

Uncertainty quantification (UQ) is essential for deploying deep neural networks in safety-critical settings. Although methods like Deep Ensembles achieve strong UQ performance, their high computational and memory costs hinder scalability to large models. We introduce Hydra Ensembles, an efficient transformer-based ensemble that prunes attention heads to create diverse members and merges them via a new multi-head attention with grouped fully-connected layers. This yields a compact model with inference speed close to a single network, matching or surpassing Deep Ensembles in UQ performance without retraining from scratch. We also provide an in-depth analysis of pruning, showing that naive approaches can harm calibration, whereas Hydra Ensembles preserves robust uncertainty. Experiments on image and text classification tasks, with various architectures, show consistent gains over Deep Ensembles. Remarkably, in zero-shot classification on ImageNet-1k, our approach surpasses state of the art methods, even without requiring additional training.