Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers

AI summary

Overview

Research area: Uncertainty quantification (UQ) for deep neural networks, specifically efficient ensembling of transformer architectures (ViT, BERT, OpenCLIP-ViT).

Technical level: Intermediate — the paper assumes familiarity with Transformer attention blocks, ensemble methods, calibration metrics (ECE, Brier, NLL), and OOD detection metrics (AUROC, FPR95, AUPR), plus second-order loss analysis for the theoretical proposition.

Scope: The paper proposes Hydra Ensembles, a method that prunes attention heads from a single pre-trained transformer to create diverse ensemble members and fuses them into one model via Grouped Fully-Connected layers, delivering Deep-Ensembles-level uncertainty estimates at close to single-model inference cost.

What This Paper Is About

Deep Ensembles is the most reliable known approach for uncertainty quantification, but it requires training, storing, and running several independent models, which is prohibitively expensive for large foundation models such as CLIP and BERT. The authors ask how pruning can be used to build transformer ensembles that stay diverse and well-calibrated while running at near single-model cost, and they show both theoretically and empirically that naive pruning can actively harm calibration rather than help it.

Key Contributions

1. An analysis of pruning's effect on uncertainty. The authors prove (Proposition 1) and demonstrate empirically that classical pruning can degrade performance more strongly on noisy/corrupted data than on clean data, so pruned networks cannot simply be reused for UQ.

2. Hydra Ensembles, described as the first pruning framework specifically designed for UQ in transformer-based large-scale models. It prunes attention heads to create diverse subnetworks from one pre-trained backbone and merges them into a single model using a fused multi-head attention with Grouped Fully-Connected (GFC) layers.

3. Two membership-construction strategies, Hydra Ensembles (Taylor), which uses classical structured Taylor pruning (Molchanov et al., 2019) when no uncertainty validation set is available, and Hydra Ensembles (Circuit), which uses the Headmap algorithm (Wang et al., 2025) to extract circuits when such a set exists.

4. Empirical validation across three tasks — supervised image classification (ViT-B/16 on ImageNet-1K and CIFAR-100), supervised text classification (bert-base-uncased on SST-2), and zero-shot image classification (OpenCLIP-ViT/B-32 on nine datasets) — showing uncertainty estimates comparable to Deep Ensembles with much lower training and inference cost, and state-of-the-art zero-shot ImageNet-1k OOD results without any additional training.

Main Findings

• Pruning can hurt calibration under noise. Under Assumption 1 (the model minimizes both training and clean-test loss, the pruning perturbation is non-negatively correlated with the noisy-loss gradient) and if the Hessian difference H^n − H^t is positive definite, Proposition 1 states that ΔL(θ) ≤ ΔL(θ + δθ), i.e. pruning degrades performance more severely on the noisy dataset than on the clean one.

• Motivating question answered by head-level structured pruning. Pruning attention heads rather than MLPs is chosen because MoE methods already exploit MLP specialization, head-level pruning is simpler for model merging, and adding heads is cheaper than merging full MLPs. In ViT-B/16 the paper notes MLPs contain 4.7M parameters versus 2.3M for MHA.

• Merging keeps diversity without degradation. At each MLP layer, weights and biases are averaged across the M members; queries, keys, values, and the output projection are computed jointly with Grouped Fully-Connected layers, with inputs reshaped from MT×d to T×Md so fused MHA processes the same number of tokens as a standard MHA. The paper reports that applying MLP fusion post training does not reduce diversity or degrade ID/OOD metrics (Appendix B.7).

• Supervised ImageNet-1K results. On ImageNet-1K with 8 heads per MHA block (8x3), Hydra Ensembles (Circ) reaches 80.88 accuracy versus 82.19 for Deep Ensembles (a −1.3% gap) and 80.67 for the Single model, while improving average OOD performance over Deep Ensembles by AUROC +0.8% and AUPR +0.4%. Hydra Ensembles (Taylor) reaches 81.20 accuracy with OOD AUROC 85.36, FPR95 50.50, AUPR 60.75.

• Supervised CIFAR-100 results. With 10 heads per block (10x3), Hydra Ensembles shows its largest UQ advantage over Deep Ensembles on average OOD metrics: AUROC +3.4%, FPR95 −2.2%, AUPR +1.2%. Hydra Ens (Taylor) reports 92.00 accuracy, AUROC 88.89, FPR95 39.57, AUPR 95.46; Hydra Ens (Circ) reports 92.11, AUROC 89.43, FPR95 36.44, AUPR 95.17.

• Inference cost close to a single model. In bfloat16 on ImageNet-1k, the Hydra Ensembles/Single ratio is 1.07× for both whole-test runtime and per-batch inference, whereas Deep Ensembles is almost three times slower and has more than twice as many parameters.

• Text classification on SST-2. With 6 heads per attention block (6x3), Hydra-Ens (Circ) matches Deep Ensembles (D) accuracy (92.55 versus 93) and calibration while improving OOD detection over Deep Ensembles (D) by AUROC +2.8%, FPR95 −7.6%, and AUPR +2.2%. Hydra-Ens (Circ) records AUROC 77.6, FPR95 55.06, AUPR 84.16, versus 74.81, 62.69, 84.9 for Deep Ensembles (D).

• Zero-shot ImageNet-1k OOD beats a trained state of the art. With both encoders pruned (10x3) and no training, Hydra Ensembles achieves the best accuracy (74.00), ECE (3.49), and aECE (3.35) and the second-best NLL (0.93), and surpasses ViLU on OOD by +1.3 AUROC (76.82 vs 75.38), −3.5 FPR95 (68.05 vs 71.59), and +4.0 AUPR (47.85 vs 43.81).

• Naive and unstructured pruning are comparatively weak. Taylor alone (8 heads) on ImageNet-1K reaches 80.68 accuracy but AUPR 59.46 and FPR95 54.51, and the unstructured OBA baseline (Sun et al., 2025) reaches 78.52 accuracy with OOD AUROC 82.61. In the zero-shot setting, Taylor drops to 60.75 accuracy with ECE 13.76, while CircAvg reaches 71.91 accuracy and AUROC 76.88.

• Fine-tuning choice. Supervised image and text classification models are fine-tuned; zero-shot image classification is not. All experiments use M = 3 ensemble members, described as a standard choice in Deep Ensembles.

Methodology in Plain English

The authors start from one already-trained transformer rather than from scratch. They make several copies of that backbone and prune a different set of attention heads from each copy, so each copy becomes a slightly different model that still knows the task. Two ways of choosing which heads to remove are used: a classical importance-based method (Taylor) when no validation data is available, and circuit extraction (Headmap) when noisy or uncertainty-focused validation data exists. The pruned copies are then stitched into one network. The MLP parts are simply averaged across copies, while the attention parts keep each copy's surviving heads separate and compute them together in a single grouped matrix operation, with the token dimension reshaped so the attention still processes the same number of tokens. Because the surviving heads come from differently pruned or differently tuned copies, the merged model retains ensemble-like disagreement between "members," and because everything runs in one forward pass, the cost stays near that of a single model. The authors test this on three settings: fine-tuned ViT-B/16 on ImageNet-1K and CIFAR-100 (with OpenOOD splits for OOD), fine-tuned BERT on SST-2, and untouched OpenCLIP-ViT/B-32 for zero-shot classification over nine datasets.

Why This Matters

Impact on research. The paper challenges the assumption that accuracy-preserving pruning is harmless, giving a formal condition under which pruning worsens behaviour on noisy data, and reframes pruning as a tool for building diverse, uncertainty-aware ensembles rather than merely for compression.

Real-world applications.

• Safety-critical vision systems such as medical imaging or autonomous driving, where models must signal when they are unreliable.
• Large vision-language deployments (for example CLIP-style retrieval or classification) where running several full models is too expensive.
• Text classification and sentiment pipelines built on BERT-scale models that need calibrated confidence scores.
• Any deployment constrained by memory or latency budgets, such as edge devices, where a memory-efficient alternative to Deep Ensembles is needed.

Industry relevance. The 1.07× inference ratio in bfloat16 relative to a single model, combined with no need to retrain members from scratch, makes ensemble-grade uncertainty substantially cheaper to serve, which matters for companies currently paying roughly three times the inference cost for Deep Ensembles.

Future Directions

• Scaling Hydra Ensembles beyond M = 3 members: the paper provides a cost/benefit study for more members in Appendix B.6, but the main experiments fix M = 3.
• Understanding circuit-level representations for UQ more deeply, since Appendix B.2 reports that certain attention heads are highly specialized for this task and that pruning them is particularly harmful.
• Extending the approach to other modalities and foundation models, given that the current evaluation covers image classification, text classification, and zero-shot CLIP classification.
• Deciding systematically when fine-tuning pruned members is necessary versus when zero-shot use is sufficient, which the paper ties to Proposition 1 and to whether circuit-based selection is used.

Target Audience

Researchers and engineers working on uncertainty quantification, model compression, and efficient inference for transformers, plus practitioners deploying CLIP- or BERT-scale models in latency- or memory-constrained settings. Readers with a background in deep learning but not in UQ will still follow the empirical sections; the theoretical proposition and the fused-attention derivation require comfort with loss expansions and attention mechanics.

Authors’ 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.