Reproduce “Ensembling Pruned Attention Heads For Uncertainty-Aware Efficient Transformers” on OpenResearch

About this paper

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

Hydra Ensembles constructs compact, uncertainty-aware transformer models by generating diverse subnetworks from a single pre-trained model through structured attention head pruning and then merging them. This method delivers uncertainty quantification performance comparable to Deep Ensembles, while significantly reducing computational cost and achieving near single-model inference speed.

Key insights

• Naive pruning of neural networks can theoretically and empirically degrade performance on noisy data and impair uncertainty estimation, even if predictive accuracy is maintained on clean data.
• Specific attention heads within transformers specialize in distinguishing in-distribution from out-of-distribution data; preserving these heads is crucial for robust UQ.
• Ensemble diversity, primarily driven by variations in the selected attention heads through different pruning seeds, is key to achieving robust out-of-distribution detection in the proposed merged architecture.

Summary by alphaXiv.