Composing flow transformations

[minaskar]

Hi,

Is there currently a way to "stack" and train different flows at the same time? If not, are there any plans to include this type of functionality? What I mean by that is to create composite flows (e.g. a few layers of MAF + a few layers of NSF).

Also, what is the recommended way to get both the Jacobian and the forward/inverse transform of a flow of some samples?

[francois-rozet]

Hello @minaskar, thank you for your questions!

Is there currently a way to "stack" and train different flows at the same time?

In Zuko, a flow is an instance of the Flow class and has two attributes: transform and base. To create a flow, you can either instantiate one of the built-in ones (MAF, NSF, ...) or compose your own as a Flow subclass. For instance, if you want one affine autoregressive transformation and one spline coupling transformation:

import torch
import zuko

class MyFlow(zuko.flows.Flow):
    def __init__(self, features: int, context: int = 0):
        transforms = [
            zuko.flows.MaskedAutoregressiveTransform(
                features,
                context,
                univariate=zuko.transforms.MonotonicAffineTransform,
                shapes=[(), ()],  # shapes of the transformation's parameters
            ),
            zuko.flows.MaskedAutoregressiveTransform(
                features,
                context,
                order=torch.arange(features).flip(0),  # flip the features' order
                passes=2,  # coupling
                univariate=zuko.transforms.MonotonicRQSTransform,
                shapes=[(8,), (8,), (7,)],  # 8 bins
            ),
        ]

        base = zuko.flows.Unconditional(
            zuko.distributions.DiagNormal,
            torch.zeros(features),
            torch.ones(features),
            buffer=True,
        )

        super().__init__(transforms, base)

flow = MyFlow(5, context=16)

Another way is to instantiate built-in flows and borrow their layers:

maf = zuko.flows.MAF(5, context=16, transforms=2)
nsf = zuko.flows.NSF(5, context=16, transforms=1, passes=2)

flow = zuko.flows.Flow((*maf.transforms, *nsf.transforms), maf.base)

What is the recommended way to get both the Jacobian and the forward/inverse transform of a flow of some samples?

With plain PyTorch transformations, it is not possible to get the transformed sample and the log absolute determinant of its Jacobian (LADJ) in a single call. You would first transform x and then get its LADJ.

z = transform(x)
ladj = transform.log_abs_det_jacobian(x, z)

Zuko introduces a new call_and_ladj method that returns both at once:

z, ladj = transform.call_and_ladj(x)

The full transformation defined by a flow is accessible via its transform attribute, once it has been conditioned on a context y. So, in your case,

z, ladj = flow(y).transform.call_and_ladj(x)

[minaskar]

Thanks a lot, that was very helpful!