SOAP-Bubbles: Structured Weight Uncertainty for Neural Networks

Deep learning models often benefit from weight uncertainty—representing model parameters as probability distributions rather than point estimates—for tasks like model merging, reinforcement learning, and improved uncertainty estimation. However, existing scalable variational methods are typically restricted to diagonal Gaussian approximations, which fail to capture inter-parameter correlations. The authors seek to develop a scalable, easy-to-implement method for learning expressive, non-diagonal posterior distributions in large-scale neural networks.

The authors propose SOAP-Bubbles, a family of structured posterior distributions defined by rotating a diagonal Gaussian into the eigenspace of the SOAP optimizer’s preconditioner. To optimize these distributions, they introduce Eigenspace-VON (EVON). EVON modifies the SOAP optimizer by running the Variational Online Newton (IVON) method within the eigenspace spanned by the Kronecker-factored preconditioners. By performing variational inference in this transformed coordinate system, the method effectively learns a non-diagonal covariance in the original weight space without the prohibitive memory or compute overheads associated with traditional structured variational inference.

EVON demonstrates significant improvements over existing diagonal-covariance variational methods. In theoretical benchmarks, such as binary logistic regression, EVON successfully recovers the exact Gaussian posterior. In large-scale experiments, specifically pretraining NanoGPT models on the FineWeb-1B dataset, EVON yields better validation loss compared to IVON. Furthermore, the authors show that ensembling models sampled from SOAP-Bubbles provides superior performance compared to ensembling from diagonal posteriors, confirming that the method successfully captures a richer and more useful representation of weight uncertainty.