Variational Families

AbstractVariationalFamily

AbstractVariationalFamily(posterior)

Bases: Module, Generic[L]

Abstract base class used to represent families of distributions that can be used within variational inference.

__call__

__call__(*args, **kwargs)

Evaluate the variational family's density.

For a given set of parameters, compute the latent function's prediction under the variational approximation.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

predict abstractmethod

predict(*args, **kwargs)

Predict the GP's output given the input.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

AbstractVariationalGaussian

AbstractVariationalGaussian(
    posterior, inducing_inputs, jitter=1e-06
)

Bases: AbstractVariationalFamily[L]

The variational Gaussian family of probability distributions.

num_inducing property

num_inducing

The number of inducing inputs.

__call__

__call__(*args, **kwargs)

Evaluate the variational family's density.

For a given set of parameters, compute the latent function's prediction under the variational approximation.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

predict abstractmethod

predict(*args, **kwargs)

Predict the GP's output given the input.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

VariationalGaussian

VariationalGaussian(
    posterior,
    inducing_inputs,
    variational_mean=None,
    variational_root_covariance=None,
    jitter=1e-06,
)

Bases: AbstractVariationalGaussian[L]

The variational Gaussian family of probability distributions.

The variational family is $q(f(\cdot)) = \int p(f(\cdot)\mid u) q(u) \mathrm{d}u$, where $u = f(z)$ are the function values at the inducing inputs $z$ and the distribution over the inducing inputs is $q(u) = \mathcal{N}(\mu, S)$. We parameterise this over $\mu$ and $sqrt$ with $S = sqrt sqrt^{\top}$.

num_inducing property

num_inducing

The number of inducing inputs.

__call__

__call__(*args, **kwargs)

Evaluate the variational family's density.

For a given set of parameters, compute the latent function's prediction under the variational approximation.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

prior_kl

prior_kl()

Compute the prior KL divergence.

Compute the KL-divergence between our variational approximation and the Gaussian process prior.

For this variational family, we have

$$\begin{align} \operatorname{KL}[q(f(\cdot))\mid\mid p(\cdot)] & = \operatorname{KL}[q(u)\mid\mid p(u)]\\ & = \operatorname{KL}[ \mathcal{N}(\mu, S) \mid\mid N(\mu z, \mathbf{K}_{zz}) ], \end{align}$$

where $u = f(z)$ and $z$ are the inducing inputs.

Returns:

• ScalarFloat ( ScalarFloat ) –

  The KL-divergence between our variational approximation and the GP prior.

predict

predict(test_inputs)

Compute the predictive distribution of the GP at the test inputs t.

This is the integral $q(f(t)) = \int p(f(t)\mid u) q(u) \mathrm{d}u$, which can be computed in closed form as:

$$\mathcal{N}\left(f(t); \mu t + \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} (\mu - \mu z), \mathbf{K}_{tt} - \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} \mathbf{K}_{zt} + \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} S \mathbf{K}_{zz}^{-1} \mathbf{K}_{zt}\right).$$

Parameters:

• test_inputs (Float[Array, 'N D']) –

  The test inputs at which we wish to make a prediction.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The predictive distribution of the low-rank GP at the test inputs.

GraphVariationalGaussian

GraphVariationalGaussian(
    posterior,
    inducing_inputs,
    variational_mean=None,
    variational_root_covariance=None,
    jitter=1e-06,
)

Bases: VariationalGaussian[L]

A variational Gaussian defined over graph-structured inducing inputs.

This subclass adapts the :class:VariationalGaussian family to the case where the inducing inputs are discrete graph node indices rather than continuous spatial coordinates.

The main differences are

• Inducing inputs are integer node IDs.
• Kernel matrices are ensured to be dense and 2D.

num_inducing property

num_inducing

The number of inducing inputs.

__call__

__call__(*args, **kwargs)

Evaluate the variational family's density.

For a given set of parameters, compute the latent function's prediction under the variational approximation.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

predict

predict(test_inputs)

Compute the predictive distribution of the GP at the test inputs t.

This is the integral $q(f(t)) = \int p(f(t)\mid u) q(u) \mathrm{d}u$, which can be computed in closed form as:

$$\mathcal{N}\left(f(t); \mu t + \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} (\mu - \mu z), \mathbf{K}_{tt} - \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} \mathbf{K}_{zt} + \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} S \mathbf{K}_{zz}^{-1} \mathbf{K}_{zt}\right).$$

Parameters:

• test_inputs (Float[Array, 'N D']) –

  The test inputs at which we wish to make a prediction.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The predictive distribution of the low-rank GP at the test inputs.

prior_kl

prior_kl()

Compute the prior KL divergence.

Compute the KL-divergence between our variational approximation and the Gaussian process prior.

For this variational family, we have

$$\begin{align} \operatorname{KL}[q(f(\cdot))\mid\mid p(\cdot)] & = \operatorname{KL}[q(u)\mid\mid p(u)]\\ & = \operatorname{KL}[ \mathcal{N}(\mu, S) \mid\mid N(\mu z, \mathbf{K}_{zz}) ], \end{align}$$

where $u = f(z)$ and $z$ are the inducing inputs.

Returns:

• ScalarFloat ( ScalarFloat ) –

  The KL-divergence between our variational approximation and the GP prior.

WhitenedVariationalGaussian

WhitenedVariationalGaussian(
    posterior,
    inducing_inputs,
    variational_mean=None,
    variational_root_covariance=None,
    jitter=1e-06,
)

Bases: VariationalGaussian[L]

The whitened variational Gaussian family of probability distributions.

The variational family is $q(f(\cdot)) = \int p(f(\cdot)\mid u) q(u) \mathrm{d}u$, where $u = f(z)$ are the function values at the inducing inputs $z$ and the distribution over the inducing inputs is $q(u) = \mathcal{N}(Lz \mu + mz, Lz S Lz^{\top})$. We parameterise this over $\mu$ and $sqrt$ with $S = sqrt sqrt^{\top}$.

num_inducing property

num_inducing

The number of inducing inputs.

__call__

__call__(*args, **kwargs)

Evaluate the variational family's density.

For a given set of parameters, compute the latent function's prediction under the variational approximation.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

prior_kl

prior_kl()

Compute the KL-divergence between our variational approximation and the Gaussian process prior.

For this variational family, we have

$$\begin{align} \operatorname{KL}[q(f(\cdot))\mid\mid p(\cdot)] & = \operatorname{KL}[q(u)\mid\mid p(u)]\\ & = \operatorname{KL}[N(\mu , S)\mid\mid N(0, I)]. \end{align}$$

Returns:

• ScalarFloat ( ScalarFloat ) –

  The KL-divergence between our variational approximation and the GP prior.

predict

predict(test_inputs)

Compute the predictive distribution of the GP at the test inputs t.

This is the integral q(f(t)) = \int p(f(t)\midu) q(u) du, which can be computed in closed form as

$$\mathcal{N}\left(f(t); \mu t + \mathbf{K}_{tz} \mathbf{L}z^{\top} \mu , \mathbf{K}_{tt} - \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} \mathbf{K}_{zt} + \mathbf{K}_{tz} \mathbf{L}z^{\top} S \mathbf{L}z^{-1} \mathbf{K}_{zt} \right).$$

Parameters:

• test_inputs (Float[Array, 'N D']) –

  The test inputs at which we wish to make a prediction.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The predictive distribution of the low-rank GP at the test inputs.

NaturalVariationalGaussian

NaturalVariationalGaussian(
    posterior,
    inducing_inputs,
    natural_vector=None,
    natural_matrix=None,
    jitter=1e-06,
)

Bases: AbstractVariationalGaussian[L]

The natural variational Gaussian family of probability distributions.

The variational family is $q(f(\cdot)) = \int p(f(\cdot)\mid u) q(u) \mathrm{d}u$, where $u = f(z)$ are the function values at the inducing inputs $z$ and the distribution over the inducing inputs is $q(u) = N(\mu, S)$. Expressing the variational distribution, in the form of the exponential family, $q(u) = exp(\theta^{\top} T(u) - a(\theta))$, gives rise to the natural parameterisation $\theta = (\theta_{1}, \theta_{2}) = (S^{-1}\mu, -S^{-1}/2)$, to perform model inference, where $T(u) = [u, uu^{\top}]$ are the sufficient statistics.

num_inducing property

num_inducing

The number of inducing inputs.

__call__

__call__(*args, **kwargs)

Evaluate the variational family's density.

For a given set of parameters, compute the latent function's prediction under the variational approximation.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

prior_kl

prior_kl()

Compute the KL-divergence between our current variational approximation and the Gaussian process prior.

For this variational family, we have

$$\begin{align} \operatorname{KL}[q(f(\cdot))\mid\mid p(\cdot)] & = \operatorname{KL}[q(u)\mid\mid p(u)] \\ & = \operatorname{KL}[N(\mu, S)\mid\mid N(mz, \mathbf{K}_{zz})], \end{align}$$

with $\mu$ and $S$ computed from the natural parameterisation $\theta = (S^{-1}\mu , -S^{-1}/2)$.

Returns:

• ScalarFloat ( ScalarFloat ) –

  The KL-divergence between our variational approximation and the GP prior.

predict

predict(test_inputs)

Compute the predictive distribution of the GP at the test inputs $t$.

This is the integral $q(f(t)) = \int p(f(t)\mid u) q(u) \mathrm{d}u$, which can be computed in closed form as

$$\mathcal{N}\left(f(t); \mu t + \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} (\mu - \mu z), \mathbf{K}_{tt} - \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} \mathbf{K}_{zt} + \mathbf{K}_{tz} \mathbf{K}_{zz}^{-1} S \mathbf{K}_{zz}^{-1} \mathbf{K}_{zt} \right),$$

with $\mu$ and $S$ computed from the natural parameterisation $\theta = (S^{-1}\mu , -S^{-1}/2)$.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  A function that accepts a set of test points and will return the predictive distribution at those points.

ExpectationVariationalGaussian

ExpectationVariationalGaussian(
    posterior,
    inducing_inputs,
    expectation_vector=None,
    expectation_matrix=None,
    jitter=1e-06,
)

Bases: AbstractVariationalGaussian[L]

The natural variational Gaussian family of probability distributions.

The variational family is $q(f(\cdot)) = \int p(f(\cdot)\mid u) q(u) \mathrm{d}u$, where $u = f(z)$ are the function values at the inducing inputs $z$ and the distribution over the inducing inputs is $q(u) = \mathcal{N}(\mu, S)$. Expressing the variational distribution, in the form of the exponential family, $q(u) = exp(\theta^{\top} T(u) - a(\theta))$, gives rise to the natural parameterisation $\theta = (\theta_{1}, \theta_{2}) = (S^{-1}\mu , -S^{-1}/2)$ and sufficient statistics $T(u) = [u, uu^{\top}]$. The expectation parameters are given by $\nu = \int T(u) q(u) \mathrm{d}u$. This gives a parameterisation, $\nu = (\nu_{1}, \nu_{2}) = (\mu , S + uu^{\top})$ to perform model inference over.

num_inducing property

num_inducing

The number of inducing inputs.

__call__

__call__(*args, **kwargs)

Evaluate the variational family's density.

For a given set of parameters, compute the latent function's prediction under the variational approximation.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

prior_kl

prior_kl()

Evaluate the prior KL-divergence.

Compute the KL-divergence between our current variational approximation and the Gaussian process prior.

For this variational family, we have

$$\begin{align} \operatorname{KL}(q(f(\cdot))\mid\mid p(\cdot)) & = \operatorname{KL}(q(u)\mid\mid p(u)) \\ & =\operatorname{KL}(\mathcal{N}(\mu, S)\mid\mid \mathcal{N}(m_z, K_{zz})), \end{align}$$

where $\mu$ and $S$ are the expectation parameters of the variational distribution and $m_z$ and $K_{zz}$ are the mean and covariance of the prior distribution.

Returns:

• ScalarFloat ( ScalarFloat ) –

  The KL-divergence between our variational approximation and the GP prior.

predict

predict(test_inputs)

Evaluate the predictive distribution.

Compute the predictive distribution of the GP at the test inputs $t$.

This is the integral $q(f(t)) = \int p(f(t)\mid u)q(u)\mathrm{d}u$, which can be computed in closed form as which can be computed in closed form as

$$\mathcal{N}(f(t); \mu_t + \mathbf{K}_{tz}\mathbf{K}_{zz}^{-1}(\mu - \mu_z), \mathbf{K}_{tt} - \mathbf{K}_{tz}\mathbf{K}_{zz}^{-1}\mathbf{K}_{zt} + \mathbf{K}_{tz}\mathbf{K}_{zz}^{-1}\mathbf{S} \mathbf{K}_{zz}^{-1}\mathbf{K}_{zt})$$

with $\mu$ and $S$ computed from the expectation parameterisation $\eta = (\mu, S + uu^\top)$.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The predictive distribution of the GP at the test inputs $t$.

CollapsedVariationalGaussian

CollapsedVariationalGaussian(
    posterior, inducing_inputs, jitter=1e-06
)

Bases: AbstractVariationalGaussian[GL]

Collapsed variational Gaussian.

Collapsed variational Gaussian family of probability distributions. The key reference is Titsias, (2009) - Variational Learning of Inducing Variables in Sparse Gaussian Processes.

num_inducing property

num_inducing

The number of inducing inputs.

__call__

__call__(*args, **kwargs)

Evaluate the variational family's density.

For a given set of parameters, compute the latent function's prediction under the variational approximation.

Parameters:

• *args (Any, default: () ) –

  Arguments of the variational family's predict method.

• **kwargs (Any, default: {} ) –

  Keyword arguments of the variational family's predict method.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The output of the variational family's predict method.

predict

predict(test_inputs, train_data)

Compute the predictive distribution of the GP at the test inputs.

Parameters:

• test_inputs (Float[Array, 'N D']) –

  The test inputs $t$ at which to make predictions.

• train_data (Dataset) –

  The training data that was used to fit the GP.

Returns:

• GaussianDistribution ( GaussianDistribution ) –

  The predictive distribution of the collapsed variational Gaussian process at the test inputs $t$.