Bijectors, LazyMeta, Parameters, etc.

[francois-rozet]

Hi 👋

I am interested in contributing to FlowTorch, but I find the structure of the package hard to grasp. It seems like every class is a subclass of LazyMeta instead of the usual nn.Module. What is a LazyMeta? What problem does it solve and how?
For bijectors, I guess they replace transforms, but when looking at their implementation I am confused because they don't all implement the same interface (some define a _forward method but others define forward). Why? Similarly, I don't understand the difference between sample and rsample in flows. I also don't get why autoregressive networks are specified as "parameters". And why a Tensor is a child of the class Parameter?
Finally, since the distributions are not modules and the bijectors are not either, how are the learnable parameters handeled? And the non-learnable ones (buffers)? Can they be transfered to CUDA easily (with .to(device) for instance)?

[stefanwebb]

Hi @francois-rozet, thanks for your interest in contributing to FlowTorch! There's a lot to unpack here - I'll do my best :)

flowtorch.lazy.LazyMeta is a metaclass that handles delayed initialization. It solves the problem of being able to create a "plan" for a normalizing flow before you know all of the information, e.g. the shape of the base distribution. When you try to instantiate a Bijector or Parameters object and haven't specified all the arguments, it will return a flowtorch.lazy.Lazy object that records the class type and the arguments/values you've specified so far. Then when you call .__init__ another time and fill in the remaining arguments it will convert into an actual Bijector or Parameters object.

flowtorch.bijectors.Bijector is the equivalent of torch.distributions.transforms.Transform or tfp.bijectors.Bijector. They will always have .forward, ._forward, .backward, and ._backward methods, although they may be inherited from parent classes. .forward is the user-facing method and will handle things like caching that are common to all bijectors, whereas ._forward is the developer-facing implementation of the bijector's operation, and is called by .forward.

flowtorch.distributions.Flow inherits from torch.nn.Module. It will instantiate the bijector that is passed to it using the shape from the base distribution, and the bijector pass this shape information on to the parameters and instantiate them. You can get the learnable parameters simply by flow.parameters() since it's an nn.Module, and also do flow.to(device) as well (one of the design goals has been to "play nice" with existing PyTorch interfaces).

The difference between .sample and .rsample (as well as .normalize and .rnormalize) in flowtorch.distributions.Flow is that the r version can have an autograd taken through it (the same convention as torch.distributions.Distribution). The r stands for reparameterized.

flowtorch.parameters.Parameters is a general interface for specifying the parameters for a normalizing flow, whether they are a tensor, autoregressive NN, dense NN, etc. There is also a basic type system on them to make sure that only parameters that are compatible with a given bijector are used (e.g. you have to use autoregressive parameters for an autoregressive bijector).

[stefanwebb]

I think your higher level question may be, how does one implement a new Bijector or modify an existing one? This requires a few steps:

• Fork the repo, clone the fork on your local system, and create a branch on this.
• Decide whether your Bijector should inherit from Fixed, Elementwise, Autoregressive, or some new abstract Bijector class that doesn't exist yet (e.g. Coupling, Continuous).
• If you can use an existing one, then continue to create your class, defining ._forward, ._inverse, and ._log_abs_det_jacobian.
• Run scripts/generate_imports.py so the new Bijector is added to flowtorch/bijectors/__init__.py
• Make sure all linting, type checks, and unit tests pass
• Commit your changes, and push to a new branch on the fork
• Create a PR from the branch on your fork

[francois-rozet]

Thank you for the detailed answer! It's much clearer now!
So, if I understood correctly, the "big" difference with usual NN in PyTorch is that the layers, here bijectors, don't hold their own parameters. Instead the flow is responsible for holding them in the format of Parameter objects. Aren't you afraid that the class Tensor would be confusing with respect to torch.Tensor?

I forgot to ask, but how is handled the conditioning of flows ? In pyro there is a distinction between conditional and unconditional distributions/transforms which adds a lot of complexity to the interface (and actually makes it very slow). In nflows, all distributions and transforms are conditional, but some do not vary with respect to the context (they simply ignore it), which reduces both the complexity and the overhead.

Concerning .sample and .normalize, if they always block the gradients, you might consider adding the decorator @torch.no_grad() to them instead of having a context with torch.no_grad(): in them. It is both nicer and more explicit.

[stefanwebb]

That's almost right- the bijectors hold their parameters after they have been initialized in self._params. Although, Bijectors are not nn.Module's. The class distribution.Flow is however, and it stores a reference to the parameters so you can do flow.parameters().

That's an excellent point about flowtorch.parameters.Tensor and torch.Tensor. We'll have to give that some more thought. I think either renaming it in FlowTorch or else making it more of an extension to torch.Tensor (subclassing it?) might be good solutions!

Conditional flows are not yet implemented - I had to take them out when we did a major refactoring a while back. The interface for conditional flows will be much more like nflows rather than pyro. It seems like it's the most requested feature so will be completed soon.

That's an excellent suggestion about using @torch.no_grad(). The custom documentation system we developed for Docusaurus v2 picks up on decorators so if we make this change you would then it would be apparent from the docs that it doesn't pass a gradient

[francois-rozet]

The class distribution.Flow is however, and it stores a reference to the parameters so you can do flow.parameters().

Then can you save and load the dict of parameters as any nn.Module (with .state_dict and .load_state_dict) ? Depending how load_state_dict is implemented, it might break the link between the weights held by the flow and the weights held by the bijectors.

[stefanwebb]

I need to write some docs on how these classes work and a tutorial on creating a new Bijector/Parameters class - I'll post these in the near future!