Free-form Flows: Make Any Architecture a Normalizing Flow

TL;DR

This paper introduces a training method that enables any neural network to function as a normalizing flow, broadening the design space and improving performance in molecule generation and inverse problems.

Contribution

It presents a novel training procedure that removes the invertibility constraint, allowing arbitrary neural networks to be used as normalizing flows for generative modeling.

Findings

Achieved state-of-the-art results in molecule generation benchmarks.

Demonstrated competitive performance in inverse problem benchmarks.

Enabled flexible architecture choices for normalizing flows.

Abstract

Normalizing Flows are generative models that directly maximize the likelihood. Previously, the design of normalizing flows was largely constrained by the need for analytical invertibility. We overcome this constraint by a training procedure that uses an efficient estimator for the gradient of the change of variables formula. This enables any dimension-preserving neural network to serve as a generative model through maximum likelihood training. Our approach allows placing the emphasis on tailoring inductive biases precisely to the task at hand. Specifically, we achieve excellent results in molecule generation benchmarks utilizing $E(n)$-equivariant networks. Moreover, our method is competitive in an inverse problem benchmark, while employing off-the-shelf ResNet architectures.