Smoothing, Volume, and Samples

Take a look at the data below. What do you see?

Despite the sparsity and noise, a pattern is visible, and one often wants to infer the underlying smooth function which produced the data. A common misconception is that it is necessary to assume a degree of smoothness in order to avoid overfitting the noise.

This notebook illustrates why that is not the case. It turns out we can jointly infer the function and its smoothness in a sensible way using Bayes' theorem and Gaussian processes. This problem teaches us the importance of parameter volume in high dimensions, explains why the posterior maximum often looks nothing whatsoever like a posterior sample, and demonstrates that uncertainty quantification is valuable even if one does not need uncertainties.

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
plt.rcParams['figure.dpi'] = 200
rng = jr.key(123)

# Example data
x = jnp.linspace(0, 1, 100)
y = jnp.sin(30*x) * jnp.exp(-3*x)
k1, k2 = jr.split(jr.key(2))
data_x = jr.uniform(k2, (100,))
data_err = 0.1 * jnp.ones_like(data_x)
data_y = jnp.interp(data_x, x, y) + data_err * jr.normal(k1, data_x.shape)
plt.errorbar(data_x, data_y, yerr=data_err, fmt='o', c='k')
plt.show()

Before diving in, consider a lower-dimensional problem. Suppose we observe a single data point, $d=1$, which we think is drawn from a Gaussian distribution with mean $\mu$ and variance $v$. Bayes' theorem says

$$ P(\mu,v|d) = \frac{P(d|\mu,v)\,P(\mu,v)}{P(d)} $$

For the prior let us choose wide Gaussians for $\mu$ and $\log v$, to guarantee positive variance. Since we are only in two dimensions, we can grid the posterior explicitly. Take a moment to reason through the plots below.

As you may have noticed, the posterior maximum (MAP) represents a narrow Gaussian with finely-tuned mean, increasing the likelihood but overfitting the data. Posterior samples, on the other hand, are typically wider Gaussians with a larger range in mean. Any given point in this region has lower probability density, but there is a larger total "volume" of acceptable parameter space, leading more random samples to come from this region! The interplay between volume and density is key to understanding Bayes in high dimensions.

# Data and parameters
d = 1
mu, logvar = jnp.mgrid[-5:5:300j, -5:5:300j]
var = jnp.exp(logvar)

# Bayes theorem
prior = jnp.exp(-1/2 * (mu**2 / 5)) * jnp.exp(-1/2 * (logvar**2 / 5))
likelihood = jnp.exp(-1/2 * ((mu - d)**2 / var)) / jnp.sqrt(2 * jnp.pi * var)
posterior = prior * likelihood

# Posterior samples
samples_idx = jr.choice(jr.key(0), a=posterior.size, p=jnp.ravel(posterior), shape=(30,))
samples_mu, samples_logvar = jnp.ravel(mu)[samples_idx], jnp.ravel(logvar)[samples_idx]

# Posterior maximum (MAP)
map_idx = jnp.argmax(posterior)
map_mu, map_logvar = jnp.ravel(mu)[map_idx], jnp.ravel(logvar)[map_idx]

# Plot!
fig, axes = plt.subplots(ncols=3, figsize=(9,3), sharey=True)
axes[0].pcolormesh(mu, logvar, prior, cmap='inferno')
axes[1].pcolormesh(mu, logvar, likelihood, cmap='inferno')
axes[2].pcolormesh(mu, logvar, posterior, cmap='inferno')
axes[2].scatter(samples_mu, samples_logvar, color='deepskyblue', edgecolor='black', s=8, lw=0.5, label='samples')
axes[2].scatter(map_mu, map_logvar, marker='X', facecolor='tomato', edgecolor='black', s=30, lw=0.5, label='MAP')
axes[2].contour(mu, logvar, posterior, [0.005], colors='white', linewidths=0.5)
for i, ax in enumerate(axes):
    ax.set(box_aspect=1, title=['Prior', 'Likelihood', 'Posterior'][i], xlabel=r'$\mu$', ylabel=r'$\log v$' if i==0 else '')
axes[2].legend(fontsize=8, frameon=False, labelcolor='white', handletextpad=0.0)
plt.show()

To make this result more obvious, we can integrate over $\mu$ to obtain the "marginal" posterior for $\log v$. Due to the volume effect, the joint MAP is quite a bit lower than "most" of the posterior and also the maximum of this marginal posterior. One might object that this whole story depends on the parameterization, and that's exactly the point! MAP is coordinate-dependent, whereas the act of sampling from the posterior is the same no matter which parameterization we choose. Sometimes the parameterization is friendly and MAP works fine, but sometimes not. That's what we will see next with Gaussian processes.

marginal_posterior = jnp.sum(posterior, axis=0)
marginal_prior = jnp.sum(prior, axis=0)

plt.figure(figsize=(6,4))
plt.plot(logvar[0], marginal_prior / jnp.max(marginal_prior), c='C0', ls=':', label='Prior')
plt.plot(logvar[0], marginal_posterior / jnp.max(marginal_posterior), c='C0', label='Posterior')
plt.axvline(map_logvar, color='red', linestyle='--', label='Joint MAP')
plt.gca().set(xlabel=r'$\log v$', ylabel='Probability (unnormalized)', yticks=[])
plt.legend()
plt.show()

A Gaussian process is a "stochastic process (a collection of random variables indexed by time or space), such that every finite collection of those random variables has a multivariate normal distribution" (Wikipedia). The key ingredient is the covariance between any two points in space or time. A stationary Gaussian process is one where the covariance depends only on some kind of distance between the points, resulting in a field which is (statistically) homogeneous and isotropic.

Gaussian processes are great for representing distributions of smooth functions. For example, take the simple covariance kernel below. By sampling from a mean-zero GP with different values of $s$, we can obtain functions of varying smoothness. Remember, all we are doing here is sampling from one big multivariate Gaussian.

$$ k(r) = \exp\left(-\frac{r^2}{2s^2} \right) $$

# Forward model
def kernel(logscale):
    r = jnp.abs(x[:, None] - x[None, :])
    K = jnp.exp(-0.5 * (r / jnp.exp(logscale))**2)
    K *= jnp.where(r == 0, 1 + 1e-14, 1)
    return K

def field(logscale, xi):
    K = kernel(logscale)
    L = jnp.linalg.cholesky(K)
    return L @ xi

# Random samples
rng, k1 = jr.split(rng)
xi = jr.normal(k1, (x.shape[0],))

fig, axes = plt.subplots(ncols=3, figsize=(12,4))
for ax, s in zip(axes, [0.002, 0.02, 0.2]):
    f = field(jnp.log(s), xi)
    ax.plot(x, f)
    ax.set(title=f"s={s}", box_aspect=1)
plt.show()

What happens if we jointly infer the smoothing scale $s$ along with the Gaussian process parameters which define the field?

As a technical aside, high-dimensional problems like this can become intractable very quickly. However, in the case of a fixed kernel, the Gaussian prior and Gaussian likelihood make for a Gaussian posterior which can be calculated analytically. We can then afford to build an explicit grid over $s$, the only troublesome nonlinear parameter, and compute things like the marginal posterior, posterior samples, and posterior maximum in standardized coordinates. These are all great introductory exercises so they are left to the reader with the code below as a hint :)

Take a look at the results and see if you can understand them in terms of the intuition from the 2D problem above.

# Set up problem
logscale_grid = jnp.linspace(-7, -1, 500)
logscale_prior = (-4, 1)

def response(logscale, xi):
    f = field(logscale, xi)
    return jnp.interp(data_x, x, f)

def logprior(logscale):
    return -0.5 * ((logscale - logscale_prior[0])/logscale_prior[1])**2

def loglike(logscale):
    R = jax.jacobian(lambda xi: response(logscale, xi))(jnp.zeros_like(x))
    S = R @ R.T + jnp.diag(data_err**2)
    _, logdet = jnp.linalg.slogdet(S)
    return -0.5 * (data_y.T @ jnp.linalg.solve(S, data_y) + logdet)

# Scale marginal posterior
prior = jnp.exp(jax.vmap(logprior)(logscale_grid))
posterior = jnp.exp(jax.vmap(loglike)(logscale_grid)) * prior

# Posterior samples
def draw_sample(logscale, noise):
    R = jax.jacobian(lambda xi: response(logscale, xi))(jnp.zeros_like(x))
    D_inv = jnp.eye(len(x)) + (R.T / data_err**2) @ R
    m = jnp.linalg.solve(D_inv, R.T @ (data_y / data_err**2))
    L = jnp.linalg.cholesky(D_inv)
    return field(logscale, m + jnp.linalg.solve(L.T, noise))

n_samples = 8
rng, k1, k2 = jr.split(rng, 3)
logscale_samples = jr.choice(k1, logscale_grid, p=posterior, shape=(n_samples,))
noise = jr.normal(k2, (n_samples, len(x)))
samples = jax.vmap(draw_sample)(logscale_samples,noise)

# Maximum a posteriori
def max_logprob(logscale):
    R = jax.jacobian(lambda xi: response(logscale, xi))(jnp.zeros_like(x))
    D_inv = jnp.eye(len(x)) + (R.T / data_err**2) @ R
    m = jnp.linalg.solve(D_inv, R.T @ (data_y / data_err**2))
    logprob = -0.5 * jnp.sum(((data_y - R @ m) / data_err)**2) - 0.5 * jnp.sum(m**2) + logprior(logscale)
    return logprob, field(logscale, m)

logprobs, fields = jax.vmap(max_logprob)(logscale_grid)
map_logscale = logscale_grid[jnp.argmax(logprobs)]
map_field = fields[jnp.argmax(logprobs)]

# Plot!
fig, axes = plt.subplots(nrows=2, figsize=(6, 6), dpi=300)
axes[0].plot(logscale_grid, prior / jnp.max(prior), c='C0', ls=':', label='Prior')
axes[0].plot(logscale_grid, posterior / jnp.max(posterior), c='C0', label='Posterior')
axes[0].axvline(map_logscale, color='C3', label='MAP')
axes[0].set(yticks=[], xlabel='log(scale)', ylabel='Probability (unnormalized)')
axes[0].legend()
axes[1].plot(x, y, c='k', alpha=0.8)
axes[1].plot(x, jnp.mean(samples, axis=0), c='C0')
axes[1].plot(x, samples.T, c='C0', alpha=0.3)
axes[1].plot(x, map_field, c='C3', zorder=0)
axes[1].errorbar(data_x, data_y, yerr=data_err, fmt='o', c='k', ms=4, zorder=-1)
axes[1].set(xlabel='$x$', ylabel='$f(x)$', yticks=[])
fig.tight_layout()
plt.show()

As before, the MAP overfits the noise with a rough function which is carefully tuned to match the data. On the other hand, posterior samples come from the smooth region of parameter space which fits less well but has a much larger volume of acceptable configurations. The difference between the two is so large because volume becomes more important as dimensionality increases. If one wants a principled way to learn both the function and its degree of smoothness, accounting for volume is mandatory.

There are a million more things to say here. Is this effect obvious or profound? Is the assumption of stationarity good or bad? What about the MAP in correlated coordinates? How does this relate to entropy and statistical mechanics? How can we implement this in practice for large, nonlinear models? What other types of problems require accounting for parameter volume? I won't answer these questions here, but fortunately there are many excellent resources that delve into these topics in more depth. Enjoy!

Further reading:

• A Conceptual Introduction to Hamiltonian Monte Carlo: a vivid discussion of volume, the "typical set," and why HMC works
• Information Field Theory: a wealth of resources on Bayesian inference of fields using Gaussian processes
• Metric Gaussian Variational Inference: a volume-aware algorithm which scales nearly-linearly to massive problems

Acknowledgement: This notebook was created by Benjamin Dodge and is available online here.