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.

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\).