Augmented Neural ODE Flows
Abstract
Normalizing flows based on neural ODEs, as implemented in FFJORD, provide a powerful theoretical framework for density estimation and data generation. While the neural ODE formulation 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 augmented transformation dynamics. To derive the Jacobian determinant of the general augmented form, we generalise 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 and CIFAR-10, show that AFFJORD outperforms the baseline FFJORD through the improved flexibility of the underlying vector field.
Domains
Machine Learning [cs.LG]
Origin : Files produced by the author(s)