Variational Families

AbstractVariationalFamily

AbstractVariationalFamily(
    posterior: AbstractPosterior[P, L],
)

Bases: Module, Generic[L]

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

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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: Any, **kwargs: Any) -> GaussianDistribution

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: AbstractPosterior[P, L],
    inducing_inputs: Union[
        Int[Array, "N D"], Float[Array, "N D"], Real
    ],
    jitter: ScalarFloat = 1e-06,
)

Bases: AbstractVariationalFamily[L]

The variational Gaussian family of probability distributions.

num_inducing property

num_inducing: int

The number of inducing inputs.

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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: Any, **kwargs: Any) -> GaussianDistribution

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: AbstractPosterior[P, L],
    inducing_inputs: Union[
        Int[Array, "N D"], Float[Array, "N D"]
    ],
    variational_mean: Union[
        Float[Array, "N 1"], None
    ] = None,
    variational_root_covariance: Union[
        Float[Array, "N N"], None
    ] = None,
    jitter: ScalarFloat = 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: int

The number of inducing inputs.

prior_kl

prior_kl() -> ScalarFloat

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: Union[
        Int[Array, "N D"], Float[Array, "N D"]
    ],
) -> GaussianDistribution

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.

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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.

GraphVariationalGaussian

GraphVariationalGaussian(
    posterior: AbstractPosterior[P, L],
    inducing_inputs: Int[Array, "N D"],
    variational_mean: Union[
        Float[Array, "N 1"], None
    ] = None,
    variational_root_covariance: Union[
        Float[Array, "N N"], None
    ] = None,
    jitter: ScalarFloat = 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: int

The number of inducing inputs.

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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: Union[
        Int[Array, "N D"], Float[Array, "N D"]
    ],
) -> GaussianDistribution

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() -> ScalarFloat

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: AbstractPosterior[P, L],
    inducing_inputs: Union[
        Int[Array, "N D"], Float[Array, "N D"]
    ],
    variational_mean: Union[
        Float[Array, "N 1"], None
    ] = None,
    variational_root_covariance: Union[
        Float[Array, "N N"], None
    ] = None,
    jitter: ScalarFloat = 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: int

The number of inducing inputs.

prior_kl

prior_kl() -> ScalarFloat

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: Float[Array, "N D"],
) -> GaussianDistribution

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.

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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.

NaturalVariationalGaussian

NaturalVariationalGaussian(
    posterior: AbstractPosterior[P, L],
    inducing_inputs: Float[Array, "N D"],
    natural_vector: Union[Float[Array, "M 1"], None] = None,
    natural_matrix: Union[Float[Array, "M M"], None] = None,
    jitter: ScalarFloat = 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: int

The number of inducing inputs.

prior_kl

prior_kl() -> ScalarFloat

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: Float[Array, "N D"],
) -> GaussianDistribution

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.

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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.

ExpectationVariationalGaussian

ExpectationVariationalGaussian(
    posterior: AbstractPosterior[P, L],
    inducing_inputs: Float[Array, "N D"],
    expectation_vector: Union[
        Float[Array, "M 1"], None
    ] = None,
    expectation_matrix: Union[
        Float[Array, "M M"], None
    ] = None,
    jitter: ScalarFloat = 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: int

The number of inducing inputs.

prior_kl

prior_kl() -> ScalarFloat

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: Float[Array, "N D"],
) -> GaussianDistribution

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

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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.

CollapsedVariationalGaussian

CollapsedVariationalGaussian(
    posterior: AbstractPosterior[P, GL],
    inducing_inputs: Float[Array, "N D"],
    jitter: ScalarFloat = 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: int

The number of inducing inputs.

predict

predict(
    test_inputs: Float[Array, "N D"], train_data: Dataset
) -> GaussianDistribution

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

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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.

VariationalGaussianInit dataclass

VariationalGaussianInit(
    inducing_inputs: Union[
        Int[Array, "N D"], Float[Array, "N D"]
    ],
    variational_mean: Union[
        Float[Array, "N 1"], None
    ] = None,
    variational_root_covariance: Union[
        Float[Array, "N N"], None
    ] = None,
)

Initialization parameters for a variational Gaussian distribution.

HeteroscedasticPrediction

Bases: NamedTuple

Mean and variance of the signal and noise latent processes.

HeteroscedasticVariationalFamily

HeteroscedasticVariationalFamily(
    posterior: HP,
    inducing_inputs: Union[
        Int[Array, "N D"], Float[Array, "N D"]
    ] = None,
    inducing_inputs_g: Union[
        Int[Array, "M D"], Float[Array, "M D"], None
    ] = None,
    variational_mean_f: Union[
        Float[Array, "N 1"], None
    ] = None,
    variational_root_covariance_f: Union[
        Float[Array, "N N"], None
    ] = None,
    variational_mean_g: Union[
        Float[Array, "M 1"], None
    ] = None,
    variational_root_covariance_g: Union[
        Float[Array, "M M"], None
    ] = None,
    jitter: ScalarFloat = 1e-06,
    signal_init: Optional[VariationalGaussianInit] = None,
    noise_init: Optional[VariationalGaussianInit] = None,
)

Bases: AbstractVariationalFamily[HL]

Variational family for two independent latent processes f and g.

__call__

__call__(*args: Any, **kwargs: Any) -> GaussianDistribution

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.