Haiku integration

[MilesCranmer]

Hi all,

Very nice package. Just starting to port my model to it.

I was wondering what you'd recommend as the best way to integrate a flow with a larger Haiku model? To give a bit of context: I am using a custom encoder to encode a high-dimensional dataset, and then using a flow conditioned on the encoded vector to model a distribution.

I noticed the trainer class seems like it might be able to handle general loss functions - would you recommend using that for my larger Haiku model? (and pass the loss_fn) Or instead, would it be better I use the flow.apply functions with the subset of parameters used in the flow?

Thanks!!
Miles

[C-J-Cundy]

Hi Miles,
Giving my two cents as a user, this is what I use to integrate a flow from NuX into a bigger model:

    flow = nux.Flow(create_flow, rand_key, train_inputs, batch_axes=(0,))

    def sample_flow(params, state, rng_key, n):
        # Sample from base distribution, i.e. normal
        samples = rnd.normal(rng_key, shape=(n, input_shape))
        # samples = rnd.sample(...)
        output = flow.stateful_apply(
            rng_key, {"x": samples}, params, None, sample=True, reconstruction=True
        )
        return output[0]["x"], output[0]["log_px"], output[1]  # I *think* this is right 

Then sample_flow is a pure function mapping params, state, rng_key to samples and log probs from the model.
Based on these samples, I evaluate a loss and get derivatives of the loss, update params using that gradient, etc.
(The way I'm using the flow is a bit different to the docs example which assume you're trying to eval log probs for data, wheras I'm in a variational workflow where I want to sample from the flow and evaluate the samples' log-probs)

I think it would be a bit more 'jax-y' if the stateful_apply functions were exposed more, maybe something for future updates

[MilesCranmer]

Awesome, thanks for the tip!! This looks like a good way to handle things.

Cheers,
Miles

[EddieCunningham]

The Flow class is actually a wrapper around a lower level object that has the functionality you're looking for. I added a new notebook that shows how to use this for class conditional modeling.

Chris, I'm happy to hear that you are able to use NuX in a larger model! One note about the Flow class is that it can handle doubly-batched inputs with a jax.scan loop using the scan_apply function. This can be useful in your application because you can do importance sampling in a few lines of code:

n_importance_samples = 10000
make_doubly_batched = lambda x: jnp.broadcast_to(x[None,...], (n_importance_samples,) + x.shape)
doubly_batched_inputs = jax.tree_map(make_doubly_batched, inputs)
outputs = flow.scan_apply(key, doubly_batched_inputs)
log_px = logsumexp(outputs["log_px"], axis=0) - jnp.log(n_importance_sample)

This will run your model on "inputs" 10000 times using a really fast jax.scan loop with a different random key each time. Its not necessary to use the Flow class or the FlowTrainer class for training, but they can definitely save a bit of time writing code that runs fast!

[EddieCunningham]

Hi Miles,

Unfortunately I'm not sure how easy it will be to use NuX in a Haiku model. The reason is that NuX is built on Haiku but adds a bit of extra functionality at a pretty low level that might make combining Haiku and NuX difficult. While Haiku keeps track of parameters and state, NuX also tracks compile time constants behind the scenes. This is mainly used to save off the shapes of inputs at initialization so that at runtime, the flow layers know how inputs are batched and can deal with batches automatically.

I have not done any testing to see if you can combine NuX layers and Haiku modules, but it might be possible? That being said, the Layer class in NuX is a wrapped version of the Haiku module, so it should be straight forward to rewrite your Haiku model as a NuX Layer. Check out nux.networks for examples of this.

I also highly recommend that you use the trainer class to define a new loss function. You can do this by creating a new class that inherits from FlowTrainer and write a new loss function. The max likelihood and classification trainer are examples of this. I'll write a tutorial with more details over the next couple of days. Until then, if you give the implementation a shot I'll be happy to answer any questions you have.