diff --git a/docs/examples/dual_svgp.py b/docs/examples/dual_svgp.py index 42728f7f5..f2f3e7153 100644 --- a/docs/examples/dual_svgp.py +++ b/docs/examples/dual_svgp.py @@ -19,59 +19,42 @@ # # 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. +# This is the applied companion to the +# [natural gradients notebook](natural_gradients.py). That notebook derives +# the whole geometry this one puts to work: the exponential-family view of +# $q(\mathbf{u})$, the Fisher identity that makes the natural gradient free +# to compute, the site — or *dual* — reparameterisation of +# {cite:t}`adam2021dual`, its EP heritage, the two silent convention traps in +# the source material, the tied update and why it never inverts anything, +# the positive-semidefinite cone-safety argument for the site branch, and +# the claims table for `dual_elbo` versus `elbo` as an M-step objective. +# **Read that one first.** Nothing below re-derives any of it; this notebook +# assumes it and asks instead: does GPJax's implementation actually deliver +# what the theory promises, on real models, with real numbers? A sibling +# notebook, [natgrads.py](natgrads.py), runs the same kind of check for the +# moment-storage branch, `VariationalGaussian`. +# +# Four checks, in order. A conjugate regression where `DualVariationalGaussian` +# plus one `natural_gradient_step` at $\rho=1$ reproduces the +# {cite:t}`titsias2009` optimum exactly. A non-conjugate classification +# problem where the site branch and the moment branch are driven through +# matched steps and compared directly, which locates precisely where the +# $\rho=\gamma$ identity's log-concavity condition fails and `beta_floor` +# starts to matter. The banana benchmark from the natural-gradients notebook, +# run again with all three optimisers — Adam alone, natural gradients on +# $(\mathbf{m},\mathbf{L})$, and t-SVGP's dual branch — with timings. And, +# last, the M-step claims table exercised in practice: a bound slice showing +# where `dual_elbo` dominates `elbo` and where it does not, and a full +# variational-EM training loop comparing the two as M-step objectives. +# +# **One notational reminder, carried over from the natural-gradients +# notebook.** $\boldsymbol{\theta}$ is the kernel hyperparameters, not a +# natural parameter; the natural parameter of $q(\mathbf{u})$ is +# $\boldsymbol{\eta}$, its expectation parameter is $\boldsymbol{\mu}$, and +# $\boldsymbol{\lambda} = (\boldsymbol{\lambda}_1, \boldsymbol{\Lambda}_2)$ +# is the site — the pair `DualVariationalGaussian` stores as `dual_vector` +# and `dual_matrix`. $\mathbf{a}_i = \mathbf{K}_{zz}^{-1}\mathbf{k}_z(x_i)$ +# throughout. # %% # Enable Float64 for more stable matrix inversions. @@ -82,7 +65,6 @@ 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 @@ -125,169 +107,36 @@ def negative_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 +# ## What `DualVariationalGaussian` stores +# +# In brief, because the natural-gradients notebook has the derivation: $q$ +# is the prior reweighted by an unnormalised Gaussian site, +# $q(\mathbf{u}) \propto p_{\boldsymbol{\theta}}(\mathbf{u})\, +# \exp(\boldsymbol{\lambda}_1^\top\tilde{\mathbf{u}} - +# \tfrac12\tilde{\mathbf{u}}^\top\boldsymbol{\Lambda}_2\tilde{\mathbf{u}})$ +# with $\tilde{\mathbf{u}} = \mathbf{u} - \boldsymbol{\mu}_z$, giving moments +# $\mathbf{S} = (\mathbf{K}_{zz}^{-1} + \boldsymbol{\Lambda}_2)^{-1}$, +# $\mathbf{m} = \boldsymbol{\mu}_z + \mathbf{S}\boldsymbol{\lambda}_1$. +# `DualVariationalGaussian` stores $\boldsymbol{\lambda}_1$ as `dual_vector` +# and $\boldsymbol{\Lambda}_2$ as `dual_matrix`, both zero by default — so +# $q = p$ at initialisation — and exposes the pair above via `.moments()`. +# GPJax stores the **flanked, precision** convention (never the un-flanked +# sums, never $-\tfrac12\boldsymbol{\Lambda}_2$); that is the rule the +# natural-gradients notebook derives and states, and every demo below +# depends on it silently, the way any user of the API does. +# +# Where do $\boldsymbol{\lambda}_1$ and $\boldsymbol{\Lambda}_2$ come from at +# the optimum? One local likelihood approximation per data point, EP-style, +# with the site values read off Bonnet's and Price's theorems rather than +# fitted by moment matching: # # $$\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`. +# so a single `jax.grad` of the likelihood's existing +# `expected_log_likelihood` gives both, for closed-form and quadrature +# likelihoods alike. Here is that call pattern, checked against the +# Gaussian closed form $\alpha_i = (y_i-m_i)/\sigma^2$, +# $\beta_i = 1/\sigma^2$. # %% key, alpha_beta_key = jr.split(key) @@ -326,41 +175,11 @@ def total_expected_log_likelihood(mean, variance): 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. +# Both match the closed form to machine precision. This is the ingredient +# every demo below is built from: `expected_log_likelihood` plus one +# `jax.grad` call, tied across data points into the $M$- and +# $M\times M$-sized `dual_vector`/`dual_matrix` update the natural-gradients +# notebook derives in full. # %% [markdown] # ## Conjugate models: one step is enough @@ -369,20 +188,18 @@ def total_expected_log_likelihood(mean, variance): # # $$\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: +# The $m_i$ cancels, so the tied update's target does not move either, and +# $\rho=1$ lands on the fixed point 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. +# 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,\mathbf{S}^\star)$ +# is the {cite:t}`titsias2009` optimal $q(\mathbf{u})$ verbatim — the site +# instantiation of the general "$\gamma=1$ is exact" argument the +# natural-gradients notebook proves for any storage convention. The mean +# function below is deliberately non-zero: the sites act on the *centred* +# process, so $\mathbf{y}-\boldsymbol{\mu}_x$ is where that matters. # %% num_data = 200 @@ -408,6 +225,7 @@ def conjugate_model(lengthscale): prior = gpx.gps.Prior( mean_function=gpx.mean_functions.Constant(jnp.array(prior_constant)), kernel=jk.RBF(lengthscale=lengthscale), + jitter=regression_jitter, ) return prior * gpx.likelihoods.Gaussian(obs_stddev=noise_stddev) @@ -532,175 +350,59 @@ def exact_sites(lengthscale, inducing_inputs, dataset): ) # %% [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 +# 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 numbers above confirm both. +# +# 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 eight 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}" -) +# $c(\boldsymbol{\theta})$ is minus the Titsias trace term — 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, flanked site GPJax 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] -# 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. +# ## Locating where $\rho=\gamma$ stops holding +# +# The natural-gradients notebook proves the site step and the moment step +# are the *same iteration* whenever the computed $\beta_i$ stay +# non-negative, and shows this on the banana classification problem: a +# single point crosses into $\beta_i<0$ at step five of six matched +# $\rho=\gamma=0.8$ steps, the `beta_floor` clip engages, and the two +# branches part company at $\mathcal{O}(10^{-3})$ in $(\mathbf{m},\mathbf{S})$ +# — while agreeing to the float64 noise floor everywhere before that. We +# reproduce that check here rather than take it on faith, because it is the +# load-bearing claim behind everything that follows: the banana benchmark +# and the VEM loop below both interleave many such steps, so knowing exactly +# when and how far the two branches can diverge tells us how much of any +# difference in their trajectories to attribute to the E-step versus the +# M-step. +# +# The next cell is the data-generating function and inducing-point layout +# from the natural-gradients notebook, reproduced character for character — +# same function body, same `jr.key(42)`, same 2000 points, same $10\times5$ +# grid — so the problem here is the same problem, point for point, as the +# one there. # %% @@ -723,11 +425,6 @@ def make_banana(key, num_points): 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) @@ -757,29 +454,22 @@ def make_banana_moment_family(): ) -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"train / test : {banana_train.n} / {banana_data.n - banana_train.n}") +print(f"class balance : {float(banana_data.y.mean()):.3f}") 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 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 + + 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) @@ -855,47 +545,37 @@ def six_matched_steps(beta_floor): 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. +# The table matches the natural-gradients notebook's: for the first four +# steps every $\beta_i$ is positive, the clip is inert, and the branches +# agree to $10^{-13}$. At step five a single training point out of 1600 +# crosses into $\beta_i<0$ — a label-$0$ point whose marginal mean of +# $+2.427$ sits almost exactly on the mirrored $+2.44$ threshold the +# cone-safety section derives for GPJax's clipped `inv_probit` — +# `beta_floor` engages, and from that step on the gap jumps +# to $\mathcal{O}(10^{-3})$ and compounds. Disabling the clip brings the same +# six steps back to the noise floor, which is the control that pins the +# cause down to the clip alone. +# +# For the two demos ahead, the practical reading is: whatever difference +# shows up between the dual and moment branches beyond the low-$10^{-3}$ +# level located here is not coming from the E-step being a different +# search direction. It is either wall-clock, or the M-step objective. + +# %% [markdown] +# ## The banana benchmark: three optimisers # -# 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. +# All three models below start at $q=p$, that is $\mathbf{m}=\mathbf{0}$, +# $\mathbf{S}=\mathbf{K}_{zz}$ — where a dual family with zero sites already +# is — over the same inducing grid just built. # %% +banana_dual_family = DualVariationalGaussian( + model=banana_model, + inducing_inputs=banana_inducing, +) +natgrad_family = make_banana_moment_family() +adam_family = make_banana_moment_family() + # The log-linear ramp of the natural-gradients notebook: 1e-4 -> 1e-1 over K = 100. num_iterations = 1000 batch_size = 256 @@ -1014,135 +694,66 @@ def smooth(history): ) # %% [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. +# 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}$ located a section 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. +# ## The M-step in practice +# +# The natural-gradients notebook's claims table is the reference for what +# follows; we do not restate its proof here, only exercise it. In brief, +# `elbo` freezes the *whole* natural parameter of $q$ at its E-step optimum +# while $\boldsymbol{\theta}$ moves; `dual_elbo` freezes only the +# data-derived sites and lets the prior half of $q$ track +# $\mathbf{K}_{zz}(\boldsymbol{\theta})$. The two agree in value and +# gradient at a converged E-step, by the envelope theorem, and can disagree +# substantially away from it — which in practice is always, since nobody +# runs an E-step to convergence between Adam steps. Two questions follow. +# Does the dual bound actually dominate the standard one, as the paper's +# figure suggests? And does that translate into a better fitted model at the +# end of a real VEM loop? # %% log_offsets = jnp.linspace(-1.2, 0.6, 61) @@ -1175,25 +786,6 @@ def bound_slice(inducing_inputs, dataset, sites, moments, offsets): 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)") @@ -1230,120 +822,38 @@ def bound_slice(inducing_inputs, dataset, sites, moments, offsets): 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 +# Left panel: both bounds pass through the same point at +# $\boldsymbol{\theta}_t$, which is the value-equality row of the claims +# table. To the *right*, at longer lengthscales, $l$ falls off a cliff — a +# $q$ chosen for one kernel is a bad approximation under another — while +# $\bar l$ has barely moved, since only the data-dependent half of it was +# frozen. To the *left* the two collapse together instead, since the sparse +# approximation itself degrades there regardless of which $q$ is frozen. +# 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. +# Right panel: at $M=20$ the dual bound dominates everywhere probed, up to +# float64 noise at the crossing point (the value-equality point itself). At +# $M=5$ and $M=10$ the gap dips below zero for real — printed above — so the +# dominance guarantee, which the claims table restricts to $\mathbf{Z}=\mathbf{X}$, +# genuinely does not extend to every sparse configuration. (At +# $\mathbf{Z}=\mathbf{X}$ itself the flanked sites reduce to +# $\boldsymbol{\lambda}_1=(\mathbf{y}-\boldsymbol{\mu})/\sigma^2$, +# $\boldsymbol{\Lambda}_2=\mathbf{I}/\sigma^2$ — genuinely +# $\boldsymbol{\theta}$-free — and dominance holds everywhere; that is a +# corollary of the claims table, not a new demo, so we do not reproduce it +# here.) # %% [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. +# 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, and so does the claims table. +# 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 @@ -1493,80 +1003,87 @@ def test_metrics(model): 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") +# The last round the lead is non-positive, so we can report where it becomes +# permanent rather than where it first (and only briefly) turns positive. +last_nonpositive_round = int(jnp.max(jnp.where(bound_lead <= 0.0, rounds, 0))) +if last_nonpositive_round == 0: + print("dual_elbo leads for the whole run") else: + remaining = vem_rounds - last_nonpositive_round 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" + f"last round with a non-positive lead: {last_nonpositive_round}; " + f"the lead stays positive for all {remaining} rounds after that" ) # %% [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". +# 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 plots the +# difference of the two bounds rather than the two bounds themselves, +# because on a negative ELBO of a few hundred a lead of a nat is invisible +# on the natural scale. +# +# The sign is not settled early. For the first two rounds the dual bound +# leads by a few tenths of a nat, then falls behind for rounds three through +# six — both lengthscales are still moving fast and $q$ is far from +# converged at every round, exactly the regime where the two M-step +# objectives are proven to differ — bottoming out at $-1.15$ nats. The round +# printed above is the *last* one with a non-positive lead; every round +# after it is positive, rising to a peak around $+1.27$ nats and easing back +# to the final value printed above. So the honest reading is not "the +# dual M-step is uniformly ahead"; early on it is not. 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 metrics agree closely. 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}$. +# 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. +# * **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. +# * **$\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 +# broke the $\rho=\gamma$ identity above, by $\sim\!5\times10^{-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; never write a test against +# $\boldsymbol{\Lambda}_2$ directly — see the natural-gradients notebook +# for the numerical demonstration of why. # * **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). +# $(\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}$. # %% [markdown] # ## System configuration diff --git a/docs/examples/natgrads.py b/docs/examples/natgrads.py index 8c2741af8..8171d2576 100644 --- a/docs/examples/natgrads.py +++ b/docs/examples/natgrads.py @@ -15,43 +15,33 @@ # --- # %% [markdown] -# # Natural Gradients +# # Natural Gradients in Practice # # 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. +# This notebook assumes the +# [natural gradients notebook](natural_gradients.py) throughout: the +# exponential-family view of $q(\mathbf{u})$, the Fisher=Jacobian identity +# that makes the natural gradient free to compute, the mirror-descent +# reading of the step, the "one step is enough" theorem for conjugate +# models, and the cone-safety theorem with its proof. None of that is +# re-derived here — this notebook connects that geometry to the GPJax API +# instead: `gpx.fit_natgrads`, the lower-level `natural_gradient_step`, +# `partition_variational`, and `WhitenedVariationalGaussian`, on two real +# training runs, checking the theory's predictions against what actually +# happens on this machine. # # 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. +# 1. **demo (i)** — a conjugate 1D regression where a single $\gamma=1$ step +# lands on the exact variational optimum, while Adam is still crawling +# after two thousand iterations; +# 2. **demo (ii)** — a mini-batched Bernoulli classification benchmark, +# comparing natural gradients + Adam against Adam alone, per iteration +# *and* per second; +# 3. the failure mode: what a large $\gamma$ does to the +# $\boldsymbol{\Theta}_2$ cone, 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) @@ -82,7 +72,6 @@ from gpjax.natural_gradients import ( expectation_from_moments, moments_from_expectation, - moments_from_natural, natural_from_moments, natural_gradient_step, partition_variational, @@ -102,232 +91,21 @@ def negative_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. +# ## Demo (i): conjugate regression +# +# Recall from the natural gradients notebook that for a conditionally +# conjugate model — here, a Gaussian likelihood — the ELBO is affine in the +# expectation parameters, so the step collapses to +# $\boldsymbol{\theta}_{\text{new}} = (1-\gamma)\,\boldsymbol{\theta} + \gamma\,\boldsymbol{\lambda}$ +# for a fixed $\boldsymbol{\lambda}$ that does not depend on $q$. At +# $\gamma=1$ this is not an approximation to the optimum, it *is* the +# optimum: $\boldsymbol{\theta}_{\text{new}} = \boldsymbol{\lambda} = \boldsymbol{\theta}^\star$, +# reached in one step from any starting point (the "one step is enough" +# theorem there, after Sato 2001; for the SVGP it recovers the +# {cite:t}`titsias2009` optimum). We watch that happen, then race it against +# Adam from the same bad start. # %% -# 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 @@ -344,8 +122,9 @@ def expectation_of_natural(flat): 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. +# 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), @@ -368,13 +147,13 @@ def expectation_of_natural(flat): # 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}$. +# $\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 as $q(\mathbf{u})$, 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 @@ -382,6 +161,10 @@ def expectation_of_natural(flat): # # $$\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 .$$ +# +# This is used only as a reference value below, computed once with plain +# linear algebra so that the natural-gradient step has something exact to be +# checked against. # %% unwrapped_initial = paramax.unwrap(initial_family) @@ -406,7 +189,8 @@ def expectation_of_natural(flat): 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. +# 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, @@ -418,7 +202,10 @@ def expectation_of_natural(flat): print(f"ELBO at the closed-form optimum: {reference_elbo:.6f}") # %% -# One natural-gradient step at gamma = 1. +# One natural-gradient step at gamma = 1. `partition_variational` splits the +# family into the pytree leaves the step is allowed to touch (the +# variational parameters) and everything else (the hyperparameters); the +# step is exactly what `fit_natgrads` calls once per iteration. variational_partition, hyper_partition = partition_variational(initial_family) stepped_partition, loss_before = natural_gradient_step( variational_partition, @@ -467,20 +254,23 @@ def expectation_of_natural(flat): ) # %% [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 +# One step, from a random initialisation more than 3000 ELBO nats away, lands +# on the closed-form optimum to $\sim10^{-13}$ in both the mean and the +# covariance — the float64 noise floor for a problem of this size — and the +# ELBO itself matches to all six printed decimal places. A second step moves +# the mean by the same $\sim10^{-14}$, confirming the fixed point. +# +# Notice the `map_jitter=0.0` keyword: 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}$. +# It defaults to zero in `fit_natgrads` for exactly 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. +# 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 at all. # %% exact_posterior = paramax.unwrap(regression_model).condition(regression_data) @@ -526,25 +316,34 @@ def expectation_of_natural(flat): clean_legend(ax) axes[0].set_ylabel(r"$f(x)$") +data_range_mask = (test_inputs[:, 0] >= -3.0) & (test_inputs[:, 0] <= 3.0) print( - "max |sparse mean - exact mean| : " + "max |sparse mean - exact mean|, full grid [-3.2,3.2] : " f"{jnp.max(jnp.abs(unwrapped_stepped(test_inputs).mean - exact_mean)):.3e}" ) +print( + "max |sparse mean - exact mean|, data range [-3,3] : " + f"{jnp.max(jnp.abs((unwrapped_stepped(test_inputs).mean - exact_mean)[data_range_mask])):.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. +# 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 two printed +# maxima confirm it quantitatively: restricted to the data range $[-3,3]$ +# the sparse and exact means agree about fifteen times more closely than +# they do on the full test grid, where the largest gap sits at the grid's +# edge, past the last inducing input. Both 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` — applied to `hyper_partition`, the half of the +# pytree `partition_variational` carved off as *not* the natural gradient's +# business — so that both methods solve 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( @@ -595,9 +394,19 @@ def expectation_of_natural(flat): ) clean_legend(axes[0]) -axes[1].plot(iteration_index + 1, adam_gap, color=cols[0], label="Adam on $(m, L)$") +axes[1].plot( + iteration_index + 1, + jnp.maximum(adam_gap, 1e-16), + 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" + [1], + [max(natgrad_gap, 1e-16)], + 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)" @@ -605,68 +414,73 @@ def expectation_of_natural(flat): 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. +# 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: it takes 867 iterations merely to come within ten nats of the +# optimum, 1800 to come within one nat, and it never gets within a tenth of +# a nat across the full 2000 — the ELBO gap is still $6.0\times10^{-1}$ +# nats and still shrinking, while the single natural-gradient step closed +# the gap to zero at double precision. 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$ +# ## Demo (ii): non-conjugate banana classification # -# 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 +# 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 +# $\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)". +# +# `natgrad_lr` accepts any Optax schedule — that is the API surface for this +# whole recommendation. 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? Recall the mirror-descent +# reading from the natural gradients notebook: the step is always a convex +# 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. +# — the "when natural gradients fail" section below writes its second block +# out explicitly. The $N/B$ rescaling inside the ELBO keeps the stochastic +# gradient unbiased at every $\gamma$, including $\gamma=1$; what degrades +# is *variance*. Outside conjugacy $\boldsymbol{\theta}^{\text{tgt}}$ is not +# a fixed optimum — it is where one fixed-point iteration from *here* would +# land, and it moves with both $q$ and the mini-batch. 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, which is where the variance +# reduction comes from. # # Time for a harder problem. @@ -720,7 +534,8 @@ def make_banana(key, num_points): clean_legend(ax) # %% -# Two identical models, built from the same arrays, so the comparison is fair. +# 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) @@ -746,7 +561,8 @@ def 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. +# The log-linear ramp, 1e-4 -> 1e-1 over K = 100 iterations, as an Optax +# schedule handed straight to `natgrad_lr`. num_iterations = 1000 batch_size = 256 natgrad_schedule = ox.exponential_decay( @@ -806,11 +622,14 @@ def timed_fit(run): ) # %% [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. +# 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 — `gpx.fit_natgrads` alternates a `natural_gradient_step` on those +# with an ordinary `gpx.fit`-style Adam step on everything else; `gpx.fit` +# moves everything with Adam. 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 @@ -827,8 +646,8 @@ def smooth(history): 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. +# 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())) @@ -857,8 +676,8 @@ def smooth(history): 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". +# 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)) @@ -881,17 +700,21 @@ def smooth(history): ) # %% [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. +# 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: it reaches Adam's thousand-iteration bound of $335.31$ by +# iteration $117$, and finishes at $323.98$ against Adam's $335.31$. 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 $1.5$–$1.7\times$ the per-iteration cost of Adam alone across +# repeated runs — see the timings printed above, which are what your +# machine actually measured. Salimbeni et al. report a comparable ratio of +# about $1.5\times$ on their own hardware, 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 @@ -931,8 +754,9 @@ def predictive_probability(model, inputs, num_chunks=8): 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. + # 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( @@ -968,65 +792,47 @@ def predictive_probability(model, inputs, num_chunks=8): 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. +# 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: $93.50\%$ accuracy +# and $0.166$ NLPD for natural gradients against $93.25\%$ and $0.169$ for +# Adam alone. 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 +# above, and both models then classify the held-out points about equally +# well. Both runs also 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 +# Recall the update 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 +# and that $\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 automatically. The +# natural gradients notebook's **cone-safety theorem** proves — via Price's +# theorem applied to the ELBO's data-fit term — that for any log-concave +# likelihood and any starting point, $\gamma\in[0,1]$ keeps +# $\boldsymbol{\Theta}_2^{\text{new}}$ inside that cone, mini-batching +# included; the full statement and proof are there, and are not repeated +# here. What escapes the guarantee is $\gamma>1$, which extrapolates past +# the target, and likelihoods that are not log-concave *as computed* rather +# than 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. 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. # %% @@ -1053,7 +859,8 @@ def banana_loss_of_expectation(expectation): 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. +# 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) @@ -1064,24 +871,26 @@ def banana_loss_of_expectation(expectation): 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`. +# Read that table as a statement about *this initialisation*, not about +# $\gamma=2$ in general. 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: the sign flips between $\gamma=1$ ($-4.67$) and +# $\gamma=2$ ($+40.66$), and interpolating those two rows puts the crossing +# at $\gamma\approx1.10$. 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` and by `natural_gradient_step` directly. # %% print("gamma = 100 from the over-confident initialisation") @@ -1108,41 +917,46 @@ def banana_loss_of_expectation(expectation): ) # %% [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. +# 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. +# * **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 cone-safety theorem's 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)$. +# 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. +# $\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) is the applied +# notebook for the other storage convention this geometry admits — the +# site, or dual, parameterisation of {cite:t}`adam2021dual` — and, like this +# one, it assumes the [natural gradients notebook](natural_gradients.py) +# throughout rather than re-deriving anything. # %% [markdown] # ## System configuration diff --git a/docs/examples/natural_gradients.py b/docs/examples/natural_gradients.py new file mode 100644 index 000000000..4d7cf7f2d --- /dev/null +++ b/docs/examples/natural_gradients.py @@ -0,0 +1,1468 @@ +# --- +# 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}`natural_gradients.ipynb` +# +# This notebook is the prerequisite for two others: the +# [natural gradients](natgrads.py) notebook and the +# [dual sparse GP](dual_svgp.py) notebook. Between them those two cover the +# GPJax API — how to call `gpx.fit_natgrads`, how to choose a variational +# family, how the method behaves on real training runs — and they assume +# everything below already: the exponential-family geometry of +# $q(\mathbf{u})$, the identity that makes the natural gradient free to +# compute, the site/dual reparameterisation of that same geometry, and the +# guarantees and failure modes attached to both. Read this one first; the +# other two will not re-derive any of it. +# +# 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, or, as we will see, a pair of +# site parameters — 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. `gpjax.fit_natgrads` implements the +# recipe of {cite:t}`salimbeni2018`, alternating a natural-gradient step on +# $q$ with an ordinary gradient step, taken with any Optax optimiser, on the +# kernel hyperparameters. Every call it makes at the $q$-step is a call to +# the lower-level primitive `natural_gradient_step`, which is what every demo +# below calls directly, so that we can inspect one step at a time. +# +# The remarkable practical point, developed in the first half of this +# notebook, is that for a Gaussian variational family the natural gradient +# costs *no* Fisher matrix at all: the Fisher information turns out to be the +# Jacobian between two standard exponential-family coordinate systems, so the +# natural gradient with respect to one of them is the plain gradient with +# respect to the other. The second half turns the same geometry over and +# looks at it from a different storage convention — the *site*, or *dual*, +# parameterisation of {cite:t}`adam2021dual` — and asks what changes and what +# provably does not. +# +# The route is: +# +# 1. the exponential-family view of $q(\mathbf{u})$, its two canonical +# coordinate systems, and the third one GPJax actually stores; +# 2. the Fisher information is the Jacobian between them, checked +# numerically; +# 3. the mirror-descent reading of the step, and why $\gamma \le 1$ is +# special; +# 4. conjugate models, where one step at $\gamma=1$ is the exact answer — a +# single shared demo that reaches the same optimum through both the +# moment/whitened storage and the site/dual storage; +# 5. the site, or dual, reparameterisation of the same $q$, its EP heritage, +# and the two silent convention traps that wait in the source material; +# 6. the tied natural-gradient update in site coordinates, and why it never +# needs to invert anything; +# 7. cone-safety in both storage conventions — a negative-definite cone for +# the moments, a positive-semidefinite cone for the sites — with the +# numerical checks that locate exactly where each guarantee ends; +# 8. the two hyperparameter objectives, `elbo` and `dual_elbo`, and precisely +# what is proven about the gap between them, versus what is only +# measured; +# 9. practical guidance spanning both storage conventions. +# +# 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. + +# %% +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 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.objectives import dual_elbo, elbo + from gpjax.parameters import LowerTriangular, Real + from gpjax.variational_families import ( + DualVariationalGaussian, + VariationalGaussian, + WhitenedVariationalGaussian, + ) + +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] +# ## 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. Later in this notebook a fourth system joins them: the *site*, or +# *dual*, coordinates that `DualVariationalGaussian` stores instead of +# $(\mathbf{m},\mathbf{L})$. + +# %% [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, on a small non-conjugate model (Bernoulli likelihood, +# $M=3$). + +# %% +key, input_key, label_key, mean_key, root_key = jr.split(key, 5) + +fisher_inputs = jr.uniform(input_key, (30, 1), minval=-2.0, maxval=2.0) +fisher_labels = ( + jr.uniform(label_key, (30, 1)) < jax.nn.sigmoid(2.0 * fisher_inputs) +).astype(jnp.float64) +fisher_data = gpx.Dataset(X=fisher_inputs, y=fisher_labels) + +fisher_model = ( + gpx.gps.Prior(mean_function=gpx.mean_functions.Zero(), kernel=jk.RBF()) + * gpx.likelihoods.Bernoulli() +) + +num_fisher_inducing = 3 +fisher_mean = 0.5 * jr.normal(mean_key, (num_fisher_inducing, 1)) +fisher_factor = 0.5 * jr.normal(root_key, (num_fisher_inducing, num_fisher_inducing)) +fisher_root = jnp.linalg.cholesky( + fisher_factor @ fisher_factor.T + jnp.eye(num_fisher_inducing) +) +fisher_family = gpx.variational_families.VariationalGaussian( + model=fisher_model, + inducing_inputs=jnp.linspace(-2.0, 2.0, num_fisher_inducing).reshape(-1, 1), + variational_mean=fisher_mean, + variational_root_covariance=fisher_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), + fisher_family, + (Real(variational_mean), LowerTriangular(variational_root_covariance)), + ) + return negative_elbo(paramax.unwrap(trial), fisher_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_fisher_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_fisher_inducing)) + ) + + +def expectation_of_natural(flat): + """The map whose Jacobian is the Fisher information.""" + moments = moments_from_natural(*unpack(flat, num_fisher_inducing)) + return pack(*expectation_from_moments(*moments)) + + +flat_natural = pack(*natural_from_moments(fisher_mean, fisher_root)) +flat_expectation = pack(*expectation_from_moments(fisher_mean, fisher_root)) + +fisher_matrix = jax.jacfwd(expectation_of_natural)(flat_natural) +natural_gradient = jnp.linalg.solve( + fisher_matrix, 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_matrix - fisher_matrix.T)):.3e}" +) +print( + f"smallest eigenvalue of F : {jnp.min(jnp.linalg.eigvalsh(fisher_matrix)):.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 — the asymmetry is at +# float64 noise and the smallest eigenvalue is a healthy $0.67$ — and the +# natural gradient obtained by solving with it agrees with the plain +# gradient in expectation coordinates to $3.6\times10^{-15}$, machine +# precision for a problem of this size. 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. Every demo from +# here on uses the right-hand side of the identity, so that this +# $\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 the next +# section shows concretely, 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 — a fact this notebook returns to twice, once +# for each storage convention, in the cone-safety section below. + +# %% [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. Nothing in that argument refers to how $q$ +# is stored — it is a statement about the $(\boldsymbol{\theta}, +# \boldsymbol{\eta})$ geometry itself — so it has to hold equally for +# whatever storage convention we hand the step. We check that directly, on +# one shared problem, with two storage conventions at once. The second of +# them, `DualVariationalGaussian`, is not yet defined — that is the subject +# of the rest of this notebook — but it needs nothing more here than to be +# treated as a black box that also implements `natural_gradient_step`. +# +# The problem is a 1D conjugate regression, with a deliberately non-zero +# mean function so that neither branch gets a free pass on that front. + +# %% +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) +test_inputs = jnp.linspace(-3.2, 3.2, 300).reshape(-1, 1) + +# A conjugate SVGP: prior * likelihood, exactly as `Prior.__mul__` builds it. +regression_model = gpx.gps.Prior( + mean_function=gpx.mean_functions.Constant(jnp.array(prior_constant)), + kernel=jk.RBF(lengthscale=regression_lengthscale), + jitter=regression_jitter, +) * gpx.likelihoods.Gaussian(obs_stddev=noise_stddev) + +unwrapped_regression_model = paramax.unwrap(regression_model) +regression_kernel = unwrapped_regression_model.prior.kernel +regression_mean_function = unwrapped_regression_model.prior.mean_function + +# %% [markdown] +# The Titsias optimum, in the original (non-whitened) coordinates at the +# inducing points, is the reference both branches are checked against: +# +# $$\boldsymbol{\Lambda}_{\text{Tit}} = \mathbf{K}_{zz} + \sigma^{-2}\mathbf{K}_{zx}\mathbf{K}_{xz}, \qquad \mathbf{m}^\star = \boldsymbol{\mu}_z + \sigma^{-2}\mathbf{K}_{zz}\boldsymbol{\Lambda}_{\text{Tit}}^{-1}\mathbf{K}_{zx}(\mathbf{y}-\boldsymbol{\mu}_x), \qquad \mathbf{S}^\star = \mathbf{K}_{zz}\boldsymbol{\Lambda}_{\text{Tit}}^{-1}\mathbf{K}_{zz} .$$ + +# %% +Kzz = regression_kernel.gram(regression_inducing).as_matrix() +Kzz = Kzz + regression_jitter * jnp.eye(num_inducing) +Lz = jnp.linalg.cholesky(Kzz) +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) + +# %% [markdown] +# **Branch A: moment storage, whitened.** We use the whitened family, 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`. Numerically it helps a great deal, because +# $\mathbf{m}_w$ and $\mathbf{S}_w$ are $\mathcal{O}(1)$ regardless of the +# kernel scale. We start it from a deliberately bad initialisation and take +# one step at $\gamma=1$. + +# %% +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)) + +whitened_initial = WhitenedVariationalGaussian( + model=regression_model, + inducing_inputs=regression_inducing, + variational_mean=bad_mean, + variational_root_covariance=bad_root, +) +unwrapped_whitened_initial = paramax.unwrap(whitened_initial) + +whitened_variational, whitened_hyper = partition_variational(whitened_initial) +whitened_stepped_partition, whitened_loss_before = natural_gradient_step( + whitened_variational, + whitened_hyper, + regression_data, + negative_elbo, + 1.0, + map_jitter=0.0, +) +whitened_stepped = paramax.unwrap( + eqx.combine(whitened_stepped_partition, whitened_hyper) +) + +# Un-whiten to compare against the Titsias optimum in the original (u) space. +m_w = whitened_stepped.variational_mean +L_w = whitened_stepped.variational_root_covariance +S_w = L_w @ L_w.T +mu_z = regression_mean_function(regression_inducing) +whitened_mean_in_u = mu_z + Lz @ m_w +whitened_covariance_in_u = Lz @ S_w @ Lz.T + +# %% [markdown] +# **Branch B: site storage.** `DualVariationalGaussian` starts at +# $\boldsymbol{\lambda}=\mathbf{0}$, i.e. $q=p$ — there is no analogue of +# "deliberately bad" to choose, since every initialisation of this family +# is $q=p$. We take one step at $\rho=1$, the site branch's name for the +# same step size, and read $(\mathbf{m},\mathbf{S})$ off directly with +# `.moments()`; no un-whitening is needed here, because the sites are always +# stored relative to the un-whitened prior. + +# %% +dual_initial = DualVariationalGaussian( + model=regression_model, inducing_inputs=regression_inducing +) +dual_variational, dual_hyper = partition_variational(dual_initial) +dual_stepped_partition, dual_loss_before = natural_gradient_step( + dual_variational, dual_hyper, regression_data, negative_dual_elbo, 1.0 +) +dual_stepped = paramax.unwrap(eqx.combine(dual_stepped_partition, dual_hyper)) +dual_mean, dual_covariance = dual_stepped.moments() + +print(f"ELBO before the whitened step : {-whitened_loss_before:12.6f}") +print( + "ELBO after the whitened step : " + f"{float(elbo(whitened_stepped, regression_data)):12.6f}" +) +print( + "dual_elbo after the dual step : " + f"{float(dual_elbo(dual_stepped, regression_data)):12.6f}" +) +print( + "max |m_whitened - m*| (Titsias) : " + f"{jnp.max(jnp.abs(whitened_mean_in_u - optimal_mean)):.3e}" +) +print( + "max |S_whitened - S*| (Titsias) : " + f"{jnp.max(jnp.abs(whitened_covariance_in_u - optimal_covariance)):.3e}" +) +print( + "max |m_dual - m*| (Titsias) : " + f"{jnp.max(jnp.abs(dual_mean - optimal_mean)):.3e}" +) +print( + "max |S_dual - S*| (Titsias) : " + f"{jnp.max(jnp.abs(dual_covariance - optimal_covariance)):.3e}" +) +print( + "max |m_whitened - m_dual| : " + f"{jnp.max(jnp.abs(whitened_mean_in_u - dual_mean)):.3e}" +) +print( + "max |S_whitened - S_dual| : " + f"{jnp.max(jnp.abs(whitened_covariance_in_u - dual_covariance)):.3e}" +) + +# %% [markdown] +# One step from two completely different starting points and two completely +# different storage conventions — a whitened mean and Cholesky factor on one +# side, a pair of site parameters on the other — land on the same point to +# $3\times10^{-12}$ in the mean and $3\times10^{-13}$ in the covariance, both +# measured against the closed-form Titsias optimum, and to +# $1\times10^{-13}$ against *each other* directly. That gap is the float64 +# noise floor for a problem of this size; both ELBOs agree to the printed six +# decimal places. This is the cleanest statement this notebook can make +# about what "two coordinate systems for the same geometry" means: not an +# analogy, but the same arithmetic answer, reached two different ways. +# +# The plot makes the same point visually — the initial, deliberately absurd +# $q$ on the left, and the two stepped posteriors overlaid on the exact GP +# posterior on the right, indistinguishable from it and from each other. + +# %% +exact_posterior = unwrapped_regression_model.condition(regression_data) +exact_predictive = exact_posterior(test_inputs) +exact_mean = exact_predictive.mean +exact_stddev = jnp.sqrt(exact_predictive.variance) +whitened_predictive = whitened_stepped(test_inputs) +dual_predictive = dual_stepped(test_inputs) + +fig, axes = plt.subplots(ncols=2, figsize=(10, 3.0), sharey=True) +init_predictive = unwrapped_whitened_initial(test_inputs) +for ax, mean_curve, stddev_curve, title in [ + ( + axes[0], + init_predictive.mean, + jnp.sqrt(init_predictive.variance), + "Initialisation", + ), + ( + axes[1], + whitened_predictive.mean, + jnp.sqrt(whitened_predictive.variance), + "After one step (both branches)", + ), +]: + ax.scatter( + regression_inputs, + regression_outputs, + alpha=0.15, + 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, mean_curve, color=cols[1], label="Variational $q$ (whitened)") + ax.fill_between( + test_inputs.flatten(), + mean_curve - 2 * stddev_curve, + mean_curve + 2 * stddev_curve, + alpha=0.3, + color=cols[1], + ) + ax.set(xlabel=r"$x$", title=title, ylim=(-3.0, 3.0)) +axes[1].plot( + test_inputs, + dual_predictive.mean, + color=cols[2], + linestyle=":", + linewidth=2, + label="Variational $q$ (dual)", +) +clean_legend(axes[0]) +clean_legend(axes[1]) +axes[0].set_ylabel(r"$f(x)$") + +# %% [markdown] +# The rest of this notebook is about the second branch: what it stores, +# where the update in that section came from, and exactly when — not if, +# *when* — the two branches stop being the same iteration. +# +# **One notational break, from here on.** Above, $\boldsymbol{\theta}$ was +# the natural parameter of $q(\mathbf{u})$ and $\boldsymbol{\eta}$ was the +# expectation parameter. From here $\boldsymbol{\theta}$ is reserved for the +# kernel hyperparameters, which enter for the first time in the +# hyperparameter-learning section near the end. The natural parameter of +# $q(\mathbf{u})$ becomes $\boldsymbol{\eta}$, the expectation parameter +# becomes $\boldsymbol{\mu}$, and $\boldsymbol{\lambda}$ is the *site* — +# not the fixed conjugate-likelihood vector of the derivation just above, +# which will not be needed again. In these letters the Fisher identity reads +# $\tilde\nabla_{\boldsymbol{\eta}}\mathcal{L} = \partial\mathcal{L}/\partial\boldsymbol{\mu}$, +# and it is restated that way below. + +# %% [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 half of the 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 — exactly the +# $\boldsymbol{\mu}_z$ that made the shared demo above need a non-zero mean +# function to be a fair test. `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 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? +# {cite:t}`adam2021dual` 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. For a Gaussian +# likelihood $\alpha_i = (y_i - m_i)/\sigma^2$ and $\beta_i = 1/\sigma^2$; +# here are both, by autodiff. + +# %% +key, alpha_beta_key = jr.split(key) +ep_response = jr.normal(alpha_beta_key, (5, 1)) +ep_mean = jnp.linspace(-1.0, 1.0, 5) +ep_variance = jnp.linspace(0.2, 0.9, 5) +ep_stddev = 0.37 +ep_likelihood = gpx.likelihoods.Gaussian(obs_stddev=ep_stddev) + + +def total_expected_log_likelihood(mean, variance): + """Summed variational expectation, as a function of the marginal moments.""" + return jnp.sum( + ep_likelihood.expected_log_likelihood( + ep_response, mean[:, None], variance[:, None] + ) + ) + + +bonnet_alpha, price_derivative = jax.grad( + total_expected_log_likelihood, argnums=(0, 1) +)(ep_mean, ep_variance) +price_beta = -2.0 * price_derivative + +closed_form_alpha = (ep_response.squeeze(-1) - ep_mean) / ep_stddev**2 +closed_form_beta = jnp.full_like(ep_mean, 1.0 / ep_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 / {ep_stddev}^2)") + +# %% [markdown] +# Both match the closed form to $10^{-14}$. 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 silent convention traps +# +# 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 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}_{zz}\boldsymbol{\lambda}_1$, +# $\bar{\boldsymbol{\Lambda}}_2 = \mathbf{K}_{zz}\boldsymbol{\Lambda}_2\mathbf{K}_{zz}$. +# Both conventions describe the same $q$ and give the same $\mathbf{R}$, but +# they are not interchangeable for our purposes: the additive, +# hyperparameter-free split of the previous section holds *exactly* only in +# the flanked convention, since in the un-flanked one +# $\boldsymbol{\eta}_1 = \mathbf{K}_{zz}^{-1}\bar{\boldsymbol{\lambda}}_1$ +# moves with $\boldsymbol{\theta}$ too. 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 across its own equations: a natural-parameter one +# ($-\tfrac12\beta_i$) and a precision one ($g_{2,i} = \beta_i$). The dense +# limit settles it — with $\mathbf{Z}=\mathbf{X}$, +# $\mathbf{S}^{-1} = \mathbf{K}_{ff}^{-1} + \operatorname{diag}(\boldsymbol{\beta})$ +# 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. That +# is dramatic entrywise, invisible in everything anyone actually reads off +# the model, and cheap to measure. Put the inducing inputs on the data +# themselves, $\mathbf{Z}=\mathbf{X}$, for the first forty points of the +# conjugate regression above — the worst case for +# $\operatorname{cond}(\mathbf{K}_{zz})$, and the one configuration in +# which the optimal sites are known on paper: +# $\boldsymbol{\lambda}_1 = (\mathbf{y}-\boldsymbol{\mu}_x)/\sigma^2$ and +# $\boldsymbol{\Lambda}_2 = \mathbf{I}/\sigma^2$. One $\rho=1$ step lands +# on that fixed point, exactly as in the shared demo, and what the family +# then *stores* can be compared entrywise against the paper answer. + +# %% +dense_count = 40 +dense_inputs = regression_inputs[:dense_count] +dense_data = gpx.Dataset(X=dense_inputs, y=regression_outputs[:dense_count]) + +dense_variational, dense_hyper = partition_variational( + DualVariationalGaussian(model=regression_model, inducing_inputs=dense_inputs) +) +dense_variational, _ = natural_gradient_step( + dense_variational, dense_hyper, dense_data, negative_dual_elbo, 1.0 +) +dense_fitted = paramax.unwrap(eqx.combine(dense_variational, dense_hyper)) + +dense_gram = regression_kernel.gram(dense_inputs).as_matrix() + ( + regression_jitter * jnp.eye(dense_count) +) +dense_centred = dense_data.y - regression_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_fitted.dual_vector - 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_fitted.dual_matrix - 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)" +) + +# %% [markdown] +# The stored matrix is wrong by $8.4$ entrywise — against an analytic +# answer whose every diagonal entry is $1/\sigma^2 \approx 11.1$ — because +# forming it routes through $\mathbf{K}_{zz}^{-1}$ at a condition number of +# $10^{9}$. Yet the flanked quantity +# $\mathbf{K}_{zz}\boldsymbol{\lambda}_1$, which is what everything +# downstream actually consumes, is right to $2.6\times10^{-9}$ in relative +# error: the error lives in the near-null space of $\mathbf{K}_{zz}$ and is +# annihilated on the way back out. A measurement in one configuration, not +# a theorem — and the rule it teaches is the one to keep: **GPJax stores +# the flanked, precision convention; never test $\boldsymbol{\Lambda}_2$ +# entrywise — test $\mathbf{R}$, the moments, or the predictions instead.** + +# %% [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 Fisher identity from the first half of this +# notebook, restated in the new letters — 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. GPJax supplies it. +# +# 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$ from the first half of this 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; the shared conjugate demo two sections back already +# used `1.0` for both. The cone-safety section below checks the two +# branches against each other step by step at $\rho=\gamma=0.8$. +# +# It is worth being concrete about what the site step does *not* do. A +# natural-gradient step in the moment 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 site 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 directly; +# 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 | site (dual) | 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 site side of the table. +# Whether that turns into wall-clock depends on how large a share of the +# iteration the saved work was; the [dual sparse GP notebook](dual_svgp.py) +# measures it on a real training run. What is certain is the direction of any difference: since +# the iterates are the same either way (subject to the cone-safety condition +# below), 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 the caveat +# that their "implementation is not as optimized as SVGP in GPflow." + +# %% [markdown] +# ## Cone-safety +# +# The step is +# $\boldsymbol{\theta}\leftarrow\boldsymbol{\theta} - \gamma\,\partial\ell/\partial\boldsymbol{\eta}$ +# (moment coordinates) or +# $\boldsymbol{\lambda} \leftarrow (1-\rho)\boldsymbol{\lambda} + \rho\,\nabla_{\boldsymbol{\mu}}\mathcal{L}_{\text{ell}}$ +# (site coordinates), and each has a cone it must not leave: +# $\boldsymbol{\Theta}_2 \prec 0$ because it is $-\tfrac12$ a covariance +# inverse, and $\boldsymbol{\Lambda}_2 \succeq 0$ because it is a precision. +# Nothing in either update enforces that automatically. We take the two +# cones in turn. +# +# **Moment coordinates: the negative-definite cone.** 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 the step: +# +# $$\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 (moments).** 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 *as computed* — GPJax's `inv_probit` clips its output into +# $[10^{-3},1-10^{-3}]$, which flattens $\log p$ enough to give it a +# positive second derivative for $f\lesssim-2.44$, so even the Bernoulli +# model below leaves the guaranteed regime once a point is confidently +# mislabelled. We sweep $\gamma$ from an over-confident starting point, +# $\mathbf{S}_0=10^{-2}\mathbf{I}$, sharper than the target — precisely the +# regime where extrapolation bites — on the two-dimensional "banana" +# classification problem used again 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) +num_train = 1600 +banana_train = gpx.Dataset(X=banana_inputs[:num_train], y=banana_labels[:num_train]) + +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() +) + +overconfident_family = 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) + +gamma_values = jnp.array([0.1, 0.5, 1.0, 2.0, 5.0, 10.0]) +largest_eigenvalues = jnp.array( + [ + jnp.max(jnp.linalg.eigvalsh(natural_matrix - g * matrix_gradient)) + for g in gamma_values + ] +) +print("gamma max eig(Theta2_new) status") +for gamma, largest in zip(gamma_values, largest_eigenvalues, strict=True): + status = "negative definite" if largest < 0 else "*** LEFT THE CONE ***" + print(f"{float(gamma):6.2f} {float(largest):+18.5f} {status}") + +fig, ax = plt.subplots(figsize=(5.5, 3.2)) +ax.plot(gamma_values, largest_eigenvalues, marker="o", color=cols[1]) +ax.axhline(0.0, color="black", linestyle="--", linewidth=1) +ax.axvline(1.0, color="gray", linestyle=":", linewidth=1) +ax.set( + xlabel=r"$\gamma$", + ylabel=r"$\lambda_{\max}(\boldsymbol{\Theta}_2^{\text{new}})$", + title="Where this initialisation leaves the cone", +) + +# %% [markdown] +# Read that as a statement about *this initialisation*, not about +# $\gamma=2$ in general. 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 little room to extrapolate +# into: the sign flips between $\gamma=1$ ($-4.67$) and $\gamma=2$ +# ($+40.66$), and linear interpolation puts the crossing at +# $\gamma\approx1.10$. What the theorem actually guarantees is +# $\gamma\in[0,1]$, for any log-concave likelihood and any starting point, +# and nothing whatsoever beyond that. +# +# 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` 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`. This backoff is specific +# to the moment branch — the site branch's affine update needs no such +# rescue, for the reason the rest of this section develops. +# +# **Site coordinates: the positive-semidefinite cone.** The tied update +# derived above is affine in $\boldsymbol{\lambda}$, so for +# $\rho\in[0,1]$ a convex combination of $\boldsymbol{\Lambda}_2\succeq0$ +# and a PSD target stays PSD automatically — no Cholesky-validity check is +# needed, because there is nothing to fail. The one place this can still go +# wrong is upstream of the convex combination: the target itself is built +# from Price's curvature $\beta_i$, and $\beta_i\ge0$ needs the same +# log-concavity condition as the moment branch's cone-safety theorem — as +# *computed*, not as written. GPJax's dual step guards this with a floor, +# `beta_floor` (default $10^{-8}$), clipping $\boldsymbol{\beta}$ from below +# before it enters $\boldsymbol{\Lambda}_2$. It is $\boldsymbol{\beta}$ that +# is clipped, never $\boldsymbol{\Lambda}_2$ itself, so the update stays +# affine and `jit`/`scan`-safe. +# +# Since $\rho=\gamma$, the two branches should step *identically* — the same +# $(\mathbf{m},\mathbf{S})$ at every iteration — for as long as the computed +# $\beta_i$ stay non-negative, and only then. We check this directly: six +# matched $\rho=\gamma=0.8$ steps on the banana problem, both branches +# started at $q=p$, comparing $(\mathbf{m},\mathbf{S})$ after every step and +# recording Price's curvature just before it. + + +# %% +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 + + +def make_banana_moment_family(): + """A fresh SVGP over the banana data, at q = p.""" + banana_gram = paramax.unwrap(banana_model).prior.kernel.gram( + banana_inducing + ).as_matrix() + 1e-6 * jnp.eye(num_banana_inducing) + return VariationalGaussian( + model=banana_model, + inducing_inputs=banana_inducing, + variational_mean=jnp.zeros((num_banana_inducing, 1)), + variational_root_covariance=jnp.linalg.cholesky(banana_gram), + ) + + +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 + + +floored_rows = six_matched_steps(1e-8) +print("step |(m, S) gap| beta < 0 min beta its marginal mean") +for step, (gap, negative_count, smallest, mean_there) in enumerate( + floored_rows, 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 floored_rows) +unfloored_rows = six_matched_steps(-jnp.inf) +unfloored_gap = max(gap for gap, _, _, _ in unfloored_rows) +print(f"\nworst gap, default beta_floor = 1e-8 : {banana_gap:.3e}") +print(f"worst gap, clip disabled (-inf) : {unfloored_gap:.3e}") + +fig, ax = plt.subplots(figsize=(5.5, 3.2)) +steps = jnp.arange(1, 7) +ax.plot( + steps, + jnp.array([g for g, _, _, _ in floored_rows]), + marker="o", + color=cols[1], + label="default beta_floor", +) +ax.plot( + steps, + jnp.array([g for g, _, _, _ in unfloored_rows]), + marker="x", + color=cols[0], + label="clip disabled", +) +ax.set( + xlabel="Step", + ylabel=r"$\max|(\mathbf{m}, \mathbf{S})_{\text{site}} - (\mathbf{m}, \mathbf{S})_{\text{moment}}|$", + yscale="log", + title="The two branches, step by step", +) +clean_legend(ax) + +# %% [markdown] +# For the first four steps every $\beta_i$ is positive, the clip does +# nothing, and the two branches agree to $\sim10^{-13}$ — the float64 noise +# floor. At step five a single training point out of 1600 crosses into +# $\beta_i<0$ — a label-$0$ point, whose log-likelihood is the +# $f\mapsto-f$ mirror of the $y=1$ case, so the $-2.44$ threshold derived +# above sits at $+2.44$ for it, and its marginal mean has just reached +# $+2.427$ — the `beta_floor` clip engages, and from +# that step the gap jumps to $\mathcal{O}(10^{-3})$ and compounds at step +# six. Disabling the clip (`beta_floor=-jnp.inf`, the second line above) +# brings the same six steps back to the noise floor — $8.9\times10^{-13}$ +# — which pins the cause down precisely: it is neither conditioning nor the +# cancellation in $\mathbf{H}_2=\mathbf{S}+\mathbf{m}\mathbf{m}^\top$ that +# the moment branch has to undo, since disabling the one thing that differs +# between the branches removes the discrepancy entirely. The condition the +# $\rho=\gamma$ identity needs — $\beta_i\ge0$, log-concavity as computed — +# is real, and this is exactly where and how it fails, on the same problem +# and the same clip that guarantees the site branch never leaves its own +# cone. The residual is still far below anything visible in the ELBO, +# which is the number either optimiser is steering by. + +# %% [markdown] +# ## The M-step objective: `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}$ | +# | $\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 first row is free: any Gaussian $q$, whatever produced it, gives a +# valid ELBO. The second and third rows are the envelope theorem doing real +# work — at a stationary $q$ the *implicit* dependence of +# $\boldsymbol{\eta}_0(\boldsymbol{\theta})$'s contribution on +# $\boldsymbol{\theta}$ contributes nothing to the total derivative, so it +# does not matter whether the prior half of $q$ is allowed to move with +# $\boldsymbol{\theta}$ or not; away from stationarity it matters a great +# deal, since nobody runs an E-step to convergence between Adam steps in +# practice. The fourth and fifth rows need the sites to carry no implicit +# $\boldsymbol{\theta}$-dependence, which is only exactly true without +# sparsity: with $\mathbf{Z}\neq\mathbf{X}$ the flanked sites still route +# through $\mathbf{K}_{zx}(\boldsymbol{\theta})$, so freezing them at +# $\boldsymbol{\theta}_t$ can make $\bar l$ sub-optimal, and in fact +# non-dominant, elsewhere. The sixth row is honestly labelled: the paper +# measures a speed-up and does not derive one. +# +# The claim we *can* check directly here is rows two and three: value and +# gradient equality, and how quickly they set in as the E-step converges. +# Two matched families — sites and moments, started at the same $q$ — take +# an increasing number of $\rho=\gamma=0.8$ E-steps before we read off the +# hyperparameter gradient of each bound. + +# %% +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"dual_elbo - elbo at the shared q = p init : {shared_bound:.3e}") + + +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("\nE-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] +# At the shared initialisation the two bounds already agree to +# $4.3\times10^{-5}$ nats — both are the ELBO at $q=p$, so the only source +# of disagreement is the jitter each objective's Cholesky picks up +# differently, not a real difference in value. The gradient row is the one +# that matters: at zero E-steps the two hyperparameter gradients disagree by +# as much as their own magnitude ($39.0$ against a norm of $39.5$), and as +# the E-step is allowed to run longer the disagreement collapses +# geometrically — $0.80$, then $0.048$, then $9.1\times10^{-4}$, down to +# $2.9\times10^{-14}$ by 60 steps — exactly the envelope-theorem prediction +# that the two gradients coincide once, and only once, $q$ has actually +# stationarised. Away from that limit they are not close: they are +# different vectors, of comparable size, pointing the M-step in different +# directions. This is the fourth-row caveat made concrete on a specific +# model, not a claim that one direction is better; the +# [dual sparse GP notebook](dual_svgp.py) runs a full variational-EM loop +# on both objectives and checks which one actually gets further, which a +# static gradient comparison cannot answer. + +# %% [markdown] +# ## Practical guidance +# +# Guidance that applies to **either** storage convention: +# +# * **Conjugate and full batch: use $\gamma=\rho=1$.** One iteration is the +# exact solution — the shared demo above reached it from two different +# starting points and two different storage conventions, both to +# $\sim10^{-12}$ — and further iterations are fixed points. +# * **Never exceed a step size of $1$.** The convex-combination guarantee +# stops there for both branches. On the moment side the backoff exists to +# catch mistakes, not to enable them; on the site side there is no backoff +# at all, because the update never needs rescuing within $[0,1]$. +# * **Non-log-concave likelihoods have no guarantee at all, on either +# branch — as *computed*, not as written.** GPJax's `inv_probit` clips +# its output, which flattens $\log p$ enough to break log-concavity past +# $f\approx-2.44$; a Student-$t$ likelihood is not log-concave anywhere +# near that mild. The moment branch's target can leave the +# negative-definite cone; the site branch's computed $\beta_i$ can go +# negative and rely on `beta_floor` to stay safe. Neither is a defect in +# the optimiser — both are a property of the likelihood. +# * **The natural gradient buys optimiser speed, not a better $q$ at the +# same $\boldsymbol{\theta}$.** Wherever the cone-safety condition holds, +# the two branches are the *same iteration*; the only thing either +# storage convention can change is how fast that iteration is computed, +# and how the M-step behaves once $q$ moves. A more sharply peaked $q$ +# only appears if the underlying variational optimum actually is one. +# +# Guidance specific to **moment storage** ($(\mathbf{m},\mathbf{L})$, via +# `VariationalGaussian` or `WhitenedVariationalGaussian`): +# +# * **Non-conjugate or mini-batched: ramp $\gamma$.** Salimbeni et al. +# recommend starting around $10^{-4}$ and reaching $\approx10^{-1}$ +# "sufficiently quickly ($<1000$ iterations)"; `natgrad_lr` accepts any +# Optax schedule. +# * **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 only when fighting an ill-conditioned +# $\mathbf{S}$. +# * **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. +# +# Guidance specific to **site storage** +# (`DualVariationalGaussian`): +# +# * **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 +# would need to be re-derived for a multi-output model. +# * **Flanked storage squares $\operatorname{cond}(\mathbf{K}_{zz})$.** +# Benign 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. +# * **`beta_floor` is not a no-op for Bernoulli.** It is what breaks the +# $\rho=\gamma$ identity once a point is confidently mislabelled, and it +# is doing exactly its job when it does — keeping the update inside the +# PSD cone rather than letting a negative $\beta_i$ push it out. + +# %% [markdown] +# ## System configuration + +# %% +# %reload_ext watermark +# %watermark -n -u -v -iv -w -a 'Thomas Pinder' diff --git a/docs/index.md b/docs/index.md index 54d657efb..5de596478 100644 --- a/docs/index.md +++ b/docs/index.md @@ -152,6 +152,7 @@ examples/intro_to_kernels examples/regression examples/classification examples/poisson +examples/natural_gradients ``` ```{toctree}