Reproduce “Variational Deep Learning via Implicit Regularization” on OpenResearch

About this paper

Jonathan Wenger, Beau Coker, Juraj Marusic, John P. Cunningham

IBVI trains a distribution over neural-network predictions by minimizing expected data loss, relying on initialization, parameterization, and stochastic gradient descent (SGD) rather than an explicit prior penalty. In overparameterized linear regression, its zero-loss Gaussian is closest to the prior in 2-Wasserstein distance. On tested image-corruption benchmarks, uncertainty metrics are competitive, while deep ensembles have lower error; memory overhead is about 10%.

Key insights

• The implicit bias of SGD in overparametrized linear models, when minimizing only the expected loss for a Gaussian variational distribution, performs generalized variational inference that minimizes the 2-Wasserstein distance to the prior.
• This implicit regularization leads to a variational posterior where uncertainty collapses at training data points, but off the data manifold, the predictive uncertainty reflects the prior's influence.
• The variational distribution learned by IBVI for overparametrized linear models is theoretically equivalent to the distribution over weights of an ensemble of independently trained models initialized from the prior.

Summary by alphaXiv.