Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers

Firas Gabetni, Giuseppe Curci, Andrea Pilzer, Subhankar Roy, Elisa Ricci, Gianni Franchi

TL;DR

This work tackles uncertainty quantification in transformers by addressing the high cost of Deep Ensembles. It introduces Hydra Ensembles, a framework that prunes attention heads to create diverse subnetworks and fuses them into a single transformer via grouped fully-connected layers, achieving near-single-model inference with competitive UQ performance. Across image, text, and zero-shot tasks, Hydra Ensembles delivers strong calibration and OOD detection, often surpassing training-free baselines and narrowing or matching the Deep Ensembles gap. The approach demonstrates that careful, circuit-informed head pruning can preserve or enhance uncertainty estimates while dramatically reducing computational demands, making scalable UQ feasible for large-scale language-vision models.

Abstract

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.

Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers

TL;DR

This work tackles uncertainty quantification in transformers by addressing the high cost of Deep Ensembles. It introduces Hydra Ensembles, a framework that prunes attention heads to create diverse subnetworks and fuses them into a single transformer via grouped fully-connected layers, achieving near-single-model inference with competitive UQ performance. Across image, text, and zero-shot tasks, Hydra Ensembles delivers strong calibration and OOD detection, often surpassing training-free baselines and narrowing or matching the Deep Ensembles gap. The approach demonstrates that careful, circuit-informed head pruning can preserve or enhance uncertainty estimates while dramatically reducing computational demands, making scalable UQ feasible for large-scale language-vision models.

Abstract

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.
Paper Structure (56 sections, 1 theorem, 46 equations, 7 figures, 22 tables, 1 algorithm)

This paper contains 56 sections, 1 theorem, 46 equations, 7 figures, 22 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. On Pruning and Uncertainty Quantification
  4. Preliminaries
  5. Pruning under Noisy Data Conditions
  6. Efficient Ensembles with Hydra Ensembles
  7. Transformer preliminaries.
  8. Proposed Method
  9. How to construct ensemble members
  10. Experimental results
  11. Supervised Image Classification
  12. Supervised Text Classification
  13. Zero-shot Image Classification
  14. Conclusions
  15. Reproducibility Statement.
  16. ...and 41 more sections

Key Result

Proposition 1

Suppose Assumption 1 holds. Let $\mathbf{H}^t$ and $\mathbf{H}^n$ denote the Hessians of the loss on $\mathcal{D}^t$ and $\mathcal{D}^n$, respectively. If $\mathbf{H}^n - \mathbf{H}^t \succ 0,$ then the loss gap after pruning satisfies $\Delta\mathcal{L}(\boldsymbol{\theta}) \;\leq\; \Delta\mathcal{

Figures (7)

  • Figure 1: Hydra Ensembles. We start from a single transformer backbone and prune its attention heads to create multiple diverse subnetworks. These subnetworks are then combined at the head level into a Fused Multi-Head Attention (MHA), and then also merged at the MLP level, as described in Section \ref{['sec:Proposed_Method_imp']}. The pruned heads and MLPs can either be fine-tuned or kept frozen. Transformer heads are shown in matrix form for illustration only.
  • Figure 2: Illustration of Fused MHA.
  • Figure 3: Inference costs of Deep Ensembles, Single, and Hydra Ensembles with 8 heads on ImageNet-1k. We report runtime per batch under BF16/FP32 (left), as well as parameter count (middle) and multiply-adds (right).
  • Figure 4: Internal representation of the six most important (left) and least important (right) attention heads for OOD and ID data using PCA. The representation on the left shows a clearer separation between the two distributions, while on the right it illustrates how the distributions sometimes even collapse onto each other.
  • Figure 5: Effect of pruning different number of heads count using Taylor and CircAvg for ViT-B/16 on ImageNet-1K (Zero shot)
  • ...and 2 more figures

Theorems & Definitions (2)

  • Proposition 1: Pruning is Worse under Noise
  • proof