WhitenedVariationalGaussian#

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

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

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)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:

The predictive distribution of the low-rank GP at

the test inputs.

Return type:

GaussianDistribution

prior_kl()[source]#

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}[N(\mu , S)\mid\mid N(0, I)]. \end{aligned}\end{split}\]

Against a standard normal prior the divergence has a closed form that needs no matrix factorisation at all. Writing \(S = LL^{\top}\) for the stored triangular root \(L\), and using \(\operatorname{tr}[S] = \lVert L\rVert_F^2\) and \(\log\lvert S\rvert = 2\sum_i \log\lvert L_{ii}\rvert\),

\[\operatorname{KL}[\mathcal{N}(\mu, S)\mid\mid\mathcal{N}(0, I)] = \tfrac{1}{2}\left( \lVert\mu\rVert^2 + \lVert L\rVert_F^2 - m - 2\sum_i \log\lvert L_{ii}\rvert \right),\]

where \(m\) is the number of inducing points.

Returns:

The KL-divergence between our variational

approximation and the GP prior.

Return type:

ScalarFloat