NaturalGuide + ConjugateVI diverges on non-conjugate likelihoods at every damping
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- Avg merge
- 19h 2m
- Merged PRs (30d)
- 16
Description
## Summary
`NaturalGuide` trained with `ConjugateVI.step` diverges on a non-conjugate problem, at every damping value I tried. The other guides (`MeanFieldGuide`, `FullRankGuide`, `WhitenedGuide`, `DeltaGuide`) all converge sensibly on the same prior, likelihood and data, so this looks specific to the natural-parameter update rather than to the problem being hard.
## Observed
Latent GP + GEV likelihood, 107 inducing points, `GaussHermiteIntegrator(order=20)`. After 30–40 CVI steps the ELBO is not merely poor, it is astronomically large:
| damping | ELBO loss after training |
|---|---|
| 1.0 | 1.5e10 |
| 0.3 | 3.7e9 |
| 0.1 | 1.7e9 |
| 0.05 | 1.2e9 |
On a smaller variant (40 inducing points, Gumbel likelihood) `damping=0.3` reached ~1e250 and `damping=0.1` / `0.05` raised `EquinoxRuntimeError` instead. For scale, the converged guides on the same problem sit at ELBO losses of roughly 350–730.
Because there is no damping at which this is merely inaccurate, I left `NaturalGuide` out of a published benchmark rather than report a misleading score.
## Reproduction
pyrox-gp 0.1.0, gaussx 0.0.18, jax 0.10.2, Python 3.13.
```python
import jax, jax.numpy as jnp, numpy as np
jax.config.update("jax_enable_x64", True)
from gaussx import GaussHermiteIntegrator, DenseSolver
import pyrox_gp as P
import numpyro.distributions as dist
S = 40
rng = np.random.default_rng(0)
X = jnp.asarray(rng.normal(size=(S, 2)))
y = jnp.asarray(rng.normal(40.0, 2.0, size=S))
solver = DenseSolver()
prior = P.SparseGPPrior(
kernel=P.Matern(nu=1.5, init_variance=9.0, init_lengthscale=0.9),
Z=X, mean_fn=lambda Xq: jnp.full(Xq.shape[0], 40.0), solver=solver,
)
lik = P.DistLikelihood(lambda f: dist.Gumbel(loc=f, scale=2.0))
integ = GaussHermiteIntegrator(order=20)
for damping in (0.3, 0.1, 0.05):
g = P.NaturalGuide.init(S, scale=1.0, solver=solver)
cvi = P.ConjugateVI(damping=damping, integrator=integ)
for _ in range(30):
g = cvi.step(prior, g, lik, X, y)
print(damping, -float(P.svgp_elbo(prior, g, lik, X, y, integrator=integ)))
```
Swapping `P.NaturalGuide` + `ConjugateVI` for `P.WhitenedGuide` + an optax loop on the same `svgp_elbo` converges normally, which is what makes me think the problem is in the natural-parameter update and not the setup.
## Possible directions
I have not read the CVI implementation closely enough to propose a fix, but the shape of the failure (blow-up rather than slow progress, insensitive to damping) suggests one of:
- a sign or factor-of-two convention mismatch between the natural parameters `(θ₁, θ₂) = (Σ⁻¹μ, -½Σ⁻¹)` and what `natural_update` expects;
- the projected site precisions not being constrained to keep `Σ` positive-definite, so one bad step inverts the curvature and every later step amplifies it;
- the non-conjugate expected-statistics path (with an integrator) differing in convention from the conjugate path that presumably has test coverage.
A regression test that runs a handful of CVI steps against a Gaussian likelihood — where the closed form is available and the answer is known — would catch whichever of these it is.
## Note
Worth pairing with jejjohnson/gaussx#229: `safe_cholesky`'s `while_loop` blocks reverse-mode AD, which independently makes `MeanFieldGuide` and `FullRankGuide` untrainable through `svgp_elbo`. Between the two, three of the five guides are currently unusable as shipped for gradient-based training on a non-conjugate likelihood.
Found while benchmarking `pyrox-gp` inference methods for a latent GP with a GEV likelihood.
Contributor guide
Research direction
Start with ConjugateVI.step, NaturalGuide, and the natural_update path, then compare them with the svgp_elbo and WhitenedGuide behavior in the reproduction. Check the natural-parameter conventions and projected precisions during several non-conjugate steps. Add a Gaussian-likelihood regression test with a few CVI steps and verify that the result remains finite and approaches the known closed-form answer.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, testing-qa
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100