Variational fit (ADVI) - initialisation

As you are discovering, the AVDI codebase is in need of a loving refactor.

Everything ultimately gets set up deep in the bowels of the MeanFieldGroup class, here. To understand what you need to give, it is good to look at the _prepare_start function.

For optimization problems, PyMC has a class called DictToArrayBijection that lets you freely flip-flop between a dictionary with parameter name keys and n-dimensional value-array values, and a flat array of all raveled parameter arrays concatenated. The later is what is expected by e.g. scipy.optimize.minimize.

Edit: I brought up DictToArrayBijection because there is a comment in the linked method that mentions it, but staring at the code, it doesn’t appear to be required. I think standard dictionaries will be sufficient.

So what it appears you need to provide is:

  • For the mean, a dictionary of parameter_name: ndarray. These will be piped into make_initial_point_fn, which in turn will handle this bijection business for you (I think!)
  • For the sigmas, based on the start_sigmas.get here, it looks like you also want a dictionary of parameter_name: ndarray.

One note: the names you will need to provide are not necessarily the ones your expect. They will be the names you see in model.continuous_value_vars, which attaches a suffix based on the bijective transformation internally used to allow samplers/optimizers to propose values in \mathbb{R}^N, rather than in constrained subspaces. For example, if you make a variable sigma = pm.HalfNormal('sigma'), the name you will find is sigma_log__.

1 Like