Feat: Implement Full-Rank VI

[gil2rok]

Implement variational inference (VI) with full-rank Gaussian approximation

While mean-field VI learns a Gaussian approximation $N(\mu, \sigma I)$ with diagonal covariance $\sigma I$, full-rank VI learns a Gaussian approximation $N(\mu, \Sigma)$ with full-rank covariance $\Sigma$.

We use the Cholesky decomposition $\Sigma = C C^T$ parameterized by the log standard deviation $\rho$ and lower triangle matrix $L$ such that $C = \exp(\rho) + L$. This (1) ensures $\Sigma$ remains symmetric and positive-definite during optimization ¹ and (2) admits better sampling and log density computation (with improved time complexity, space complexity, and numerical stability)². Thus the full-rank Gaussian is parameterized by $(\mu, \rho, L)$.

To-Do:

• Support arbitrary PyTrees instead of just arrays
• Decide on format for lower triangular matrix $L$ ($n$-length flattened array vs dense $2D$ array for $n = d(d+1)/2$)
• Confirm initialization strategy from standard normal
• Write tests (current test fails)

Footnotes

1. Automatic Differentiation Variational Inference Section 2.4 ↩

2. Lecture notes ↩

[gil2rok]

@junpenglao I am trying to better grok JAX's PyTrees. I would love specific feedback on how to support PyTrees instead of just JAX arrays in my code. I believe this mostly affects my init(), _sample(), and generate_fullrank_logdensity() functions.

For example, consider my sampling implementation vs that in meanfield_vi.py:

def _sample(rng_key, mu, rho, L, num_samples):
    cholesky = jnp.tril(L, k=-1) + jnp.diag(jnp.exp(L))
    eps = jax.random.normal(rng_key, (num_samples,) + mu.shape)
    return mu + eps @ cholesky.T

def _sample(rng_key, mu, rho, num_samples):
    sigma = jax.tree.map(jnp.exp, rho)
    mu_flatten, unravel_fn = jax.flatten_util.ravel_pytree(mu)
    sigma_flat, _ = jax.flatten_util.ravel_pytree(sigma)
    flatten_sample = (
        jax.random.normal(rng_key, (num_samples,) + mu_flatten.shape) * sigma_flat
        + mu_flatten
    )
    return jax.vmap(unravel_fn)(flatten_sample)

How + why would I change my code? Any resource recommendations would be greatly appreciated.

[junpenglao]

PyTree with covariate matrix is a headache to deal with. I suggest you take a look at how pathfinder deal with it internally: basically all the state parameter is represented as flatten array, and you only unflatten at the end. Considering what you have right now already assume everything is a flatten array, you just need to add the flatten and unflatten part.

[junpenglao]

Why not use multivariate_normal.logpdf from JAX?

[gil2rok]

I wanted to avoid computing the inverse and log determinant of the covariance matrix $\Sigma$ by using the Cholesky factor when computing the logpdf.

Does jax.random.multivariate_normal.logpdf take Cholesky factors as input? I want to avoid needing to compute the covariance matrix $\Sigma = C C^T$, and then pass it into JAX's multivariate normal which separates it back into the Cholesky factor $C$.

From https://jax.readthedocs.io/en/latest/_autosummary/jax.random.multivariate_normal.html it appears the multivariate normal log density only accepts the covariance as a dense matrix!

See jax-ml/jax#11386. Thoughts on tradeoff btwn readability (with JAX's multivariate normal) and speed (custom implementation)?

[gil2rok]

A couple notes about the Cholesky factor:

1. I believe my implementation is incorrect. I compute the $d \times d$ Cholesky factor as $C = \exp(\rho) + L$ for log standard deviation $\rho$ and lower triangle matrix $L$. However, I store $L$ as a $d \times d$ matrix instead of just the lower triangle:

L = jax.tree.map(lambda x: jnp.zeros((*x.shape, x.shape)), position)

During optimization, this means we may learn non-zero values for the diagonal and upper triangle in $L$! Thus I instead will store $L$ as list of length $n$ representing the number of elements in the lower triangle of a $d \times d$ matrix excluding the diagonal. Thus $n = \sum_{i=1}^{d-1} i = (d - 1)(d) / 2$. We now construct the Cholesky factor as:

def unflatten_lower_triangular(tril_flat):
    n = tril_flat.size  # Number of elements in the lower triangular part
    d = int(jnp.sqrt(1 + 8 * n) - 1) // 2  # Dimension of the original matrix
    
    lower_tri_matrix = jnp.zeros((d, d))
    indices = jnp.tril_indices(d, k=-1)  # Indices for the lower triangle (excluding the diagonal)
    return lower_tri_matrix.at[indices].set(tril_flat)

C = jnp.diag(jnp.exp(rho)) + unflatten_lower_triangular(L)

2. Given my full-rank Gaussian parameterization as $(\mu, \rho, L)$, how should I flatten/unflatten to work with arbitrary Pytrees? Following @junpenglao 's advice, I see Pathfinder's approximate function:

def approximate(...):
    initial_position_flatten, unravel_fn = ravel_pytree(initial_position)
    ...
    unravel_fn_mapped = jax.vmap(unravel_fn)
    pathfinder_result = PathfinderState(
        elbo,
        unravel_fn_mapped(position),
        unravel_fn_mapped(grad_position),
        alpha,
        beta,
        gamma,
    )
    ...

How would I apply this to my _sample() function? I'm struggling to understand why flattening helps us handle arbitrary PyTrees, and how this works for the lower triangle $L$ vs mean $\mu$ or log std $\rho$. Is this correct or is more flattening needed?

def _sample(rng_key, mu, rho, L, num_samples):
    cholesky = jnp.diag(jnp.exp(rho)) + unflatten_lower_triangular(L)
    eps = jax.random.normal(rng_key, (num_samples,) + mu.shape)
    return mu + eps @ cholesky.T

[gil2rok]

1. Largely disregard my previous comments about the Cholesky factor in Feat: Implement Full-Rank VI #720 (comment). I would love some feedback on my current implementation, which more closely follows the original pull request Implement Fullrank vi #479.
2. For simplicity, I changed the code to use jax.random.multivariate_normal.logpdf even though it redundantly recommits the Cholesky decomposition internally. Do you think I should keep it or revert back to my more efficient implementation?

[junpenglao]

You have a good point about computation efficiency, let's keep the cholesky version of the logpdf. Could you rewrite the API to similar to multivariate_normal.logpdf, and use function partial when you call it?

[gil2rok]

Can you specify the API you imagine my implementation would have and how function partial would be applied to it? Feel free to reference my code I just committed.

[junpenglao]

I see, you are following the pattern in meanfield VI, I think there are a few places need refactoring there. Let me send out a PR so you see what i meant.

[gil2rok]

I see, you are following the pattern in meanfield VI, I think there are a few places need refactoring there. Let me send out a PR so you see what i meant.

Once you finish refactoring Meanfield VI, I'm happy to adapt the new style. Let me know when you're done!

[gil2rok]

Update: need to figure out why the full covariance matrix is not being recovered. May need to come back to this in a few weeks b/c of some deadlines.