ExpectationVariationalGaussian#

class gpjax.variational_families.ExpectationVariationalGaussian(posterior, inducing_inputs, expectation_vector=None, expectation_matrix=None, jitter=1e-06)[source]#

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.

Parameters:
predict(test_inputs)[source]#

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:

The predictive distribution of the GP at the

test inputs \(t\).

Return type:

GaussianDistribution

Parameters:

test_inputs (Float[jaxlib._jax.Array, 'N D'] | Float[ndarray, 'N D'])

prior_kl()[source]#

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{split}\begin{aligned} \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{aligned}\end{split}\]

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:

The KL-divergence between our variational approximation and

the GP prior.

Return type:

ScalarFloat