Numpyro: Simplify training phase of SVI

[fehiepsi]

Currently, to train SVI, users need to know either

• jax.lax.fori_loop
• jax.lax.scan
• or jax.jit

Either way, users code will look like

svi = SVI(model, guide, ...)
svi_state = svi.init(rng_key)
def body_fn:
    ...
svi_state = fori_loop/scan/python loop
params = svi.get_params(svi_state)

To me, this is unfortunate and in most of the cases (given the fact that we already support subsampling via numpyro.subsample), the model input will be constant. I think it is better to support a high-level api for the above job. For example, we can add a method svi.run (which mimics mcmc.run) which do all the job, kind of

svi = SVI(model, guide, ...)
svi.run(rng_key, num_steps, *args, progress_bar=True, **kwargs)
params = svi.get_params()
losses = svi.get_losses()

With this, we

• need not to know what svi_state is
• need not to know how to use those jax transforms
• switch between with progress_bar and no progress_bar easily
• can see that the code is more intuitively

[neerajprad]

I think it is better to support a high-level api for the above job. For example, we can add a method svi.run (which mimics mcmc.run) which do all the job, kind of

A higher level API is fine for convenience, but note that if we cannot subsample in memory, i.e. have to load the data into memory in mini-batches, I think we will need finer grained control than what svi.run will provide.

[eb8680]

As an alternative, are there any widely used JAX libraries we could offload this to? A lot of people seem to use PyTorch Lightning for training their PyTorch models, is there something similar in JAX we could either integrate or recommend?

[fehiepsi]

we cannot subsample in memory

Yeah, it is a good point! For subsampling, currently, we only slice over subsample_indices so I guess the memory requirement is likely just the same as creating a data loader under JAX (e.g. the current implementation in numpyro.examples.load_dataset.get_batch). But for very large data that does not fit into memory, we should use a custom dataloader (not compiled in JAX) and init, update pattern as in vae example. Actually, for SVI training, there are also more complicated scenariors such as distributed training, which I am not familiar about.

I am not sure how we can deliver those technical issues for users. Probably, we can mention in svi_run that for large data, it is better to follow the pattern in vae example.

JAX libraries we could offload this to

I believe @ahmadsalim has put a lot of thought into this at https://github.com/pyro-ppl/numpyro/pull/649, which implements callbacks, following both Keras and PyTorch Lightning designs. I found that it is very convenient to use callbacks in both Pyro/NumPyro because it supports diagnosing training procedure, early stopping, epoch evaluation,... and the api is easy to use (like Keras and Lightning). It would be great if we have a package, which works with both Pyro/NumPyro SVI api, to do the job.

[neerajprad]

A lot of people seem to use PyTorch Lightning for training their PyTorch models, is there something similar in JAX we could either integrate or recommend?

Interesting! I have not heard about a similar higher level library for JAX, though both PyTorch or TensorFlow dataset loaders can be used generically.

Probably, we can mention in svi_run that for large data, it is better to follow the pattern in vae example.

Note that none of our examples deal with large enough datasets and therefore our loaders are trivial - they load the entire dataset into memory. Since part of the reason to use SVI is scalable inference, I think it is important to not overfit the interface to the small data examples that we currently have.

I believe @ahmadsalim has put a lot of thought into this at #649

Interesting, I'll take a look at the PR to learn more about these patterns.

[fehiepsi]

FYI, here are links to callbacks in keras or lightning. I don't think we should integrate it in SVI because it will violate the lightweight of the current SVI interface, but it deserves to be on a separate package (or in numpyro.contrib - but I think the api will work for Pyro's SVI too) for its convenience to make a complex training procedure.

it is important to not overfit the interface

Yes, I agree. Let's see what we will come up with in https://github.com/pyro-ppl/numpyro/pull/799. I still have not found a good solution yet, but will make a proposal soon.

[ahmadsalim]

Hi!

I have been a bit ill lately, but @OlaRonning who is a co-author of EinStein VI and PhD fellow in our group is interested in integrating the interface. He is an awesome person who knows a lot about the project, and knows what to do.

[OlaRonning]

Hi,
@ahmadsalim thanks for the kind words :D

The base class that is used in EinStein VI for both SteinVI and SVI is here. It implements train with callbacks (checkpoints etc.), SVI overloads evaluate and predict with the expected semantics. I've used it for training a DMM, and found it convenient.

Aside #800 is implemented here

[fritzo]

God bedring @ahmadsalim!

[ahmadsalim]

Tak Fritz! 😄