Example: Uncertainty Sampling

Hide code cell source

 
import matplotlib.pyplot as plt
%matplotlib inline
import matplotlib_inline
matplotlib_inline.backend_inline.set_matplotlib_formats('svg')
import seaborn as sns

# Uncomment the next two lines if running for the book
import warnings
warnings.filterwarnings("ignore")

from functools import partial
import jax
from jax import tree, vmap, jit, grad, value_and_grad
import jax.numpy as jnp
import jax.random as jr
from jax.flatten_util import ravel_pytree
from tinygp import GaussianProcess, kernels, transforms
import equinox as eqx
import optax
from jax.scipy.optimize import minimize
import scipy.stats.qmc as qmc
import numpy as np
from typing import NamedTuple

jax.config.update("jax_enable_x64", True)
np.random.seed(0)
key = jr.PRNGKey(np.random.randint(0, 10000))

Example: Uncertainty Sampling#

We will use uncertainty sampling, which maximizes the expected information gain under the Gaussian assumptions developed in the preceding section (MacKay, 1992).

1D example#

Here is the data-generating process. (Yes, it’s just a simple python function. But pretend it’s something much more expensive to evaluate, like a physical experiment or a complex simulation.)

def data_generating_process_1d(x):
    """Stand-in for an expensive simulator or physical experiment."""
    x = x[0]  # Unpack the 1D array
    return np.sin(10*x**2) + x + 0.1*np.random.normal(size=x.shape)

Start with just one data point:

X_1d = jr.uniform(key, (1, 1))
y_1d = data_generating_process_1d(X_1d)

It is not a good idea to train the model when we have so few data. So, we will simply fix the hyperparameters of the GP to something reasonable.

def build_gp(params, X):
    """Build a 1D Gaussian process with RBF kernel."""
    k = params['amplitude']*kernels.ExpSquared(params['lengthscale'])
    return GaussianProcess(k, X, diag=params['noise_stdev']**2, mean=params['mean'])

params = {
    'mean': 0.0,
    'amplitude': 4.0,
    'lengthscale': 0.15,
    'noise_stdev': 0.1
}

def eval_gp(Xq, X, y, params):
    gp = build_gp(params, X)
    _, cond_gp = gp.condition(y, Xq)
    return cond_gp

Xq_plt = jnp.linspace(0, 1, 1000)[:, None]
cond_gp = eval_gp(Xq_plt, X_1d, y_1d, params)

And let’s visualize the uncertainty in the model’s prediction:

def plot_posterior_uncertainty_1d(X, y, params, ax=None, legend=True):
    X_pred = jnp.linspace(0, 1, 100)[:, None]
    cond_gp = eval_gp(X_pred, X, y, params)
    mean = cond_gp.mean
    std = jnp.sqrt(cond_gp.variance)
    samples = cond_gp.sample(key, shape=(5,))
    if ax is None:
        fig, ax = plt.subplots(figsize=FIGURE_SIZES["half_standard"])
    ax.plot(X_pred, mean, label="Predictive mean" if legend else None)
    ax.fill_between(X_pred.squeeze(), mean - 2*std, mean + 2*std, alpha=0.2, label=r"$\pm 2\sigma$" if legend else None)
    ax.scatter(X, y, label="Observations" if legend else None)
    ax.plot(X_pred, samples.T, color='tab:red', lw=0.5, alpha=0.5)
    if legend:
        # ax.legend(loc='upper right')
        ax.legend(loc='best')
    ax.set_ylim(-2.0, 4.0)
    ax.set_xlabel(r"$x$")
    ax.set_ylabel(r"$y$")
    finalize_axes(keep_box=False)
    return ax

plot_posterior_uncertainty_1d(X_1d, y_1d, params);
One-dimensional Gaussian-process predictive mean, two-standard-deviation band, posterior samples, and the initial observation; uncertainty is smallest near the observed input.

Uncertainty is lower near the data point and high everywhere else. Let’s find the “best” next point to evaluate. Since our goal is to create the “best” model, we’ll select the point where there is maximum predictive uncertainty.

Hide code cell source

 
def get_next_x(X_obs, y_obs, params, n_start=1000):
    """Gets the next input x to evaluate. It does this by doing a random search to find the point of maximum predictive uncertainty."""
    Xs = qmc.LatinHypercube(X_obs.shape[1]).random(n=n_start)
    variances = eval_gp(Xs, X_obs, y_obs, params).variance.squeeze()
    next_x = Xs[jnp.argmax(variances)]
    return next_x

next_x = get_next_x(X_1d, y_1d, params)
ax = plot_posterior_uncertainty_1d(X_1d, y_1d, params, legend=False)
ax.axvline(next_x, color="red", linestyle="--", label=r"Next $x$")
ax.legend(loc='upper right');
The initial one-dimensional Gaussian-process posterior with a vertical line marking the next query at maximum predictive uncertainty near the boundary.

The “best” point is near the boundary. Let’s generate a data value at this point and plot a new fitted GP:

key, key_data, key_train = jr.split(key, 3)
next_y = data_generating_process_1d(next_x)
X_1d = jnp.concatenate([X_1d, next_x[None, :]], axis=0)
y_1d = jnp.concatenate([y_1d, next_y[None]], axis=0)

ax = plot_posterior_uncertainty_1d(X_1d, y_1d, params, legend=False)
The updated one-dimensional Gaussian-process posterior after adding the boundary observation, with uncertainty reduced around both observed inputs.

Let’s repeat this process a few more times:

Hide code cell source

 
for i in range(8):
    next_x = get_next_x(X_1d, y_1d, params)
    ax = plot_posterior_uncertainty_1d(X_1d, y_1d, params, legend=False)
    ax.axvline(next_x, color="red", linestyle="--", label=r"Next $x$")
    if i == 0:
        ax.legend(loc='upper right')
    key, key_data, key_train = jr.split(key, 3)
    next_y = data_generating_process_1d(next_x)
    X_1d = jnp.concatenate([X_1d, next_x[None, :]], axis=0)
    y_1d = jnp.concatenate([y_1d, next_y[None]], axis=0)
Eight successive one-dimensional active-learning posteriors, each marking the next maximum-uncertainty query as observations progressively cover the domain. Eight successive one-dimensional active-learning posteriors, each marking the next maximum-uncertainty query as observations progressively cover the domain. Eight successive one-dimensional active-learning posteriors, each marking the next maximum-uncertainty query as observations progressively cover the domain. Eight successive one-dimensional active-learning posteriors, each marking the next maximum-uncertainty query as observations progressively cover the domain. Eight successive one-dimensional active-learning posteriors, each marking the next maximum-uncertainty query as observations progressively cover the domain. Eight successive one-dimensional active-learning posteriors, each marking the next maximum-uncertainty query as observations progressively cover the domain. Eight successive one-dimensional active-learning posteriors, each marking the next maximum-uncertainty query as observations progressively cover the domain. Eight successive one-dimensional active-learning posteriors, each marking the next maximum-uncertainty query as observations progressively cover the domain.

In a realistic scenario, we would keep going until the epistemic uncertainty is low enough or we run out of computational budget.

2D example#

Let’s do the same thing in 2D. The data-generating process is:

@jit
def data_generating_process_2d(x):
    """Stand-in for an expensive simulator or physical experiment."""
    x1, x2 = x[0], x[1]
    f1 = 1/2*jnp.sin(5/2*x1 + 2/3*x2)**2 + 2/3*jnp.exp(-x1*(x2 - 0.5)**2)*jnp.cos(4*x1 + x2)**2
    f2 = x1*(1 - x1)*x2*(1 - x2)*10
    alpha = jax.nn.sigmoid(10*(x2 - x1))
    return alpha*f1 + (1 - alpha)*f2

Again, pretend data_generating_process_2d is actually some expensive process.

Let’s start with a few data points:

X_2d = qmc.LatinHypercube(d=2).random(n=5)
y_2d = vmap(data_generating_process_2d)(X_2d)

Build a GP:

def build_gp(params, X):
    """Build a 2D Gaussian process with RBF kernel."""
    amp = params['amplitude']
    ell = params['lengthscales']
    sigma = params['noise_stdev']
    k = amp*transforms.Linear(1/ell, kernels.ExpSquared())  # Must be constructed this way if ell is a vector
    return GaussianProcess(k, X, diag=sigma**2, mean=params['mean'])

# Fixed hyperparameters
params = {
    'mean': 0.0,
    'amplitude': 4.0,
    'lengthscales': jnp.array([0.3, 0.3]),  # Set these to different values for an anisotropic GP
    'noise_stdev': 0.01
}

def eval_gp(Xq, X, y, params):
    gp = build_gp(params, X)
    _, cond_gp = gp.condition(y, Xq)
    return cond_gp

X1, X2 = jnp.meshgrid(jnp.linspace(0, 1, 30), jnp.linspace(0, 1, 30))
Xq_plt = jnp.stack([X1.ravel(), X2.ravel()], axis=1)
cond_gp = eval_gp(Xq_plt, X_2d, y_2d, params)

Find the best next point at which to collect data:

next_x = get_next_x(X_2d, y_2d, params)

And plot:

Hide code cell source

 
def plot_posterior_uncertainty_2d(X, y, params, unc_clim=None):
    X1_plt, X2_plt = jnp.meshgrid(jnp.linspace(0, 1, 50), jnp.linspace(0, 1, 50))
    X_plt = jnp.stack([X1_plt.ravel(), X2_plt.ravel()], axis=1)
    
    fig, ax = plt.subplots(1, 3, figsize=FIGURE_SIZES["half_tall"])

    # Ground truth
    Y_gt_plt = vmap(data_generating_process_2d)(X_plt).reshape(X1_plt.shape)
    p = ax[0].pcolormesh(X1_plt, X2_plt, Y_gt_plt, cmap="viridis")
    p.set_edgecolor('face')
    ax[0].set_aspect("equal")

    # Evaluate GP on a grid
    cond_gp = eval_gp(X_plt, X=X, y=y, params=params)
    Y_mean_plt = cond_gp.mean.reshape(X1_plt.shape)
    Y_variance_plt = cond_gp.variance.reshape(X1_plt.shape)

    # Plot the mean
    p = ax[1].pcolormesh(X1_plt, X2_plt, Y_mean_plt, cmap="viridis", vmin=p.get_clim()[0], vmax=p.get_clim()[1])
    p.set_edgecolor('face')
    ax[1].scatter(X[:,0], X[:,1], color="C1", label="Observation")
    ax[1].set_aspect("equal")

    # Plot the variance
    p = ax[2].pcolormesh(X1_plt, X2_plt, Y_variance_plt, cmap="viridis", vmin=unc_clim[0] if unc_clim is not None else None, vmax=unc_clim[1] if unc_clim is not None else None)
    p.set_edgecolor('face')
    ax[2].scatter(X[:,0], X[:,1], color="C1", label="Observation")
    ax[2].set_aspect("equal")
    ax[2].set_xticks([0, 0.5, 1])
    ax[2].set_yticks([0, 0.5, 1])

    finalize_axes(keep_box=True)
    return ax, p.get_clim()

ax, unc_clim = plot_posterior_uncertainty_2d(X_2d, y_2d, params)
ax[1].plot(next_x[0], next_x[1], "rx", ms=6, mew=2, label="Next $x$")
ax[2].plot(next_x[0], next_x[1], "rx", ms=6, mew=2, label="Next $x$")
ax[1].legend(loc='upper right');
Three maps for the initial two-dimensional active-learning step: the true response, Gaussian-process mean with observations, and predictive variance with the next maximum-uncertainty query marked.

The highest uncertainty is at the edge of the domain, so that is where the next best point is.

Let’s do a few more iterations of active learning:

plot_iter = [0, 1, 2, 3, 4, 5, 10, 20]
for i in range(plot_iter[-1] + 1):
    next_x = get_next_x(X_2d, y_2d, params)
    
    if i in plot_iter:
        ax, _ = plot_posterior_uncertainty_2d(X_2d, y_2d, params, unc_clim=unc_clim)
        ax[1].plot(next_x[0], next_x[1], "rx", ms=6, mew=2, label="Next $x$")
        ax[2].plot(next_x[0], next_x[1], "rx", ms=6, mew=2, label="Next $x$")
        if i == 0:
            ax[1].legend(loc='upper right')

    key, key_data, key_train = jr.split(key, 3)
    next_y = data_generating_process_2d(next_x)
    X_2d = jnp.concatenate([X_2d, next_x[None, :]], axis=0)
    y_2d = jnp.concatenate([y_2d, next_y[None]], axis=0)
Eight stages of two-dimensional active learning, each showing the true response, Gaussian-process mean, predictive variance, accumulated observations, and next query. Eight stages of two-dimensional active learning, each showing the true response, Gaussian-process mean, predictive variance, accumulated observations, and next query. Eight stages of two-dimensional active learning, each showing the true response, Gaussian-process mean, predictive variance, accumulated observations, and next query. Eight stages of two-dimensional active learning, each showing the true response, Gaussian-process mean, predictive variance, accumulated observations, and next query. Eight stages of two-dimensional active learning, each showing the true response, Gaussian-process mean, predictive variance, accumulated observations, and next query. Eight stages of two-dimensional active learning, each showing the true response, Gaussian-process mean, predictive variance, accumulated observations, and next query. Eight stages of two-dimensional active learning, each showing the true response, Gaussian-process mean, predictive variance, accumulated observations, and next query. Eight stages of two-dimensional active learning, each showing the true response, Gaussian-process mean, predictive variance, accumulated observations, and next query.

We have created a model in an optimally efficient manner (in terms of reducing epistemic uncertainty).