SOAP-Bubbles: Structured Weight Uncertainty for Neural Networks

2026-06-22

AI Summary

Although structured weight uncertainty can enhance deep learning performance, its efficient estimation remains challenging. This work proposes the EVON optimizer combined with the SOAP-Bubbles posterior, which introduces IVON variational inference into the feature space of the SOAP optimizer and leverages a preconditioner to effectively extend diagonal covariance to a non-diagonal form. This approach enables efficient estimation of structured posteriors with negligible additional training overhead. The method recovers exact Gaussian covariances in logistic regression and significantly outperforms existing diagonal-covariance approaches in large language model pretraining.

Abstract

Structured weight-uncertainty can improve many aspects of deep learning, but it remains costly to estimate and difficult to implement. Here, we show that these issues can be addressed by adapting the SOAP optimizer. Our key idea is to run IVON, an existing diagonal-covariance variational method, in the eigenspace of SOAP's preconditioner and then use the preconditioner to transform the diagonal estimate into a non-diagonal covariance. The resulting method has costs similar to those of SOAP and requires no drastic changes to training pipelines. We call the posteriors obtained in this way SOAP-Bubbles and our new optimizer Eigenspace-VON (EVON). We show that, for logistic regression, EVON recovers the exact Gaussian covariance and that, for language model pretraining, it yields significantly better results than existing diagonal-covariance methods. Our work makes it easier to estimate more expressive posterior distributions for deep learning at scale.

Problem

• structured weight uncertainty
• deep learning
• covariance estimation
• posterior distribution
• scalability

Innovation

• structured weight uncertainty
• SOAP optimizer
• non-diagonal covariance
• variational inference
• EVON