[Paper Review] Structured and Efficient Variational Deep Learning with Matrix Gaussian Posteriors

Christos Louizos, Max Welling|arXiv (Cornell University)|Mar 15, 2016

Gaussian Processes and Bayesian InferenceComputer Science27 references95 citations

TL;DR

This paper proposes a variational Bayesian neural network using matrix variate Gaussian posteriors to model correlations between weight matrix rows and columns, enabling efficient uncertainty estimation and improved generalization. By leveraging the local reparameterization trick and pseudo-data, the model achieves Gaussian Process-like behavior, outperforming fully factorized posteriors and dropout in regression and classification tasks with better uncertainty calibration and lower error rates.

ABSTRACT

We introduce a variational Bayesian neural network where the parameters are governed via a probability distribution on random matrices. Specifically, we employ a matrix variate Gaussian \cite{gupta1999matrix} parameter posterior distribution where we explicitly model the covariance among the input and output dimensions of each layer. Furthermore, with approximate covariance matrices we can achieve a more efficient way to represent those correlations that is also cheaper than fully factorized parameter posteriors. We further show that with the "local reprarametrization trick" \cite{kingma2015variational} on this posterior distribution we arrive at a Gaussian Process \cite{rasmussen2006gaussian} interpretation of the hidden units in each layer and we, similarly with \cite{gal2015dropout}, provide connections with deep Gaussian processes. We continue in taking advantage of this duality and incorporate "pseudo-data" \cite{snelson2005sparse} in our model, which in turn allows for more efficient sampling while maintaining the properties of the original model. The validity of the proposed approach is verified through extensive experiments.

Motivation & Objective

• To address the limitations of fully factorized weight posteriors in Bayesian neural networks, which fail to capture correlations among parameters and lead to poor uncertainty estimation.
• To develop a scalable variational inference method that models joint correlations across input and output dimensions of weight matrices using matrix variate Gaussians.
• To enhance model flexibility and uncertainty calibration by incorporating pseudo-data while preserving the Gaussian Process interpretation via the local reparameterization trick.
• To empirically validate the method on regression and classification benchmarks, demonstrating improved predictive performance and uncertainty quantification over existing approaches.

Proposed method

• Uses a matrix variate Gaussian distribution p(W) = MN(M, U, V) to model the posterior over weight matrices, where U captures row (output feature) correlations and V captures column (input feature) correlations.
• Applies the local reparameterization trick to the matrix variate Gaussian posterior, enabling efficient backpropagation and enabling a Gaussian Process interpretation of hidden units in each layer.
• Introduces pseudo-data pairs for each layer to maintain GP-like properties and improve posterior flexibility without increasing model capacity.
• Employs variational inference to minimize the KL divergence between the approximate and true posterior, using a tractable lower bound on the marginal likelihood.
• Utilizes Kronecker product structure (V ⊗ U) to parameterize the full covariance of vec(W), reducing the number of variance parameters compared to full covariance matrices.
• Employs a hybrid approach combining matrix variate posteriors with variational dropout, where the dropout rate is constrained via a variational upper bound.

Experimental results

Research questions

• RQ1 Can modeling correlations between input and output features of weight matrices improve uncertainty estimation in Bayesian neural networks?
• RQ2 Does the matrix variate Gaussian posterior enable a more efficient and accurate approximation of the true posterior than fully factorized posteriors?
• RQ3 To what extent does the Gaussian Process interpretation via the local reparameterization trick enhance model flexibility and predictive performance?
• RQ4 How does the introduction of pseudo-data affect the uncertainty calibration and generalization of the model?
• RQ5 Can the proposed method outperform state-of-the-art approaches like variational dropout and PBP in both regression and classification tasks?

Key findings

• The model achieves significantly better RMSE and predictive log-likelihood than VI, PBP, and dropout uncertainty on the same regression task, demonstrating improved predictive accuracy and uncertainty calibration.
• On the permutation-invariant MNIST classification task, the model achieves lower error rates than state-of-the-art methods, especially for small network architectures, indicating enhanced generalization.
• The matrix Gaussian posterior with pseudo-data maintains strong GP-like properties, allowing the model to remain flexible even with limited capacity, countering the performance drop typically seen in small networks.
• The fully factorized Gaussian posterior with local reparameterization underperforms on the toy regression task due to severe underfitting, highlighting the importance of structured weight posteriors.
• The predictive distribution from the matrix Gaussian model provides more realistic confidence intervals than dropout, which exhibits overconfidence in data-sparse regions.
• The method reduces the number of variance parameters compared to full covariance posteriors, resulting in more efficient and stable uncertainty estimation.