Fullrank VI implementation

[xidulu]

In mean-field VI, the variational posterior for all latent variables are independent therefore we can directly utilize tree_map to compute the density https://github.com/blackjax-devs/blackjax/blob/main/blackjax/vi/meanfield_vi.py#L127

However, for fullrank VI, the variational posterior is a multivariate Gaussian over all latent variables, so I assume we have to first flatten all the latent variables in order to compute log_q, is that correct?

[rlouf]

Indeed, look at what we do in /mcmc/metrics.py to compute the kinetic energy when the mass matrix is diagonal/dense: https://github.com/blackjax-devs/blackjax/blob/main/blackjax/mcmc/metrics.py

[xidulu]

@rlouf
Can I directly copy-and-paste code from NumPyro's util function?
https://github.com/pyro-ppl/numpyro/blob/master/numpyro/distributions/util.py#L324
I need it to convert unnormalized vector to lower-Cholesky factor

[junpenglao]

I would prefer we implement it ourselves. There are many ways to implement this, for example, concatenating the original vector with itself, reshape to square, then take the upper or lower triangular part:

vect = np.arange(1, 15+1)
m = len(vect)
# m = n(n+1)/2., solve n = .5 * (-1 + sqrt(1 + 8 m))
n = int(.5 * (-1 + np.sqrt(1 + 8 * m)))
vect_extended = np.concatenate([vect, vect[n:][::-1]])
mat = np.reshape(vect_extended, [n, n])
mat = np.triu(mat, k=0)

[junpenglao]

Here is another way using indexing:

vect = np.arange(1, 15+1)
m = len(vect)
# m = n(n+1)/2., solve n = .5 * (-1 + sqrt(1 + 8 m))
n = int(.5 * (-1 + np.sqrt(1 + 8 * m)))

mat2 = np.zeros([n, n], dtype=vect.dtype)
mat2[np.triu_indices(n)] = vect

[xidulu]

Thanks!

[xidulu]

@rlouf @junpenglao
Another question, when implementing the FullrankVIState, there is going to two members: The mean and the lower Cholesky factor, the mean can be a pytree but the cholesky factor cannot:

class FullrankVIState(NamedTuple):
    mu: PyTree
    L: Array

What do you think?

[rlouf]

That's not a problem, in the sense that we already have the same situation with inverse mass matrices in HMC algorithms.