Normalizing flow neural networks by JKO scheme

TL;DR

This paper introduces JKO-iFlow, a neural ODE-based normalizing flow model inspired by the JKO scheme, which simplifies training and reduces computational costs while maintaining competitive performance.

Contribution

The paper develops a novel neural ODE flow network based on the JKO scheme, enabling efficient block-wise training without SDE sampling or score matching.

Findings

Achieves competitive performance with existing flow and diffusion models.

Reduces computational and memory costs significantly.

Enables efficient end-to-end training of flow networks.

Abstract

Normalizing flow is a class of deep generative models for efficient sampling and likelihood estimation, which achieves attractive performance, particularly in high dimensions. The flow is often implemented using a sequence of invertible residual blocks. Existing works adopt special network architectures and regularization of flow trajectories. In this paper, we develop a neural ODE flow network called JKO-iFlow, inspired by the Jordan-Kinderleherer-Otto (JKO) scheme, which unfolds the discrete-time dynamic of the Wasserstein gradient flow. The proposed method stacks residual blocks one after another, allowing efficient block-wise training of the residual blocks, avoiding sampling SDE trajectories and score matching or variational learning, thus reducing the memory load and difficulty in end-to-end training. We also develop adaptive time reparameterization of the flow network with a progressive refinement of the induced trajectory in probability space to improve the model accuracy further. Experiments with synthetic and real data show that the proposed JKO-iFlow network achieves competitive performance compared with existing flow and diffusion models at a significantly reduced computational and memory cost.