Skip to content

NormalizingFlow

Jorge Prado requested to merge github/fork/RasmusOrsoe/final_normalizing_flow into main

Created by: RasmusOrsoe

This PR adds compatibility for normalizing flows to graphnet. Specifically, this PR does the following:

  1. Adds NormalizingFlow in graphnet.models.flows as a new generic model class which all normalizing flows are intended to inherit from. These models are assumed to output a single tensor containing the transformed coordinates and the jacobian associated with the transformation. For book-keeping reasons, the column indices that identifies which columns are transformed coordinates and jacobians, these models have properties coordinate_columns and jacobian_columns. These properties are used to split the tensor for loss computation (see diagram).
  2. Adds a normalizing flow model called INGA which is similar to what is presented in https://arxiv.org/abs/1906.04032 . It relies on splines and the spline functions are currently put in graphnet.models.flows.spline_blocks and can be used both inside and outside the context of NormalizingFlows (although likely most relevant there).
  3. Adds an example of training INGA under examples/04_training/05_train_normalizing_flow.py. Here the syntax of training the normalizing flow is shown. The example trains INGA to learn a mapping from the rather complicated distribution of the pulse attributes (xyz, time, charge, etc) into a multivariate gaussian distribution.
  4. Introduces slight changes to the inheritance structure of our loss functions: the generic base-class LossFunction no longer has an abstract _forward method; this is left for implementation of further problem specific subclasses. This is introduced here because not all machine learning paradigms follow our usual prediction, target argument structure; which is indeed the case for normalizing flows.
  5. Created a new LossFunction subclass called StandardLossFunction which all of our current loss functions now inherit from. These loss functions all follow the prediction, target argument structure.
  6. Created a new LossFunction subclass FlowLossFunction which follow the prediction, jacobian argument structure.
  7. Created MultivariateGaussianFlowLoss, subclassed of FlowLossFunction, which calculates the nllh of a multivariate Gaussian.
  8. Slight changes to StandardModel that ties everything together. While working on this PR i noticed a few list comprehensions were hard to read; so I took the liberty of writing them out explicitly and adding a few comments here and there.

A diagram showing the new inheritance structure for loss functions is shown here:

subclasses

A diagram showing the overall data flow and usage of the column indices for normalizing flows can be seen here:

inga drawio

After Christmas, I'll make a new PR that adds the ability use NormalizingFlows in a generative sense and an example of how this is done. I decided not to include it here to avoid making this PR larger than it already is.

Merge request reports