jejjohnson / jejjohnson/pyrox

NaturalGuide + ConjugateVI diverges on non-conjugate likelihoods at every damping

Open
#219 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.