Initializing batch of parameters for SteinVI

[OlaRonning]

To support re-initializable guides, we suggest the WrappedGuide interface, which requires implementing a function find_params that accepts a list of RNG keys in addition to the arguments for the guide and returns a set of freshly initialized parameters for each RNG key.

class ReinitGuide(ABC):
   @abstractmethod
   def init_params(self):
       raise NotImplementedError

   @abstractmethod
   def find_params(self, rng_keys, *args, **kwargs):
       raise NotImplementedError

The WrappedGuide class provides a guide written as a function. WrappedGuide makes a callable guide re-initializable. It works by running the provided guide multiple times and reinitializing the parameters using NumPyro's interface as follows:

1. WrappedGuide runs the guide transforming each parameter to unconstrained space.
2. For each parameter in the guide, WrappedGuide sets an RNG key.
3. It replaces the values of the parameters with values provided by a NumPyro initialization strategy (e.g., init_to_uniform, init_to_median, or init_to_noise).
4. It saves the parameter values for each particle and the required inverse transformations to constrained space to run the model correctly.

To avoid reinitializing specific parameters, WrappedGuide implements reinit_hide_fn to filter the parameters. SteinVI allows parameters without reinitialization to support neural network libraries like Stax, which provide bespoke initializers.

An alternative solution is to lift parameters to sample sites as suggested in #655 (and again by @fehiepsi below). Related to #833.

[fehiepsi]

Hi @OlaRonning, could you add some tests and let's see if we can come up with a solution that does not change how init strategy works? Currently, those strategies are intended to be applied for latent sites, not param sites. I think we can either:

• lift params in your model/guide to sample sites. Like ReinitGuide, this solution is kind of "hack" because modifies user programs (e.g. even though users already define init_value=1. for those params, Stein VI will still provide initial random values for them)
• explicitly requires users to provide random statements in param for init_value.

numpyro.param("a", lambda rng: ...)

I think this is a better solution for some reasons:

• The user program can be interpreted as-is, i.e. constant init values will be constant init values, random init values will be random init values.
• We don't need to introduce ways (like lift or ReinitGuide) to modify user programs
• We can leverage many init strategies in other jax nn frameworks.

[OlaRonning]

Hi @fehiepsi,

I’ll add some details to the PR and update with some tests tonight. It should be possible to keep the existing API, but need to sketch your ideas to get feeling for which is superior.

[fehiepsi]

Another way to not touch init strategies is to change the site type, fn, is observed keys for this site. (You can use ImproperUniform for fn.)

[fehiepsi]

SteinVI allows parameters without reinitialization to support neural network libraries like Stax

Though the solution in my above comment untouched the current api of stein vi, I still think that letting users define their own initial strategy for each parameter is better. So we don't need to introduce notions like reinit_hide_fn. But it depends on you. I think that addressing that issue has minor priority.

[OlaRonning]

I don't particularly like the reinit_hide_fn as it requires the user to list parameters twice. A solution that eliminates it is superior, in my opinion.

[OlaRonning]

I have been thinking about your suggestion with numpyro.param("a", lambda rng: ...) and I agree it's the better (more flexible) solution. The SteinVI.init will instantiate the batch (particle) dimension for all parameters.

For AutoGuides we would still use the init_loc_fn with a different init for custom guides and AutoGuides. I'll sketch the solution in the tests and we can discuss extensions/alternatives.

[OlaRonning]

Hi @fehiepsi, I rewrote the test with batching using your suggestions. The init strategy is now unaffected.
I changed the semantics slightly. Now only loc params (for AutoGuides) and user annotated random params are stochastic.

ELBO-within-Stein does not spread the variational distributions based on overlap (I think I have a fix for this coming), and SVGD is known to recover well from poor instantiations. So no need to take distribution shape into account (yet) when placing the particles with an AutoGuide.

[OlaRonning]

Found a bug: callable(site['value']) is always false, so all particles were initiziated to the same value for custom guides. Instead, I now vmap

def _reinit(seed):
  with handlers.seed(rng_seed=seed):
    return site_fn(*site_args)

I understand the above to produce constraint samples when the site_fn is not identity. Is this consistent with current usage? So for example

def model():
  numpyro.param("b", lambda rng_key: Normal(0, 1.).sample(rng_key), constraint=positive)

is invalid.

[fehiepsi]

LGTM, thanks @OlaRonning!

Is this consistent with current usage?

Yup. lambda key: Normal... might produce invalid samples for positive constraint.

only loc params (for AutoGuides) and user annotated random params are stochastic.

I like this behavior! init_loc_fn will initialize values for the latent variables and for parameters, it depends on users.

[OlaRonning]

Thanks for the review and suggestions @fehiepsi! C: