Variational Inference for Bayesian Logistic Regression

In this notebook we approximate the posterior distribution of a Bayesian logistic regression in three different ways and compare them:

1. Gaussian Variational Inference (VI) — fit a multivariate Gaussian $q(\beta) = \mathcal{N}(m, LL^\top)$ to the posterior by maximizing the evidence lower bound (ELBO).

2. Laplace approximation — fit a Gaussian centered at the posterior mode (the MAP), with covariance equal to the inverse negative Hessian of the log posterior at that mode.

3. Random-walk Metropolis–Hastings (MH) — a long MCMC run that we treat as the ground truth posterior.

We work two examples:

• A small 1-D simulation where we can also compute the exact posterior on a fine grid and see how VI and Laplace each compare to truth.

• A real dataset, frogs.csv (presence/absence of frogs at 212 sites with 7 environmental covariates), where the posterior lives in 8 dimensions and exhaustive enumeration is no longer possible.

The big-picture questions:

• How well does a Gaussian approximation work for a logistic regression posterior?

• When VI and Laplace disagree, which one is closer to the true posterior?

• How do the answers change when the design matrix has near-collinear predictors that produce a wide, elongated posterior?

1. Warm-up: a 1-D Bayesian logistic regression¶

We start with a tiny problem so that we can compute the exact posterior on a fine grid and use it as the gold standard.

Model. With $i = 1, \ldots, n$,

with $\sigma$ the logistic function. We simulate $n=30$ data points with $\theta_{\text{true}} = 2$ and prior scale $\tau = 2$.

Three approximations to the posterior:

• Exact posterior by grid quadrature. Evaluate the unnormalized log posterior on a fine grid in $\theta$, normalize with logsumexp, and read off the mean, sd, and mode by Riemann sums. This is the truth in 1-D.

• Gaussian variational inference. Restrict $q(\theta)$ to a univariate Gaussian with parameters $(\mu, s)$ and maximize the ELBO

  $$\mathrm{ELBO}(\mu, s) \;=\; \mathbb{E}_{q}\big[\log p(\theta, y)\big] \;+\; \mathcal{H}\big(q\big),$$

  (2)

  where the entropy of $\mathcal{N}(\mu, s^2)$ is $\tfrac{1}{2}\log(2\pi e\, s^2)$. We compute the expectation by Gauss–Hermite quadrature (numpy.polynomial.hermite.hermgauss): with $\theta = \mu + \sqrt{2}\, s\, z$ the integral against $q$ becomes $\sum_k w_k\, \log p(\theta_k, y) / \sqrt{\pi}$. Optimization is done by BFGS over $(\mu, \log s)$ — taking the log of $s$ keeps it positive.

• Laplace approximation. Find the posterior mode $\hat\theta_{\text{MAP}}$ by minimizing the negative log posterior, then approximate the posterior by $\mathcal{N}(\hat\theta_{\text{MAP}}, [-H(\hat\theta_{\text{MAP}})]^{-1})$, where $H$ is the second derivative of the log posterior. In 1-D the Hessian is just a number and the Laplace sd is $1/\sqrt{-H}$.

The cell below builds all three and overlays them on a single plot.

What we see in the 1-D example¶

The three densities — exact (solid), VI (dashed), and Laplace (dot-dashed) — are essentially indistinguishable on the plot, and the printed summaries agree to two or three decimals.

Two reasons this comparison is so flattering:

• The posterior is smooth, unimodal, and roughly bell-shaped, so a Gaussian is a good approximation almost regardless of how you tune it. In 1-D with a proper prior, the posterior is also automatically log-concave for logistic regression, so VI and Laplace are both well-posed.

• VI and Laplace are different fits in principle: Laplace is a local quadratic approximation at the MAP, whereas VI minimizes the global KL divergence $\mathrm{KL}(q \,\Vert\, p)$. When the posterior is nearly Gaussian, those two objectives have nearly the same optimum — that’s exactly what we see here. But even though the means are nearly the same, there is a strict difference between the mean learned by VI and the Laplace approximation.

The interesting question is how the two methods fare in higher dimensions on a real dataset, where the posterior may be skewed, heavy-tailed, or stretched along directions of weak identification. That’s what the rest of the notebook investigates.

2. The frogs dataset¶

We now switch to a real dataset. frogs.csv records presence/absence (pres.abs) of the southern corroboree frog at 212 sites in Australia, along with environmental covariates: altitude, distance to the nearest extant population, number of nearby pools and sites, average rainfall, and mean min/max temperatures.

Building the design matrix¶

Three pre-processing choices:

• Log-transform distance and NoOfPools. Both are highly right-skewed count-like variables; on the log scale they are roughly symmetric, which makes the Gaussian approximation more reasonable on this scale.

• Standardize each predictor to mean 0, sd 1. This is purely for numerical stability of the optimizer — it puts every coefficient on the same scale, so Adam’s single learning rate can move them all sensibly. We will undo this transformation at the end and report results on the original predictor scale too.

• Add an explicit intercept column of ones. The total dimension is therefore $p = 8$ (intercept plus 7 covariates), with $n = 212$ observations.

X_t and y_t are torch tensors used by the VI optimizer.

3. Variational inference for the multivariate posterior¶

We now do Gaussian VI in $p = 8$ dimensions. The variational family is the full-covariance multivariate normal,

parameterized by a mean vector $m \in \mathbb{R}^p$ and a lower-triangular Cholesky factor $L$. We will maximize the ELBO using the reparameterization trick so that we can backpropagate through Monte Carlo samples.

The log-likelihood, vectorized¶

For a batch of $S$ samples $\beta^{(1)}, \ldots, \beta^{(S)}$ each of shape $p$, this returns the log-likelihood at each sample. We use torch.nn.functional.softplus(eta) = log(1 + exp(eta)), which is the numerically stable form of the log-partition. The Bernoulli log-likelihood written this way is

Variational parameters and the Cholesky parameterization¶

Two trainable tensors:

• m is the variational mean, a free vector in $\mathbb{R}^p$.

• raw_L is an unconstrained $p \times p$ matrix. We turn it into a valid Cholesky factor via make_L:

  • take its lower-triangular part,

  • exponentiate the diagonal so the diagonal entries are guaranteed positive.

  This trick — optimizing the log of the diagonal — keeps $L$ in the space of valid Cholesky factors throughout training without any constrained optimization, exactly as we did with log s in the 1-D problem.

We use Adam with lr=0.003. Adam is a sensible default for noisy stochastic ELBO gradients because it adapts per-parameter step sizes.

The ELBO via the reparameterization trick¶

Direct Monte Carlo of $\mathbb{E}_q[\log p(y, \beta)]$ would not let us take gradients with respect to the parameters of $q$ — the sampling distribution would change with the parameters and stop being differentiable. The reparameterization trick fixes this by writing

so that $\beta$ is a deterministic function of $(m, L)$ and the random noise $z$. Then

which we estimate by drawing num_samples independent $z$’s and averaging. PyTorch’s autograd does the rest.

Entropy term. For $q = \mathcal{N}(m, LL^\top)$,

i.e. the standard multivariate-Gaussian entropy with $\log|\Sigma|^{1/2} = \sum_i \log L_{ii}$. The diagonal of $L$ is positive by construction, so the log is safe.

Running the optimizer¶

10,000 Adam steps, each estimating the ELBO with 5,000 reparameterized Monte Carlo samples. Printing every 500 iterations lets us watch the negative ELBO descend.

Optimization summary. The negative ELBO drops sharply in the first ~500 steps (from 262 down to 111) and then keeps creeping down, settling around 99 by step ~5000. The remaining wiggle (between roughly 98.99 and 99.06) is just Monte Carlo noise from the 5,000 samples per step — not a sign of further progress. So 10K iterations is more than enough.

A useful sanity check: the log-posterior at the MLE that we’ll compute below is $\approx -98.81$, and the ELBO is a lower bound on the log marginal likelihood. Our converged ELBO of $\approx -99$ being just a hair below -98.81 tells us the Gaussian approximation is doing very well in absolute terms — there isn’t much room for an even tighter Gaussian fit.

ELBO trace¶

The trace is the visual version of the printout above: a fast initial drop followed by a noisy plateau. If the curve had still been declining steadily at iteration 10,000 we would have wanted to keep training.

Extracting the variational mean and covariance¶

After training, m is the variational mean and $LL^\top$ is the variational covariance. The marginal sds are the square roots of its diagonal.

Reading the VI table. Two coefficients stand out: altitude (sd $\approx 4.91$) and meanmax (sd $\approx 5.36$). Their posterior means are relatively small (about -0.22 and -3.03) but their uncertainty is huge compared to the others. That is a classic symptom of near-collinearity: altitude and the temperature variables (meanmin, meanmax) carry overlapping information, so the data have very little to say about any one of them individually, even though their sum (or some linear combination) is well-identified.

By contrast, log_distance, log_NoOfPools, and meanmin look genuinely identified: nontrivial means with sds an order of magnitude smaller, and the mean for meanmin (+3.60) is roughly 3.7 posterior sds away from zero.

4. Laplace approximation for the same model¶

For comparison we now build the Laplace approximation on the standardized scale. Two equivalent things happen:

• statsmodels fits an ordinary logistic regression by IRLS, returning the MLE $\hat\beta$ and the fitted probabilities $\hat p_i$.

• The Hessian of the log-likelihood at the MLE is $-X^\top W X$ where $W = \mathrm{diag}\big(\hat p_i (1-\hat p_i)\big)$. Under a uniform prior, the Laplace mean is the MLE and the Laplace covariance is $(X^\top W X)^{-1}$ — i.e. exactly the asymptotic frequentist covariance of the MLE.

So under a flat prior, “Laplace approximation” and “MLE plus standard errors” are the same thing, just relabeled in Bayesian language. This is why the Laplace_mean_MLE column below is just the GLM coefficients.

VI vs Laplace, side by side. Most coefficients agree to two or three decimals: the intercept, log_distance, log_NoOfPools, NoOfSites, and even avrain are nearly identical between VI and Laplace. The two do disagree noticeably on a few:

These are the same coefficients with the largest sds — the ones least identified by the data, where the posterior is broad and the two objectives (KL minimization vs local quadratic fit at the mode) start to disagree. Whether VI or Laplace is closer to the truth here is exactly what the MH run later will let us check.

The sds, on the other hand, agree very tightly between VI and Laplace across all 8 coefficients — VI sds are slightly larger by roughly 1–2 %, which is the expected mode-seeking behavior of $\mathrm{KL}(q \,\Vert\, p)$ on a smooth, log-concave target.

Visual comparison: forest plot¶

Each row shows the posterior mean $\pm 2$ posterior sds for one coefficient, under VI (circles) and Laplace (squares), on the standardized scale. Coefficients whose interval crosses the dashed zero line are not “significant” in the loose sense; the others are.

5. Putting the answers back on the original predictor scale¶

We standardized predictors only for numerical convenience. To report results in interpretable units we need to undo that.

Let $\gamma_j$ be the coefficient on the standardized predictor and $\beta_j$ the coefficient on the original predictor, with predictor mean $\bar x_j$ and sd $s_j$. Then

This is a linear map $\beta_{\text{orig}} = A\, \beta_{\text{std}}$, so under any Gaussian approximation,

No re-fitting needed — the transformation is exact.

Sanity check: refit the GLM on the original predictors¶

For Laplace on the original scale, we could apply the same linear map. As an independent check, we instead refit the logistic regression directly on the un-standardized predictors and read off the MLE and standard errors. The results should match (and do match, up to floating-point noise).

Note on the intercept. On the original scale the intercept has mean $\approx 26$ and sd $\approx 135$ — both numerically large. That isn’t something pathological, it’s just what happens when the predictors are measured in their natural units (altitude in meters, rainfall in millimeters). The intercept is an extrapolated value at $x = 0$, far outside the observed range of the data, so its scale is dominated by the $\sum_j \gamma_j\, \bar x_j / s_j$ correction. The slopes themselves become tiny on the original scale (e.g. altitude ≈ -0.0017 per meter), which is also unsurprising because altitude varies by hundreds of meters across sites. This is exactly why the standardized scale is more comfortable for diagnostics, even when the original scale is what we want to report.

6. Ground truth: a long Metropolis–Hastings run¶

VI and Laplace are both Gaussian approximations. To know how good they are, we need a method that converges to the true posterior in the limit of infinite computation. We use random-walk Metropolis–Hastings, run long enough that Monte Carlo error is well below the differences we care about. The resulting samples will be our gold-standard reference.

Numpy log-posterior¶

We need a fast pure-numpy version of the log-posterior to call inside the MH loop. The prior is uniform, so log posterior = log likelihood (up to an additive constant that cancels in the MH ratio).

A quick sanity check: the log-likelihood at the MLE (-98.81) is much higher than at $\beta = 0$ (-146.95), so the data are very informative — moving to the MLE earns us about 48 nats of likelihood.

The Metropolis–Hastings sampler¶

Three implementation choices worth flagging:

• Proposal covariance. We use $\Sigma_{\text{prop}} = c^2\, \Sigma_{\text{Laplace}}$ with $c = 1.1$. Setting the proposal shape to match the (estimated) target shape is the standard way to handle the geometry of a multivariate target — without it the chain would have a hard time exploring the long, narrow direction defined by the altitude/temperature collinearity.

• Pre-computed Cholesky factor. prop_chol = chol(prop_cov) lets us generate proposals as $\beta_{\text{curr}} + L_{\text{prop}} z$ with $z \sim \mathcal{N}(0, I)$ — much faster than calling np.random.multivariate_normal inside the loop.

• Initialization at the MLE. Starting at the Laplace mode means burn-in is essentially free. With 1,000,000 iterations, even a generous 20K burn-in is a tiny fraction of the chain.

We work with log-acceptance ratios and compare against log(uniform) to avoid numerical underflow when the proposal is very poor.

Acceptance rate ≈ 0.17.

Diagnostics: ESS and trace plots¶

Two checks of MCMC quality:

• Effective sample size (ESS). Even though we have 980K samples, they are autocorrelated. ESS $= n / (1 + 2\sum_k \rho_k)$ measures the equivalent number of independent draws. We approximate this with the “initial positive sequence” of autocorrelations (truncated when $\rho_k$ first drops below 0.05). The Monte Carlo standard error of the posterior mean is then $\hat\sigma / \sqrt{\mathrm{ESS}}$.

• Trace plots. Visually a healthy trace looks like a “fuzzy caterpillar”: no trends, no sticky periods, the chain wandering freely around the posterior mean. Trends or extended flat stretches would indicate poor mixing.

Reading the diagnostics. All eight ESS values are between roughly $3.5 \times 10^4$ and $3.7 \times 10^4$ — about 3.7% of the raw 980K samples. That’s reasonable for a vanilla RWMH chain. The Monte Carlo standard errors on the posterior means are at most $\approx 0.03$ (for altitude and meanmax, the two collinear ones with the widest posteriors) and well below 0.01 for everything else. So when we compare VI and Laplace to the MH means below, any difference larger than $\approx 0.05$ is real and not Monte Carlo noise.

7. Three-way comparison: MH (truth) vs VI vs Laplace¶

With MH as ground truth, how close are VI and Laplace? We compute the absolute error of each approximation relative to MH for both the posterior mean and posterior sd, coefficient by coefficient.

The headline numbers (from the printed output):

So Laplace edges out VI on means while VI clearly beats Laplace on sds (roughly half the total error).

Looking coefficient-by-coefficient:

• Six well-identified coefficients (intercept, log_distance, log_NoOfPools, NoOfSites, avrain, meanmin): both methods are excellent. Mean errors are below 0.07 and sd errors below 0.04.

• Two badly identified coefficients (altitude, meanmax): both approximations show their largest errors here, but with opposite tendencies.

  • VI’s mean for altitude is -0.22 vs MH’s -0.58 (error 0.36). Laplace gets -0.83, which overshoots by 0.25.

  • VI’s mean for meanmax is -3.03 vs MH’s -3.41 (error 0.39). Laplace gets -3.45, which is essentially spot-on (error 0.03).

  The bias in VI is the well-known “mode-seeking” tendency of $\mathrm{KL}(q\Vert p)$: when the posterior is wide along a direction, $q$ would rather sit slightly inside the bulk than stretch out to cover the tails, so its marginal mean shrinks toward zero. Laplace, anchored at the MAP, doesn’t have this pull — but it pays for that with sd errors that are 2–3× worse on these same two coefficients (0.16 and 0.18 vs VI’s 0.08).

• VI sds are uniformly closer to MH for every single coefficient. The intuition: Laplace is a local quadratic at the mode, which underestimates variance whenever the true posterior has heavier-than-Gaussian tails. VI optimizes the global KL and so accounts for some of that tail mass.

Visual check: posterior marginals¶

We overlay the MH histogram (treated as truth) with the VI Gaussian (solid) and the Laplace Gaussian (dashed) for each of the 8 coefficients. Look for where the dashed and solid Gaussians peel away from the histogram — that is, where Gaussian-ness fails the data.

What the histograms reveal. For most coefficients the histogram is visibly Gaussian and both VI and Laplace track it nearly perfectly. The exceptions are again altitude and meanmax: their histograms are slightly shifted relative to the Gaussian fits — VI is centered a bit too far right, Laplace too far left — and their tails are a touch heavier than either Gaussian. That is a more honest picture than the table alone: the true posterior in those directions is not quite Gaussian, and that’s the fundamental obstacle no Gaussian approximation can fully overcome on this dataset.

8. Original-scale comparison¶

We repeat the three-way comparison after transforming everything back to the original predictor scale. The transformation is the same linear map $A$ as before — applied to the MH samples directly (then re-summarized), to the VI mean and covariance analytically, and to Laplace via the refit GLM.

Because the transformation is linear, the ranking of approximation errors is preserved up to scale: any coefficient where VI was closer to MH on the standardized scale is still closer on the original scale, and likewise for Laplace. The final two lines tally the wins.