Problem
NaturalVariationalGaussian.prior_kl (gpjax/variational_families.py:~522) computes a triangular root of S via the reversed-Cholesky trick, then throws it away:
sqrt_inv = jnp.swapaxes(
jnp.linalg.cholesky(S_inv[..., ::-1, ::-1])[..., ::-1, ::-1], -2, -1
)
sqrt = jsp.linalg.solve_triangular(sqrt_inv, jnp.eye(m), lower=True)
S = lx.MatrixLinearOperator(sqrt @ sqrt.T) # <- root discarded into a dense S
...
return qu.kl_divergence(pu) # <- generic KL re-factorises it
This is the same pattern fixed for VariationalGaussian / WhitenedVariationalGaussian in #665 (see #708), which deliberately left this family alone to stay scoped.
Measured cost
jnp.linalg.cholesky calls per prior_kl() on main, m = 5:
| family |
current |
minimum |
NaturalVariationalGaussian |
5 |
2 |
ExpectationVariationalGaussian |
4 |
2 |
The minimum of 2 for the natural family is the reversed Cholesky it needs anyway (to form μ = S θ₁) plus one for Kzz.
Note the counts above are against main; #707 reduces the generic _kl_divergence from 4 factorisations to 2, which will absorb part of this on its own. The remaining waste — discarding sqrt and rebuilding S densely — is specific to this family and needs the same closed-form treatment #708 applied:
KL = ½( ‖Lz⁻¹(m_z − μ)‖² + ‖Lz⁻¹ L‖²_F − m + 2Σ log[Lz]ᵢᵢ − 2Σ log|Lᵢᵢ| )
with L = sqrt already in hand.
ExpectationVariationalGaussian is not the same bug
Worth stating explicitly so it isn't "fixed" by mistake: ExpectationVariationalGaussian builds S = η₂ − η₁η₁ᵀ, which is genuinely dense with no root available, so factorising it once is legitimate. Its surplus over the minimum is entirely the generic-KL double-factorisation that #707 addresses, not a variational-families problem.
Acceptance criteria
NaturalVariationalGaussian.prior_kl down to 2 factorisations.
- Value and gradient equivalence to the current implementation, pinned with literals captured beforehand — gradients w.r.t.
natural_vector and natural_matrix specifically, since those are what fit optimises.
- A factorisation-count test that fails before the change.
- jit / grad / vmap clean.
Related caveat
#666 notes these families store their matrices as unconstrained Real and can leave the PD cone under plain Adam. That's a separate correctness concern; this issue is purely about redundant factorisations and should not be taken as a sign the family is otherwise ready for use.
Found while verifying #665 (see #708).
Problem
NaturalVariationalGaussian.prior_kl(gpjax/variational_families.py:~522) computes a triangular root ofSvia the reversed-Cholesky trick, then throws it away:This is the same pattern fixed for
VariationalGaussian/WhitenedVariationalGaussianin #665 (see #708), which deliberately left this family alone to stay scoped.Measured cost
jnp.linalg.choleskycalls perprior_kl()onmain, m = 5:NaturalVariationalGaussianExpectationVariationalGaussianThe minimum of 2 for the natural family is the reversed Cholesky it needs anyway (to form
μ = S θ₁) plus one forKzz.Note the counts above are against
main; #707 reduces the generic_kl_divergencefrom 4 factorisations to 2, which will absorb part of this on its own. The remaining waste — discardingsqrtand rebuildingSdensely — is specific to this family and needs the same closed-form treatment #708 applied:with
L = sqrtalready in hand.ExpectationVariationalGaussianis not the same bugWorth stating explicitly so it isn't "fixed" by mistake:
ExpectationVariationalGaussianbuildsS = η₂ − η₁η₁ᵀ, which is genuinely dense with no root available, so factorising it once is legitimate. Its surplus over the minimum is entirely the generic-KL double-factorisation that #707 addresses, not a variational-families problem.Acceptance criteria
NaturalVariationalGaussian.prior_kldown to 2 factorisations.natural_vectorandnatural_matrixspecifically, since those are whatfitoptimises.Related caveat
#666 notes these families store their matrices as unconstrained
Realand can leave the PD cone under plain Adam. That's a separate correctness concern; this issue is purely about redundant factorisations and should not be taken as a sign the family is otherwise ready for use.Found while verifying #665 (see #708).