fit_natgrads#

gpjax.fit.fit_natgrads(*, model, objective, train_data, optim, natgrad_lr=1e-1, key=jr.key(42), num_iters=100, batch_size=-1, map_jitter=0.0, backoff=0.5, max_backoff=5, beta_floor=1e-8, log_rate=10, verbose=True, unroll=1, safe=True)[source]#

Train a variational family by alternating natural-gradient and Optax steps.

Implements the NGD+Adam scheme of Salimbeni, Eleftheriadis and Hensman (2018), arXiv:1803.09151. Each iteration takes one natural-gradient step on the exponential-family coordinates of the variational distribution, then one step of the supplied Optax optimiser on everything else – kernel and likelihood hyperparameters, the mean function, and the inducing inputs, which the paper counts as hyperparameters.

The natural gradient with respect to the natural parameters \(\boldsymbol\theta\) is the ordinary gradient with respect to the expectation parameters \(\boldsymbol\eta\), so the update \(\boldsymbol\theta\leftarrow\boldsymbol\theta -\gamma\,\partial\ell/\partial\boldsymbol\eta\) needs no Fisher matrix. For a conjugate (Gaussian-likelihood) model on the full batch, natgrad_lr=1.0 reaches the exact optimal \(q\) in a single iteration.

Example

>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> import optax as ox
>>> import gpjax as gpx
>>>
>>> xtrain = jnp.linspace(0, 1, 20).reshape(-1, 1)
>>> ytrain = jnp.sin(xtrain)
>>> D = gpx.Dataset(X=xtrain, y=ytrain)
>>>
>>> meanf = gpx.mean_functions.Constant()
>>> kernel = gpx.kernels.RBF()
>>> likelihood = gpx.likelihoods.Gaussian()
>>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
>>> posterior = prior * likelihood
>>>
>>> z = jnp.linspace(0, 1, 5).reshape(-1, 1)
>>> q = gpx.variational_families.VariationalGaussian(
...     model=posterior, inducing_inputs=z
... )
>>>
>>> negative_elbo = lambda p, d: -gpx.objectives.elbo(p, d)
>>> trained_model, history = gpx.fit_natgrads(
...     model=q, objective=negative_elbo, train_data=D,
...     optim=ox.adam(0.01), natgrad_lr=1.0, num_iters=10, verbose=False,
... )
Parameters:

  • model (Model) – The variational family to be optimised.

  • objective (Objective) – The loss to minimise, e.g. lambda q, d: -gpjax.objectives.elbo(q, d).

  • train_data (Dataset) – The training data used to evaluate the objective.

  • optim (GradientTransformation) – The Optax optimiser applied to the hyperparameter partition.

  • natgrad_lr (float | int | jax.Array | optax.Schedule) – The natural-gradient step size \(\gamma\in(0,1]\), or an Optax schedule mapping the iteration number to a step size. Defaults to 1e-1, the value Salimbeni et al. recommend in the stochastic, non-conjugate regime; natgrad_lr=1.0 is optimal only when the model is conjugate and the batch is full. Adam, Chang, Khan and Solin (2021) write this step size \(\rho\) for the dual parameterisation; it is the same quantity, and started from the same \(q\) the two branches produce identical iterates – provided the dual branch’s computed \(\boldsymbol\beta\) stays non-negative, so that its beta_floor never engages. GPJax’s clipped probit link breaks that in the far tails. On a DualVariationalGaussian a value above \(1\) is rejected, because the site update is a convex combination towards its target – for a schedule this is checked over the whole num_iters-long trajectory, not just at construction.

  • key (KeyArray) – The random key used for mini-batch selection. Defaults to jr.key(42).

  • num_iters (int) – The number of alternating iterations to run. Defaults to 100.

  • batch_size (int) – The size of the mini-batch to use. Defaults to -1 (i.e. full batch). The same batch feeds both sub-steps of an iteration.

  • map_jitter (float) – Jitter added inside the \(\boldsymbol\theta\leftrightarrow\boldsymbol\xi\) maps. Defaults to 0.0 and is deliberately not inherited from the model’s Prior.jitter: a non-zero value biases the recovered covariance by \(\approx\varepsilon\lVert\mathbf S\rVert^2\) regardless of conditioning, which destroys the exactness of the conjugate one-step solution. Raise it to \(10^{-12}\)–\(10^{-10}\) only when fighting an ill-conditioned \(\mathbf S\), and note that a non-zero value also shifts every entry of history by \(\mathcal O(\varepsilon)\), because the logged loss is read off the differentiated \(\boldsymbol\eta\) closure.

  • backoff (float) – Multiplicative shrink factor applied to \(\gamma\) when a step would leave the negative-definite cone. Defaults to 0.5.

  • max_backoff (int) – The number of shrink attempts after the first, so \(\gamma\) can fall by \(\beta^{K}\). Defaults to 5.

  • beta_floor (float) – Lower clip on the expected negative curvature \(\beta\) in the dual step, which keeps \(\boldsymbol\Lambda_2\) inside the positive semi-definite cone for likelihoods that are not log-concave. The Salimbeni-family step ignores it. Defaults to 1e-8.

  • log_rate (int) – How frequently the objective value should be printed. Defaults to 10.

  • verbose (bool) – Whether to display the training progress bar. Defaults to True.

  • unroll (int) – The number of unrolled steps to use for the optimisation. Defaults to 1.

  • safe (bool) – Whether to validate inputs before optimisation. Defaults to True.

Returns:

A tuple of the optimised model and a 1-D history of

length num_iters.

Return type:

tuple[Model, jax.Array]

Notes

Step ordering. Within one iteration the natural-gradient step runs first and the Optax step second, on the already-updated \(q\). Salimbeni et al. describe the reverse order and explicitly allow either; natgrad-first is chosen here because the forward pass that produces \(\partial\ell/\partial\boldsymbol\eta\) also yields \(\ell(\boldsymbol\xi_t,\boldsymbol\phi_t)\) for free, which is exactly fit()’s history[t] convention, and because it decouples a bad hyperparameter step from the Cholesky factorisations of the natural-gradient step by one iteration. The ordering changes traces bit-for-bit, so do not reverse it casually.

Choice of family. The step differentiates the loss through \(\boldsymbol\xi(\boldsymbol\eta)\), which subtracts \(\boldsymbol\eta_1\boldsymbol\eta_1^\top\) from \(\mathbf H_2\). When \(\lVert\mathbf m\rVert^2\gg\lVert\mathbf S\rVert\) that cancellation loses digits quietly – finite, unguarded and increasingly wrong – so prefer WhitenedVariationalGaussian, whose \(q(\mathbf v)\) stays close to \(\mathcal N(\mathbf 0,\mathbf I)\), in that regime. DualVariationalGaussian is immune to this particular cancellation for a different reason: its step is affine in the stored sites and takes no \(\boldsymbol\xi(\boldsymbol\eta)\) round trip at all, so \(\boldsymbol\eta_1\boldsymbol\eta_1^\top\) is never formed. It buys that with a second \(M\times M\) factorisation per objective evaluation and a step size capped at \(1\).

Expand for references to gpjax.fit.fit_natgrads

fit_natgrads