SOAP-Bubbles: Effective and Scalable Variational Learning with Structured Covariances

Structured posteriors are expected to be better than mean-field posteriors, but recent variational learning methods for large deep networks only use Gaussians with diagonal covariance. A reason is that diagonal covariances can be estimated by simple modifications of existing optimizer's implementations, for instance, the IVON optimizer closely follows Adam's code. Unfortunately, no such alternatives for Gaussians with structured covariances are effective and scalable. Here, we fill this gap and show that block-diagonal covariances can be obtained by adapting the SOAP optimizer. Specifically, our Eigenspace Variational Online Newton (EVON) method modifies SOAP to run IVON instead of Adam in the Eigenspace of the preconditioner. We refer to the posteriors obtained this way as 'SOAP-Bubbles' and show that for logistic regression they recover the optimal posterior approximation among full Gaussians. For language model pretraining we get significant improvements in validation loss over IVON without any increase in cost, and ensembling models drawn from SOAP-Bubbles reduces the test loss more than ensembling using IVON's posterior. Our work shows that there is a fundamental connection between second-order optimization and Gaussian posteriors, which can be used to improve the accuracy of variational learning.