Search papers, labs, and topics across Lattice.
This paper introduces a novel preconditioned sampling strategy that leverages curvature estimates from adaptive optimizers like AdamW to enhance the exploration of high-dimensional posterior distributions in Bayesian neural networks. By utilizing optimizer-derived geometries, the method significantly reduces or eliminates the need for a lengthy burn-in phase during sampling, leading to improved numerical stability and predictive performance. The approach is validated across multiple datasets and network architectures, demonstrating its effectiveness without incurring additional computational costs.
Curvature estimates from adaptive optimizers can transform Bayesian sampling efficiency, slashing burn-in times while boosting predictive accuracy.
Sampling-based methods offer a principled approach to uncertainty quantification in Bayesian neural networks. Their practical use, however, is often challenged by the computational cost of exploring high-dimensional and multimodal posterior distributions. To overcome these difficulties, Bayesian Deep Ensembles, i.e., warmstarting the sampling from several optimized solutions, have proven to be an effective strategy. In this paper, we demonstrate that curvature estimates computed during the warmstart as a byproduct in adaptive optimizers such as AdamW can inform the sampling phase at negligible additional cost. Specifically, our proposed preconditioned sampling strategy based on optimizer-derived geometries can substantially reduce or even eliminate the need for a lengthy sampling burn-in phase and leads to greater numerical stability. This approach consistently maintains or improves predictive performance and uncertainty quantification without any additional computational costs. We confirm the consistency of our findings across various datasets and network architectures.