Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers

TL;DR

Hydra Ensembles is a novel, efficient transformer ensemble method that prunes attention heads and merges them to achieve strong uncertainty quantification with minimal computational overhead.

Contribution

We propose Hydra Ensembles, a pruning and merging technique for transformers that improves uncertainty quantification efficiently without retraining from scratch.

Findings

Hydra Ensembles match or surpass Deep Ensembles in UQ performance.

The method achieves inference speeds close to a single network.

It outperforms state-of-the-art methods in zero-shot ImageNet classification.

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.