[Paper Review] The k-tied Normal Distribution: A Compact Parameterization of Gaussian Mean Field Posteriors in Bayesian Neural Networks

Jakub Świątkowski, Kevin A. Roth+8 more|arXiv (Cornell University)|Feb 7, 2020

Gaussian Processes and Bayesian Inference51 references22 citations

TL;DR

This paper proposes the k-tied Normal distribution, a compact parameterization for Gaussian mean-field variational posteriors in Bayesian neural networks by enforcing low-rank structure on the posterior standard deviation matrix. By tying parameters via low-rank SVD, the method reduces the number of variational parameters, improves gradient signal-to-noise ratio, and accelerates training without sacrificing predictive performance across multiple models.

ABSTRACT

Variational Bayesian Inference is a popular methodology for approximating posterior distributions over Bayesian neural network weights. Recent work developing this class of methods has explored ever richer parameterizations of the approximate posterior in the hope of improving performance. In contrast, here we share a curious experimental finding that suggests instead restricting the variational distribution to a more compact parameterization. For a variety of deep Bayesian neural networks trained using Gaussian mean-field variational inference, we find that the posterior standard deviations consistently exhibit strong low-rank structure after convergence. This means that by decomposing these variational parameters into a low-rank factorization, we can make our variational approximation more compact without decreasing the models' performance. Furthermore, we find that such factorized parameterizations improve the signal-to-noise ratio of stochastic gradient estimates of the variational lower bound, resulting in faster convergence.

Motivation & Objective

• To investigate whether simpler, compact parameterizations of the variational posterior can outperform richer ones in Bayesian neural networks.
• To address the high parameter count and noisy gradients in standard Gaussian mean-field variational inference.
• To explore whether low-rank structure in posterior standard deviations enables effective parameter tying without performance loss.
• To develop a parameterization that reduces memory and computation while maintaining predictive accuracy and convergence speed.

Proposed method

• The k-tied Normal distribution parameterizes the posterior standard deviation matrix as a low-rank SVD decomposition, reducing the number of free parameters.
• It enforces a tied structure across weights using a rank-k factorization of the standard deviation matrix, where k is much smaller than the original dimension.
• The method applies this factorization to the diagonal covariance matrix in mean-field variational inference, maintaining the independence assumption but reducing parameter count.
• The variational parameters are optimized via stochastic gradient descent on a differentiable approximation of the ELBO, with reduced gradient variance due to parameter tying.
• The approach is compatible with both fully connected and convolutional layers, and can be combined with existing variance reduction techniques.
• The method is shown to be distinct from matrix variate Gaussians, as it supports rank >1 tying, which cannot be represented in that framework.

Experimental results

Research questions

• RQ1 Can a low-rank parameterization of the posterior standard deviation matrix maintain or improve model performance in Bayesian neural networks?
• RQ2 Does reducing the number of variational parameters through parameter tying improve the signal-to-noise ratio of ELBO gradient estimates?
• RQ3 Is the low-rank structure in learned posterior standard deviations a consistent phenomenon across diverse Bayesian neural network architectures?
• RQ4 Can a compact parameterization outperform standard mean-field variational inference in terms of training speed and generalization?
• RQ5 Is the k-tied Normal distribution compatible with existing training techniques and scalable to deeper models?

Key findings

• The posterior standard deviation matrices in trained Bayesian neural networks consistently exhibit strong low-rank structure, enabling effective parameter tying.
• The k-tied Normal distribution reduces the number of variational parameters significantly while maintaining or improving ELBO and predictive performance.
• Parameter tying via low-rank SVD improves the signal-to-noise ratio of stochastic gradient estimates, leading to faster convergence during training.
• The method outperforms standard mean-field variational inference in terms of training efficiency and generalization across multiple model types.
• The k-tied approach is not representable within the matrix variate Gaussian framework when rank >1, establishing a distinct and more expressive class of parameterizations.
• The method is compatible with existing variance reduction techniques and can be combined with them to further improve training stability.