Variational Bayes with the Reparameterisation Trick

Approximate the posterior of a Bayesian linear regression with variational inference, and compare it with the exact answer.

The model has a Gaussian prior on the weights and Gaussian noise:

w ~ Normal(0, prior_std^2)     for each of the d weights
y ~ Normal(X @ w, noise_std^2)

For this model the exact posterior is known, which makes it a good test. Variational inference does not use that formula. It picks a simple family, here a diagonal Gaussian q(w) = Normal(mu, sigma), and fits mu and sigma by maximising the evidence lower bound:

ELBO = E_q[ log p(y | X, w) ] - KL(q || prior)

Estimate the expectation with samples drawn by q.rsample(). Those samples are mu + sigma * eps with eps ~ Normal(0, 1), so gradients flow back to mu and sigma. That is the reparameterisation trick. Get the KL term from torch.distributions.kl_divergence, which knows the closed form for two Gaussians.

Return the fitted q as a torch.distributions.Normal.

Function signature

def fit_variational(X, y, noise_std, prior_std, steps=1000):

Example

Input

X: (50, 3) with columns that are not correlated, y = X @ [1.5, -2, 0.5] + noise
q = fit_variational(X, y, noise_std=0.5, prior_std=1.0)

Output

q.mean    ->  close to the exact posterior mean, about [1.5, -2.0, 0.5]
q.stddev  ->  close to the exact posterior standard deviations

Reasoning

When the columns of X are not correlated, the exact posterior has a diagonal covariance, so a diagonal Gaussian can match it exactly. The ELBO is largest at that match.

Constraints and notes

• X is float32 (n, d) and y is float32 (n,). Keep sigma positive by training an unconstrained tensor and applying softplus.
• Use rsample, not sample. sample cuts the gradient.
• Average over several samples each step, and lower the learning rate near the end, so that the noise in the estimate settles.
• When the columns are correlated, the exact posterior is not diagonal and a diagonal q cannot match it. The tests check what the optimal diagonal q is in that case.
• Return q built from detached tensors.