Enhanced Distribution Modelling via Augmented Architectures For Neural
  ODE Flows
Etrit Haxholli Marco Lorenzi
While the neural ODE formulation of normalizing flows such as in FFJORD
enables us to calculate the determinants of free form Jacobians in O(D) time,
the flexibility of the transformation underlying neural ODEs has been shown to
be suboptimal. In this paper, we present AFFJORD, a neural ODE-based
normalizing flow which enhances the representation power of FFJORD by defining
the neural ODE through special augmented transformation dynamics which preserve
the topology of the space. Furthermore, we derive the Jacobian determinant of
the general augmented form by generalizing the chain rule in the continuous
sense into the Cable Rule, which expresses the forward sensitivity of ODEs with
respect to their initial conditions. The cable rule gives an explicit
expression for the Jacobian of a neural ODE transformation, and provides an
elegant proof of the instantaneous change of variable. Our experimental results
on density estimation in synthetic and high dimensional data, such as MNIST,
CIFAR-10 and CelebA 32x32, show that AFFJORD outperforms the baseline FFJORD
through the improved flexibility of the underlying vector field.