Skip to content

NaturalVariationalGaussian.prior_kl discards the triangular root it already computes #710

Description

@thomaspinder

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).

Metadata

Metadata

Assignees

No one assigned

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions