Reduced rank VI implementation

Hi Everyone,

Is there a way to custom specify the reduced rank covariance matrix of the ADVI inference? I have a multivariate survival model with splines that works quite well with mean-field ADVI (with underestimated variance, of course) but the full-rank has trouble due to its attempt to model the correlation between the spline coefficients and the frailty parameter (these correlations can actually be ignored for our model). I’d like to be able to make the covariance matrix sparse for the corresponding entries in the VI approximating distribution. Is there some way to pass a masked matrix?



Hi David, there are couple of ways you can try. I am not sure which way is the best, as we are still experimenting with them.

  1. using Normalizing Flow to introduce a low rank correlation, details here:

  2. using the Grouped Approximation
    The idea here is to put part of the latent RVs into a FullRank Group, and the rest into a meanfield Group

cc @ferrine

Thanks so much. I’ve been using the normalizing flows approach but with limited success. As I understand it, passing the string “'scale-hh*10-loc” would model the dependencies of 10 variables in the covariance matrix (please correct me if that’s incorrect). If so, how does the NFVI choose which variables it will model?

The grouped approximation sounds exactly like what I need though. I’ll be testing it out as well.

Thanks again,


“'scale-hh*10-loc” would model the dependencies of 10 variables in the covariance matrix

Intuition is very close to what happens with covariance, yes. But still I think Groups are your choice as @junpenglao mentioned.