diff --git a/CHANGELOG.md b/CHANGELOG.md index be4d59b7c..a85086a53 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,8 +8,122 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- **Dense joint predictive covariance for state-space GPs.** + `StateSpacePrior.predict`, `StateSpaceConjugateModel.predict`, and + `StateSpacePosterior.__call__` (`gpjax.state_space`) now accept + `covariance="dense"`, returning the full joint covariance across test + points rather than marginal variances only + ([#651](https://github.com/thomaspinder/GPJax/issues/651)). For the + unconditioned prior this is just the kernel's own dense gram — the + state-space SDE is an exact representation with no training data to + marginalise out. For the conditioned (smoothed) posterior it is built from + the RTS smoother's cross-covariance recursion (Särkkä & Solin 2019 §12.2): + `rts_smoother` gains an opt-in `return_gains=True` that exposes its + already-computed per-step smoother gains, and a new + `gpjax.state_space.prediction._dense_smoothed_test_covariance` chains them + into the `M x M` test-point covariance. This keeps the state-space + formulation's linear-in-`N` cost — no `N x N` gram over the training set is + ever formed — with the cross-covariance work landing at `O(M^2 d^3)`, + independent of `N` and no larger than the `O(M^2)` already required to + store the dense output. `StateSpacePosterior.filtered` / + `StateSpaceConjugateModel.predict_filter` (the *causal* predictive) keep + raising `NotImplementedError` for `covariance="dense"`: each of their test + points conditions on a different information set (observations up to its + own timestamp), so a "joint" filtered covariance is not the dense + conjugate-predictive-shaped object the smoothed path now matches, and + extending it is left to a future issue if there is demand for it. + +- **`gpjax.fit_natgrads` and the `gpjax.natural_gradients` module.** Trains a variational + family by alternating one natural-gradient step on the variational distribution with + one step of a supplied Optax optimiser on everything else — kernel and likelihood + hyperparameters, the mean function, and the inducing inputs — following Salimbeni, + Eleftheriadis and Hensman (2018), [arXiv:1803.09151](https://arxiv.org/abs/1803.09151). + `VariationalGaussian` and `WhitenedVariationalGaussian` are supported. For a conjugate + model on the full batch, `natgrad_lr=1.0` reaches the optimal `q` in one iteration. + +- **`DualVariationalGaussian` and `gpjax.objectives.dual_elbo`.** The dual (t-SVGP) + parameterisation of Adam, Chang, Khan and Solin (2021), + [arXiv:2111.03412](https://arxiv.org/abs/2111.03412). Instead of the moments of + `q(u)`, the family stores an unnormalised Gaussian *site* on the centred inducing + outputs — `dual_vector` is the site's first natural parameter and `dual_matrix` its + precision — from which `q(u)` is recovered through the working matrix + `R = Kzz + Kzz Lambda_2 Kzz`. Because the stored coordinates are an affine image of + the natural parameters, a natural-gradient step is a convex combination of the + current sites with a closed-form target, so no expectation-to-natural round trip is + needed and the KL is never differentiated. `fit_natgrads` dispatches on the family + and takes that step; the step size means the same thing in both branches, and from + the same starting `q` the dual and Salimbeni E-steps produce identical iterates — + provided the dual branch's computed per-point curvature stays non-negative, so that + its `beta_floor` never engages. That holds for a genuinely log-concave likelihood; + GPJax's `inv_probit` clips its probabilities away from 0 and 1, which breaks + log-concavity in the far tails, and there the two branches diverge. The + dual branch restricts the step size to the interval from zero to one, since the + update is a convex combination. + `dual_elbo` has the same *value* as `elbo` at the implied moments, for any sites and + any hyperparameters, but a different *hyperparameter gradient from `elbo` evaluated on + the matched `VariationalGaussian`*: the prior part of `q` tracks the kernel while the + data-dependent sites stay frozen. The difference is between the two + parameterisations, not between the two functions — calling `elbo` directly on a + `DualVariationalGaussian` returns the same value and the same gradients as + `dual_elbo`, which is simply the batched-marginals fast path for that family. That + frozen-site gradient is what gives the M-step its reported behaviour, so `Kzz` must + not be detached and the implied moments must not be cached on the family. + `DualVariationalGaussian` also works with plain `gpjax.fit`, where it is ordinary + gradient descent in the dual coordinates. The `VariationalParametrisationSuite` ASV + benchmark gains a `dual` axis value. + +### Changed (breaking) + +- **`OILMMPosterior.__call__`/`.predict` default to `covariance="diagonal"`** + (was `"dense"`) ([#682](https://github.com/JaxGaussianProcesses/GPJax/issues/682)). + The dense joint covariance costs `O(m n^2 p^2)` — an `np x np` matrix built + from `m` dense `n x n` latent covariances — which forfeits the + `O(mn^3 + nmp)` scaling that OILMM exists to provide, and is unaffordable + well before the mean/marginal-variance query is. Marginal variances are the + common case and are unaffected by the mixing matrix's off-diagonal + structure, so the cheap path is now the default; pass + `covariance="dense"` explicitly for the joint covariance. + +### Removed + +- **`NaturalVariationalGaussian` and `ExpectationVariationalGaussian`.** These were + parameterisation-only classes with no optimiser attached: they stored the natural or + expectation coordinates of `q(u)` but offered no way to take a natural-gradient step + in them. Natural-gradient geometry belongs to the optimiser — the Fisher matrix *is* + the Jacobian dη/dθ, so a natural-gradient step in the natural parameters θ is exactly + an ordinary gradient step in the expectation parameters η, and either coordinate system + can be recovered on the fly from whatever the family happens to store. `fit_natgrads` + therefore operates directly on `VariationalGaussian` and `WhitenedVariationalGaussian`, + which store constraint-respecting coordinates. Users of the removed classes should + switch to `VariationalGaussian` with `gpjax.fit_natgrads`. + The `VariationalParametrisationSuite` ASV benchmark loses its `natural` and + `expectation` axis values; previously recorded results for those two arms are orphaned. + ### Fixed +- **`gpx.kernels.RBF()` was a type error, and `White()` carried a phantom + trainable lengthscale** + ([#695](https://github.com/JaxGaussianProcesses/GPJax/issues/695)). Pyright + synthesises `__init__` signatures from dataclass fields for any kernel that + inherits its `__init__` (e.g. `RBF`, `Matern12/32/52`), and those fields had + no defaults, so the canonical `RBF()` call was flagged as missing arguments + while the nonsensical `RBF(name="xyz")` type-checked cleanly (raising + `TypeError` at runtime). Kernel fields now carry real defaults matching + their `__init__` defaults, and `name` is a `ClassVar` rather than a + dataclass field, so the synthesised and hand-written signatures agree. + Separately, `White` hardcoded `lengthscale=1.0` into + `StationaryKernel.__init__` even though `White.__call__` never reads it, + so every `White` kernel carried a real, trainable `PositiveReal` leaf with + zero gradient that showed up in optimiser state and MCMC traces; `White` + now has its own minimal `__init__` and no longer carries a lengthscale at + all (`White().lengthscale is None`, and it is absent from + `jax.tree_util.tree_flatten`). The stale `_compute_base_init` workaround in + `kernels/base.py`, whose docstring claimed "equinox modules are frozen + after `super().__init__()`" -- no longer true under the pinned Equinox + version -- was removed in favour of a plain `super().__init__(...)` call. + - **`Zero` mean function is trainable and drifts away from zero.** Fitting a model with the default `Zero()` mean function moved its constant towards the data mean (0.0 → 5.09 on a dataset with mean 5), silently changing the diff --git a/CLAUDE.md b/CLAUDE.md index c9e78a79c..2df8f3efe 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -78,17 +78,18 @@ Functions `(model, Dataset) -> scalar`: - `conjugate_mll` / `conjugate_loocv` -- for `ConjugatePosterior` - `log_posterior_density` (alias `non_conjugate_mll`) -- for `NonConjugatePosterior` - `elbo` / `collapsed_elbo` -- for variational families +- `dual_elbo` -- for `DualVariationalGaussian` (t-SVGP; same value as `elbo`, different hyperparameter gradient) - `heteroscedastic_elbo` -- for heteroscedastic models Optimise by negating: `nmll = lambda p, d: -conjugate_mll(p, d)` ### Fitting (`gpjax/fit.py`) -Three optimisers: `fit()` (Optax gradient descent with scan), `fit_scipy()` (SciPy L-BFGS-B), `fit_lbfgs()` (Optax L-BFGS with `while_loop`). All handle the constrained/unconstrained bijection automatically: `paramax.unwrap(model)` is called inside the loss function, and `eqx.partition`/`eqx.combine` with `eqx.is_array` manage trainable vs static parts. +Four optimisers: `fit()` (Optax gradient descent with scan), `fit_scipy()` (SciPy L-BFGS-B), `fit_lbfgs()` (Optax L-BFGS with `while_loop`), `fit_natgrads()` (natural-gradient steps on a variational family, alternated with Optax steps on the hyperparameters). All handle the constrained/unconstrained bijection automatically: `paramax.unwrap(model)` is called inside the loss function, and `eqx.partition`/`eqx.combine` with `eqx.is_array` manage trainable vs static parts. ### Variational inference (`gpjax/variational_families.py`) -`VariationalGaussian`, `WhitenedVariationalGaussian`, `NaturalVariationalGaussian`, `ExpectationVariationalGaussian`, `CollapsedVariationalGaussian`, `GraphVariationalGaussian`, `HeteroscedasticVariationalFamily`. All inherit from `AbstractVariationalFamily` and implement `predict()` + `prior_kl()`. +`VariationalGaussian`, `WhitenedVariationalGaussian`, `DualVariationalGaussian`, `CollapsedVariationalGaussian`, `GraphVariationalGaussian`, `HeteroscedasticVariationalFamily`. All inherit from `AbstractVariationalFamily` and implement `predict()` + `prior_kl()`. ### NumPyro integration (`gpjax/numpyro_extras.py`) diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 000000000..8d7b643eb --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,89 @@ +# GPJax domain glossary + +The ubiquitous language of this codebase. Code, docs, tests, and reviews use +these terms exactly; when a concept is missing, add it here in the same PR +that introduces it. + +## Design principle + +**Maths-first, comprehensible to the maths-familiar non-expert.** API names +follow the textbook equation unless the textbook word is jargon a +practitioner would not own. When the two conflict, choose the word a +scikit-learn/PyMC user already speaks and state the maths once in the +docstring (e.g. the class is `JointModel`, and its docstring says "the joint +distribution p(f, y)"). + +## Terms + +**Prior** — the Gaussian process prior p(f), pairing a kernel with a mean +function. Queryable at any inputs: `prior(x)` returns the prior predictive. +Owns the model's single numerical-stabilisation knob, `jitter`. + +**Likelihood** — the conditional distribution p(y | f). A *pure* conditional: +it holds no priors, no dataset facts, and no data sizes. (Both were tried and +removed at v1.0 — see ADR-0001.) + +**JointModel** — the joint distribution p(f, y) = p(y | f) · p(f), created by +`prior * likelihood`. The *trainable* object: `gpx.fit` optimises its +hyperparameters. It carries state that must exist before conditioning +(hyperparameters; the non-conjugate latent; the heteroscedastic noise-process +model) and no derived quantities. Concrete kinds: `ConjugateModel`, +`NonConjugateModel`, `HeteroscedasticModel` — each holds state the others +cannot. + +**condition** — the operation p(f, y) + 𝒟 → p(f | 𝒟), spelled +`model.condition(D)` or the operator form `model | D` (read: "f given D"). + +The signature is `condition(train_data)` uniformly, on every conditionable +object, with `train_data` required. It is universal across the exact, latent, +sparse, collapsed and state-space modes. Where the maths does not consume the +data — the uncollapsed variational families, which already carry the fitted +q(u) — the argument is still accepted, for interface uniformity, and the +docstring says so plainly. A signature that varied by object would be a +worse API than one argument occasionally ignored. + +The one documented exclusion is the heteroscedastic path +(`HeteroscedasticModel` and `HeteroscedasticVariationalFamily`), which has no +closed-form conditioned process: it carries two latent processes, signal and +noise, so there is no single p(f | 𝒟) to return. Both raise +`NotImplementedError` naming the alternative — inference runs through +`HeteroscedasticVariationalFamily` and the `heteroscedastic_elbo` objective, +and prediction through `predict` / `predict_latents`, or by conditioning the +`signal_variational` and `noise_variational` components individually. + +`prior_kl` is deliberately *not* part of this contract: it keeps a per-family +signature, because only the collapsed family's KL is a function of the data. + +**Posterior** — the conditioned process p(f | 𝒟) returned by `condition`. An +*immutable* pytree: the training-covariance factorisation is computed once and +cached; every query is a view of it. The uniform query surface: +`posterior(xtest, covariance="dense"|"diagonal")`, plus per-mode views — +`log_marginal_likelihood` / `loo` / `sample_approx` on the exact mode, +`log_posterior_density` on the latent mode. Users never name the concrete +implementations behind the interface. + +**variational family** — a trainable approximate posterior over inducing +values: to sparse GPs what JointModel is to exact ones. It carries the joint +model in its `model` field, and `.condition(D)` yields a Posterior like any +other; `elbo`-style objectives are its training criteria. The model's +`Prior.jitter` is the only stabilisation knob — families carry none of their +own. + +**evidence / log marginal likelihood** — p(𝒟), the normalising constant of +conditioning, exposed as `posterior.log_marginal_likelihood`. "Evidence" and +"marginal likelihood" are the same quantity; the attribute uses the +GP-community's term. + +**sugar** — a documented one-line composition kept for ergonomics, never a +second implementation. `model.predict(x, D)` and `model(x, D)` are sugar for +`model.condition(D)(x)`; `model | D` is sugar for `model.condition(D)`. + +**objective** — a scalar function `(model, Dataset) -> ScalarFloat` consumed +by `gpx.fit`. Objectives are thin: `conjugate_mll` is the evidence view of +the conditioned posterior, not a second derivation. + +**Dataset** — the data container. `n_total` records the full-dataset size +when the object is a minibatch view (stamped by `get_batch`); the minibatch +ELBO scale is derived from it, never supplied by hand. Read it through +`full_size`, which falls back to `n` for a whole dataset — production code +uses `data.full_size / data.n` and never re-spells the fallback inline. diff --git a/benchmarks/README.md b/benchmarks/README.md index 5fa8ee86e..859c2fb36 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -77,6 +77,6 @@ Prints a regression report with significance markers. No branches touched. ## Notes on bench-check `bench-check` runs `asv check`, which builds an env for the commit -referenced by `asv.conf.json`'s `branches: ["main"]`. It only succeeds -once `benchmarks/` and `asv-constraints.txt` have landed on `main`. After -that, it's safe to add to the `all-tests` poe sequence. +referenced by `asv.conf.json`'s `branches: ["main"]`. It is wired into +the `all-tests` poe sequence, so it runs on every `uv run poe all-tests` +alongside `lint`, `docstrings` and `test`. diff --git a/benchmarks/compile.py b/benchmarks/compile.py index 0601c57ba..accecb63e 100644 --- a/benchmarks/compile.py +++ b/benchmarks/compile.py @@ -39,10 +39,10 @@ def setup(self): kernel = gpx.kernels.RBF() mean = gpx.mean_functions.Zero() prior = gpx.gps.Prior(kernel=kernel, mean_function=mean) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=n) + likelihood = gpx.likelihoods.Gaussian() self.posterior = prior * likelihood self.q = VariationalGaussian( - posterior=self.posterior, inducing_inputs=X[:M_INDUCING] + model=self.posterior, inducing_inputs=X[:M_INDUCING] ) self.jitted_mll = jax.jit(objectives.conjugate_mll) diff --git a/benchmarks/objectives.py b/benchmarks/objectives.py index 5689eab81..60f9e5806 100644 --- a/benchmarks/objectives.py +++ b/benchmarks/objectives.py @@ -9,9 +9,9 @@ both M (inducing count) and batch_size; VfeElboSuite (collapsed analytic ELBO) varies M only — the collapsed objective requires the full dataset. -VariationalParametrisationSuite compares the four variational Gaussian -parameterisations (standard, whitened, natural, expectation) at one -fixed (n, M) so users can see the per-step cost of the choice. +VariationalParametrisationSuite compares the three variational Gaussian +parameterisations (standard, whitened, dual) at one fixed (n, M) so users +can see the per-step cost of the choice. HeteroscedasticElboSuite and OilmmPredictSuite are independent — they do not participate in the alignment because their model structure @@ -26,9 +26,8 @@ from gpjax.likelihoods import HeteroscedasticGaussian from gpjax.variational_families import ( CollapsedVariationalGaussian, - ExpectationVariationalGaussian, + DualVariationalGaussian, HeteroscedasticVariationalFamily, - NaturalVariationalGaussian, VariationalGaussian, WhitenedVariationalGaussian, ) @@ -50,7 +49,7 @@ def _conjugate_posterior(n: int): kernel = gpx.kernels.RBF() mean = gpx.mean_functions.Zero() prior = gpx.gps.Prior(kernel=kernel, mean_function=mean) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=n) + likelihood = gpx.likelihoods.Gaussian() return prior * likelihood, data @@ -86,7 +85,7 @@ class SvgpElboSuite: def setup(self, n, m, batch_size): posterior, data, Z = _sparse_setup(n, m=m) - self.q = VariationalGaussian(posterior=posterior, inducing_inputs=Z) + self.q = VariationalGaussian(model=posterior, inducing_inputs=Z) self.batch = Dataset(X=data.X[:batch_size], y=data.y[:batch_size]) realise(objectives.elbo(self.q, self.batch)) @@ -107,7 +106,7 @@ class VfeElboSuite: def setup(self, n, m): posterior, self.data, Z = _sparse_setup(n, m=m) - self.q = CollapsedVariationalGaussian(posterior=posterior, inducing_inputs=Z) + self.q = CollapsedVariationalGaussian(model=posterior, inducing_inputs=Z) realise(objectives.collapsed_elbo(self.q, self.data)) def time_collapsed_elbo(self, n, m): @@ -117,17 +116,25 @@ def time_collapsed_elbo(self, n, m): _VARIATIONAL_FAMILIES = { "standard": VariationalGaussian, "whitened": WhitenedVariationalGaussian, - "natural": NaturalVariationalGaussian, - "expectation": ExpectationVariationalGaussian, + "dual": DualVariationalGaussian, } class VariationalParametrisationSuite: - """Per-step ELBO cost across the four variational Gaussian families. - - All four parameterise the same q(u); the differences are in how the - KL term and predictive moments are computed. Holding (n, M) fixed - isolates the parameterisation cost. + """Per-step ELBO cost across the variational Gaussian parameterisations. + + All three parameterise the same q(u); the differences are in how the + KL term and predictive moments are computed. The dual (t-SVGP) family + stores sites rather than moments and recovers q(u) through + R = Kzz + Kzz Lambda_2 Kzz, so every predict and every KL factorises + both Kzz and R, against one factorisation for the standard family and + none in the whitened KL. Holding (n, M) fixed isolates that cost. + + The arm deliberately times the generic ``elbo`` on all three families + rather than ``dual_elbo`` on the dual one: the point of the comparison + is the parameterisation, so the objective has to be held fixed. The + dual family's own fast path is ``dual_elbo``, which replaces the + per-point ``predict`` with one batched ``marginals`` call. """ params = (list(_VARIATIONAL_FAMILIES),) @@ -136,7 +143,7 @@ class VariationalParametrisationSuite: def setup(self, family): n = 1000 posterior, self.data, Z = _sparse_setup(n) - self.q = _VARIATIONAL_FAMILIES[family](posterior=posterior, inducing_inputs=Z) + self.q = _VARIATIONAL_FAMILIES[family](model=posterior, inducing_inputs=Z) realise(objectives.elbo(self.q, self.data)) def time_elbo(self, family): @@ -157,11 +164,11 @@ def setup(self, n): noise_prior = gpx.gps.Prior( kernel=noise_kernel, mean_function=gpx.mean_functions.Zero() ) - likelihood = HeteroscedasticGaussian(num_datapoints=n, noise_prior=noise_prior) + likelihood = HeteroscedasticGaussian(noise_prior=noise_prior) posterior = signal_prior * likelihood Z = data.X[:M_INDUCING] self.q = HeteroscedasticVariationalFamily( - posterior=posterior, inducing_inputs=Z, inducing_inputs_g=Z + model=posterior, inducing_inputs=Z, inducing_inputs_g=Z ) self.data = data realise(objectives.heteroscedastic_elbo(self.q, self.data)) @@ -183,7 +190,7 @@ def setup(self, m): X = jnp.linspace(0, 1, N).reshape(-1, 1) y = jr.normal(key, (N, P)) dataset = gpx.Dataset(X=X, y=y) - self.posterior = model.condition_on_observations(dataset) + self.posterior = model.condition(dataset) self.X_test = jnp.linspace(0.1, 0.9, 20).reshape(-1, 1) realise(self.posterior.predict(self.X_test)) diff --git a/benchmarks/state_space.py b/benchmarks/state_space.py index f51a2ec77..110470c16 100644 --- a/benchmarks/state_space.py +++ b/benchmarks/state_space.py @@ -9,7 +9,7 @@ import gpjax as gpx from gpjax.state_space import StateSpacePrior, state_space_mll -from gpjax.state_space.gps import StateSpaceConjugatePosterior +from gpjax.state_space.gps import StateSpaceConjugateModel import jax.numpy as jnp import jax.random as jr @@ -35,10 +35,8 @@ def setup(self, n): mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0), ) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=n, obs_stddev=0.1) - self.posterior = StateSpaceConjugatePosterior( - prior=prior, likelihood=likelihood - ) + likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) + self.posterior = StateSpaceConjugateModel(prior=prior, likelihood=likelihood) self.data = _temporal_dataset(n) realise(state_space_mll(self.posterior, self.data)) diff --git a/docs/adr/0001-conditioning-architecture.md b/docs/adr/0001-conditioning-architecture.md new file mode 100644 index 000000000..68b3df246 --- /dev/null +++ b/docs/adr/0001-conditioning-architecture.md @@ -0,0 +1,154 @@ +# ADR-0001: The v1.0 conditioning architecture + +- **Status:** accepted +- **Date:** 2026-08-06 +- **Deciders:** Thomas Pinder (design settled in an architecture-review + + grilling session; full decision tree recorded there) + +## Context + +An architecture review (2026-08-06) found that no module owned "condition a +GP on data": the derivation *stabilise → factor → solve → predictive +moments* was written out at eleven call sites across `gps.py`, +`objectives.py`, `variational_families.py`, and `models/oilmm.py`. Fixes +reached some copies and missed others (the diagonal-predict fix landed in two +of eleven), jitter had two owners (`Prior.jitter` vs `Posterior.jitter`), so +`predict` and `conjugate_mll` could factorise *different matrices* for the +same model, and `conjugate_mll` — the oracle for the Kalman MLL and +`collapsed_elbo` — had no value-level test of its own. + +## Decision + +One deep conditioning module, with the public API as its veneer, rolled out +universally at v1.0: + +- `prior * likelihood` returns a **`JointModel`** — the joint p(f, y), the + trainable object (`ConjugateModel`, `NonConjugateModel` with a lazily-sized + latent, `HeteroscedasticModel` owning the noise-process prior). +- `model.condition(D)` (operator sugar `model | D`) returns a **`Posterior`**: + an immutable pytree caching the factorisation, with the predictive, + `log_marginal_likelihood`, `loo`, and pathwise `sample_approx` as views of + it. One abstract interface; per-mode internal implementations. +- The signature is `condition(train_data)` **uniformly**, with `train_data` + required, on every conditionable object. Objects whose maths does not + consume the data still accept it and say so in their docstring. A patchy + signature was ruled out explicitly: universal, or not at all. +- `prior.jitter` is the single stabilisation knob, applied exactly once, + inside conditioning, through `linalg.stabilised_cholesky` (the seed of the + linalg deepening). +- `predict(x, D)` survives as documented one-line sugar. Objectives become + one-line views. `return_covariance_type` is renamed `covariance`. +- Likelihoods are pure conditionals: `num_datapoints` deleted (`Dataset` + carries a static `n_total`, stamped by `get_batch`, from which the + minibatch ELBO scale is derived) and `noise_prior` moved to + `HeteroscedasticModel` (removing the `likelihoods -> gps` circular import). +- Deleted as failed deletion-tests: `AbstractPrior` (one-adapter seam), + `AbstractPosterior` (splits into JointModel/Posterior), the + `LatentPosterior` and `ChainedPosterior` markers. +- Universality is a property of the **v1.0 release**, assembled from stacked + PRs: safety net → core conditioning → variational universalisation. The + variational PR was sequenced **after** the natural-gradients stack + (#714–#730), since that stack rewrites `variational_families.py`. + +The safety net landed first, in its own PR: closed-form oracles for the MLL, +predict, and LOOCV; cross-derivation equivalence pins; and an integration +harness whose failures actually raise. + +### What "universal" means, as shipped + +`condition(train_data) -> Posterior`, with `__or__` as its operator sugar, is +implemented across every conditioning mode: + +| Mode | Object | Posterior | +|---|---|---| +| exact | `ConjugateModel` | `ExactPosterior` | +| latent | `NonConjugateModel` | `LatentPosterior` | +| sparse | `VariationalGaussian`, `WhitenedVariationalGaussian`, `DualVariationalGaussian`, `GraphVariationalGaussian` | `SparsePosterior` | +| collapsed | `CollapsedVariationalGaussian` | `CollapsedPosterior` | +| state-space | `StateSpaceConjugateModel` | `StateSpacePosterior` (Kalman-backed) | +| multi-output | `OILMMModel` | `OILMMPosterior` (M cached latent factorisations) | + +The sparse families carry q(u) internally and ignore `train_data`; the +collapsed family solves its optimal q\*(u) from it; the state-space mode +takes an additional keyword-only `observation_mask`. `StateSpacePosterior` +runs the Kalman recursions rather than inheriting the dense assembly, so +conditioning a state-space model stays O(N) rather than silently becoming +O(N³). + +**The one exclusion** is the heteroscedastic path. `HeteroscedasticModel` and +`HeteroscedasticVariationalFamily` carry two latent processes, signal and +noise, and there is no closed-form single conditioned process to return. +Both raise `NotImplementedError` with a message naming the alternative: +inference runs through `HeteroscedasticVariationalFamily` fitted with the +`heteroscedastic_elbo` objective; prediction runs through `predict` / +`predict_latents`, or by conditioning the `signal_variational` and +`noise_variational` components individually. This is documented rather than +papered over — returning some partial object from `condition` would make the +contract dishonest, which is the failure mode the uniform signature exists to +prevent. + +`gpjax.models.oilmm` is on the contract. `OILMMModel` is not a `JointModel` +— it is not built from `prior * likelihood` — but it does not need to be: +`condition` is a method returning a `Posterior`, and membership is defined by +that, not by ancestry. `model.condition(D)` and `model | D` return an +`OILMMPosterior`, now a `Posterior` subclass holding the `M` conditioned +latent processes; `condition_on_observations` and `predict(return_full_cov=)` +survive as deprecated aliases. + +Two things fell out of that change rather than being designed into it. The +old `OILMMPosterior` held `M` *unconditioned* `ConjugateModel`s alongside `M` +datasets, deferring the real conditioning to `predict` — so every prediction +re-factorised `M` Choleskys. Holding `ExactPosterior`s instead caches them at +`condition` time (measured ~1.8x on repeated prediction at n=300). And the +docstring's stated reason for the plain class — that it "holds Dataset +objects which are not JAX pytree nodes" — was simply untrue: `Dataset` is a +registered pytree, and `ExactPosterior` has held one as a field all along. + +## Rejected alternatives + +- **Dropping the JointModel object** ("just `prior.condition(data, + likelihood)`"): the trainable/derived split is load-bearing. `fit` needs a + named pytree of everything trainable; the Posterior carries derived cache + that a gradient step must never touch. Any container invented to make + `fit` ergonomic *is* the JointModel under another name. +- **Staged or partial public rollout**: a patchy API (some objects + conditioning, others not) was ruled out; stacking PRs into one v1.0 release + achieves universality without a big-bang branch. +- **Prior inside the likelihood**: the codebase already ran this experiment — + `HeteroscedasticGaussian.noise_prior` produced circular ownership across + three modules. Priors live in models; likelihoods stay conditional + families. +- **Keeping `num_datapoints`**: only four real reads existed; nothing + validated the value, so a wrong one silently mis-scaled the ELBO; no major + GP library puts dataset size on likelihoods (research: + `plans/2026-08-06-drop-num-datapoints-research.md`). +- **Naming the joint `Model`** (rejected in favour of `JointModel`): the + prefix self-documents the maths and answers "why can't I predict with this + yet?"; "model" remains in the name for practitioners. + +## Consequences + +- Locality: one home for the algebra; the two-owner jitter bug is + structurally impossible; `sample_approx` reuses the same factor as + `predict` and refuses multi-output loudly instead of broadcasting wrongly. +- The evidence is cached with the factorisation, so repeated + predict-then-score workflows stop re-factorising. +- Breaking changes at v1.0 are recorded in `docs/migration.md`; the + vocabulary lives in `CONTEXT.md`. +- Variational universalisation (landed as the stack's final PR): every + Gaussian-output family conditions through the module's sparse/collapsed + modes, deleting the five copied predictive derivations and the duplicated + Titsias bound. The family-side `jitter` knob is gone — `Prior.jitter` is + applied inside conditioning for families exactly as for joint models — so + the collapsed-ELBO/MLL equivalence test that was strictly xfailed at + non-default jitter now passes, and the families' `posterior` field is + renamed `model` (it holds the joint, not a posterior). +- Follow-ups tracked for the stack: the linalg structure-preserving + deepening, one training loop with stepper adapters, the compute-engine + seam, sharing the `K_zz` factor between `prior_kl` and `condition(D)` + within a single ELBO step (today each factorises its own; XLA CSE merges + them under `jit`), and the + heteroscedastic family's NamedTuple predict (candidate-3 debt — it still + derives its predictive through its two sub-families rather than a + conditioning mode of its own, which is why it is the documented exclusion + above rather than a sixth mode). diff --git a/docs/conf.py b/docs/conf.py index 2e16ea00c..6435d6b8f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -72,6 +72,7 @@ exclude_patterns = [ "_build", + "adr/*", # ADRs are in-repo records, not (yet) part of the docs site "Thumbs.db", ".DS_Store", "conf.py", # this config module is not a document @@ -181,8 +182,14 @@ # Everything else (broken xrefs, bad anchors, malformed directives) stays fatal # on deploy, which it was not when that build ran without `-W` at all. # The PR gate suppresses nothing. +# codeautolink cannot match doctest blocks carrying `# doctest: +SKIP` markers +# against their rendered HTML (a matcher limitation, not a doc defect — xdoctest +# validates the examples). Suppressed on every path; predates the v1.0 stack but +# first surfaced when this workflow ran cold post-Sphinx-migration. +suppress_warnings = ["codeautolink.match_block"] + if os.environ.get("GPJAX_DOCS_RESILIENT") == "1": - suppress_warnings = [ + suppress_warnings = suppress_warnings + [ "mystnb.exec", # execution failure + "traceback saved in:" follow-up "mystnb.glue", # a glue key that never got produced by a failed notebook # A notebook that fails to execute renders with no outputs, and MyST-NB diff --git a/docs/examples/backend.py b/docs/examples/backend.py index 89b874eb5..8ddf0ea9a 100644 --- a/docs/examples/backend.py +++ b/docs/examples/backend.py @@ -157,7 +157,7 @@ def glue(*args, **kwargs): prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) -likelihood = gpx.likelihoods.Gaussian(100) +likelihood = gpx.likelihoods.Gaussian() posterior = likelihood * prior print(posterior) diff --git a/docs/examples/barycentres.py b/docs/examples/barycentres.py index 591ec61c5..453e82546 100644 --- a/docs/examples/barycentres.py +++ b/docs/examples/barycentres.py @@ -132,7 +132,7 @@ def glue(*args, **kwargs): # ## Dataset # # We'll simulate five datasets and develop a Gaussian process -# [posterior](#gpjax.gps.ConjugatePosterior) before +# [posterior](#gpjax.gps.ConjugateModel) before # identifying the Gaussian process barycentre at a set of test points. Each dataset # will be a sine function with a different vertical shift, periodicity, and quantity # of noise. @@ -184,7 +184,7 @@ def fit_gp(x: jax.Array, y: jax.Array) -> gpx.distributions.GaussianDistribution y = y.reshape(-1, 1) D = gpx.Dataset(X=x, y=y) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=n) + likelihood = gpx.likelihoods.Gaussian() posterior = ( gpx.gps.Prior( mean_function=gpx.mean_functions.Constant(), kernel=gpx.kernels.RBF() diff --git a/docs/examples/classification.py b/docs/examples/classification.py index 74edeabdb..88df10b97 100644 --- a/docs/examples/classification.py +++ b/docs/examples/classification.py @@ -105,7 +105,7 @@ kernel = gpx.kernels.RBF() meanf = gpx.mean_functions.Constant() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) -likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n) +likelihood = gpx.likelihoods.Bernoulli() # %% [markdown] # We construct the posterior through the product of our prior and likelihood. @@ -144,7 +144,7 @@ ) # %% [markdown] -# From which we can [make predictions](#gpjax.gps.NonConjugatePosterior.predict) at +# From which we can [make predictions](#gpjax.gps.JointModel.predict) at # novel inputs, as illustrated in {numref}`fig-classification-map-predictive`. # %% mystnb={"figure": {"caption": "The MAP predictive mean and its one-sigma band over the binary observations.", "name": "fig-classification-map-predictive"}} diff --git a/docs/examples/collapsed_vi.py b/docs/examples/collapsed_vi.py index 5e8b92b05..f22c2716b 100644 --- a/docs/examples/collapsed_vi.py +++ b/docs/examples/collapsed_vi.py @@ -121,7 +121,7 @@ # %% meanf = gpx.mean_functions.Constant() kernel = gpx.kernels.RBF() # 1-dimensional inputs -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) posterior = prior * likelihood @@ -135,7 +135,7 @@ # %% q = gpx.variational_families.CollapsedVariationalGaussian( - posterior=posterior, inducing_inputs=z + model=posterior, inducing_inputs=z ) # %% [markdown] @@ -170,7 +170,7 @@ # %% mystnb={"figure": {"caption": "Predictive mean and two-standard-deviation band of the sparse posterior, shown against the observations and the latent function, with the optimised inducing point locations overlaid.", "name": "fig-collapsed-vi-predictions"}} latent_dist = opt_posterior(xtest, train_data=D) -predictive_dist = opt_posterior.posterior.likelihood(latent_dist) +predictive_dist = opt_posterior.model.likelihood(latent_dist) inducing_points = opt_posterior.inducing_inputs.unwrap() @@ -242,7 +242,7 @@ # %% full_rank_model = gpx.gps.Prior( mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.RBF() -) * gpx.likelihoods.Gaussian(num_datapoints=D.n) +) * gpx.likelihoods.Gaussian() nmll = jit(lambda: -gpx.objectives.conjugate_mll(full_rank_model, D)) # %timeit nmll().block_until_ready() diff --git a/docs/examples/constructing_new_kernels.py b/docs/examples/constructing_new_kernels.py index 88da1e2f9..a2d3b3e2f 100644 --- a/docs/examples/constructing_new_kernels.py +++ b/docs/examples/constructing_new_kernels.py @@ -287,7 +287,7 @@ def __call__( # Define polar Gaussian process PKern = Polar() meanf = gpx.mean_functions.Zero() -likelihood = gpx.likelihoods.Gaussian(num_datapoints=n) +likelihood = gpx.likelihoods.Gaussian() circular_posterior = gpx.gps.Prior(mean_function=meanf, kernel=PKern) * likelihood # Optimise GP's marginal log-likelihood using BFGS diff --git a/docs/examples/deep_kernels.py b/docs/examples/deep_kernels.py index 2a382e44c..6003285cc 100644 --- a/docs/examples/deep_kernels.py +++ b/docs/examples/deep_kernels.py @@ -196,7 +196,7 @@ def __call__(self, x: jax.Array) -> jax.Array: kernel = DeepKernelFunction(network=forward_linear, base_kernel=base_kernel) meanf = gpx.mean_functions.Zero() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() posterior = prior * likelihood # %% [markdown] # ### Optimisation diff --git a/docs/examples/dual_svgp.py b/docs/examples/dual_svgp.py new file mode 100644 index 000000000..42728f7f5 --- /dev/null +++ b/docs/examples/dual_svgp.py @@ -0,0 +1,1576 @@ +# --- +# jupyter: +# jupytext: +# cell_metadata_filter: -all +# custom_cell_magics: kql +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.19.1 +# kernelspec: +# display_name: Python 3 +# language: python +# name: python3 +# --- + +# %% [markdown] +# # Dual Parameterisation of Sparse GPs (t-SVGP) +# +# Download this notebook: {nb-download}`dual_svgp.ipynb` +# +# A sparse variational GP stores its approximate posterior $q(\mathbf{u})$ as a mean +# and a Cholesky factor. That is a choice, not a law. The +# {cite:t}`adam2021dual` *dual* +# parameterisation stores something else: the **likelihood sites**, one per data +# point, tied down to the inducing points. The distribution is the same object; what +# changes is which numbers are held in memory, and therefore what the optimiser is +# allowed to hold fixed. +# +# Two things follow from that change, and this notebook is about establishing how +# much each one is worth. +# +# 1. **The natural-gradient step becomes an explicit convex combination of the stored +# parameters.** No conversion to natural parameters, no conversion back, and the +# KL term is never differentiated. The step is the same *iteration* as the one in +# the [natural gradients notebook](natgrads.py) — +# we check that below to $10^{-15}$ — **provided the computed per-point curvature +# $\beta_i$ stays non-negative**. That holds for a genuinely log-concave +# likelihood. GPJax's probit link clips its probabilities, which breaks +# log-concavity in the far tails; where that bites, the dual branch's `beta_floor` +# engages and the two branches really do differ. We locate that below and measure +# it. Away from it, the difference is wall-clock only, never accuracy. +# 2. **The hyperparameter objective changes.** Because the sites are, in this +# convention, free of the kernel hyperparameters, letting $\mathbf{K}_{zz}$ move +# while the sites stay put gives a *different* function of $\boldsymbol{\theta}$ +# from the usual "freeze $(\mathbf{m},\mathbf{S})$" bound: the same value at the +# current hyperparameters, and a gradient that coincides with the standard one +# there only once the E-step has converged — which, between optimiser steps, it +# never has. That gap is the mechanism. It is worth more than the first point and +# is harder to pin down; the last third of the notebook is spent being precise +# about what is proven and what is merely measured. +# +# The route is: the dual coordinates and their EP heritage; the tying that restores +# $\mathcal{O}(M^2)$ memory; the two storage conventions and which one GPJax picked; +# the tied update and why it needs no round trip; a conjugate model where one step at +# $\rho=1$ is the exact answer; the $\rho=\gamma$ check that discharges claim 1; the +# banana classification benchmark from the natural-gradients notebook, run again with +# all three optimisers; and finally hyperparameter learning, where `dual_elbo` and +# `elbo` part company. +# +# This notebook assumes the natural-gradients notebook. Read that one first: it +# derives the exponential-family view of $q(\mathbf{u})$, the identity that the +# natural gradient in one of its two canonical coordinate systems is the ordinary +# gradient in the other, and the mirror-descent reading of the step size, all of which +# are assumed here. +# +# **One notational break from it.** That notebook writes the natural parameter of +# $q(\mathbf{u})$ as $\boldsymbol{\theta}$ and the expectation parameter as +# $\boldsymbol{\eta}$. Here $\boldsymbol{\theta}$ is reserved for the kernel +# hyperparameters, so the natural parameter is $\boldsymbol{\eta}$, the expectation +# parameter is $\boldsymbol{\mu}$, and $\boldsymbol{\lambda}$ is the site — not that +# notebook's conjugate likelihood parameter. In these letters its identity reads +# $\tilde\nabla_{\boldsymbol{\eta}}\mathcal{L} = \partial\mathcal{L}/\partial\boldsymbol{\mu}$, +# and it is restated that way where it is used below. + +# %% +# Enable Float64 for more stable matrix inversions. +import time + +import equinox as eqx +import jax +from jax import config +import jax.numpy as jnp +import jax.random as jr +import jax.tree_util as jtu +from jaxtyping import install_import_hook +import matplotlib as mpl +import matplotlib.pyplot as plt +import optax as ox +import paramax +from utils import clean_legend, use_mpl_style + +config.update("jax_enable_x64", True) + + +with install_import_hook("gpjax", "beartype.beartype"): + import gpjax as gpx + import gpjax.kernels as jk + from gpjax.natural_gradients import ( + natural_gradient_step, + partition_variational, + ) + from gpjax.objectives import dual_elbo, elbo + from gpjax.parameters import Real + from gpjax.variational_families import ( + DualVariationalGaussian, + VariationalGaussian, + ) + +key = jr.key(123) + +# set the default style for plotting +use_mpl_style() +cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"] + + +def negative_elbo(model, data): + """The loss for a family that stores moments; GPJax optimisers descend.""" + return -elbo(model, data) + + +def negative_dual_elbo(model, data): + """The loss for a family that stores sites.""" + return -dual_elbo(model, data) + + +# %% [markdown] +# ## From natural to dual coordinates +# +# Write the natural parameter of $q(\mathbf{u}) = \mathcal{N}(\mathbf{m},\mathbf{S})$ +# as $\boldsymbol{\eta} = (\mathbf{S}^{-1}\mathbf{m},\ -\tfrac12\mathbf{S}^{-1})$, and +# the natural parameter of the prior +# $p(\mathbf{u}) = \mathcal{N}(\mathbf{0},\mathbf{K}_{zz})$ as +# $\boldsymbol{\eta}_0(\boldsymbol{\theta}) = (\mathbf{0},\ -\tfrac12\mathbf{K}_{zz}^{-1})$. +# Their difference is the object this notebook stores: +# +# $$\boldsymbol{\eta} = \underbrace{\left(\mathbf{0},\ -\tfrac12\mathbf{K}_{zz}^{-1}\right)}_{\boldsymbol{\eta}_0(\boldsymbol{\theta})\ \text{prior}} \;+\; \underbrace{\left(\boldsymbol{\lambda}_1,\ -\tfrac12\boldsymbol{\Lambda}_2\right)}_{\boldsymbol{\lambda}\ \text{sites}} .$$ +# +# The decomposition is additive, and — in this convention — the second half carries no +# dependence on the kernel hyperparameters $\boldsymbol{\theta}$ at all. Equivalently, +# $q$ is the prior reweighted by an unnormalised Gaussian *site*, +# +# $$t(\tilde{\mathbf{u}}) = \exp\!\left(\boldsymbol{\lambda}_1^\top\tilde{\mathbf{u}} - \tfrac12\tilde{\mathbf{u}}^\top\boldsymbol{\Lambda}_2\tilde{\mathbf{u}}\right), \qquad q(\mathbf{u}) \propto p_{\boldsymbol{\theta}}(\mathbf{u})\,t(\tilde{\mathbf{u}}),$$ +# +# from which the moments follow by completing the square, +# +# $$\mathbf{S} = \left(\mathbf{K}_{zz}^{-1} + \boldsymbol{\Lambda}_2\right)^{-1}, \qquad \tilde{\mathbf{m}} = \mathbf{S}\boldsymbol{\lambda}_1, \qquad \mathbf{m} = \boldsymbol{\mu}_z + \tilde{\mathbf{m}} .$$ +# +# Here $\tilde{\mathbf{u}} = \mathbf{u} - \boldsymbol{\mu}_z$ are the inducing outputs +# centred on the prior mean function, so a non-zero mean function needs no special +# case anywhere below. `DualVariationalGaussian` stores $\boldsymbol{\lambda}_1$ as +# `dual_vector` ($M\times1$) and $\boldsymbol{\Lambda}_2$ as `dual_matrix` +# ($M\times M$), both defaulting to zero — which sets $q = p$ and makes the KL vanish +# at initialisation. +# +# Nothing here is ever inverted. Every quantity the family needs routes through +# +# $$\mathbf{R} := \mathbf{K}_{zz} + \mathbf{K}_{zz}\boldsymbol{\Lambda}_2\mathbf{K}_{zz} = \mathbf{K}_{zz}\mathbf{S}^{-1}\mathbf{K}_{zz},$$ +# +# which satisfies $\mathbf{R} \succeq \mathbf{K}_{zz} \succ 0$ whenever +# $\boldsymbol{\Lambda}_2 \succeq 0$. So $\operatorname{chol}(\mathbf{R})$ cannot fail, +# and — this is the point — it is *better* conditioned than +# $\operatorname{chol}(\boldsymbol{\Lambda}_2)$ would be, which is rank deficient at +# initialisation and whenever the batch is smaller than $M$. Two Cholesky +# factorisations per iteration, $\mathbf{L}_K$ and $\mathbf{L}_R$, and no more. + +# %% [markdown] +# ## The EP connection +# +# Where do $\boldsymbol{\lambda}_1$ and $\boldsymbol{\Lambda}_2$ come from? Adam et +# al. show that the ELBO-optimal $q$ has the site form +# +# $$q^*(\mathbf{u}) \;\propto\; p_{\boldsymbol{\theta}}(\mathbf{u})\prod_{i=1}^{N} t_i^*(\mathbf{u}), \qquad t_i^*(\mathbf{u}) = \exp\!\left(\langle\boldsymbol{\lambda}_i^*,\ \mathbf{T}(\mathbf{a}_i^\top\mathbf{u})\rangle\right),$$ +# +# with $\mathbf{T}(v) = (v, v^2)$ the Gaussian sufficient statistics and +# $\mathbf{a}_i = \mathbf{K}_{zz}^{-1}\mathbf{k}_z(x_i)$. Each $t_i$ is a +# **two-dimensional** object acting on the scalar projection +# $\mathbf{a}_i^\top\mathbf{u}$: one local likelihood approximation per data point, +# exactly as in expectation propagation. The difference from EP is where the site +# values come from. EP computes them by matching moments against a cavity +# distribution; here they are read straight off the first two derivatives of the +# expected log likelihood. With $q(f_i) = \mathcal{N}(m_i, v_i)$, Bonnet's and Price's +# theorems give +# +# $$\alpha_i = \frac{\partial}{\partial m_i}\,\mathbb{E}_{q(f_i)}\!\left[\log p(y_i\mid f_i)\right], \qquad \beta_i = -2\,\frac{\partial}{\partial v_i}\,\mathbb{E}_{q(f_i)}\!\left[\log p(y_i\mid f_i)\right],$$ +# +# so a single `jax.grad` of the likelihood's existing `expected_log_likelihood` +# suffices. No second derivatives, and it works for closed-form and quadrature +# likelihoods alike. +# +# Stored naively that is $\mathcal{O}(N)$ memory, which would be a poor trade. But +# every $t_i$ enters $q(\mathbf{u})$ only through the rank-one projection +# $\mathbf{a}_i^\top\mathbf{u}$, so the $N$ sites can be **tied**: summed into two +# inducing-space objects of size $M$ and $M\times M$. Writing +# $g_{1,i} = \alpha_i + \beta_i\,(m_i - \mu(x_i))$ and $g_{2,i} = \beta_i$, the tied +# values at a converged full-batch E-step are +# +# $$\boldsymbol{\lambda}_1 = \sum_{i=1}^{N}\mathbf{a}_i g_{1,i} = \mathbf{A}\mathbf{g}_1, \qquad \boldsymbol{\Lambda}_2 = \sum_{i=1}^{N}g_{2,i}\,\mathbf{a}_i\mathbf{a}_i^\top = \mathbf{A}\operatorname{diag}(\mathbf{g}_2)\mathbf{A}^\top .$$ +# +# Memory is back to $\mathcal{O}(M^2)$, the same as standard SVGP. Two warnings. The +# tying introduces a bias — the paper says so, and reports that it "does not seem to +# affect convergence in practice". And these sums are the *fixed point*, not the value +# at a general iterate: during training the stored pair is a running convex +# combination of such targets, and is never computed by evaluating the sum. + +# %% [markdown] +# ## Two conventions, and which one GPJax stores +# +# Two traps wait for anyone reading the paper alongside the code, and both are +# silent: each produces a valid-looking $q$ that is simply not the one intended. +# +# **The flanking trap.** The paper's main text (its Eq. 21) stores the *un-flanked* +# sums, built from $\mathbf{k}_z(x_i)$ rather than +# $\mathbf{a}_i = \mathbf{K}_{zz}^{-1}\mathbf{k}_z(x_i)$: +# +# $$\bar{\boldsymbol{\lambda}}_1 = \mathbf{K}_{zx}\mathbf{g}_1 = \mathbf{K}_{zz}\boldsymbol{\lambda}_1, \qquad \bar{\boldsymbol{\Lambda}}_2 = \mathbf{K}_{zx}\operatorname{diag}(\mathbf{g}_2)\mathbf{K}_{xz} = \mathbf{K}_{zz}\boldsymbol{\Lambda}_2\mathbf{K}_{zz}.$$ +# +# Both conventions describe the same $q$, and $\mathbf{R}$ is literally the same +# matrix in each. They are not interchangeable for our purposes, though: in the +# un-flanked form *both* halves of $\boldsymbol{\eta} - \boldsymbol{\eta}_0$ move with +# $\boldsymbol{\theta}$, since +# $\boldsymbol{\eta}_1 = \mathbf{K}_{zz}^{-1}\bar{\boldsymbol{\lambda}}_1$. The +# additive, hyperparameter-free split that the whole second half of this notebook +# rests on holds *exactly* only in the flanked convention. The paper flags the choice +# in a single sentence and calls the flanked form "an alternative tying method"; GPJax +# stores the alternative. +# +# **The $-\tfrac12$ trap.** The paper uses $\lambda_2$ with two incompatible meanings: +# the natural-parameter one ($-\tfrac12\beta_i$, following its Eq. 13) and the +# precision one ($g_{2,i} = \beta_i$, in its Eq. 21 and Algorithm 2). The dense limit +# settles it: with $\mathbf{Z} = \mathbf{X}$ we have $\mathbf{a}_i = \mathbf{e}_i$ and +# $\mathbf{S}^{-1} = \mathbf{K}_{ff}^{-1} + \operatorname{diag}(\boldsymbol{\beta})$, +# which forces $\boldsymbol{\Lambda}_2 = \operatorname{diag}(\boldsymbol{\beta})$, +# positive. **GPJax stores $\boldsymbol{\Lambda}_2$ in the precision convention**: +# positive semi-definite, no $-\tfrac12$. +# +# The flanked convention is not free. Storing $\boldsymbol{\Lambda}_2$ rather than +# $\bar{\boldsymbol{\Lambda}}_2$ means a round trip through $\mathbf{K}_{zz}^{-1}$ and +# back, which squares its condition number. We measure the consequence in the +# $\mathbf{Z}=\mathbf{X}$ demo below, where it is dramatic entrywise and invisible in +# everything anyone actually reads off the model. The practical rule that falls out: +# never test $\boldsymbol{\Lambda}_2$ entrywise — test $\mathbf{R}$, the moments, the +# bound, or the predictions. + +# %% [markdown] +# ## The tied natural-gradient update +# +# Now the payoff. Split the ELBO into its two terms, with $\boldsymbol{\mu}$ the +# expectation parameter of $q$: +# +# $$\mathcal{L}(\boldsymbol{\eta}) = \mathcal{L}_{\text{ell}}(\boldsymbol{\eta}) - \operatorname{KL}\left[q_{\boldsymbol{\eta}}\,\|\,p_{\boldsymbol{\eta}_0}\right], \qquad \mathcal{L}_{\text{ell}} = \frac{N}{B}\sum_{i\in\mathcal{B}}\mathbb{E}_{q(f_i)}\!\left[\log p(y_i\mid f_i)\right].$$ +# +# For an exponential family the KL between two of its members is +# $\langle\boldsymbol{\eta}-\boldsymbol{\eta}_0,\boldsymbol{\mu}\rangle - A(\boldsymbol{\eta}) + A(\boldsymbol{\eta}_0)$, +# and $\nabla_{\boldsymbol{\eta}}A = \boldsymbol{\mu}$, so the two Jacobian terms +# cancel and +# +# $$\nabla_{\boldsymbol{\mu}}\operatorname{KL}\left[q_{\boldsymbol{\eta}}\,\|\,p_{\boldsymbol{\eta}_0}\right] = \boldsymbol{\eta} - \boldsymbol{\eta}_0 = \boldsymbol{\lambda} .$$ +# +# **The KL's gradient is the stored parameter itself.** Since the natural gradient in +# $\boldsymbol{\eta}$ is the ordinary gradient in $\boldsymbol{\mu}$, the ascent step +# $\boldsymbol{\eta} \leftarrow \boldsymbol{\eta} + \rho\,\nabla_{\boldsymbol{\mu}}\mathcal{L}$ +# collapses to +# +# $$\boldsymbol{\lambda} \;\leftarrow\; (1-\rho)\,\boldsymbol{\lambda} \;+\; \rho\,\nabla_{\boldsymbol{\mu}}\mathcal{L}_{\text{ell}},$$ +# +# a convex combination between where the sites are and where this mini-batch wants +# them. The KL never has to be differentiated at all. Chaining +# $\nabla_{\boldsymbol{\mu}}\mathcal{L}_{\text{ell}}$ through the marginals and +# converting out of the $-\tfrac12$ convention gives the update in stored +# coordinates, +# +# $$\boldsymbol{\lambda}_1 \leftarrow (1-\rho)\boldsymbol{\lambda}_1 + \rho\,\frac{N}{B}\,\mathbf{A}_{\mathcal{B}}\mathbf{g}_1^{\mathcal{B}}, \qquad \boldsymbol{\Lambda}_2 \leftarrow (1-\rho)\boldsymbol{\Lambda}_2 + \rho\,\frac{N}{B}\,\mathbf{A}_{\mathcal{B}}\operatorname{diag}\!\left(\mathbf{g}_2^{\mathcal{B}}\right)\mathbf{A}_{\mathcal{B}}^\top .$$ +# +# The $N/B$ factor is not in the paper's printed update; without it the sites converge +# to $B/N$ of their correct value, since a mini-batch sum is $B/N$ of the full sum in +# expectation. The reference implementation supplies it, and so does GPJax. +# +# Two consequences worth stating separately. First, the update is **affine in the +# stored parameters**, so for $\rho\in[0,1]$ and $\beta_i\ge0$ it can never leave the +# positive semi-definite cone: a convex combination of PSD matrices is PSD. Second, +# $\rho$ **is** the Salimbeni step size $\gamma$ of the natural-gradients notebook, +# not a separate damping coefficient — the display above is +# $\boldsymbol{\eta}\leftarrow\boldsymbol{\eta}+\rho\nabla_{\boldsymbol{\mu}}\mathcal{L}$ +# written out. GPJax accordingly uses one keyword, `natgrad_lr`, for both dispatch +# branches. We check that claim numerically two sections from now. +# +# First, the ingredients. For a Gaussian likelihood $\alpha_i = (y_i - m_i)/\sigma^2$ +# and $\beta_i = 1/\sigma^2$; here are both, by autodiff through +# `expected_log_likelihood`. + +# %% +key, alpha_beta_key = jr.split(key) +check_response = jr.normal(alpha_beta_key, (5, 1)) +check_mean = jnp.linspace(-1.0, 1.0, 5) +check_variance = jnp.linspace(0.2, 0.9, 5) +check_stddev = 0.37 +check_likelihood = gpx.likelihoods.Gaussian(obs_stddev=check_stddev) + + +def total_expected_log_likelihood(mean, variance): + """Summed variational expectation, as a function of the marginal moments.""" + return jnp.sum( + check_likelihood.expected_log_likelihood( + check_response, mean[:, None], variance[:, None] + ) + ) + + +bonnet_alpha, price_derivative = jax.grad( + total_expected_log_likelihood, argnums=(0, 1) +)(check_mean, check_variance) +price_beta = -2.0 * price_derivative + +closed_form_alpha = (check_response.squeeze(-1) - check_mean) / check_stddev**2 +closed_form_beta = jnp.full_like(check_mean, 1.0 / check_stddev**2) + +print( + "max |alpha - (y - m) / sigma^2| : " + f"{jnp.max(jnp.abs(bonnet_alpha - closed_form_alpha)):.3e}" +) +print( + "max |beta - 1 / sigma^2| : " + f"{jnp.max(jnp.abs(price_beta - closed_form_beta)):.3e}" +) +print(f"beta : {price_beta[0]:.6f} (= 1 / {check_stddev}^2)") + +# %% [markdown] +# ## No round trip needed +# +# It is worth being concrete about what the dual step does *not* do. A +# natural-gradient step in the stored parameterisation $(\mathbf{m},\mathbf{L})$ has +# to convert to $\boldsymbol{\eta}$, differentiate the whole ELBO — Cholesky of +# $\mathbf{K}_{zz}$, the conditional, and the KL — apply a Jacobian, and then convert +# back through $\boldsymbol{\theta}$, which costs an inverse and a fresh Cholesky. In +# dual coordinates none of that happens, for two structural reasons: the stored +# coordinates *are* an affine image of $\boldsymbol{\eta}$, so the step is an affine +# step on them; and the target $\nabla_{\boldsymbol{\mu}}\mathcal{L}_{\text{ell}}$ has +# a closed form whose only dependence on $q$ is through the marginals $(m_i, v_i)$, +# which the ELBO computes anyway. +# +# | stage | dual (t-SVGP) | natural gradients on $(\mathbf{m},\mathbf{L})$ | +# |---|---|---| +# | $\operatorname{chol}(\mathbf{K}_{zz})$ | $\mathcal{O}(M^3)$ | $\mathcal{O}(M^3)$ | +# | $\mathbf{A}_{\mathcal{B}} = \mathbf{K}_{zz}^{-1}\mathbf{K}_{zb}$ | $\mathcal{O}(M^2B)$ | $\mathcal{O}(M^2B)$ | +# | covariance factor | $\operatorname{chol}(\mathbf{R})$, $\mathcal{O}(M^3)$ | $\mathbf{S} = \mathbf{L}\mathbf{L}^\top$, $\mathcal{O}(M^3)$ | +# | marginals $(m_i, v_i)$ | $\mathcal{O}(M^2B)$ | $\mathcal{O}(M^2B)$ | +# | $(\boldsymbol{\alpha},\boldsymbol{\beta})$ | one `jax.grad` of a scalar in two $B$-vectors | the same, but inside the full AD tape | +# | gradient assembly | two `einsum`s, $\mathcal{O}(M^2B)$ | reverse-mode AD through chol / conditional / **KL**, plus a Jacobian | +# | $\boldsymbol{\eta}\to\boldsymbol{\xi}$ round trip | **none** | inverse + Cholesky, $\mathcal{O}(M^3)$ | +# +# Same asymptotics, with strictly less work on the dual side of the table. Whether +# that turns into wall-clock depends on how large a share of the iteration the saved +# work was, and we measure it on the banana below rather than assert it here. What is +# certain is the direction of any difference: since the iterates are the same either +# way, the E-step can only differ in time, never in accuracy. Adam et al. measure about +# $5\times$ on MNIST ($N = 70{,}000$, $M = 100$, $B = 200$, ten latent GPs) against +# GPflow's SVGP with natural gradients, with their own caveat that "our implementation +# is not as optimized as SVGP in GPflow". Their sweep over $M$ they describe as "a +# constant factor caused by our computationally cheaper E-step; the effect is +# substantial in most practical settings where $m$ is set below 250". Those are their +# numbers on their hardware. Everything printed below is ours, on the CPU that +# rendered this page. + +# %% [markdown] +# ## Conjugate models: one step is enough +# +# With a Gaussian likelihood the site targets do not depend on $q$ at all: +# +# $$\alpha_i = \frac{y_i - m_i}{\sigma^2}, \quad \beta_i = \frac{1}{\sigma^2} \qquad\Longrightarrow\qquad g_{1,i} = \alpha_i + \beta_i\left(m_i - \mu(x_i)\right) = \frac{y_i - \mu(x_i)}{\sigma^2}, \quad g_{2,i} = \frac{1}{\sigma^2} .$$ +# +# The $m_i$ cancels. So the update is an affine contraction towards a fixed point that +# does not move, and $\rho=1$ lands on it from anywhere in one step: +# +# $$\boldsymbol{\lambda}_1^\star = \frac{1}{\sigma^2}\mathbf{K}_{zz}^{-1}\mathbf{K}_{zx}(\mathbf{y}-\boldsymbol{\mu}_x), \qquad \boldsymbol{\Lambda}_2^\star = \frac{1}{\sigma^2}\mathbf{K}_{zz}^{-1}\mathbf{K}_{zx}\mathbf{K}_{xz}\mathbf{K}_{zz}^{-1},$$ +# +# whereupon +# $\mathbf{R}^\star = \mathbf{K}_{zz} + \sigma^{-2}\mathbf{K}_{zx}\mathbf{K}_{xz}$ is +# exactly the inverse of Titsias' $\boldsymbol{\Sigma}$, and +# +# $$\mathbf{m}^\star = \boldsymbol{\mu}_z + \frac{1}{\sigma^2}\mathbf{K}_{zz}\boldsymbol{\Sigma}\mathbf{K}_{zx}(\mathbf{y}-\boldsymbol{\mu}_x), \qquad \mathbf{S}^\star = \mathbf{K}_{zz}\boldsymbol{\Sigma}\mathbf{K}_{zz}, \qquad \boldsymbol{\Sigma} = \left(\mathbf{K}_{zz} + \sigma^{-2}\mathbf{K}_{zx}\mathbf{K}_{xz}\right)^{-1},$$ +# +# the {cite:t}`titsias2009` optimal $q(\mathbf{u})$ verbatim. The mean function below is +# deliberately non-zero: the sites act on the *centred* process, and the +# $\mathbf{y}-\boldsymbol{\mu}_x$ above is where that shows up. + +# %% +num_data = 200 +noise_stddev = 0.3 +observation_variance = noise_stddev**2 +prior_constant = 0.4 +regression_lengthscale = 0.5 +regression_jitter = 1e-8 + +key, input_key, noise_key = jr.split(key, 3) +regression_inputs = jr.uniform(input_key, (num_data, 1), minval=-3.0, maxval=3.0) +regression_outputs = jnp.sin(2.0 * regression_inputs) + noise_stddev * jr.normal( + noise_key, (num_data, 1) +) +regression_data = gpx.Dataset(X=regression_inputs, y=regression_outputs) + +num_inducing = 20 +regression_inducing = jnp.linspace(-3.0, 3.0, num_inducing).reshape(-1, 1) + + +def conjugate_model(lengthscale): + """The conjugate joint model (prior * likelihood) at a given RBF lengthscale.""" + prior = gpx.gps.Prior( + mean_function=gpx.mean_functions.Constant(jnp.array(prior_constant)), + kernel=jk.RBF(lengthscale=lengthscale), + ) + return prior * gpx.likelihoods.Gaussian(obs_stddev=noise_stddev) + + +def site_family(lengthscale, inducing_inputs, sites=None): + """A dual family, optionally carrying a frozen pair of sites.""" + family = DualVariationalGaussian( + model=conjugate_model(lengthscale), + inducing_inputs=inducing_inputs, + ) + if sites is None: + return family + return eqx.tree_at( + lambda tree: (tree.dual_vector, tree.dual_matrix), + family, + (Real(sites[0]), Real(sites[1])), + ) + + +def moment_family(lengthscale, inducing_inputs, moments): + """A moment family carrying a frozen $(m, S)$.""" + mean, covariance = moments + return VariationalGaussian( + model=conjugate_model(lengthscale), + inducing_inputs=inducing_inputs, + variational_mean=mean, + variational_root_covariance=jnp.linalg.cholesky(covariance), + ) + + +def exact_sites(lengthscale, inducing_inputs, dataset): + """One rho = 1 conjugate step from lambda = 0: the exactly optimal sites.""" + variational, hyper = partition_variational( + site_family(lengthscale, inducing_inputs) + ) + variational, _ = natural_gradient_step( + variational, hyper, dataset, negative_dual_elbo, 1.0 + ) + fitted = paramax.unwrap(eqx.combine(variational, hyper)) + return (fitted.dual_vector, fitted.dual_matrix), fitted.moments() + + +# %% +# The Titsias optimum in closed form, against the same jittered K_zz the family uses. +initial_dual = site_family(regression_lengthscale, regression_inducing) +regression_prior = paramax.unwrap(initial_dual).model.prior +regression_kernel = regression_prior.kernel +regression_mean_function = regression_prior.mean_function + +Kzz = regression_kernel.gram(regression_inducing).as_matrix() +Kzz = Kzz + regression_jitter * jnp.eye(num_inducing) +Kzx = regression_kernel.cross_covariance(regression_inducing, regression_inputs) +centred_outputs = regression_outputs - regression_mean_function(regression_inputs) + +titsias_precision = Kzz + Kzx @ Kzx.T / observation_variance +optimal_mean = ( + regression_mean_function(regression_inducing) + + Kzz + @ jnp.linalg.solve(titsias_precision, Kzx @ centred_outputs) + / observation_variance +) +optimal_covariance = Kzz @ jnp.linalg.solve(titsias_precision, Kzz) + +# The collapsed (Titsias) bound, which the dual ELBO must reproduce at that optimum. +nystrom = Kzx.T @ jnp.linalg.solve(Kzz, Kzx) +marginal_covariance = nystrom + observation_variance * jnp.eye(num_data) +_, marginal_logdet = jnp.linalg.slogdet(marginal_covariance) +marginal_quadratic = centred_outputs.squeeze(-1) @ jnp.linalg.solve( + marginal_covariance, centred_outputs.squeeze(-1) +) +prior_variance_diagonal = jnp.diag( + regression_kernel.gram(regression_inputs).as_matrix() +) +sparsity_gap = jnp.sum(prior_variance_diagonal - jnp.diag(nystrom)) / ( + 2 * observation_variance +) +collapsed_bound = ( + -0.5 * (num_data * jnp.log(2 * jnp.pi) + marginal_logdet + marginal_quadratic) + - sparsity_gap +) + +# %% +# One dual natural-gradient step at rho = 1, from lambda = 0. +dual_variational, dual_hyper = partition_variational(initial_dual) +stepped_variational, loss_before = natural_gradient_step( + dual_variational, dual_hyper, regression_data, negative_dual_elbo, 1.0 +) +stepped_dual = paramax.unwrap(eqx.combine(stepped_variational, dual_hyper)) +stepped_mean, stepped_covariance = stepped_dual.moments() + +# A second step must be a no-op. +twice_stepped_variational, _ = natural_gradient_step( + stepped_variational, dual_hyper, regression_data, negative_dual_elbo, 1.0 +) +twice_stepped_dual = paramax.unwrap(eqx.combine(twice_stepped_variational, dual_hyper)) +twice_stepped_mean, twice_stepped_covariance = twice_stepped_dual.moments() + +stepped_bound = dual_elbo(stepped_dual, regression_data) + +print(f"ELBO before the step : {-loss_before:12.6f}") +print(f"dual_elbo after one step : {float(stepped_bound):12.6f}") +print(f"Titsias collapsed bound : {float(collapsed_bound):12.6f}") +print( + f"max |m - m*| : {jnp.max(jnp.abs(stepped_mean - optimal_mean)):.3e}" +) +print( + "max |S - S*| : " + f"{jnp.max(jnp.abs(stepped_covariance - optimal_covariance)):.3e}" +) +print( + "max |m_2 - m_1| (fixed pt) : " + f"{jnp.max(jnp.abs(twice_stepped_mean - stepped_mean)):.3e}" +) +print( + "max |S_2 - S_1| (fixed pt) : " + f"{jnp.max(jnp.abs(twice_stepped_covariance - stepped_covariance)):.3e}" +) +print(f"collapsed bound - dual_elbo : {float(collapsed_bound - stepped_bound):.12e}") +print( + "N * jitter / (2 sigma^2) : " + f"{num_data * regression_jitter / (2 * observation_variance):.12e}" +) + +# %% [markdown] +# One step from $\boldsymbol{\lambda}=\mathbf{0}$ reproduces the Titsias optimum to +# around $10^{-12}$ in the mean and $10^{-13}$ in the covariance, and a second step +# moves nothing. +# +# The last two printed lines deserve a sentence, because the residual between +# `dual_elbo` and the analytic collapsed bound is not noise — it is +# $N\varepsilon/(2\sigma^2)$ to nine significant figures, where $\varepsilon$ is the +# model's `Prior.jitter`, which is why both are printed to twelve. +# The conditioned sparse posterior adds that jitter to every marginal variance it +# returns, so `elbo` carries the inflation too, and the dual family reproduces it +# deliberately: matching the two objectives to machine precision is worth more than +# matching either to a formula on paper. Lower the jitter and the gap falls +# proportionally. +# +# One thing this demo does *not* show, contrary to a remark in the paper that is easy +# to over-read: the constant $c(\boldsymbol{\theta})$ relating the dual ELBO to +# $\log\mathcal{Z}(\boldsymbol{\theta})$ is **not** zero here. Its value depends on +# which site convention $\mathcal{Z}$ is taken against, and the two have to be paired +# consistently. Against the *normalised projected* site +# $t_i(\mathbf{u}) = \mathcal{N}(y_i \mid \mathbf{a}_i^\top\mathbf{u}, \sigma^2)$, +# $c(\boldsymbol{\theta})$ is minus the Titsias trace term, that is the negated +# `sparsity_gap` computed above, and it vanishes only when +# $\mathbf{Z} = \mathbf{X}$ — the non-sparse case the +# paper's remark actually covers. Against the unnormalised site of the previous +# section, the one this notebook stores, it picks up the site normaliser as well and is +# a different and much larger constant. Either way it is non-zero and +# $\boldsymbol{\theta}$-dependent, which is why GPJax evaluates the bound as +# (variational expectation $-$ KL) rather than as a log-partition function. + +# %% [markdown] +# ## $\rho$ is $\gamma$ +# +# The claim from the tied-update section was that the dual E-step and the +# natural-gradient E-step of the previous notebook are the *same iteration*, not two +# algorithms that happen to converge to the same place. Started from the same $q$, +# with the same rate and the same batches, they should produce the same +# $(\mathbf{m},\mathbf{S})$ at every step, to floating-point noise. There is one +# condition on that, which the derivation left implicit and which the second demo below +# violates: the dual branch clips the per-point curvature $\beta_i$ at `beta_floor`, so +# the identity needs the computed $\beta_i$ to be non-negative. We check the clean case +# first and then go looking for the exception. +# +# Testing that needs a non-conjugate problem — in the conjugate case both branches +# jump to the same optimum at $\rho=1$, which proves nothing about the path — and +# matched initialisations. `DualVariationalGaussian` starts at +# $\boldsymbol{\lambda}=\mathbf{0}$, i.e. $q = p$, so the `VariationalGaussian` here is +# built at $\mathbf{m}=\mathbf{0}$, $\mathbf{S}=\mathbf{K}_{zz}$ rather than at its +# default $\mathbf{S}=\mathbf{I}$. + +# %% +num_logit_data = 200 +num_logit_inducing = 8 +logit_jitter = 1e-8 + +key, logit_input_key, logit_label_key = jr.split(key, 3) +logit_inputs = jr.uniform(logit_input_key, (num_logit_data, 1), minval=-2.0, maxval=2.0) +logit_labels = ( + jr.uniform(logit_label_key, (num_logit_data, 1)) + < jax.nn.sigmoid(3.0 * jnp.sin(2.0 * logit_inputs)) +).astype(jnp.float64) +logit_data = gpx.Dataset(X=logit_inputs, y=logit_labels) +logit_inducing = jnp.linspace(-2.0, 2.0, num_logit_inducing).reshape(-1, 1) + +logit_model = ( + gpx.gps.Prior( + mean_function=gpx.mean_functions.Zero(), + kernel=jk.RBF(lengthscale=0.5, variance=1.7), + ) + * gpx.likelihoods.Bernoulli() +) + +logit_dual = DualVariationalGaussian( + model=logit_model, inducing_inputs=logit_inducing +) +logit_gram = paramax.unwrap(logit_model).prior.kernel.gram( + logit_inducing +).as_matrix() + logit_jitter * jnp.eye(num_logit_inducing) +logit_moments = VariationalGaussian( + model=logit_model, + inducing_inputs=logit_inducing, + variational_mean=jnp.zeros((num_logit_inducing, 1)), + variational_root_covariance=jnp.linalg.cholesky(logit_gram), +) + +shared_bound = float( + dual_elbo(paramax.unwrap(logit_dual), logit_data) + - elbo(paramax.unwrap(logit_moments), logit_data) +) +print(f"cond(K_zz) : {jnp.linalg.cond(logit_gram):.3e}") +print(f"dual_elbo - elbo at the shared init : {shared_bound:.3e}") + + +# %% +def implied_moments(family): + """Return $(m, S)$ for either parameterisation.""" + unwrapped = paramax.unwrap(family) + if isinstance(unwrapped, DualVariationalGaussian): + return unwrapped.moments() + root = unwrapped.variational_root_covariance + return unwrapped.variational_mean, root @ root.T + + +print("rate max |(m, S) gap| over six full-batch steps") +for rate in [0.3, 0.8, 1.0]: + site_partition, site_hyper = partition_variational(logit_dual) + moment_partition, moment_hyper = partition_variational(logit_moments) + worst_gap = 0.0 + for _ in range(6): + site_partition, _ = natural_gradient_step( + site_partition, site_hyper, logit_data, negative_dual_elbo, rate + ) + moment_partition, _ = natural_gradient_step( + moment_partition, moment_hyper, logit_data, negative_elbo, rate + ) + site_mean, site_covariance = implied_moments( + eqx.combine(site_partition, site_hyper) + ) + moment_mean, moment_covariance = implied_moments( + eqx.combine(moment_partition, moment_hyper) + ) + worst_gap = max( + worst_gap, + float(jnp.max(jnp.abs(site_mean - moment_mean))), + float(jnp.max(jnp.abs(site_covariance - moment_covariance))), + ) + print(f"{rate:5.2f} {worst_gap:.3e}") + +# %% +# The same statement one level up, through `fit_natgrads`, with the hyperparameters +# held still by a zero-learning-rate optimiser so that only the E-steps move. +frozen_hyperparameters = dict( + train_data=logit_data, + optim=ox.sgd(0.0), + natgrad_lr=0.8, + num_iters=50, + key=jr.key(1), + verbose=False, +) +_, site_history = gpx.fit_natgrads( + model=logit_dual, objective=negative_dual_elbo, **frozen_hyperparameters +) +_, moment_history = gpx.fit_natgrads( + model=logit_moments, objective=negative_elbo, **frozen_hyperparameters +) +print(f"negative ELBO after 50 E-steps, sites : {float(site_history[-1]):.10f}") +print(f"negative ELBO after 50 E-steps, moments: {float(moment_history[-1]):.10f}") +print( + "max gap over the whole trace : " + f"{jnp.max(jnp.abs(site_history - moment_history)):.3e}" +) + +# %% [markdown] +# The two traces are the same trace. Whatever else is true of the dual +# parameterisation, it is not a different approximation: at $\rho=\gamma$ the E-steps +# coincide, so any difference in a fitted model has to come from somewhere else — the +# M-step, or the one modelling assumption the identity rests on. On this problem every +# $\beta_i$ stays positive, which is that assumption; the banana model in the next +# section drives a point far enough into the tail that GPJax's *computed* $\beta_i$ +# turns negative, the `beta_floor` guard engages, and the two branches then genuinely +# part company. We locate that point and measure the consequence. + +# %% [markdown] +# ## The banana, again +# +# The next cell is the data-generating function from the +# [natural gradients notebook](natgrads.py), +# reproduced character for character — same function body, same `jr.key(42)`, same +# 2000 points — so the problem here is the same problem, point for point, as the one +# there. The initialisation of $q$ differs, for a reason given below. + + +# %% +def make_banana(key, num_points): + """Two-class banana problem with a curved Bayes-optimal boundary.""" + key_latent, key_label = jr.split(key) + latent = jr.uniform(key_latent, (num_points, 2), minval=-3.0, maxval=3.0) + decision = latent[:, 1] - (0.7 * latent[:, 0] ** 2 - 1.5) + probability = jax.nn.sigmoid(3.0 * decision) + labels = (jr.uniform(key_label, (num_points,)) < probability).astype(jnp.float64) + return latent, labels[:, None] + + +banana_key = jr.key(42) +banana_inputs, banana_labels = make_banana(banana_key, 2000) +banana_data = gpx.Dataset(X=banana_inputs, y=banana_labels) + +num_train = 1600 +train_inputs, test_inputs_2d = banana_inputs[:num_train], banana_inputs[num_train:] +train_labels, test_labels = banana_labels[:num_train], banana_labels[num_train:] +banana_train = gpx.Dataset(X=train_inputs, y=train_labels) + +print(f"train / test : {banana_train.n} / {banana_data.n - banana_train.n}") +print(f"class balance : {float(banana_data.y.mean()):.3f}") + +# %% +# Three models over the same inducing grid, all started from q = p. +num_banana_inducing = 50 +inducing_grid = jnp.meshgrid(jnp.linspace(-2.8, 2.8, 10), jnp.linspace(-2.8, 2.8, 5)) +banana_inducing = jnp.stack([axis.ravel() for axis in inducing_grid], axis=1) +banana_jitter = 1e-6 + +banana_model = ( + gpx.gps.Prior( + mean_function=gpx.mean_functions.Zero(), + kernel=jk.RBF(active_dims=[0, 1]), + ) + * gpx.likelihoods.Bernoulli() +) + +banana_gram = paramax.unwrap(banana_model).prior.kernel.gram( + banana_inducing +).as_matrix() + banana_jitter * jnp.eye(num_banana_inducing) +banana_prior_root = jnp.linalg.cholesky(banana_gram) + + +def make_banana_moment_family(): + """A fresh SVGP over the banana data, at q = p.""" + return VariationalGaussian( + model=banana_model, + inducing_inputs=banana_inducing, + variational_mean=jnp.zeros((num_banana_inducing, 1)), + variational_root_covariance=banana_prior_root, + ) + + +banana_dual_family = DualVariationalGaussian( + model=banana_model, + inducing_inputs=banana_inducing, +) +natgrad_family = make_banana_moment_family() +adam_family = make_banana_moment_family() + +print(f"inducing inputs : {banana_inducing.shape}") +print(f"cond(K_zz) : {jnp.linalg.cond(banana_gram):.3e}") + +# %% [markdown] +# All three start at $q = p$, that is $\mathbf{m}=\mathbf{0}$ and +# $\mathbf{S}=\mathbf{K}_{zz}$, which is where a dual family with zero sites already +# is. The natural-gradients notebook used `VariationalGaussian`'s own default +# $\mathbf{S}=\mathbf{I}$ instead, so the curves below start from a slightly different +# place than the ones there; matched initialisations matter more within a comparison +# than across notebooks. + + +# %% +# The rho = gamma check again, on a harder model, this time step by step and carrying +# the diagnostic that explains what happens: Price's curvature beta_i, which the site +# update needs to be non-negative. +def price_curvature(family, data): + """Return the marginal means and $\\beta_i=-2\\,\\partial_{v_i}E_q[\\log p]$.""" + marginal_mean, marginal_variance = family.marginals(data.X) + + def total_expectation(variance): + return jnp.sum( + family.model.likelihood.expected_log_likelihood( + data.y, marginal_mean[:, None], variance[:, None] + ) + ) + + return marginal_mean, -2.0 * jax.grad(total_expectation)(marginal_variance) + + +def six_matched_steps(beta_floor): + """Six rho = 0.8 steps in both branches, from the shared q = p start.""" + site_partition, site_hyper = partition_variational( + DualVariationalGaussian( + model=banana_model, + inducing_inputs=banana_inducing, + ) + ) + moment_partition, moment_hyper = partition_variational(make_banana_moment_family()) + rows = [] + for _ in range(6): + # Measured before the step, at the q both branches currently share. + marginal_mean, curvature = price_curvature( + paramax.unwrap(eqx.combine(site_partition, site_hyper)), banana_train + ) + site_partition, _ = natural_gradient_step( + site_partition, + site_hyper, + banana_train, + negative_dual_elbo, + 0.8, + beta_floor=beta_floor, + ) + moment_partition, _ = natural_gradient_step( + moment_partition, moment_hyper, banana_train, negative_elbo, 0.8 + ) + site_mean, site_covariance = implied_moments( + eqx.combine(site_partition, site_hyper) + ) + moment_mean, moment_covariance = implied_moments( + eqx.combine(moment_partition, moment_hyper) + ) + rows.append( + ( + max( + float(jnp.max(jnp.abs(site_mean - moment_mean))), + float(jnp.max(jnp.abs(site_covariance - moment_covariance))), + ), + int(jnp.sum(curvature < 0)), + float(jnp.min(curvature)), + float(marginal_mean[jnp.argmin(curvature)]), + ) + ) + return rows + + +print("step |(m, S) gap| beta < 0 min beta its marginal mean") +for step, (gap, negative_count, smallest, mean_there) in enumerate( + six_matched_steps(1e-8), start=1 +): + print( + f"{step:4d} {gap:12.3e} {negative_count:4d}/{banana_train.n}" + f" {smallest:+8.4f} {mean_there:+8.3f}" + ) + +banana_gap = max(gap for gap, _, _, _ in six_matched_steps(1e-8)) +unfloored_gap = max(gap for gap, _, _, _ in six_matched_steps(-jnp.inf)) +print(f"\nworst gap, default beta_floor = 1e-8 : {banana_gap:.3e}") +print(f"worst gap, clip disabled (-inf) : {unfloored_gap:.3e}") + +# %% [markdown] +# The two branches part company, and the table says exactly when and why. It is not +# conditioning, and it is not the cancellation in +# $\mathbf{H}_2 = \mathbf{S} + \mathbf{m}\mathbf{m}^\top$ that the moment branch has +# to undo. Disabling the clip — the last line above — brings the same six steps back +# to the noise floor, which rules both of those out: they are unchanged by the value +# of `beta_floor`. +# +# What the $\rho=\gamma$ identity actually needs is a condition the derivation left +# implicit. The site target is built from Price's curvature +# $\beta_i = -2\,\partial_{v_i}\mathbb{E}_{q}\!\left[\log p(y_i\mid f_i)\right]$, and +# the dual branch clips it at `beta_floor` before it enters +# $\boldsymbol{\Lambda}_2$ while the Salimbeni branch never sees it at all. So long as +# $\beta_i \ge 0$ the clip is inert and the two are the same iteration. $\beta_i \ge 0$ +# is guaranteed by log-concavity of $\log p(y\mid f)$ — and GPJax's Bernoulli +# likelihood is not quite log-concave, *as computed*. `inv_probit` squashes its output +# into $[10^{-3},\,1-10^{-3}]$ so that the log stays finite, and that floor flattens +# the tail: $\log p$ as computed has *positive* second derivative for +# $f \lesssim -2.44$, where the exact probit log-likelihood would still be concave. +# A point the model has become confident is mislabelled sits in that region and +# contributes $\beta_i < 0$. Mind the sign: for a $y_i = 0$ point the log-likelihood is +# $\log\Phi(-f_i)$, so the quantity that has to fall below $-2.44$ is $-m_i$, and the +# table's offending point — marginal mean $+2.43$, label $0$ — is exactly on that +# threshold at step five and past it at step six. +# +# That is what the table shows. For the first four steps every $\beta_i$ is positive, +# the clip does nothing, and the branches agree to $10^{-13}$ — the true noise floor of +# this problem. At step five a single training point out of 1600 crosses over, and from +# that step the two are stepping differently by +# $\rho\,\tfrac{N}{B}\,(\beta_{\text{floor}} - \beta_i)\, +# \mathbf{a}_i\mathbf{a}_i^\top$, with +# $\mathbf{a}_i = \mathbf{K}_{zz}^{-1}\mathbf{k}_{zi}$. One rank-one term at +# $\beta_i \approx -0.26$ is enough to move $(\mathbf{m},\mathbf{S})$ by +# $\sim\!10^{-3}$, and the gap compounds over the following step. +# +# Two things follow. The residual is still far below anything visible in the ELBO, +# which is the number either optimiser is steering by, so it does not undermine the +# comparisons below. But the "same iteration" claim is conditional, not absolute, and +# the condition is a property of the *computed* likelihood rather than of the +# mathematical one. + +# %% +# The log-linear ramp of the natural-gradients notebook: 1e-4 -> 1e-1 over K = 100. +num_iterations = 1000 +batch_size = 256 +natgrad_schedule = ox.exponential_decay( + init_value=1e-4, transition_steps=100, decay_rate=1000.0, end_value=1e-1 +) + + +def timed_fit(run): + """Run twice: the first call pays JIT compilation, the second is steady state.""" + model, history = run() + history.block_until_ready() + start = time.perf_counter() + model, history = run() + history.block_until_ready() + return model, history, time.perf_counter() - start + + +shared_settings = dict( + train_data=banana_train, + optim=ox.adam(1e-2), + batch_size=batch_size, + num_iters=num_iterations, + key=jr.key(1), + verbose=False, +) + +# %% +dual_model, dual_history, dual_seconds = timed_fit( + lambda: gpx.fit_natgrads( + model=banana_dual_family, + objective=negative_dual_elbo, + natgrad_lr=natgrad_schedule, + **shared_settings, + ) +) +natgrad_model, natgrad_history, natgrad_seconds = timed_fit( + lambda: gpx.fit_natgrads( + model=natgrad_family, + objective=negative_elbo, + natgrad_lr=natgrad_schedule, + **shared_settings, + ) +) +adam_model, adam_history, adam_seconds = timed_fit( + lambda: gpx.fit(model=adam_family, objective=negative_elbo, **shared_settings) +) + +for name, seconds in [ + ("t-SVGP (dual) + Adam", dual_seconds), + ("natural gradients + Adam", natgrad_seconds), + ("Adam only", adam_seconds), +]: + print( + f"{name:26s}: {seconds:5.2f} s " + f"({1e3 * seconds / num_iterations:.2f} ms / iteration)" + ) + +# %% +smoothing_window = 25 + + +def smooth(history): + """Trailing mean over `smoothing_window` iterations.""" + return jnp.convolve( + history, jnp.ones(smoothing_window) / smoothing_window, mode="valid" + ) + + +smoothed_iterations = jnp.arange(smoothing_window - 1, num_iterations) +curves = [ + ("t-SVGP (dual) + Adam", smooth(dual_history), dual_seconds, cols[2]), + ("Natural gradients + Adam", smooth(natgrad_history), natgrad_seconds, cols[1]), + ("Adam only", smooth(adam_history), adam_seconds, cols[0]), +] + +elbo_floor = 0.95 * min(float(curve.min()) for _, curve, _, _ in curves) +elbo_ceiling = 1.10 * max(float(curve.max()) for _, curve, _, _ in curves) + +fig, axes = plt.subplots(ncols=2, figsize=(10, 3.0), sharey=True) +for name, curve, seconds, colour in curves: + axes[0].plot(smoothed_iterations, curve, color=colour, label=name) + axes[1].plot( + jnp.linspace(0.0, seconds, num_iterations)[smoothing_window - 1 :], + curve, + color=colour, + label=name, + ) +axes[0].set(xlabel="Iteration", yscale="log", ylim=(elbo_floor, elbo_ceiling)) +axes[1].set(xlabel="Wall-clock seconds", yscale="log", ylim=(elbo_floor, elbo_ceiling)) +axes[0].set_ylabel("Negative ELBO (mini-batch)") +clean_legend(axes[0]) +clean_legend(axes[1]) + +for name, curve, seconds, _ in curves: + print(f"{name:26s}: negative ELBO {float(curve[-1]):8.2f} after {seconds:.2f} s") +print( + "max gap between the two natural-gradient curves: " + f"{float(jnp.max(jnp.abs(curves[0][1] - curves[1][1]))):.2f} nats" +) + +# Sentinel above every attainable iteration index, so "never crossed" is +# distinguishable from "crossed on the last iteration". +adam_target = float(curves[2][1][-1]) +never = num_iterations + 1 +for name, curve, seconds, _ in curves[:2]: + crossing = int(jnp.min(jnp.where(curve < adam_target, smoothed_iterations, never))) + if crossing == never: + print(f"{name:26s}: never reaches Adam's final value") + else: + print( + f"{name:26s}: reaches Adam's {num_iterations}-iteration value at " + f"iteration {crossing}, i.e. after " + f"{crossing * seconds / num_iterations:.2f} s of the {curves[2][2]:.2f} s " + "Adam spent" + ) + +# %% [markdown] +# Both natural-gradient runs leave Adam behind per iteration, and both reach Adam's +# thousand-iteration bound in a fraction of the wall-clock time Adam needed for it — +# the crossings are printed above. +# +# What the timings do *not* show is a cheaper dual iteration. On this problem the two +# natural-gradient runs cost within a few percent of each other per iteration, with +# the dual one marginally the more expensive. The round trip the dual parameterisation +# avoids is $\mathcal{O}(M^3)$ at $M=50$, which is nothing next to the +# $\mathcal{O}(BM^2)$ marginals at $B=256$, so there is little to save here in the +# first place. (GPJax's dual step also evaluates the objective once more per iteration +# than it strictly needs to, so that `history[t]` means the same thing in both +# branches — but under `jit`, which is how `fit_natgrads` always runs, XLA normally +# folds that repeat away, and nothing measured here separates the two effects.) Adam +# et al. report their gains at $M=100$ with ten latent GPs and $N=70{,}000$, where the +# constant they save is a much larger share of the total. Take the numbers above as a +# measurement of this configuration on this CPU, not as a refutation or a +# confirmation of theirs. +# +# The two natural-gradient curves do *not* lie on top of each other, and the gap is +# far too large to be the $10^{-3}$ that the `beta_floor` clip contributed a few cells +# ago. Up to that clip their E-steps are still the same iteration; what differs is +# that `fit_natgrads` interleaves an Adam step on the +# kernel hyperparameters and the inducing inputs, and the objective it differentiates +# for that step is `dual_elbo` in one run and `elbo` in the other. Those two have the +# same value and different hyperparameter gradients away from a converged E-step — and +# with $\gamma$ ramping up from $10^{-4}$, the E-step spends most of the first hundred +# iterations far from converged. From iteration 1 onwards the two runs are optimising +# the same model from different hyperparameters, and they never rejoin. +# +# Which way does the divergence go? Here the dual run ends at the *higher*, that is +# worse, negative ELBO of the two; both final values are printed above. That is one +# seed, on a mini-batch bound, with the kernel hyperparameters and all fifty inducing +# inputs moving under a ramping $\gamma$ — the two runs sit at different +# $\boldsymbol{\theta}$ from iteration 1, so this is not a controlled comparison of the +# two M-step objectives and should not be read as one, in either direction. The +# controlled version — frozen inducing inputs, one kernel hyperparameter, matched +# E-steps — is the VEM run at the end of the notebook. What this figure does establish +# is that the choice of M-step objective changes the trajectory by tens of nats, which +# is why the rest of the notebook is about that choice. + +# %% [markdown] +# ## Hyperparameter learning: `dual_elbo` versus `elbo` +# +# Variational EM alternates an E-step, which maximises the ELBO over $q$ at fixed +# $\boldsymbol{\theta}$, with an M-step, which maximises it over $\boldsymbol{\theta}$ +# at fixed $q$. "Fixed $q$" is the ambiguous part. In natural coordinates the E-step +# returns +# $\boldsymbol{\eta}^*_t = \boldsymbol{\eta}_0(\boldsymbol{\theta}_t) + \boldsymbol{\lambda}^*_t$, +# and there are two ways to hold that still: +# +# $$\text{standard:}\quad l(\boldsymbol{\theta}) = \mathcal{L}\big(\underbrace{\boldsymbol{\eta}_0(\boldsymbol{\theta}_t) + \boldsymbol{\lambda}^*_t}_{\text{all frozen}},\ \boldsymbol{\theta}\big), \qquad\qquad \text{dual:}\quad \bar l(\boldsymbol{\theta}) = \mathcal{L}\big(\boldsymbol{\eta}_0(\boldsymbol{\theta}) + \boldsymbol{\lambda}^*_t,\ \boldsymbol{\theta}\big).$$ +# +# `elbo` computes the first, because a `VariationalGaussian` stores +# $(\mathbf{m},\mathbf{L})$ and those are what stay fixed. `dual_elbo` computes the +# second, because a `DualVariationalGaussian` stores the sites, and the prior half of +# $q$ is rebuilt from $\mathbf{K}_{zz}(\boldsymbol{\theta})$ every time the bound is +# evaluated. The intuition is that the sites encode what the *data* said, which is a +# property of the likelihood and should not be re-derived when the kernel moves, +# whereas the prior contribution to $q$ *should* move with the kernel. +# +# That is also why nothing derived from $\boldsymbol{\theta}$ may be cached on the +# family. Caching $(\mathbf{m},\mathbf{S})$ would turn `dual_elbo` back into `elbo` +# under differentiation while leaving every printed value identical — a silent bug of +# the worst kind. +# +# Here is what is actually guaranteed, which is less than the headline suggests: +# +# | claim | status | +# |---|---| +# | $\bar l$ is a valid lower bound on $\log p_{\boldsymbol{\theta}}(\mathbf{y})$ everywhere | **proven** — it is the ELBO at a legitimate Gaussian $q$ | +# | $\bar l(\boldsymbol{\theta}_t) = l(\boldsymbol{\theta}_t)$ | **proven**, exactly, at a converged E-step | +# | $\nabla_{\boldsymbol{\theta}}\bar l(\boldsymbol{\theta}_t) = \nabla_{\boldsymbol{\theta}}l(\boldsymbol{\theta}_t)$ | **proven**, same condition, by the envelope theorem | +# | $\bar l(\boldsymbol{\theta}) \ge l(\boldsymbol{\theta})$ for *all* $\boldsymbol{\theta}$ | proven only when the sites are genuinely $\boldsymbol{\theta}$-free — a conjugate likelihood with its exact sites *and* $\mathbf{Z} = \mathbf{X}$, which is the regime of the second demo below | +# | $\bar l$ is a local upper bound on $l$ | proven in the conjugate case; the paper writes "we can't show this in the non-conjugate setting" | +# | faster EM convergence when non-conjugate | **empirical only** — "exact theoretical reasons behind the speed-ups are currently unknown to us" | +# +# The honest headline: the two bounds agree in value *and* gradient at a converged +# E-step, and the dual M-step objective is less sensitive to +# $\boldsymbol{\theta}_{\text{old}}$, which permits larger M-steps. It is not a +# uniformly tighter bound in the general sparse case. Take the claims in order. + + +# %% +def kernel_gradient(variational, hyper, objective, dataset): + """Gradient of `objective` with respect to the unconstrained kernel parameters.""" + + def loss(hyper): + return objective(paramax.unwrap(eqx.combine(variational, hyper)), dataset) + + gradient = eqx.filter_grad(loss)(hyper) + leaves = jtu.tree_leaves(gradient.model.prior.kernel) + return jnp.concatenate([jnp.atleast_1d(jnp.ravel(leaf)) for leaf in leaves]) + + +print("E-steps max |grad dual_elbo - grad elbo| |grad dual_elbo|") +for num_e_steps in [0, 1, 3, 6, 20, 60]: + site_partition, site_hyper = partition_variational(logit_dual) + moment_partition, moment_hyper = partition_variational(logit_moments) + for _ in range(num_e_steps): + site_partition, _ = natural_gradient_step( + site_partition, site_hyper, logit_data, negative_dual_elbo, 0.8 + ) + moment_partition, _ = natural_gradient_step( + moment_partition, moment_hyper, logit_data, negative_elbo, 0.8 + ) + site_gradient = kernel_gradient( + site_partition, site_hyper, negative_dual_elbo, logit_data + ) + moment_gradient = kernel_gradient( + moment_partition, moment_hyper, negative_elbo, logit_data + ) + print( + f"{num_e_steps:7d} " + f"{float(jnp.max(jnp.abs(site_gradient - moment_gradient))):24.3e} " + f"{float(jnp.max(jnp.abs(site_gradient))):.3e}" + ) + +# %% [markdown] +# The gradients converge onto each other as the E-step converges, which is the +# envelope theorem doing its work: at a stationary $q$ the implicit dependence of the +# prior half of $\boldsymbol{\eta}$ on $\boldsymbol{\theta}$ contributes nothing. Away +# from stationarity the difference is not a rounding effect but a different vector: at +# the shared initialisation the two gradients disagree by as much as the whole +# magnitude of either one. +# +# So the two M-step objectives can only differ when the E-step is incomplete, which in +# practice is always: nobody runs an E-step to convergence between Adam steps. The +# question is whether the difference helps. Freeze the sites at their +# $\boldsymbol{\theta}_t$ values and slide the lengthscale. + +# %% +log_offsets = jnp.linspace(-1.2, 0.6, 61) +frozen_sites, frozen_moments = exact_sites( + regression_lengthscale, regression_inducing, regression_data +) + + +def bound_slice(inducing_inputs, dataset, sites, moments, offsets): + """`dual_elbo` and `elbo` along a log-lengthscale slice, at frozen q.""" + dual_values, moment_values = [], [] + for offset in offsets: + lengthscale = regression_lengthscale * jnp.exp(offset) + dual_values.append( + dual_elbo( + paramax.unwrap(site_family(lengthscale, inducing_inputs, sites)), + dataset, + ) + ) + moment_values.append( + elbo( + paramax.unwrap(moment_family(lengthscale, inducing_inputs, moments)), + dataset, + ) + ) + return jnp.array(dual_values), jnp.array(moment_values) + + +dual_slice, moment_slice = bound_slice( + regression_inducing, regression_data, frozen_sites, frozen_moments, log_offsets +) + +# The slice is not symmetric about theta_t, so print both of its ends. +reference_index = int(jnp.argmin(jnp.abs(log_offsets))) +inducing_spacing = float(regression_inducing[1, 0] - regression_inducing[0, 0]) +shortest_lengthscale = float(regression_lengthscale * jnp.exp(log_offsets[0])) +print("delta log-l dual_elbo elbo") +for label, index in [ + ("left edge ", 0), + ("at theta_t", reference_index), + ("right edge", len(log_offsets) - 1), +]: + print( + f"{label} {float(log_offsets[index]):+5.2f} {float(dual_slice[index]):12.2f} " + f"{float(moment_slice[index]):14.2f}" + ) +print( + f"inducing spacing {inducing_spacing:.3f}, shortest lengthscale on the slice " + f"{shortest_lengthscale:.3f}" +) + +fig, axes = plt.subplots(ncols=2, figsize=(10, 3.0)) +axes[0].plot(log_offsets, dual_slice, color=cols[2], label=r"$\bar l$ (dual_elbo)") +axes[0].plot(log_offsets, moment_slice, color=cols[1], label=r"$l$ (elbo)") +axes[0].axvline(0.0, color="black", linestyle="--", linewidth=1) +axes[0].set( + xlabel=r"$\Delta\log\ell$ from $\theta_t$", + ylabel="Bound (nats)", + ylim=(float(dual_slice.min()) - 40.0, float(dual_slice.max()) + 10.0), + title=f"Sparse, $M = {num_inducing}$", +) +clean_legend(axes[0]) + +for inducing_count, colour in [(5, cols[0]), (10, cols[3]), (20, cols[2])]: + sparse_inducing = jnp.linspace(-3.0, 3.0, inducing_count).reshape(-1, 1) + sparse_sites, sparse_moments = exact_sites( + regression_lengthscale, sparse_inducing, regression_data + ) + sparse_dual, sparse_moment = bound_slice( + sparse_inducing, regression_data, sparse_sites, sparse_moments, log_offsets + ) + gap = sparse_dual - sparse_moment + axes[1].plot(log_offsets, gap, color=colour, label=f"$M = {inducing_count}$") + print( + f"M = {inducing_count:2d}: smallest gap {float(gap.min()):+.4e} nats at " + f"delta log-lengthscale {float(log_offsets[jnp.argmin(gap)]):+.3f}" + ) +axes[1].axhline(0.0, color="black", linestyle="--", linewidth=1) +axes[1].set( + xlabel=r"$\Delta\log\ell$ from $\theta_t$", + ylabel=r"$\bar l - l$ (nats)", + yscale="symlog", + title="Dominance is not uniform when sparse", +) +clean_legend(axes[1]) + +# %% [markdown] +# The left panel is the shape the paper's Fig. 2 is about, and the honest reading of it +# is asymmetric. Both bounds pass through the same point at $\boldsymbol{\theta}_t$ — +# that is the value-equality row of the table, and it holds to $7\times10^{-14}$ here, +# which is the $M=20$ minimum printed above. +# +# To the *right*, at longer lengthscales, $l$ falls off a cliff — the printed +# right-edge values put it more than $10^{5}$ nats below $\bar l$ — because a $q$ +# whose covariance was chosen for one kernel is a bad approximation under another, but +# $\bar l$ has barely moved, since only the data-dependent half of it was frozen. To +# the *left* the two collapse together instead, within a few nats of each other and +# both a couple of hundred nats below their value at $\boldsymbol{\theta}_t$. That is +# not the freezing failing but the sparse approximation itself: by the left-hand edge +# the lengthscale has dropped below half the inducing spacing — both are printed above +# — so $\mathbf{Q}_{ff}$ is a poor stand-in for $\mathbf{K}_{ff}$ and no choice of +# frozen $q$ rescues it. Since the M-step travels rightwards out of a too-short +# lengthscale, the asymmetry is the useful half: in the direction of travel an M-step +# on $\bar l$ can go much further before the bound it is climbing stops being +# informative. +# +# The right panel is the caveat. At $M=20$ the dual bound dominates everywhere we +# looked, up to float64 noise at the crossing point — the printed $M=20$ minimum is +# negative at the $10^{-14}$ level and sits at $\Delta\log\ell = 0$, which is the +# value-equality point itself. At $M=5$ and $M=10$ the gap dips below zero for real — +# the minima are printed above — and the guarantee does not hold. The reason is +# precise: with +# $\mathbf{Z}\neq\mathbf{X}$ the flanked sites +# $\boldsymbol{\lambda}_1^\star = \mathbf{K}_{zz}^{-1}\mathbf{K}_{zx}(\mathbf{y}-\boldsymbol{\mu}_x)/\sigma^2$ +# still depend on $\boldsymbol{\theta}$ through $\mathbf{K}_{zx}$, so freezing them at +# $\boldsymbol{\theta}_t$ makes +# $\boldsymbol{\eta}_0(\boldsymbol{\theta})+\boldsymbol{\lambda}^*_t$ sub-optimal +# elsewhere and the proof's hypothesis fails. Remove the sparsity and the hypothesis +# holds exactly. + +# %% +# Z = X: the sites collapse to (y - mu) / sigma^2 and I / sigma^2, free of theta. +dense_count = 40 +dense_inputs = regression_inputs[:dense_count] +dense_outputs = regression_outputs[:dense_count] +dense_data = gpx.Dataset(X=dense_inputs, y=dense_outputs) + +dense_sites, dense_moments = exact_sites( + regression_lengthscale, dense_inputs, dense_data +) +dense_prior = paramax.unwrap( + site_family(regression_lengthscale, dense_inputs) +).model.prior +dense_gram = dense_prior.kernel.gram( + dense_inputs +).as_matrix() + regression_jitter * jnp.eye(dense_count) +dense_centred = dense_outputs - dense_prior.mean_function(dense_inputs) + +exact_dual_vector = dense_centred / observation_variance +exact_dual_matrix = jnp.eye(dense_count) / observation_variance +flanked_error = jnp.max(jnp.abs(dense_gram @ (dense_sites[0] - exact_dual_vector))) +flanked_scale = jnp.max(jnp.abs(dense_gram @ exact_dual_vector)) + +print(f"cond(K_zz) at Z = X : {jnp.linalg.cond(dense_gram):.3e}") +print( + "max |Lambda_2 - I / sigma^2| : " + f"{jnp.max(jnp.abs(dense_sites[1] - exact_dual_matrix)):.3e} (never test this)" +) +print( + "relative error of K_zz lambda_1 : " + f"{float(flanked_error / flanked_scale):.3e} (test this instead)" +) + +dense_offsets = jnp.linspace(-0.6, 0.6, 41) +dense_dual_slice, dense_moment_slice = bound_slice( + dense_inputs, dense_data, dense_sites, dense_moments, dense_offsets +) + +fig, ax = plt.subplots(figsize=(5.5, 3.2)) +ax.plot(dense_offsets, dense_dual_slice, color=cols[2], label=r"$\bar l$ (dual_elbo)") +ax.plot(dense_offsets, dense_moment_slice, color=cols[1], label=r"$l$ (elbo)") +ax.axvline(0.0, color="black", linestyle="--", linewidth=1) +ax.set( + xlabel=r"$\Delta\log\ell$ from $\theta_t$", + ylabel="Bound (nats)", + yscale="symlog", + title=r"$Z = X$: dominance holds", +) +clean_legend(ax) + +dense_gap = dense_dual_slice - dense_moment_slice +print(f"smallest gap over the slice : {float(dense_gap.min()):+.3e} nats") +print(f"largest gap over the slice : {float(dense_gap.max()):+.3e} nats") + +# %% [markdown] +# With no sparsity gap the dual bound dominates over the whole slice, by several +# orders of magnitude, and stays finite where $l$ collapses through decades on a +# symlog axis. The smallest gap is float64 noise at the crossing point, not a +# violation. +# +# The two diagnostic lines above the plot are the conditioning story promised earlier, +# and they are worth reading together. At $\mathbf{Z}=\mathbf{X}$ the analytic answer +# for the stored matrix is $\boldsymbol{\Lambda}_2 = \mathbf{I}/\sigma^2$, and the +# computed one is wrong by *several units* entrywise, because forming it needs +# $\mathbf{K}_{zz}^{-1}$ twice at a condition number near $10^9$. Yet +# $\mathbf{K}_{zz}\boldsymbol{\lambda}_1$ — the flanked quantity that everything +# downstream actually consumes — is right to nine digits, and the bound plotted above +# is smooth. The error lives in the near-null space of $\mathbf{K}_{zz}$ and is +# annihilated on the way back out. That is a measurement in one configuration and not +# a theorem, which is exactly why the rule is to test $\mathbf{R}$, the moments or the +# predictions, and never $\boldsymbol{\Lambda}_2$ itself. + +# %% [markdown] +# ## The M-step in a loop +# +# Bound slices are static. The claim that actually matters is that a real VEM loop +# gets further with `dual_elbo` as its M-step objective, and that one is empirical: +# the paper says as much. So we run it. Both branches share the same E-step — the same +# iteration, up to the `beta_floor` clip located earlier — and differ only in what the +# M-step differentiates. The inducing inputs are frozen so that only the kernel moves, +# and the lengthscale starts five times too short. + +# %% +expectation_steps = 20 +maximisation_steps = 5 +vem_rounds = 40 +vem_rate = 0.5 +vem_optimiser = ox.adam(5e-2) +initial_lengthscale = 0.25 + + +def freeze_inducing(model): + """Hold the inducing inputs still, so the M-step moves only the kernel.""" + return eqx.tree_at( + lambda tree: tree.inducing_inputs, + model, + paramax.non_trainable(model.inducing_inputs), + ) + + +def vem_joint_model(lengthscale): + return ( + gpx.gps.Prior( + mean_function=gpx.mean_functions.Zero(), + kernel=jk.RBF(active_dims=[0, 1], lengthscale=lengthscale), + jitter=banana_jitter, + ) + * gpx.likelihoods.Bernoulli() + ) + + +vem_gram = paramax.unwrap(vem_joint_model(initial_lengthscale)).prior.kernel.gram( + banana_inducing +).as_matrix() + banana_jitter * jnp.eye(num_banana_inducing) + +vem_dual = freeze_inducing( + DualVariationalGaussian( + model=vem_joint_model(initial_lengthscale), + inducing_inputs=banana_inducing, + ) +) +vem_moments = freeze_inducing( + VariationalGaussian( + model=vem_joint_model(initial_lengthscale), + inducing_inputs=banana_inducing, + variational_mean=jnp.zeros((num_banana_inducing, 1)), + variational_root_covariance=jnp.linalg.cholesky(vem_gram), + ) +) + + +def run_vem(model, objective): + """Alternate `expectation_steps` E-steps with `maximisation_steps` M-steps.""" + variational, hyper = partition_variational(model) + opt_state = vem_optimiser.init(eqx.filter(hyper, eqx.is_array)) + + @eqx.filter_jit + def expectation_step(variational, hyper): + def body(carry, _): + updated, _ = natural_gradient_step( + carry, hyper, banana_train, objective, vem_rate + ) + return updated, None + + return jax.lax.scan(body, variational, None, length=expectation_steps)[0] + + @eqx.filter_jit + def maximisation_step(variational, hyper, opt_state): + def hyper_loss(hyper): + return objective( + paramax.unwrap(eqx.combine(variational, hyper)), banana_train + ) + + def body(carry, _): + hyper, opt_state = carry + loss, gradient = eqx.filter_value_and_grad(hyper_loss)(hyper) + updates, opt_state = vem_optimiser.update( + gradient, opt_state, eqx.filter(hyper, eqx.is_array) + ) + return (eqx.apply_updates(hyper, updates), opt_state), loss + + (hyper, opt_state), losses = jax.lax.scan( + body, (hyper, opt_state), None, length=maximisation_steps + ) + return hyper, opt_state, losses[-1] + + lengthscales, bounds = [], [] + for _ in range(vem_rounds): + variational = expectation_step(variational, hyper) + hyper, opt_state, loss = maximisation_step(variational, hyper, opt_state) + combined = paramax.unwrap(eqx.combine(variational, hyper)) + lengthscales.append(float(combined.model.prior.kernel.lengthscale)) + bounds.append(float(loss)) + return eqx.combine(variational, hyper), jnp.array(lengthscales), jnp.array(bounds) + + +dual_vem_model, dual_lengthscales, dual_bounds = run_vem(vem_dual, negative_dual_elbo) +moment_vem_model, moment_lengthscales, moment_bounds = run_vem( + vem_moments, negative_elbo +) + +# %% +fig, axes = plt.subplots(ncols=2, figsize=(10, 3.0)) +rounds = jnp.arange(1, vem_rounds + 1) +for name, lengthscales, colour in [ + ("M-step on dual_elbo", dual_lengthscales, cols[2]), + ("M-step on elbo", moment_lengthscales, cols[1]), +]: + axes[0].plot(rounds, lengthscales, color=colour, label=name) +axes[0].set(xlabel="VEM round", ylabel=r"Lengthscale $\ell$") +clean_legend(axes[0]) + +# The two bound traces are visually identical at this scale, so plot their difference: +# positive means the dual M-step is the further down the negative ELBO of the two. +bound_lead = moment_bounds - dual_bounds +axes[1].plot(rounds, bound_lead, color=cols[2]) +axes[1].axhline(0.0, color="black", linestyle="--", linewidth=1) +axes[1].set( + xlabel="VEM round", + ylabel="Bound lead to dual_elbo (nats)", + title="Lead of the dual M-step over the standard one", +) + + +def test_metrics(model): + """Held-out accuracy and negative log predictive density.""" + unwrapped = paramax.unwrap(model) + probability = unwrapped.model.likelihood(unwrapped(test_inputs_2d)).mean + labels = test_labels.ravel() + log_density = jnp.mean( + labels * jnp.log(probability) + (1.0 - labels) * jnp.log1p(-probability) + ) + return float(jnp.mean((probability > 0.5) == (labels > 0.5))), float(-log_density) + + +for name, model, lengthscales, bounds in [ + ("dual_elbo", dual_vem_model, dual_lengthscales, dual_bounds), + ("elbo ", moment_vem_model, moment_lengthscales, moment_bounds), +]: + accuracy, nlpd = test_metrics(model) + print( + f"M-step on {name}: lengthscale {float(lengthscales[-1]):.4f}, " + f"negative ELBO {float(bounds[-1]):8.3f}, " + f"test accuracy {accuracy:.4f}, test NLPD {nlpd:.4f}" + ) +print( + f"bound lead to dual_elbo over {vem_rounds} rounds: " + f"smallest {float(bound_lead.min()):+.3f}, largest {float(bound_lead.max()):+.3f}, " + f"final {float(bound_lead[-1]):+.3f} nats" +) +# Sentinel above every attainable round, so "never" stays distinguishable. +never_positive = vem_rounds + 1 +crossing_round = int(jnp.min(jnp.where(bound_lead > 0.0, rounds, never_positive))) +if crossing_round == never_positive: + print("the dual M-step never takes the lead") +else: + print( + f"first round with a positive lead: {crossing_round}; smallest lead from " + f"there on: {float(bound_lead[crossing_round - 1 :].min()):+.3f} nats" + ) + +# %% [markdown] +# The two lengthscale traces sit on top of each other for the first several rounds and +# then separate, with the dual branch ending the longer of the two; both final values +# are printed above. The right panel is the difference of the two bounds rather than +# the two bounds themselves, and that is deliberate: on a negative ELBO of around 294 a +# lead of a nat is invisible, so the traces would be indistinguishable and the sign of +# the difference — the whole question — unreadable. +# +# The sign changes. Over the opening rounds the dual branch is *behind*, by up to a few +# nats, while both are still far from the optimum and moving fast; it takes the lead at +# the printed crossing round and does not give it back, peaking below a nat and ending +# at the printed final value. So the honest reading is not "the dual M-step is +# uniformly ahead". It is that the two branches take different routes to the same +# place: after forty rounds the dual one has the longer lengthscale and the marginally +# better bound, and their held-out NLPDs agree to three decimal places. That is +# consistent with the theory, which promises equality at convergence and says nothing +# about the rate — "exact theoretical reasons behind the speed-ups are currently +# unknown to us". +# +# It is also a soft result on a two-dimensional problem with one kernel +# hyperparameter and fifty fixed inducing points. The regime the paper reports gains +# in — many latent GPs, large $N$, mini-batched, hyperparameters far from their +# optimum — is not this one. Read the demo as a mechanism check rather than as a +# benchmark, and if you want the mechanism in one sentence: at an incomplete E-step +# the two objectives have different hyperparameter gradients, and the dual one is the +# gradient of a function that still knows the prior depends on $\boldsymbol{\theta}$. + +# %% [markdown] +# ## Caveats +# +# * **One latent process.** Everything above assumes $L=1$. The site structure across +# multiple latent GPs is block diagonal only when the variational family is itself +# latent diagonal, and the tied projection has to be re-derived rather than reused +# for a multi-output model. `DualVariationalGaussian` targets the scalar case. +# * **$\beta_i \ge 0$ needs a log-concave likelihood — as *computed*, not as written.** +# Student-$t$ and some heteroscedastic likelihoods are not log-concave at all, and +# for those the target can push $\boldsymbol{\Lambda}_2$ out of the PSD cone. Less +# obviously, GPJax's Bernoulli joins them in the far tails: `inv_probit` clips its +# output into $[10^{-3},\,1-10^{-3}]$, which flattens $\log p$ and makes its second +# derivative positive for $f \lesssim -2.44$, so a confidently mislabelled point +# yields $\beta_i < 0$. The `beta_floor` keyword (default $10^{-8}$) clips +# $\boldsymbol{\beta}$ from below and keeps the step inside the cone. It is *not* a +# no-op for Bernoulli — it is what breaks the $\rho=\gamma$ identity on the banana +# demo above, by $\sim\!10^{-3}$ in $(\mathbf{m},\mathbf{S})$. Note that it clips +# $\boldsymbol{\beta}$, never $\boldsymbol{\Lambda}_2$: the update stays affine, so +# it stays `jit`- and `scan`-safe. +# * **$\rho \in (0,1]$.** The convex-combination guarantee stops at $1$, and beyond it +# the step extrapolates past a target that is only locally valid. `fit_natgrads` +# rejects a larger constant rate for this family at call time. +# * **Flanked storage squares $\operatorname{cond}(\mathbf{K}_{zz})$.** Benign in +# everything measured here at the level of $\mathbf{R}$, the moments and the bound, +# and visibly not benign entrywise in $\boldsymbol{\Lambda}_2$. Never write a test +# against $\boldsymbol{\Lambda}_2$ directly. +# * **The E-step is not a free lunch.** Wherever the computed $\beta_i$ stay +# non-negative it is the *same iteration* as the natural gradient step on +# $(\mathbf{m},\mathbf{L})$, and where they do not the difference is the clip above, +# not a better search direction. Whatever the dual parameterisation buys is either +# wall-clock per iteration or M-step behaviour; none of it is a better $q$ at the +# same $\boldsymbol{\theta}$. +# +# For the geometry the E-step is built on — the Fisher identity, mirror descent, the +# negative-definite cone and the step-size backoff — see the +# [natural gradients notebook](natgrads.py). + +# %% [markdown] +# ## System configuration + +# %% +# %reload_ext watermark +# %watermark -n -u -v -iv -w -a 'Thomas Pinder' diff --git a/docs/examples/graph_kernels.py b/docs/examples/graph_kernels.py index 50bdf60ee..6811745b3 100644 --- a/docs/examples/graph_kernels.py +++ b/docs/examples/graph_kernels.py @@ -174,7 +174,7 @@ def glue(*args, **kwargs): # [`fit_scipy`](#gpjax.fit.fit_scipy). # %% -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() kernel = gpx.kernels.GraphKernel(laplacian=L) prior = gpx.gps.Prior(mean_function=gpx.mean_functions.Zero(), kernel=kernel) posterior = prior * likelihood diff --git a/docs/examples/heteroscedastic_inference.py b/docs/examples/heteroscedastic_inference.py index 3ae5079f7..61f3d428c 100644 --- a/docs/examples/heteroscedastic_inference.py +++ b/docs/examples/heteroscedastic_inference.py @@ -163,8 +163,9 @@ # \mathcal{N}\!\big(y_i \mid f(x_i), \exp(g(x_i))\big), # $$ (eq-heteroscedastic-likelihood) # -# to form the posterior target that we shall approximate variationally. The product -# syntax `signal_prior * likelihood` used below constructs this augmented GP model. +# to form the posterior target that we shall approximate variationally. Because the +# joint model holds *two* priors — one per latent process — it is constructed +# directly as a `HeteroscedasticModel` rather than via the two-operand product. # %% # Signal and noise priors. @@ -176,17 +177,15 @@ mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.RBF(), ) -likelihood = HeteroscedasticGaussian( - num_datapoints=train.n, - noise_prior=noise_prior, - noise_transform=LogNormalTransform(), +likelihood = HeteroscedasticGaussian(noise_transform=LogNormalTransform()) +posterior = gpx.gps.HeteroscedasticModel( + prior=signal_prior, likelihood=likelihood, noise_prior=noise_prior ) -posterior = signal_prior * likelihood # Variational family over both processes. z = jnp.linspace(-3.2, 3.2, 25)[:, None] q = HeteroscedasticVariationalFamily( - posterior=posterior, + model=posterior, inducing_inputs=z, inducing_inputs_g=z, ) @@ -330,12 +329,10 @@ mean_function=gpx.mean_functions.Zero(), kernel=gpx.kernels.RBF(), ) -likelihood_adv = HeteroscedasticGaussian( - num_datapoints=data_adv.n, - noise_prior=noise_prior_adv, - noise_transform=SoftplusTransform(), +likelihood_adv = HeteroscedasticGaussian(noise_transform=SoftplusTransform()) +posterior_adv = gpx.gps.HeteroscedasticModel( + prior=mean_prior, likelihood=likelihood_adv, noise_prior=noise_prior_adv ) -posterior_adv = mean_prior * likelihood_adv # %% # Configure variational family @@ -349,7 +346,7 @@ q_init_g = VariationalGaussianInit(inducing_inputs=z_noise) q_sparse = HeteroscedasticVariationalFamily( - posterior=posterior_adv, + model=posterior_adv, signal_init=q_init_f, noise_init=q_init_g, ) diff --git a/docs/examples/intro_to_kernels.py b/docs/examples/intro_to_kernels.py index 945d0ed13..2bca91465 100644 --- a/docs/examples/intro_to_kernels.py +++ b/docs/examples/intro_to_kernels.py @@ -282,7 +282,7 @@ def forrester(x: Float[Array, "N"]) -> Float[Array, "N"]: # noqa: F821 prior = gpx.gps.Prior(mean_function=mean, kernel=kernel) likelihood = gpx.likelihoods.Gaussian( - num_datapoints=D.n, obs_stddev=jnp.array(1e-3) + obs_stddev=jnp.array(1e-3) ) # Our function is noise-free, so we set the observation noise's standard deviation to a very small value no_opt_posterior = prior * likelihood @@ -561,7 +561,7 @@ def plot_ribbon(ax, x, dist, color): final_kernel = gpx.kernels.SumKernel(kernels=[rbf_kernel, sum_kernel]) prior = gpx.gps.Prior(mean_function=mean, kernel=final_kernel) -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() posterior = prior * likelihood diff --git a/docs/examples/likelihoods_guide.py b/docs/examples/likelihoods_guide.py index c2f9cc9a7..ac57250dc 100644 --- a/docs/examples/likelihoods_guide.py +++ b/docs/examples/likelihoods_guide.py @@ -119,7 +119,7 @@ # argument. # %% -gpx.likelihoods.Gaussian(num_datapoints=D.n) +gpx.likelihoods.Gaussian() # %% [markdown] # ### Likelihood parameters @@ -134,7 +134,7 @@ # this as follows: # %% -gpx.likelihoods.Gaussian(num_datapoints=D.n, obs_stddev=0.5) +gpx.likelihoods.Gaussian(obs_stddev=0.5) # %% [markdown] # @@ -157,7 +157,7 @@ meanf = gpx.mean_functions.Zero() prior = gpx.gps.Prior(kernel=kernel, mean_function=meanf) -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n, obs_stddev=0.1) +likelihood = gpx.likelihoods.Gaussian(obs_stddev=0.1) posterior = prior * likelihood @@ -187,7 +187,7 @@ # Similarly, for a Bernoulli likelihood function, the samples of $y$ would be binary. # %% mystnb={"figure": {"caption": "The same latent draws passed through a Bernoulli likelihood, whose predictive samples are constrained to be binary.", "name": "fig-likelihoods-guide-bernoulli-samples"}} -likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n) +likelihood = gpx.likelihoods.Bernoulli() fig, axes = plt.subplots(ncols=3, nrows=1, figsize=(9, 2)) @@ -268,7 +268,7 @@ # %% z = jnp.linspace(-3.0, 3.0, 10).reshape(-1, 1) -q = gpx.variational_families.VariationalGaussian(posterior=posterior, inducing_inputs=z) +q = gpx.variational_families.VariationalGaussian(model=posterior, inducing_inputs=z) def q_moments(x): @@ -292,7 +292,6 @@ def q_moments(x): # %% lquad = gpx.likelihoods.Gaussian( - num_datapoints=D.n, obs_stddev=jnp.array([0.1]), integrator=gpx.integrators.GHQuadratureIntegrator(num_points=20), ) diff --git a/docs/examples/multioutput.py b/docs/examples/multioutput.py index 01260704f..2ece5d6ff 100644 --- a/docs/examples/multioutput.py +++ b/docs/examples/multioutput.py @@ -155,7 +155,7 @@ meanf = gpx.mean_functions.Zero() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) likelihood = gpx.likelihoods.MultiOutputGaussian( - num_datapoints=N, num_outputs=P, obs_stddev=1.0 + num_outputs=P, obs_stddev=1.0 ) posterior = prior * likelihood @@ -442,7 +442,7 @@ meanf_lcm = gpx.mean_functions.Zero() prior_lcm = gpx.gps.Prior(mean_function=meanf_lcm, kernel=lcm_kernel) likelihood_lcm = gpx.likelihoods.MultiOutputGaussian( - num_datapoints=N_lcm, num_outputs=P_lcm, obs_stddev=1.0 + num_outputs=P_lcm, obs_stddev=1.0 ) posterior_lcm = prior_lcm * likelihood_lcm diff --git a/docs/examples/natgrads.py b/docs/examples/natgrads.py new file mode 100644 index 000000000..8c2741af8 --- /dev/null +++ b/docs/examples/natgrads.py @@ -0,0 +1,1152 @@ +# --- +# jupyter: +# jupytext: +# cell_metadata_filter: -all +# custom_cell_magics: kql +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# jupytext_version: 1.19.1 +# kernelspec: +# display_name: Python 3 +# language: python +# name: python3 +# --- + +# %% [markdown] +# # Natural Gradients +# +# Download this notebook: {nb-download}`natgrads.ipynb` +# +# Variational inference in a sparse Gaussian process asks us to optimise a +# probability distribution $q(\mathbf{u})$, not a point in $\mathbb{R}^P$. Gradient +# descent does not know that: it moves the *storage coordinates* of $q$ — a mean +# vector and a Cholesky factor — as though they lived in flat Euclidean space, and so +# the step it takes depends on how we happened to write the distribution down. The +# natural gradient repairs this by measuring distance between distributions with the +# Fisher information metric, which makes the update invariant to the +# parameterisation. +# +# This notebook implements the recipe of +# {cite:t}`salimbeni2018`, which is what +# `gpjax.fit_natgrads` runs. The remarkable practical point is that for a Gaussian +# process the natural gradient costs *no* Fisher matrix at all: the Fisher information +# turns out to be the Jacobian $\partial\boldsymbol{\eta}/\partial\boldsymbol{\theta}$ +# between two standard coordinate systems, so the natural gradient with respect to one +# of them is the plain gradient with respect to the other. +# +# The route is: +# +# 1. write $q(\mathbf{u})$ in exponential-family form and name its two canonical +# coordinate systems, the natural parameters $\boldsymbol{\theta}$ and the +# expectation parameters $\boldsymbol{\eta}$; +# 2. show that the Fisher matrix is +# $\partial\boldsymbol{\eta}/\partial\boldsymbol{\theta}$, and check it numerically; +# 3. read the step as mirror descent, which explains why $\gamma \le 1$ is special; +# 4. **demo (i)** — a conjugate 1D regression where a single $\gamma=1$ step lands on +# the exact variational optimum (which, at the $M=20$ inducing points used there, is +# indistinguishable from the full GP posterior), while Adam is still crawling after +# two thousand; +# 5. **demo (ii)** — a mini-batched Bernoulli classification benchmark, comparing +# natural gradients + Adam against Adam alone, per iteration *and* per second; +# 6. the failure mode: what a large $\gamma$ does, and how the built-in step-size +# backoff behaves. +# +# If you have not met sparse variational GPs before, read the +# [stochastic sparse GP notebook](uncollapsed_vi.py) +# first — everything below assumes the SVGP evidence lower bound. + +# %% +# Enable Float64 for more stable matrix inversions. +import time + +import equinox as eqx +import jax +from jax import config +import jax.numpy as jnp +import jax.random as jr +from jaxtyping import install_import_hook +import matplotlib as mpl +import matplotlib.pyplot as plt +import optax as ox +import paramax +from utils import clean_legend, use_mpl_style + +config.update("jax_enable_x64", True) + + +with install_import_hook("gpjax", "beartype.beartype"): + import gpjax as gpx + import gpjax.kernels as jk + from gpjax.natural_gradients import ( + expectation_from_moments, + moments_from_expectation, + moments_from_natural, + natural_from_moments, + natural_gradient_step, + partition_variational, + ) + from gpjax.parameters import LowerTriangular, Real + +key = jr.key(123) + +# set the default style for plotting +use_mpl_style() +cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"] + + +def negative_elbo(model, data): + """The loss every fit below minimises: GPJax optimisers descend, so negate.""" + return -gpx.objectives.elbo(model, data) + + +# %% [markdown] +# ## The exponential-family view +# +# The variational distribution over the inducing outputs is +# $q(\mathbf{u}) = \mathcal{N}(\mathbf{m}, \mathbf{S})$ with $\mathbf{m}$ of shape +# $M\times 1$ and $\mathbf{S}$ of shape $M \times M$. Written as an exponential family, +# +# $$\log q(\mathbf{u};\boldsymbol{\theta}) = \log h(\mathbf{u}) + \boldsymbol{\theta}^\top \mathbf{t}(\mathbf{u}) - A(\boldsymbol{\theta}), \qquad h(\mathbf{u}) = (2\pi)^{-M/2},$$ +# +# with sufficient statistics +# $\mathbf{t}(\mathbf{u}) = [\,\mathbf{u},\ \operatorname{vec}(\mathbf{u}\mathbf{u}^\top)\,]$. +# Matching terms gives the **natural parameters** +# +# $$\boldsymbol{\theta}_1 = \mathbf{S}^{-1}\mathbf{m}, \qquad \boldsymbol{\Theta}_2 = -\tfrac{1}{2}\mathbf{S}^{-1} \prec 0,$$ +# +# so that +# $\boldsymbol{\theta}^\top\mathbf{t}(\mathbf{u}) = \mathbf{u}^\top\boldsymbol{\theta}_1 + \mathbf{u}^\top\boldsymbol{\Theta}_2\mathbf{u}$. +# The **expectation parameters** are the mean of the sufficient statistics, +# $\boldsymbol{\eta} = \mathbb{E}_q[\mathbf{t}(\mathbf{u})]$: +# +# $$\boldsymbol{\eta}_1 = \mathbf{m}, \qquad \mathbf{H}_2 = \mathbf{S} + \mathbf{m}\mathbf{m}^\top \succ 0 .$$ +# +# The log normaliser is +# +# $$A(\boldsymbol{\theta}) = -\tfrac{1}{4}\boldsymbol{\theta}_1^\top\boldsymbol{\Theta}_2^{-1}\boldsymbol{\theta}_1 - \tfrac{1}{2}\log\lvert -2\boldsymbol{\Theta}_2\rvert = \tfrac{1}{2}\mathbf{m}^\top\mathbf{S}^{-1}\mathbf{m} + \tfrac{1}{2}\log\lvert\mathbf{S}\rvert,$$ +# +# and differentiating it recovers the expectation parameters, +# $\nabla_{\boldsymbol{\theta}}A(\boldsymbol{\theta}) = \boldsymbol{\eta}$ — the +# standard duality between the two coordinate systems. +# +# There is a third coordinate system in play, the one GPJax actually *stores*: +# $\boldsymbol{\xi} = (\mathbf{m}, \mathbf{L})$ with $\mathbf{S} = \mathbf{L}\mathbf{L}^\top$ +# and $\mathbf{L}$ lower triangular with a positive diagonal. That choice keeps +# $\mathbf{S}$ positive definite under any unconstrained optimiser, but it is a +# storage convention, not a geometry. `gpjax.natural_gradients` exposes the four maps +# that connect the three systems — `expectation_from_moments`, +# `natural_from_moments`, `moments_from_expectation` and `moments_from_natural` — each +# built from Cholesky factors and triangular solves, with no explicit matrix inverse +# anywhere. + +# %% [markdown] +# ## The Fisher information is the Jacobian $\partial\boldsymbol{\eta}/\partial\boldsymbol{\theta}$ +# +# Differentiating $\log q$ twice with respect to $\boldsymbol{\theta}$ kills the +# sufficient statistics and leaves only the log normaliser, so +# +# $$\mathbf{F}_{\boldsymbol{\theta}} := -\mathbb{E}_q\!\left[\nabla^2_{\boldsymbol{\theta}}\log q\right] = \frac{\partial\boldsymbol{\eta}}{\partial\boldsymbol{\theta}} = \nabla^2_{\boldsymbol{\theta}}A(\boldsymbol{\theta}) = \operatorname{Cov}_q\!\left[\mathbf{t}(\mathbf{u})\right].$$ +# +# The Fisher information of an exponential family is simultaneously the Hessian of its +# log normaliser, the Jacobian from natural to expectation parameters, and the +# covariance of its sufficient statistics. The middle equality is the one that pays. +# Let $\ell$ be a loss (for us, the negative ELBO). The chain rule in row-gradient form +# reads +# $\partial\ell/\partial\boldsymbol{\theta} = (\partial\ell/\partial\boldsymbol{\eta})(\partial\boldsymbol{\eta}/\partial\boldsymbol{\theta})$; +# transposing to column gradients and using the self-adjointness of +# $\mathbf{F} = \mathrm{D}\boldsymbol{\eta}$ (it is a Hessian) gives +# $(\partial\ell/\partial\boldsymbol{\theta}) = \mathbf{F}(\partial\ell/\partial\boldsymbol{\eta})$, +# so that +# +# $$\tilde\nabla_{\boldsymbol{\theta}}\ell := \mathbf{F}_{\boldsymbol{\theta}}^{-1}\frac{\partial\ell}{\partial\boldsymbol{\theta}} = \frac{\partial\ell}{\partial\boldsymbol{\eta}} .$$ +# +# **The gradient with respect to the expectation parameters is the natural gradient +# with respect to the natural parameters.** No Fisher matrix is built, and no linear +# system is solved. The update is +# +# $$\boldsymbol{\theta} \leftarrow \boldsymbol{\theta} - \gamma\,\frac{\partial\ell}{\partial\boldsymbol{\eta}},$$ +# +# with $\gamma$ the step size, called `natgrad_lr` in GPJax. +# +# One technical caveat before we check this numerically. The statistic +# $\operatorname{vec}(\mathbf{u}\mathbf{u}^\top)$ has $M^2$ entries, but $q$ depends on +# $\boldsymbol{\Theta}_2$ only through its symmetric part, so in those redundant +# coordinates $\mathbf{F}$ is singular and $\mathbf{F}^{-1}$ is not defined. The fix is +# to work on the space of symmetric matrices with the trace inner product +# $\langle \mathbf{A},\mathbf{B}\rangle = \operatorname{tr}(\mathbf{A}\mathbf{B})$; +# concretely, flatten a symmetric matrix by stacking its lower triangle with the +# strictly off-diagonal entries scaled by $\sqrt{2}$. In those coordinates the +# Euclidean gradient is the correct gradient and $\mathbf{F}$ is symmetric positive +# definite. The production step never forms $\mathbf{F}$ and so never needs any of +# this; we need it only to *verify* the identity. + +# %% +# A small non-conjugate model (Bernoulli likelihood, M = 3) on which to check +# F^{-1} dl/dtheta == dl/deta directly. +key, input_key, label_key, mean_key, root_key = jr.split(key, 5) + +check_inputs = jr.uniform(input_key, (30, 1), minval=-2.0, maxval=2.0) +check_labels = ( + jr.uniform(label_key, (30, 1)) < jax.nn.sigmoid(2.0 * check_inputs) +).astype(jnp.float64) +check_data = gpx.Dataset(X=check_inputs, y=check_labels) + +check_model = ( + gpx.gps.Prior(mean_function=gpx.mean_functions.Zero(), kernel=jk.RBF()) + * gpx.likelihoods.Bernoulli() +) + +num_check_inducing = 3 +check_mean = 0.5 * jr.normal(mean_key, (num_check_inducing, 1)) +check_factor = 0.5 * jr.normal(root_key, (num_check_inducing, num_check_inducing)) +check_root = jnp.linalg.cholesky( + check_factor @ check_factor.T + jnp.eye(num_check_inducing) +) +check_family = gpx.variational_families.VariationalGaussian( + model=check_model, + inducing_inputs=jnp.linspace(-2.0, 2.0, num_check_inducing).reshape(-1, 1), + variational_mean=check_mean, + variational_root_covariance=check_root, +) + + +def symmetric_to_vector(matrix): + """Flatten a symmetric matrix isometrically: lower triangle, sqrt(2) off-diag.""" + size = matrix.shape[0] + scale = jnp.where(jnp.eye(size, dtype=bool), 1.0, jnp.sqrt(2.0)) + rows, columns = jnp.tril_indices(size) + return (matrix * scale)[rows, columns] + + +def vector_to_symmetric(vector, size): + """Invert `symmetric_to_vector`.""" + rows, columns = jnp.tril_indices(size) + lower = jnp.zeros((size, size)).at[rows, columns].set(vector) + diagonal = jnp.diag(jnp.diag(lower)) + strictly_lower = (lower - diagonal) / jnp.sqrt(2.0) + return diagonal + strictly_lower + strictly_lower.T + + +def pack(vector_part, matrix_part): + return jnp.concatenate([vector_part.ravel(), symmetric_to_vector(matrix_part)]) + + +def unpack(flat, size): + return flat[:size].reshape(-1, 1), vector_to_symmetric(flat[size:], size) + + +def loss_at_moments(variational_mean, variational_root_covariance): + trial = eqx.tree_at( + lambda family: (family.variational_mean, family.variational_root_covariance), + check_family, + (Real(variational_mean), LowerTriangular(variational_root_covariance)), + ) + return negative_elbo(paramax.unwrap(trial), check_data) + + +def loss_of_natural(flat): + """The loss as a function of the flattened natural parameters.""" + return loss_at_moments(*moments_from_natural(*unpack(flat, num_check_inducing))) + + +def loss_of_expectation(flat): + """The loss as a function of the flattened expectation parameters.""" + return loss_at_moments(*moments_from_expectation(*unpack(flat, num_check_inducing))) + + +def expectation_of_natural(flat): + """The map whose Jacobian is the Fisher information.""" + moments = moments_from_natural(*unpack(flat, num_check_inducing)) + return pack(*expectation_from_moments(*moments)) + + +flat_natural = pack(*natural_from_moments(check_mean, check_root)) +flat_expectation = pack(*expectation_from_moments(check_mean, check_root)) + +fisher = jax.jacfwd(expectation_of_natural)(flat_natural) +natural_gradient = jnp.linalg.solve(fisher, jax.grad(loss_of_natural)(flat_natural)) +expectation_gradient = jax.grad(loss_of_expectation)(flat_expectation) + +print(f"asymmetry of F : {jnp.max(jnp.abs(fisher - fisher.T)):.3e}") +print(f"smallest eigenvalue of F : {jnp.min(jnp.linalg.eigvalsh(fisher)):.4f}") +print( + "max |F^-1 dl/dtheta - dl/deta| : " + f"{jnp.max(jnp.abs(natural_gradient - expectation_gradient)):.3e}" +) + +# %% [markdown] +# $\mathbf{F}$ is symmetric and positive definite, and the natural gradient obtained by +# solving with it agrees with the plain gradient in expectation coordinates to machine +# precision. Note that the solve just performed lives in the $\operatorname{vec}_s$ +# coordinates introduced above, of dimension $P = M + \tfrac{1}{2}M(M+1)$ — nine at +# $M=3$ — and not in the $M + M^2$ coordinates, where $\mathbf{F}$ is singular. +# Everything from here on uses the right-hand side of the identity, so that +# $\mathcal{O}(P^3) = \mathcal{O}(M^6)$ Fisher solve never happens again. + +# %% [markdown] +# ## Mirror descent +# +# There is a second reading of the same update that explains the role of the step size. +# Let $\Psi = A^*$ be the convex conjugate of the log normaliser — the negative entropy +# of $q$ — so that $\boldsymbol{\theta} = \nabla\Psi(\boldsymbol{\eta})$. Mirror ascent +# on the ELBO $\mathcal{L}$ with mirror map $\Psi$ is +# +# $$\nabla\Psi(\boldsymbol{\eta}_{t+1}) = \nabla\Psi(\boldsymbol{\eta}_t) + \gamma\,\frac{\partial\mathcal{L}}{\partial\boldsymbol{\eta}}, \qquad\text{i.e.}\qquad \boldsymbol{\theta}_{t+1} = \boldsymbol{\theta}_t + \gamma\,\frac{\partial\mathcal{L}}{\partial\boldsymbol{\eta}},$$ +# +# which is precisely the natural-gradient step. The mirror-descent view is the reason +# $\gamma \le 1$ is not an arbitrary convention: as we will see in a moment, the step +# is then a *convex combination* in $\boldsymbol{\theta}$-space between where $q$ is +# and where the current data want it to be. Going beyond $\gamma = 1$ is an +# extrapolation, and extrapolation is what breaks. + +# %% [markdown] +# ## Conjugate models: one step is enough +# +# Suppose the ELBO can be written, for some fixed $\boldsymbol{\lambda}$ that does not +# depend on $q$, +# +# $$\mathcal{L}(q) = \langle\boldsymbol{\lambda},\boldsymbol{\eta}\rangle + \mathbb{H}[q] + c,$$ +# +# that is, $\mathbb{E}_q[\log p(\mathbf{y},\mathbf{u})]$ is affine in +# $\boldsymbol{\eta}$. This is exactly the conditionally-conjugate case: a Gaussian +# likelihood. Since +# $\mathbb{H}[q] = -\mathbb{E}_q[\log h] - \boldsymbol{\theta}^\top\boldsymbol{\eta} + A(\boldsymbol{\theta})$ +# and $\partial A/\partial\boldsymbol{\theta} = \boldsymbol{\eta}$, the two Jacobian +# terms cancel and $\partial\mathbb{H}/\partial\boldsymbol{\eta} = -\boldsymbol{\theta}$. +# Therefore +# +# $$\frac{\partial\mathcal{L}}{\partial\boldsymbol{\eta}} = \boldsymbol{\lambda} - \boldsymbol{\theta} \qquad\Longrightarrow\qquad \boldsymbol{\theta}_{\text{new}} = (1-\gamma)\,\boldsymbol{\theta} + \gamma\,\boldsymbol{\lambda},$$ +# +# and $\gamma = 1$ gives $\boldsymbol{\theta}_{\text{new}} = \boldsymbol{\lambda} = \boldsymbol{\theta}^\star$ +# **in one step, from any starting point**. This is Sato's (2001) observation that +# natural-gradient ascent at unit step size *is* the classical variational +# fixed-point update; for the SVGP it recovers the {cite:t}`titsias2009` optimum. +# +# Let us watch it happen. + +# %% +# Demo (i): 1D conjugate regression. +num_data = 200 +noise_stddev = 0.3 + +key, input_key, noise_key = jr.split(key, 3) +regression_inputs = jr.uniform(input_key, (num_data, 1), minval=-3.0, maxval=3.0) +regression_signal = jnp.sin(2.0 * regression_inputs) +regression_outputs = regression_signal + noise_stddev * jr.normal( + noise_key, regression_signal.shape +) +regression_data = gpx.Dataset(X=regression_inputs, y=regression_outputs) + +num_inducing = 20 +regression_inducing = jnp.linspace(-3.0, 3.0, num_inducing).reshape(-1, 1) +test_inputs = jnp.linspace(-3.2, 3.2, 300).reshape(-1, 1) + +# %% +# A conjugate SVGP, deliberately initialised a long way from its optimum. The joint +# model is prior * likelihood; the variational family approximates its posterior. +regression_model = gpx.gps.Prior( + mean_function=gpx.mean_functions.Constant(), + kernel=jk.RBF(lengthscale=0.5), + jitter=1e-8, +) * gpx.likelihoods.Gaussian(obs_stddev=noise_stddev) + +key, bad_mean_key, bad_root_key = jr.split(key, 3) +bad_mean = jr.normal(bad_mean_key, (num_inducing, 1)) +bad_factor = 0.3 * jr.normal(bad_root_key, (num_inducing, num_inducing)) +bad_root = jnp.linalg.cholesky(bad_factor @ bad_factor.T + 0.5 * jnp.eye(num_inducing)) + +initial_family = gpx.variational_families.WhitenedVariationalGaussian( + model=regression_model, + inducing_inputs=regression_inducing, + variational_mean=bad_mean, + variational_root_covariance=bad_root, +) + +# %% [markdown] +# We use the **whitened** family here, which reparameterises +# $\mathbf{u} = \boldsymbol{\mu}_z + \mathbf{L}_z\mathbf{v}$ with +# $\mathbf{L}_z\mathbf{L}_z^\top = \mathbf{K}_{zz}$ and puts a +# $\mathcal{N}(\mathbf{0},\mathbf{I})$ prior on $\mathbf{v}$. The natural-gradient +# machinery is untouched by this — $q(\mathbf{v})$ belongs to the same exponential +# family, and the whitening enters only through `prior_kl` and `predict`, which the +# loss calls polymorphically. Numerically it helps a great deal, because +# $\mathbf{m}_w$ and $\mathbf{S}_w$ are $\mathcal{O}(1)$ regardless of the kernel +# scale, and the conjugate optimum satisfies +# $\mathbf{S}_w^\star \preceq \mathbf{I}$. +# +# For the whitened family the closed-form optimum is, with +# $\mathbf{A}_w = \mathbf{K}_{xz}\mathbf{L}_z^{-\top}$ and +# $\sigma^2$ the observation variance, +# +# $$\boldsymbol{\Lambda}_w = \mathbf{I}_M + \sigma^{-2}\mathbf{A}_w^\top\mathbf{A}_w, \qquad \mathbf{b}_w = \sigma^{-2}\mathbf{A}_w^\top(\mathbf{y}-\boldsymbol{\mu}_x),$$ +# $$\mathbf{S}_w^\star = \boldsymbol{\Lambda}_w^{-1}, \qquad \mathbf{m}_w^\star = \boldsymbol{\Lambda}_w^{-1}\mathbf{b}_w .$$ + +# %% +unwrapped_initial = paramax.unwrap(initial_family) +kernel = unwrapped_initial.model.prior.kernel +mean_function = unwrapped_initial.model.prior.mean_function + +Kzz = kernel.gram(regression_inducing).as_matrix() +Kzz = Kzz + initial_family.model.prior.jitter * jnp.eye(num_inducing) +Lz = jnp.linalg.cholesky(Kzz) +Kzx = kernel.cross_covariance(regression_inducing, regression_inputs) +whitened_design = jax.scipy.linalg.solve_triangular(Lz, Kzx, lower=True).T + +observation_variance = noise_stddev**2 +whitened_precision = ( + jnp.eye(num_inducing) + whitened_design.T @ whitened_design / observation_variance +) +whitened_shift = ( + whitened_design.T + @ (regression_outputs - mean_function(regression_inputs)) + / observation_variance +) +optimal_covariance = jnp.linalg.inv(whitened_precision) +optimal_mean = jnp.linalg.solve(whitened_precision, whitened_shift) + +# The ELBO at the closed-form optimum, used below as the reference for both methods. +optimal_family = eqx.tree_at( + lambda family: (family.variational_mean, family.variational_root_covariance), + initial_family, + (Real(optimal_mean), LowerTriangular(jnp.linalg.cholesky(optimal_covariance))), +) +reference_elbo = float( + gpx.objectives.elbo(paramax.unwrap(optimal_family), regression_data) +) +print(f"ELBO at the closed-form optimum: {reference_elbo:.6f}") + +# %% +# One natural-gradient step at gamma = 1. +variational_partition, hyper_partition = partition_variational(initial_family) +stepped_partition, loss_before = natural_gradient_step( + variational_partition, + hyper_partition, + regression_data, + negative_elbo, + 1.0, + map_jitter=0.0, +) +stepped_family = eqx.combine(stepped_partition, hyper_partition) + +unwrapped_stepped = paramax.unwrap(stepped_family) +stepped_mean = unwrapped_stepped.variational_mean +stepped_root = unwrapped_stepped.variational_root_covariance +stepped_covariance = stepped_root @ stepped_root.T + +stepped_elbo = float(gpx.objectives.elbo(unwrapped_stepped, regression_data)) + +# A second step from the same place must be a fixed point. +twice_stepped_partition, _ = natural_gradient_step( + stepped_partition, + hyper_partition, + regression_data, + negative_elbo, + 1.0, + map_jitter=0.0, +) +twice_stepped_mean = paramax.unwrap( + eqx.combine(twice_stepped_partition, hyper_partition) +).variational_mean + +print(f"ELBO before the step : {-loss_before:12.6f}") +print(f"ELBO after one gamma=1 step : {stepped_elbo:12.6f}") +print(f"ELBO at the closed-form optimum: {reference_elbo:12.6f}") +print( + "max |m_1 - m*| : " + f"{jnp.max(jnp.abs(stepped_mean - optimal_mean)):.3e}" +) +print( + "max |S_1 - S*| : " + f"{jnp.max(jnp.abs(stepped_covariance - optimal_covariance)):.3e}" +) +print( + "max |m_2 - m_1| (fixed point) : " + f"{jnp.max(jnp.abs(twice_stepped_mean - stepped_mean)):.3e}" +) + +# %% [markdown] +# One step, from a random initialisation, reproduces the closed-form optimum to +# $\sim10^{-13}$ — the float64 noise floor for a problem of this size — and a second +# step moves nothing. Note the +# `map_jitter=0.0`: the jitter used inside the +# $\boldsymbol{\theta}\leftrightarrow\boldsymbol{\xi}$ maps is a *bias*, not a +# rounding effect, since +# $(\mathbf{S}^{-1}+\varepsilon\mathbf{I})^{-1} = \mathbf{S} - \varepsilon\mathbf{S}^2 + \mathcal{O}(\varepsilon^2)$. +# It defaults to zero in `fit_natgrads` for that reason, and is deliberately *not* +# inherited from the model's `Prior.jitter`, which is a different quantity applied to +# $\mathbf{K}_{zz}$. +# +# Because this model is conjugate, we can also compare the one-step posterior against +# the exact GP posterior, obtained by conditioning the joint model on the data with no +# inducing-point approximation. + +# %% +exact_posterior = paramax.unwrap(regression_model).condition(regression_data) +exact_predictive = exact_posterior(test_inputs) +exact_mean = exact_predictive.mean +exact_stddev = jnp.sqrt(exact_predictive.variance) + +fig, axes = plt.subplots(ncols=2, figsize=(10, 3.0), sharey=True) +for ax, family, title in [ + (axes[0], unwrapped_initial, "Initialisation"), + (axes[1], unwrapped_stepped, "After one $\\gamma=1$ natural-gradient step"), +]: + predictive = family(test_inputs) + predictive_mean = predictive.mean + predictive_stddev = jnp.sqrt(predictive.variance) + ax.scatter( + regression_inputs, + regression_outputs, + alpha=0.2, + s=8, + color=cols[0], + label="Observations", + ) + ax.plot( + test_inputs, exact_mean, color="black", linestyle="--", label="Exact posterior" + ) + ax.fill_between( + test_inputs.flatten(), + exact_mean - 2 * exact_stddev, + exact_mean + 2 * exact_stddev, + alpha=0.15, + color="black", + ) + ax.plot(test_inputs, predictive_mean, color=cols[1], label="Variational $q$") + ax.fill_between( + test_inputs.flatten(), + predictive_mean - 2 * predictive_stddev, + predictive_mean + 2 * predictive_stddev, + alpha=0.3, + color=cols[1], + ) + ax.set(xlabel=r"$x$", title=title, ylim=(-3.0, 3.0)) + clean_legend(ax) +axes[0].set_ylabel(r"$f(x)$") + +print( + "max |sparse mean - exact mean| : " + f"{jnp.max(jnp.abs(unwrapped_stepped(test_inputs).mean - exact_mean)):.3e}" +) + +# %% [markdown] +# The right-hand panel is the point of the whole method: a single natural-gradient step +# has taken a deliberately absurd $q$ onto the sparse variational optimum, which for +# $M=20$ inducing points on this problem is not distinguishable by eye from the exact +# posterior. The printed maximum is taken over the whole test grid $[-3.2, 3.2]$ and is +# attained at its edge, past the last inducing input; restricted to the data range +# $[-3, 3]$ the two means agree roughly ten times more closely again. Both gaps are a +# fraction of a percent of the panel height, and both are a property of the sparse +# approximation, not of the optimiser. +# +# Now the comparison. We freeze every hyperparameter with `paramax.non_trainable` — so +# that both methods are solving the *same* problem, namely finding the best +# $(\mathbf{m},\mathbf{L})$ for a fixed kernel — and run Adam on the variational +# parameters from the same bad initialisation. + +# %% +frozen_family = eqx.combine( + variational_partition, paramax.non_trainable(hyper_partition) +) + +adam_iterations = 2000 +_, adam_history = gpx.fit( + model=frozen_family, + objective=negative_elbo, + train_data=regression_data, + optim=ox.adam(1e-2), + num_iters=adam_iterations, + key=jr.key(0), + verbose=False, +) + +adam_gap = jnp.asarray(adam_history) + reference_elbo +natgrad_gap = reference_elbo - stepped_elbo +iteration_index = jnp.arange(adam_gap.size) +for tolerance in [10.0, 1.0, 0.1]: + first_hit = jnp.min(jnp.where(adam_gap < tolerance, iteration_index, adam_gap.size)) + reached = "never" if int(first_hit) == adam_gap.size else f"{int(first_hit)}" + print( + f"Adam iterations to come within {tolerance:5.1f} nats of the optimum: " + f"{reached}" + ) +print( + f"Adam ELBO gap after {adam_iterations} iterations : {float(adam_gap[-1]):.3e} nats" +) +print(f"Natural-gradient ELBO gap after 1 iteration: {natgrad_gap:.3e} nats") + +# %% +fig, axes = plt.subplots(ncols=2, figsize=(10, 3.0)) + +axes[0].plot( + iteration_index + 1, -adam_history, color=cols[0], label="Adam on $(m, L)$" +) +axes[0].axhline(reference_elbo, color="black", linestyle="--", label="Exact optimum") +axes[0].scatter( + [1], [stepped_elbo], color=cols[1], zorder=5, s=45, label="Natural gradient, 1 step" +) +axes[0].set( + xscale="log", + xlabel="Iteration", + ylabel="ELBO", + ylim=(reference_elbo - 250, reference_elbo + 25), +) +clean_legend(axes[0]) + +axes[1].plot(iteration_index + 1, adam_gap, color=cols[0], label="Adam on $(m, L)$") +axes[1].scatter( + [1], [natgrad_gap], color=cols[1], zorder=5, s=45, label="Natural gradient, 1 step" +) +axes[1].set( + xscale="log", yscale="log", xlabel="Iteration", ylabel="ELBO gap to optimum (nats)" +) +clean_legend(axes[1]) + +# %% [markdown] +# Read the right-hand panel rather than the left. On log-log axes Adam's gap barely +# bends over the first few tens of iterations and then falls faster and faster, its +# slope steepest of all over the final few hundred — the opposite of the usual "fast +# start, long crawl" picture. That shape is the optimiser's, not the problem's: Adam +# normalises its step, so each coordinate moves by at most the learning rate however +# large the gradient is, and from an initialisation this bad it is the *distance* to be +# travelled that binds, not the gradient. The printed numbers say the same thing: more +# than a thousand iterations merely to come within ten nats of the optimum, and after +# two thousand it is still several nats short and still descending, while the single +# natural-gradient step closed the gap to around $10^{-14}$ nats. Adam is converging; +# it is simply converging in coordinates that put the optimum a long way away. The +# natural gradient never travels that distance, because the Fisher metric rescales it. +# +# Two caveats before this is oversold. The hyperparameters were frozen, so this is the +# problem natural gradients are best at: a pure variational optimisation. And the +# advantage rests on conjugacy, which is what makes $\gamma=1$ a solve rather than a +# step. Neither holds in the next demo. + +# %% [markdown] +# ## Non-conjugate models: ramping $\gamma$ +# +# Outside conjugacy, $\mathbb{E}_q[\log p(\mathbf{y}\mid\mathbf{u})]$ is no longer +# affine in $\boldsymbol{\eta}$, so $\gamma=1$ is no longer a solve — it is a large +# step along a direction that was only computed locally. Salimbeni et al. find +# experimentally that "the initial natural gradient step size is a small value that is +# parameterization and likelihood dependent, but then increases to $\gamma = 1$", and +# in the stochastic setting they adopt a two-phase schedule: a log-linear ramp +# +# $$\gamma_t = \gamma_{\text{init}}\left(\frac{\gamma_{\text{final}}}{\gamma_{\text{init}}}\right)^{t/K} \quad (t < K), \qquad \gamma_t = \gamma_{\text{final}} \quad (t \ge K).$$ +# +# Their reported settings are $\gamma_{\text{init}}=10^{-4}$, +# $\gamma_{\text{final}}=10^{-1}$ with $K$ between 5 and 40 for UCI-scale problems at +# batch size 256, and $\gamma_{\text{init}}=10^{-6}$, +# $\gamma_{\text{final}}=2\times10^{-2}$, $K=2000$ for MNIST at batch size 1024, always +# with $\gamma^{\text{Adam}} = 10^{-2}$ on the hyperparameters. Their conclusion is +# that "the success of the method relies on $\gamma$ increasing to a reasonably large +# value ($\approx 0.1$) sufficiently quickly ($<1000$ iterations)". +# +# We use $K = 100$ below. Their $K$ is dataset-dependent — 5 for the smaller UCI sets, +# 40 for NAVAL, 2000 for MNIST — and $100$ buys a little extra cone headroom (see the +# last section) at this $M$ from the default $\mathbf{m}=\mathbf{0}$, +# $\mathbf{S}=\mathbf{I}$ start, while still satisfying their own $<1000$-iteration +# criterion. +# +# Why does $\gamma < 1$ help when mini-batching? The $N/B$ rescaling inside the ELBO +# makes the stochastic gradient unbiased, and because +# $\boldsymbol{\theta}_{\text{new}} = \boldsymbol{\theta} - \gamma\hat{\mathbf{g}}$ is +# affine in $\hat{\mathbf{g}}$, $\boldsymbol{\theta}_{\text{new}}$ is unbiased for the +# full-batch update at every $\gamma$, including $\gamma=1$. What degrades is +# *variance*. The step is always a combination +# $\boldsymbol{\theta}_{\text{new}} = (1-\gamma)\,\boldsymbol{\theta} + \gamma\,\boldsymbol{\theta}^{\text{tgt}}$ +# — the failure-modes section below writes its second block out explicitly — but +# outside conjugacy $\boldsymbol{\theta}^{\text{tgt}}$ is not a fixed optimum. It +# depends on the current $q$ as well as on the current mini-batch: it is where one +# fixed-point iteration from *here* would land, and it moves as $q$ moves. At +# $\gamma=1$ the step discards $\boldsymbol{\theta}_t$ entirely and jumps onto that +# noisy, moving target, so nothing averages the mini-batch noise out of it. Taking +# $\gamma<1$ makes the update an exponential moving average in $\boldsymbol{\theta}$ +# towards the target, and that is where the variance reduction comes from. A second, +# smaller effect compounds it: $\boldsymbol{\theta}\mapsto(\mathbf{m},\mathbf{S})$ is +# nonlinear, so unbiasedness in $\boldsymbol{\theta}$ does not survive the conversion +# back to moments. +# +# Time for a harder problem. + + +# %% +def make_banana(key, num_points): + """Two-class banana problem with a curved Bayes-optimal boundary.""" + key_latent, key_label = jr.split(key) + latent = jr.uniform(key_latent, (num_points, 2), minval=-3.0, maxval=3.0) + decision = latent[:, 1] - (0.7 * latent[:, 0] ** 2 - 1.5) + probability = jax.nn.sigmoid(3.0 * decision) + labels = (jr.uniform(key_label, (num_points,)) < probability).astype(jnp.float64) + return latent, labels[:, None] + + +banana_key = jr.key(42) +banana_inputs, banana_labels = make_banana(banana_key, 2000) +banana_data = gpx.Dataset(X=banana_inputs, y=banana_labels) + +num_train = 1600 +train_inputs, test_inputs_2d = banana_inputs[:num_train], banana_inputs[num_train:] +train_labels, test_labels = banana_labels[:num_train], banana_labels[num_train:] +banana_train = gpx.Dataset(X=train_inputs, y=train_labels) + +print(f"train / test : {banana_train.n} / {banana_data.n - banana_train.n}") +print(f"class balance : {float(banana_data.y.mean()):.3f}") + +# %% +boundary_inputs = jnp.linspace(-3.0, 3.0, 200) +boundary_outputs = 0.7 * boundary_inputs**2 - 1.5 + +fig, ax = plt.subplots(figsize=(5.5, 3.4)) +for label, colour, name in [(0.0, cols[0], "$y = 0$"), (1.0, cols[1], "$y = 1$")]: + mask = banana_labels.ravel() == label + ax.scatter( + banana_inputs[mask, 0], + banana_inputs[mask, 1], + s=6, + alpha=0.4, + color=colour, + label=name, + ) +ax.plot( + boundary_inputs, + boundary_outputs, + color="black", + linestyle="--", + label="Bayes-optimal boundary", +) +ax.set(xlabel=r"$x_1$", ylabel=r"$x_2$", ylim=(-3.1, 3.1), title="The banana problem") +clean_legend(ax) + +# %% +# Two identical models, built from the same arrays, so the comparison is fair. +num_banana_inducing = 50 +inducing_grid = jnp.meshgrid(jnp.linspace(-2.8, 2.8, 10), jnp.linspace(-2.8, 2.8, 5)) +banana_inducing = jnp.stack([axis.ravel() for axis in inducing_grid], axis=1) + +banana_model = ( + gpx.gps.Prior( + mean_function=gpx.mean_functions.Zero(), kernel=jk.RBF(active_dims=[0, 1]) + ) + * gpx.likelihoods.Bernoulli() +) + + +def make_banana_family(): + """A fresh SVGP over the banana data, at the default m = 0, S = I.""" + return gpx.variational_families.VariationalGaussian( + model=banana_model, inducing_inputs=banana_inducing + ) + + +natgrad_family = make_banana_family() +adam_family = make_banana_family() + +print(f"inducing inputs: {banana_inducing.shape}") + +# %% +# The log-linear ramp, 1e-4 -> 1e-1 over K = 100 iterations, as an Optax schedule. +num_iterations = 1000 +batch_size = 256 +natgrad_schedule = ox.exponential_decay( + init_value=1e-4, transition_steps=100, decay_rate=1000.0, end_value=1e-1 +) +print( + "gamma at iterations 0, 50, 100, 999: " + + ", ".join(f"{float(natgrad_schedule(t)):.2e}" for t in [0, 50, 100, 999]) +) + + +def timed_fit(run): + """Run twice: the first call pays JIT compilation, the second is steady state.""" + model, history = run() + history.block_until_ready() + start = time.perf_counter() + model, history = run() + history.block_until_ready() + return model, history, time.perf_counter() - start + + +# %% +natgrad_model, natgrad_history, natgrad_seconds = timed_fit( + lambda: gpx.fit_natgrads( + model=natgrad_family, + objective=negative_elbo, + train_data=banana_train, + optim=ox.adam(1e-2), + natgrad_lr=natgrad_schedule, + batch_size=batch_size, + num_iters=num_iterations, + key=jr.key(1), + verbose=False, + ) +) +print( + f"natural gradients + Adam : {natgrad_seconds:.2f} s " + f"({1e3 * natgrad_seconds / num_iterations:.2f} ms / iteration)" +) + +# %% +adam_model, adam_banana_history, adam_seconds = timed_fit( + lambda: gpx.fit( + model=adam_family, + objective=negative_elbo, + train_data=banana_train, + optim=ox.adam(1e-2), + batch_size=batch_size, + num_iters=num_iterations, + key=jr.key(1), + verbose=False, + ) +) +print( + f"Adam only : {adam_seconds:.2f} s " + f"({1e3 * adam_seconds / num_iterations:.2f} ms / iteration)" +) + +# %% [markdown] +# Both runs use `ox.adam(1e-2)` on the kernel hyperparameters and the inducing inputs, +# so the only difference is how $(\mathbf{m},\mathbf{L})$ move. Timings are steady +# state: each fit is called twice and only the second call is timed, so JIT +# compilation is excluded from both. They were measured on CPU while executing this +# notebook, and will differ on your machine. + +# %% +smoothing_window = 25 + + +def smooth(history): + """Trailing mean over `smoothing_window` iterations.""" + return jnp.convolve( + history, jnp.ones(smoothing_window) / smoothing_window, mode="valid" + ) + + +smoothed_iterations = jnp.arange(smoothing_window - 1, num_iterations) +smoothed_natgrad = smooth(natgrad_history) +smoothed_adam = smooth(adam_banana_history) + +# Derive the axis limits from the curves, so nothing is silently clipped on a machine +# whose run lands somewhere else. +elbo_floor = 0.95 * float(jnp.minimum(smoothed_natgrad.min(), smoothed_adam.min())) +elbo_ceiling = 1.10 * float(jnp.maximum(smoothed_natgrad.max(), smoothed_adam.max())) + +fig, axes = plt.subplots(ncols=2, figsize=(10, 3.0), sharey=True) +for ax, horizontal, xlabel in [ + (axes[0], smoothed_iterations, "Iteration"), + ( + axes[1], + jnp.linspace(0.0, natgrad_seconds, num_iterations)[smoothing_window - 1 :], + "Wall-clock seconds", + ), +]: + ax.plot( + horizontal, smoothed_natgrad, color=cols[1], label="Natural gradients + Adam" + ) + ax.set(xlabel=xlabel, yscale="log", ylim=(elbo_floor, elbo_ceiling)) +axes[0].plot(smoothed_iterations, smoothed_adam, color=cols[0], label="Adam only") +axes[1].plot( + jnp.linspace(0.0, adam_seconds, num_iterations)[smoothing_window - 1 :], + smoothed_adam, + color=cols[0], + label="Adam only", +) +axes[0].set_ylabel("Negative ELBO (mini-batch)") +clean_legend(axes[0]) +clean_legend(axes[1]) + +target_value = float(smoothed_adam[-1]) +# Sentinel above every attainable iteration index, so "never crossed" is distinguishable +# from "crossed on the last iteration". +never = num_iterations + 1 +crossing = int( + jnp.min(jnp.where(smoothed_natgrad < target_value, smoothed_iterations, never)) +) +print( + f"Adam only, negative ELBO after {num_iterations} iterations : " + f"{target_value:8.2f}" +) +if crossing == never: + print("Natural gradients, same value reached at iteration : never") +else: + print(f"Natural gradients, same value reached at iteration : {crossing}") + print( + f" i.e. {crossing * natgrad_seconds / num_iterations:.2f} s " + f"versus {adam_seconds:.2f} s" + ) +print( + "Natural gradients, negative ELBO after " + f"{num_iterations} iterations: {float(smoothed_natgrad[-1]):8.2f}" +) + +# %% [markdown] +# Both curves are mini-batch estimates and therefore noisy; they are shown as a +# 25-iteration trailing mean. Per iteration the natural-gradient run is far ahead. Per +# second it is still ahead, but by less, because each of its iterations does strictly +# more work: a natural-gradient step converts $(\mathbf{m},\mathbf{L})$ to +# $\boldsymbol{\eta}$, differentiates the loss through the inverse map, converts back +# through $\boldsymbol{\theta}$, and *then* takes the Adam step on the +# hyperparameters. On the CPU that rendered this page that came to roughly half again +# the cost per iteration — see the timings printed above, which are what your machine +# actually measured. Salimbeni et al. report a comparable ratio of about $1.5\times$, +# and their headline experiments are on datasets far larger than this one; treat the +# numbers here as a demonstration of the mechanism, not as a benchmark. + +# %% +grid_side = 64 +grid_axis = jnp.linspace(-3.1, 3.1, grid_side) +grid_x, grid_y = jnp.meshgrid(grid_axis, grid_axis) +grid_points = jnp.stack([grid_x.ravel(), grid_y.ravel()], axis=1) + + +def predictive_probability(model, inputs, num_chunks=8): + """Bernoulli success probability, evaluated in chunks to bound memory.""" + unwrapped = paramax.unwrap(model) + likelihood = unwrapped.model.likelihood + return jnp.concatenate( + [likelihood(unwrapped(chunk)).mean for chunk in jnp.split(inputs, num_chunks)] + ) + + +fig, axes = plt.subplots(ncols=2, figsize=(10, 3.6), sharey=True) +for ax, model, name, seconds in [ + (axes[0], natgrad_model, "Natural gradients + Adam", natgrad_seconds), + (axes[1], adam_model, "Adam only", adam_seconds), +]: + probability = predictive_probability(model, grid_points).reshape( + grid_side, grid_side + ) + contours = ax.contourf( + grid_x, + grid_y, + probability, + levels=jnp.linspace(0.0, 1.0, 11), + cmap="RdBu_r", + alpha=0.7, + ) + ax.contour( + grid_x, grid_y, probability, levels=[0.5], colors="black", linewidths=1.5 + ) + ax.plot( + boundary_inputs, boundary_outputs, color="black", linestyle="--", linewidth=1 + ) + # Held-out points, encoded by class in the notebook's categorical colours rather + # than in the contour colourmap, so they stay legible on top of the fill. + for label, colour, marker in [(0.0, cols[0], "o"), (1.0, cols[1], "^")]: + mask = test_labels.ravel() == label + ax.scatter( + test_inputs_2d[mask, 0], + test_inputs_2d[mask, 1], + marker=marker, + s=12, + alpha=0.9, + color=colour, + edgecolors="white", + linewidths=0.3, + ) + inducing = paramax.unwrap(model).inducing_inputs + ax.scatter(inducing[:, 0], inducing[:, 1], marker="+", s=25, color="black") + + probability_test = predictive_probability(model, test_inputs_2d, num_chunks=1) + accuracy = jnp.mean((probability_test > 0.5) == (test_labels.ravel() > 0.5)) + log_density = jnp.mean( + test_labels.ravel() * jnp.log(probability_test) + + (1.0 - test_labels.ravel()) * jnp.log1p(-probability_test) + ) + ax.set( + xlabel=r"$x_1$", + xlim=(-3.1, 3.1), + ylim=(-3.1, 3.1), + title=f"{name}\naccuracy {accuracy:.3f}, NLPD {-log_density:.3f}", + ) + print( + f"{name:26s} test accuracy {accuracy:.4f}, test NLPD {-log_density:.4f}, " + f"{seconds:.2f} s" + ) +axes[0].set_ylabel(r"$x_2$") +colourbar = fig.colorbar(contours, ax=axes, label=r"$q(y=1 \mid x)$") + +# %% [markdown] +# The solid black line is each model's $0.5$ contour and the dashed line is the +# Bayes-optimal boundary $x_2 = 0.7x_1^2 - 1.5$; crosses mark the inducing inputs after +# training. +# +# The two panels are very nearly the same picture, and the two sets of printed test +# metrics are very nearly the same numbers. That is the honest reading of this +# experiment, and it is worth stating plainly: on a densely-sampled, easily-separated +# problem the natural gradient buys *optimiser speed*, not final predictive quality. It +# reached Adam's thousand-iteration bound at the crossing iteration printed under the +# ELBO comparison above, and both models then classify the held-out points about +# equally well. Note also +# that both runs train the kernel and the inducing inputs with Adam and finish at +# different hyperparameters, so whatever small difference remains between these +# contours cannot be attributed to $\mathbf{S}$ alone. `make_banana` draws inputs +# uniformly on $[-3,3]^2$ and the plotted grid is $[-3.1,3.1]^2$, so there is no +# region here that is far from the data; a demonstration that natural gradients give +# better-calibrated *extrapolative* uncertainty would need a problem built for it. + +# %% [markdown] +# ## When natural gradients fail +# +# The step is +# $\boldsymbol{\theta}\leftarrow\boldsymbol{\theta} - \gamma\,\partial\ell/\partial\boldsymbol{\eta}$, +# and $\boldsymbol{\Theta}_2$ must stay negative definite, because +# $\boldsymbol{\Theta}_2 = -\tfrac12\mathbf{S}^{-1}$ and $\mathbf{S}$ is a covariance. +# Nothing in the update enforces that. Splitting the ELBO as +# $\mathcal{L} = \mathcal{L}_{\text{data}} - \operatorname{KL}[q\,\|\,p]$ and using +# $\partial\operatorname{KL}/\partial\mathbf{S} = \tfrac12\mathbf{K}_{zz}^{-1} - \tfrac12\mathbf{S}^{-1}$ +# gives an exact description of what happens: +# +# $$\boldsymbol{\Theta}_2^{\text{new}} = (1-\gamma)\,\boldsymbol{\Theta}_2 + \gamma\,\boldsymbol{\Theta}_2^{\text{tgt}}, \qquad \boldsymbol{\Theta}_2^{\text{tgt}} := \frac{\partial\mathcal{L}_{\text{data}}}{\partial\mathbf{S}} - \tfrac{1}{2}\mathbf{K}_{zz}^{-1}$$ +# +# (for the whitened family, replace $\mathbf{K}_{zz}^{-1}$ by $\mathbf{I}_M$). So the +# step is a convex combination in $\boldsymbol{\theta}$-space whenever +# $\gamma\in[0,1]$ — the mirror-descent reading, made concrete. +# +# **Cone-safety theorem.** If the likelihood is log-concave in $f$, then by Price's +# theorem +# ($\partial_{\mathbf{S}}\mathbb{E}_{\mathcal{N}(\mathbf{m},\mathbf{S})}[g] = \tfrac12\mathbb{E}[\nabla^2 g]$), +# +# $$\frac{\partial\mathcal{L}_{\text{data}}}{\partial\mathbf{S}} = \frac{N}{B}\sum_{n\in\mathcal{B}}\tfrac{1}{2}\,\mathbb{E}_{q(f_n)}\!\left[\frac{\partial^2\log p(y_n\mid f_n)}{\partial f_n^2}\right]\mathbf{a}_n\mathbf{a}_n^\top \preceq 0,$$ +# +# where $\mathbf{a}_n^\top$ is row $n$ of $\mathbf{A} = \mathbf{K}_{xz}\mathbf{K}_{zz}^{-1}$. +# Hence $\boldsymbol{\Theta}_2^{\text{tgt}} \prec 0$, and for $\gamma\in[0,1]$ +# $\boldsymbol{\Theta}_2^{\text{new}}$ is a convex combination of two negative-definite +# matrices, so it is negative definite. **Mini-batching does not break this**, because +# $N/B > 0$ preserves the sign. $\square$ +# +# Two things escape the theorem: $\gamma > 1$, which extrapolates past +# $\boldsymbol{\Theta}_2^{\text{tgt}}$; and likelihoods that are not log-concave +# (Student-$t$, for instance), for which +# $\partial\mathcal{L}_{\text{data}}/\partial\mathbf{S}$ can have positive eigenvalues +# and the target itself sits outside the cone. Log-concavity here is a property of the +# likelihood *as computed*, not as written: GPJax's `inv_probit` clips its output into +# $[10^{-3},\,1-10^{-3}]$, which flattens the tail of $\log p$ enough to give it a +# positive second derivative for $f \lesssim -2.44$, so even the Bernoulli model used +# below leaves the guaranteed regime once a point is confidently mislabelled. That is +# the behaviour the backoff below is really guarding. Below we sweep $\gamma$ from an +# over-confident starting point — $\mathbf{S}_0 = 10^{-2}\mathbf{I}$, sharper than the +# target — which is precisely the regime where extrapolation bites. + +# %% +overconfident_family = gpx.variational_families.VariationalGaussian( + model=banana_model, + inducing_inputs=banana_inducing, + variational_mean=jnp.zeros((num_banana_inducing, 1)), + variational_root_covariance=0.1 * jnp.eye(num_banana_inducing), +) +overconfident_mean = overconfident_family.variational_mean.unwrap() +overconfident_root = overconfident_family.variational_root_covariance.unwrap() + + +def banana_loss_of_expectation(expectation): + variational_mean, variational_root = moments_from_expectation(*expectation) + trial = eqx.tree_at( + lambda family: (family.variational_mean, family.variational_root_covariance), + overconfident_family, + (Real(variational_mean), LowerTriangular(variational_root)), + ) + return negative_elbo(paramax.unwrap(trial), banana_train) + + +cone_gradient = jax.grad(banana_loss_of_expectation)( + expectation_from_moments(overconfident_mean, overconfident_root) +) +# The matrix statistic is symmetric, so symmetrise the entrywise autodiff gradient. +matrix_gradient = 0.5 * (cone_gradient[1] + cone_gradient[1].T) +_, natural_matrix = natural_from_moments(overconfident_mean, overconfident_root) + +print("gamma max eig(Theta2_new) status") +for gamma in [0.1, 0.5, 1.0, 2.0, 5.0, 10.0]: + largest = jnp.max(jnp.linalg.eigvalsh(natural_matrix - gamma * matrix_gradient)) + status = "negative definite" if largest < 0 else "*** LEFT THE CONE ***" + print(f"{gamma:6.2f} {largest:+18.5f} {status}") + +# %% [markdown] +# Read that table as a statement about *this initialisation*, not about $\gamma=2$. +# Here $\mathbf{S}_0 = 10^{-2}\mathbf{I}$ makes $\boldsymbol{\Theta}_2 = -50\,\mathbf{I}$, +# an order of magnitude sharper than the target, so the convex combination has very +# little room to extrapolate into. Because $\boldsymbol{\Theta}_2$ is a multiple of the +# identity, $\lambda_{\max}(\boldsymbol{\Theta}_2^{\text{new}})$ is exactly linear in +# $\gamma$, and interpolating the printed $\gamma=1$ and $\gamma=2$ rows puts the +# crossing at $\gamma\approx1.1$. Where it lands is entirely a function of how far +# $\boldsymbol{\Theta}_2$ starts from $\boldsymbol{\Theta}_2^{\text{tgt}}$: in the limit +# where the two coincide, every $\gamma$ is safe. What the theorem actually guarantees +# is $\gamma\in[0,1]$, for any log-concave likelihood and any starting point, and it +# says nothing whatsoever beyond that — which is the line worth remembering. +# +# When it does go wrong, `jnp.linalg.cholesky` returns `NaN` rather than raising, +# which means validity is a *value* and the fix stays `jit`-compatible. +# `natural_gradient_step` exploits that with a backoff: it evaluates the trial steps +# $\{\gamma\beta^k\}_{k=0}^{K}$ under `vmap` and selects the first one whose Cholesky +# is finite. `backoff` ($\beta$, default $0.5$) and `max_backoff` ($K$, default $5$) +# are exposed by `fit_natgrads`. + +# %% +print("gamma = 100 from the over-confident initialisation") +overconfident_variational, overconfident_hyper = partition_variational( + overconfident_family +) +for max_backoff in [0, 3, 5, 7, 10]: + stepped, _ = natural_gradient_step( + overconfident_variational, + overconfident_hyper, + banana_train, + negative_elbo, + 100.0, + max_backoff=max_backoff, + ) + smallest_trial = 100.0 * 0.5**max_backoff + root = eqx.combine( + stepped, overconfident_hyper + ).variational_root_covariance.unwrap() + outcome = "finite" if bool(jnp.all(jnp.isfinite(root))) else "NaN" + print( + f" max_backoff = {max_backoff:2d} smallest trial gamma = " + f"{smallest_trial:7.3f} result: {outcome}" + ) + +# %% [markdown] +# The backoff is a safety net with a finite budget, not a licence to pick $\gamma$ +# carelessly: from this starting point it needs to shrink $\gamma=100$ by a factor of +# $2^7$ before the Cholesky succeeds, so the default `max_backoff=5` still returns +# `NaN`. That is the intended behaviour — a silent 32-fold reduction of a step size the +# user chose badly would be worse than a visible failure. + +# %% [markdown] +# ## Practical guidance +# +# * **Conjugate and full batch: use $\gamma = 1$.** One iteration is the exact +# solution, and further iterations are fixed points. +# * **Non-conjugate or mini-batched: ramp $\gamma$.** Salimbeni et al. recommend +# starting around $10^{-4}$ and reaching $\approx 10^{-1}$ "sufficiently quickly +# ($<1000$ iterations)"; `natgrad_lr` accepts any Optax schedule, and defaults to +# $10^{-1}$. +# * **Never exceed $\gamma = 1$.** The convex-combination guarantee stops there, and +# the backoff exists to catch mistakes, not to enable them. +# * **If a mini-batched run produces `NaN`, raise the batch size before lowering +# $\gamma$.** Small batches make +# $\boldsymbol{\Theta}_2^{\text{tgt}}$ badly conditioned, which no step size fully +# repairs. +# * **Prefer the whitened family.** The natural-gradient direction is +# parameterisation-invariant, so whitening does not change the sequence of +# distributions in exact arithmetic; it changes the *conditioning* of every map, and +# keeps $\mathbf{m}_w$, $\mathbf{S}_w$ at $\mathcal{O}(1)$. +# * **Leave `map_jitter` at $0$.** It biases $\mathbf{S}$ by +# $\approx\varepsilon\lVert\mathbf{S}\rVert^2$ independently of conditioning. Raise +# it to $10^{-12}$–$10^{-10}$ only when fighting an ill-conditioned $\mathbf{S}$. +# * **Non-log-concave likelihoods have no guarantee at all.** For a Student-$t$ +# likelihood with gross outliers the target $\boldsymbol{\Theta}_2^{\text{tgt}}$ can +# itself be outside the cone, so no positive $\gamma$ is provably safe. +# +# The companion [dual sparse GP notebook](dual_svgp.py) takes the same geometry in a +# different direction, storing the site parameters of the variational distribution +# rather than its moments. + +# %% [markdown] +# ## System configuration + +# %% +# %reload_ext watermark +# %watermark -n -u -v -iv -w -a 'Thomas Pinder' diff --git a/docs/examples/numpyro_integration.py b/docs/examples/numpyro_integration.py index a74e73014..6590bdb9c 100644 --- a/docs/examples/numpyro_integration.py +++ b/docs/examples/numpyro_integration.py @@ -131,7 +131,7 @@ # # We define a NumPyro model that samples all parameters directly using # ``numpyro.sample``, builds the GPJax -# [`ConjugatePosterior`](#gpjax.gps.ConjugatePosterior) from those samples, and +# [`ConjugateModel`](#gpjax.gps.ConjugateModel) from those samples, and # scores it with the conjugate marginal log-likelihood # ([`conjugate_mll`](#gpjax.objectives.conjugate_mll)) via ``numpyro.factor``. # No special registration step is needed -- GPJax constructors accept raw @@ -157,7 +157,7 @@ def model(X, Y, X_new=None): meanf = gpx.mean_functions.Constant() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) - likelihood = gpx.likelihoods.Gaussian(num_datapoints=N, obs_stddev=obs_noise) + likelihood = gpx.likelihoods.Gaussian(obs_stddev=obs_noise) posterior = prior * likelihood D_resid = gpx.Dataset(X=X, y=residuals) @@ -199,7 +199,7 @@ def model(X, Y, X_new=None): ) example_posterior = gpx.gps.Prior( mean_function=gpx.mean_functions.Constant(), kernel=example_kernel -) * gpx.likelihoods.Gaussian(num_datapoints=N, obs_stddev=1.0) +) * gpx.likelihoods.Gaussian(obs_stddev=1.0) parameter_priors = { "prior.kernel.kernels[0].lengthscale": dist.LogNormal(0.0, 1.0), diff --git a/docs/examples/oak.py b/docs/examples/oak.py index 1bbb3c6f0..2f610b191 100644 --- a/docs/examples/oak.py +++ b/docs/examples/oak.py @@ -273,7 +273,7 @@ def apply_flows(X_original: np.ndarray) -> jnp.ndarray: mean_function = gpx.mean_functions.Zero() prior = gpx.gps.Prior(mean_function=mean_function, kernel=oak_kernel) -likelihood = gpx.likelihoods.Gaussian(num_datapoints=num_train) +likelihood = gpx.likelihoods.Gaussian() posterior = prior * likelihood # %% @@ -286,7 +286,7 @@ def apply_flows(X_original: np.ndarray) -> jnp.ndarray: ) latent_dist = opt_posterior.predict( - X_test, train_data=train_data, return_covariance_type="diagonal" + X_test, train_data=train_data, covariance="diagonal" ) predictive_dist = opt_posterior.likelihood(latent_dist) predictive_mean = predictive_dist.mean diff --git a/docs/examples/oceanmodelling.py b/docs/examples/oceanmodelling.py index 14d5dbfa7..a4203fd49 100644 --- a/docs/examples/oceanmodelling.py +++ b/docs/examples/oceanmodelling.py @@ -334,7 +334,7 @@ def __call__( def initialise_gp(kernel, mean, dataset): prior = gpx.gps.Prior(mean_function=mean, kernel=kernel) likelihood = gpx.likelihoods.Gaussian( - num_datapoints=dataset.n, obs_stddev=jnp.array([1.0e-3], dtype=jnp.float64) + obs_stddev=jnp.array([1.0e-3], dtype=jnp.float64) ) posterior = prior * likelihood return posterior diff --git a/docs/examples/oilmm.py b/docs/examples/oilmm.py index 9f9e24626..d8d876e49 100644 --- a/docs/examples/oilmm.py +++ b/docs/examples/oilmm.py @@ -406,7 +406,7 @@ def plot_wave_output_panel( # # Before optimising any parameters, we condition with the PCA-initialised # defaults to establish a baseline. Calling -# [`condition_on_observations`](#gpjax.models.OILMMModel.condition_on_observations) +# [`condition`](#gpjax.models.OILMMModel.condition) # executes the OILMM inference algorithm: # # 1. **Project**: compute @@ -419,7 +419,7 @@ def plot_wave_output_panel( # $m$ independent posteriors. # %% -posterior = model.condition_on_observations(train_data) +posterior = model.condition(train_data) # %% [markdown] # ## Baseline predictions @@ -435,7 +435,7 @@ def plot_wave_output_panel( ) X_test = ((test_time_hours - input_mean) / input_std).reshape(-1, 1) -pre_pred = posterior.predict(X_test, return_full_cov=False) +pre_pred = posterior.predict(X_test, covariance="diagonal") pre_opt_mean_standardised = pre_pred.mean.reshape(N_test, num_outputs) pre_opt_std_standardised = jnp.sqrt(jnp.diag(pre_pred.covariance())).reshape( N_test, num_outputs @@ -542,8 +542,8 @@ def plot_wave_output_panel( # same test locations. # %% -opt_posterior = opt_model.condition_on_observations(train_data) -post_pred = opt_posterior.predict(X_test, return_full_cov=False) +opt_posterior = opt_model.condition(train_data) +post_pred = opt_posterior.predict(X_test, covariance="diagonal") post_opt_mean_standardised = post_pred.mean.reshape(N_test, num_outputs) post_opt_std_standardised = jnp.sqrt(jnp.diag(post_pred.covariance())).reshape( N_test, num_outputs @@ -674,13 +674,12 @@ def plot_wave_output_panel( for i in range(num_latent): ax = axes[i] - lat_y = opt_posterior.latent_datasets[i].y.squeeze() + latent = opt_posterior.latent_posteriors[i] + lat_y = latent.train_data.y.squeeze() ax.plot( time_hours, lat_y, "o", color=cols[i], alpha=0.4, ms=3, label="Projected data" ) - lat_pred = opt_posterior.latent_posteriors[i].predict( - X_test, train_data=opt_posterior.latent_datasets[i] - ) + lat_pred = latent.predict(X_test, covariance="diagonal") lat_mean = lat_pred.mean lat_std = jnp.sqrt(jnp.diag(lat_pred.covariance())) diff --git a/docs/examples/poisson.py b/docs/examples/poisson.py index 7c0700dc5..c36c4f859 100644 --- a/docs/examples/poisson.py +++ b/docs/examples/poisson.py @@ -135,14 +135,15 @@ kernel = gpx.kernels.RBF() meanf = gpx.mean_functions.Constant() prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel) -likelihood = gpx.likelihoods.Poisson(num_datapoints=D.n) +likelihood = gpx.likelihoods.Poisson() # %% [markdown] -# We construct the [posterior](#gpjax.gps.NonConjugatePosterior) through the product of our -# prior and likelihood. +# We construct the [model](#gpjax.gps.NonConjugateModel) through the product of our +# prior and likelihood, and initialise the whitened latent vector that MCMC +# will sample (sized by the training data). # %% -posterior = prior * likelihood +posterior = (prior * likelihood).init_latent(D.n) print(type(posterior)) # %% [markdown] diff --git a/docs/examples/regression.py b/docs/examples/regression.py index a827d4a03..6ca2c4ccb 100644 --- a/docs/examples/regression.py +++ b/docs/examples/regression.py @@ -136,11 +136,11 @@ # the evaluation of the GP's mean and covariance. # # Since we want to sample from the full posterior, we need to calculate the full covariance matrix. -# We can enforce this by including the `return_covariance_type = "dense"` attribute when predicting. +# We can enforce this by including the `covariance = "dense"` attribute when predicting. # Note this is what will be defaulted if left blank. # %% mystnb={"figure": {"caption": "Twenty function samples drawn from the zero-mean RBF prior, shown alongside the prior mean and variance band.", "name": "fig-regression-prior-samples"}} -prior_dist = prior.predict(xtest, return_covariance_type="dense") +prior_dist = prior.predict(xtest, covariance="dense") prior_mean = prior_dist.mean prior_std = prior_dist.variance @@ -181,19 +181,23 @@ # and what each one assumes about the observations. # %% -likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n) +likelihood = gpx.likelihoods.Gaussian() # %% [markdown] -# The posterior is proportional to the prior multiplied by the likelihood, written as +# The prior and likelihood together define the joint model over the latent +# function and the observations, # # $$ -# p(f(\cdot) | \mathcal{D}) \propto p(f(\cdot)) * p(\mathcal{D} | f(\cdot)). -# $$ (eq-regression-posterior) +# p(f(\cdot), \mathcal{D}) = p(\mathcal{D} | f(\cdot))\, p(f(\cdot)). +# $$ (eq-regression-joint) # -# Mimicking this construct, the posterior is established in GPJax through the `*` operator. +# Mimicking this equation, the joint model is established in GPJax through the +# `*` operator. Conditioning it on the data — `model.condition(D)`, or the +# operator form `model | D` that reads as $p(f \mid \mathcal{D})$ — yields the +# posterior process. # %% -posterior = prior * likelihood +model = prior * likelihood # %% [markdown] #