NaturalVariationalGaussian#

class gpjax.variational_families.NaturalVariationalGaussian(posterior, inducing_inputs, natural_vector=None, natural_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) = 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.

Parameters:
predict(test_inputs)[source]#

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:

A function that accepts a set of test points and will

return the predictive distribution at those points.

Return type:

GaussianDistribution

Parameters:

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

prior_kl()[source]#

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}[N(\mu, S)\mid\mid N(mz, \mathbf{K}_{zz})], \end{aligned}\end{split}\]

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

Returns:

The KL-divergence between our variational approximation and

the GP prior.

Return type:

ScalarFloat