VariationalGaussian#

class gpjax.variational_families.VariationalGaussian(posterior, inducing_inputs, variational_mean=None, variational_root_covariance=None, jitter=1e-06)[source]#

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

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).\]
Parameters:

test_inputs (Float[Array, "N D"]) – The test inputs at which we wish to make a prediction.

Returns:

The predictive distribution of the low-rank GP at

the test inputs.

Return type:

GaussianDistribution

prior_kl()[source]#

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

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

With \(S = LL^{\top}\) for the stored triangular root \(L\) and \(\mathbf{K}_{zz} = L_z L_z^{\top}\), this evaluates in closed form as

\[\tfrac{1}{2}\left( \lVert L_z^{-1}(\mu_z - \mu)\rVert^2 + \lVert L_z^{-1} L\rVert_F^2 - m + 2\sum_i \log [L_z]_{ii} - 2\sum_i \log \lvert L_{ii}\rvert \right),\]

so the Cholesky factor of \(\mathbf{K}_{zz}\) is the only factorisation required; \(S\) is never formed and never re-factorised.

Returns:

The KL-divergence between our variational

approximation and the GP prior.

Return type:

ScalarFloat