Klemens Flöge
All publications

Stein Variational Newton Neural Network Ensembles

Klemens Flöge, Mohammed Abdul Moeed, Vincent Fortuin

ICML 2024 Workshop on Structured Probabilistic Inference & Generative Modeling·workshop

Conceptual overview of the SVN method. The curvature-informed SVN updates (green) are higher quality and need fewer steps than the corresponding SVGD updates (blue).
Conceptual overview of the SVN method. The curvature-informed SVN updates (green) are higher quality and need fewer steps than the corresponding SVGD updates (blue).

Abstract

Deep neural network ensembles are powerful tools for uncertainty quantification, which have recently been re-interpreted from a Bayesian perspective. However, current methods inadequately leverage second-order information of the loss landscape, despite the recent availability of efficient Hessian approximations. We propose a novel approximate Bayesian inference method that modifies deep ensembles to incorporate Stein Variational Newton updates. Our approach uniquely integrates scalable modern Hessian approximations, achieving faster convergence and more accurate posterior distribution approximations. We validate the effectiveness of our method on diverse regression and classification tasks, demonstrating superior performance with a significantly reduced number of training epochs compared to existing ensemble-based methods, while enhancing uncertainty quantification and robustness against overfitting.

Deep ensembles are one of the most reliable tools for uncertainty quantification, and they have a clean Bayesian reading — but standard training throws away the second-order structure of the loss landscape entirely. Now that scalable Hessian approximations are practical, that is a waste.

This work modifies deep ensembles to take Stein Variational Newton steps instead of first-order ones. The curvature information sharpens the posterior approximation and cuts the number of training epochs substantially, while improving calibration and robustness to overfitting across regression and classification benchmarks.