Evaluating Prediction Uncertainty Estimates from BatchEnsemble

TL;DR

This work tackles the challenge of reliable predictive uncertainty in deep learning by evaluating BatchEnsemble as a scalable, parameter-efficient alternative to deep ensembles. It introduces GRUBE, a BatchEnsemble-enabled GRU for time-series forecasting, and demonstrates through comprehensive experiments that BatchEnsemble achieves uncertainty estimates comparable to deep ensembles while using far fewer parameters and less computation, outperforming MC dropout. Ablation studies show that adapters, GRU gates, and layer-depth all contribute to diversity and performance, with full deployment across all components delivering the strongest results. The findings support BatchEnsemble as a practical, scalable approach for calibrated uncertainty estimation across tabular and time-series domains, including under distribution shift and when selective prediction is required.

Abstract

Deep learning models struggle with uncertainty estimation. Many approaches are either computationally infeasible or underestimate uncertainty. We investigate \textit{BatchEnsemble} as a general and scalable method for uncertainty estimation across both tabular and time series tasks. To extend BatchEnsemble to sequential modeling, we introduce GRUBE, a novel BatchEnsemble GRU cell. We compare the BatchEnsemble to Monte Carlo dropout and deep ensemble models. Our results show that BatchEnsemble matches the uncertainty estimation performance of deep ensembles, and clearly outperforms Monte Carlo dropout. GRUBE achieves similar or better performance in both prediction and uncertainty estimation. These findings show that BatchEnsemble and GRUBE achieve similar performance with fewer parameters and reduced training and inference time compared to traditional ensembles.

Figures (19)

• Figure 1: GRU with BatchEnsemble (GRUBE). The architecture combines the standard GRU architecture with the BatchEnsemble formalism applied to the weights $W_z, W_f$, and $W_h$.
• Figure 2: Regression results across tabular datasets. For each metric it shows the mean $\pm$ standard error. The dotted line shows the lower limit of BatchEnsemble.
• Figure 3: Regression results across tabular datasets with distribution shift. For each metric it shows the mean $\pm$ standard error. The dotted line shows the lower limit of BatchEnsemble.
• Figure 4: Selective prediction evaluation on tabular regression datasets.
• Figure 5: Classification results across tabular datasets. The dotted line shows the lower limit (upper for accuracy) of BatchEnsemble.