Implement Variational Inference

[albcab]

Start by implementing basic transforms/change of variables: affine, exp/log, sigmoid/logistic, rotations. Either through objects that are created with log pdfs, apply transformations and can be called for its log pdf; or through functions that input a log pdf and output that log pdf with the transformation on its parameters + log abs det jac. These can be used to unconstrain variables as a preprocessing step before running HMC and shift & scale normal distributions for mean field approximations.

The vi.py would mimic the kernel implementation of HMC with VIInfo and a kernel generator that performs a step of gradient ascent, using jax.example_libraries.optimizers or optax, given a log pdf, a reference pdf and a parametrized transformation. Then a subfolder with base.py including VIState and a state generator and loss.py with ELBO implementation and potentially others.

Finally implement other transformations such as normalizing flows, polynomial chaos expansions or iterative gaussianization. Also allowing for these transformations to be stacked sequentially, one on top of the other.

Any thoughts?

[junpenglao]

VI was briefly discussed during the last meeting, while we dont have any intermediate plan to work on it, design discussion is certainly welcome.
One of the difficulty is to NOT committing to a PPL (and not end up implementing a new PPL), I think before thinking about bijection for change of variable, we should assume that transformation is done at user side (same assumption we are currently making of all the logposterior_fn function for HMC - we assume all unknown parameter are unbounded, and it is the user's responsibility to make sure that)

I think the first step for implementing VI will be a base class/function that takes multiple meta functions, for example, functions that takes approximation variables (VI variables) and generate callable that computes the variational loss.

[albcab]

So do you mean creating something like this:

class VariationalInference(NamedTuple):
    loss: Callable
    transform: Callable

vi = blackjax.vi(logprob_fn)
loss = vi.loss(vi_params)
samples = vi.transform(vi_params, iid_normal_samples)

to then create a kernel that initializes parameters and performs a gradient ascent step?