JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

Sampling weights and reasoning in function space

SGD that secretly samples, what Bayesian model averaging is really after, the leap from weight space to function space, deep kernels, and the cold-posterior puzzle.

SGD that secretly samples

Variational methods approximate the posterior; Markov-chain methods sample from it, and the workhorse of full Bayes — MCMC — is normally far too slow for deep nets. Stochastic gradient Langevin dynamics (SGLD) is the bridge: take an ordinary SGD step, but add a precisely-scaled dose of Gaussian noise each update. With a decaying step size, that noisy trajectory stops behaving like an optimiser converging to a point and starts behaving like a sampler whose visited weights are draws from the posterior.

for batch in data:                      # SGLD update
    g = grad_log_posterior(w, batch)     # gradient of log p(w|D), minibatch-scaled
    noise = sqrt(2 * lr) * randn_like(w)
    w = w + lr * g + noise               # SGD step + Langevin noise
    if past_burnin: samples.append(w.clone())   # collect posterior draws
SGLD: an SGD step plus calibrated noise turns optimisation into posterior sampling.
SGLD's noisy SGD is secretly a Markov chain wandering the weight posterior — each step is a draw, and the trajectory is your set of samples.

Interactive Markov chain hopping between states according to transition probabilities.

The beauty is that you keep SGD's minibatch scalability while gaining a chain of posterior samples; the headache is mixing — deep posteriors are multimodal, and a single chain explores only a neighbourhood of basins. Modern stochastic-gradient MCMC adds momentum, cyclical step sizes to hop between modes, and preconditioning, but honest multimodal sampling of a deep net remains genuinely hard.

What you actually want: model averaging

Whether you got your weight samples from a flow, an ensemble, Laplace, or SGLD, the payoff is the same act: Bayesian model averaging. Do not pick the single best network — average the predictions of all posterior samples, each weighted by its posterior probability. This is the Bayesian answer to overfitting to one solution: predictions hedge across every model the data leave plausible, and the hedge itself, the disagreement among samples, is your epistemic uncertainty.

p(y\mid x,\mathcal{D}) = \int p(y\mid x,\theta)\,p(\theta\mid\mathcal{D})\,d\theta \approx \frac{1}{S}\sum_{s=1}^{S} p\!\left(y\mid x,\theta^{(s)}\right),\quad \theta^{(s)}\sim p(\theta\mid\mathcal{D})

Bayesian model averaging integrates predictions over the posterior, approximated by averaging the predictions of the weight samples.

The leap to function space

Here is a quiet scandal of weight-space Bayes: we agonise over priors on weights, but nobody has intuition about what a Gaussian prior on a weight means for the functions the network can represent. Two very different weight vectors can encode the same function; a tidy prior on weights can imply a bizarre prior on functions. Function-space inference sidesteps this by placing the prior, and doing inference, directly over functions — the object you actually care about predicting.

The canonical function-space prior is the Gaussian process: a prior over functions whose behaviour is set by a kernel, and which gives exact posterior uncertainty that grows smoothly as you move away from the data — exactly the property weight-space mean-field struggles to produce. The deep link, the neural tangent kernel, shows that an infinitely-wide network trained by gradient descent is a Gaussian process with a specific kernel, making the GP not a rival to deep nets but their wide-limit shadow.

f\sim\mathcal{GP}\!\left(m(x),\,k(x,x')\right),\qquad \left(f(x_1),\dots,f(x_n)\right)\sim\mathcal{N}(\mathbf{m},\,\mathbf{K})

A Gaussian process places the prior directly in function space: any finite set of points is jointly Gaussian with covariance set by the kernel.

Deep kernels: the best of both

Gaussian processes have principled uncertainty but weak features; deep nets have rich features but shaky uncertainty. Deep kernel learning marries them: feed your inputs through a neural network, then put a GP kernel on top of the learned representation. The net learns where points are similar; the GP supplies calibrated uncertainty over that learned geometry. It is, in spirit, the neural-linear idea from Guide 3 promoted to a full nonparametric head.

The cold-posterior puzzle

Now an unsettling empirical fact. In deep nets, the true Bayesian posterior is often not the best predictor: you get better accuracy by raising the posterior to a power greater than one — sharpening it, concentrating mass near the modes. This cold posterior effect is awkward because, if Bayes were perfectly specified, the temperature-one posterior should be optimal. That it is not tells us something is misspecified.

p_T(\theta\mid\mathcal{D}) \propto \left(p(\mathcal{D}\mid\theta)\,p(\theta)\right)^{1/T},\qquad T<1\ \text{(cold)}

The tempered posterior raised to power 1/T; the cold-posterior effect is that a temperature below one often predicts better.

The leading suspects: priors over weights that are vague and badly-matched to good functions; data augmentation and label noise that effectively change the likelihood's temperature; and minibatch/curvature artefacts in how the posterior is sampled. The cold posterior is a healthy reminder that BDL is an approximation stack — prior, likelihood, and inference are all imperfect, and the temperature knob is quietly absorbing all three mismatches at once.