IMPLICIT NORMALIZING FLOWS

Abstract

Normalizing flows define a probability distribution by an explicit invertible transformation z = f (x). In this work, we present implicit normalizing flows (ImpFlows), which generalize normalizing flows by allowing the mapping to be implicitly defined by the roots of an equation F (z, x) = 0. ImpFlows build on residual flows (ResFlows) with a proper balance between expressiveness and tractability. Through theoretical analysis, we show that the function space of ImpFlow is strictly richer than that of ResFlows. Furthermore, for any ResFlow with a fixed number of blocks, there exists some function that ResFlow has a nonnegligible approximation error. However, the function is exactly representable by a single-block ImpFlow. We propose a scalable algorithm to train and draw samples from ImpFlows. Empirically, we evaluate ImpFlow on several classification and density modeling tasks, and ImpFlow outperforms ResFlow with a comparable amount of parameters on all the benchmarks.

1. INTRODUCTION

Normalizing flows (NFs) (Rezende & Mohamed, 2015; Dinh et al., 2014) are promising methods for density modeling. NFs define a model distribution p x (x) by specifying an invertible transformation f (x) from x to another random variable z. By change-of-variable formula, the model density is ln p x (x) = ln p z (f (x)) + ln |det(J f (x))| , where p z (z) follows a simple distribution, such as Gaussian. NFs are particularly attractive due to their tractability, i.e., the model density p x (x) can be directly evaluated as Eqn. (1). To achieve such tractability, NF models should satisfy two requirements: (i) the mapping between x and z is invertible; (ii) the log-determinant of the Jacobian J f (x) is tractable. Searching for rich model families that satisfy these tractability constraints is crucial for the advance of normalizing flow research. For the second requirement, earlier works such as inverse autoregressive flow (Kingma et al., 2016) and RealNVP (Dinh et al., 2017) restrict the model family to those with triangular Jacobian matrices. More recently, there emerge some free-form Jacobian approaches, such as Residual Flows (Res-Flows) (Behrmann et al., 2019; Chen et al., 2019) . They relax the triangular Jacobian constraint by utilizing a stochastic estimator of the log-determinant, enriching the model family. However, the Lipschitz constant of each transformation block is constrained for invertibility. In general, this is not preferable because mapping a simple prior distribution to a potentially complex data distribution may require a transformation with a very large Lipschitz constant (See Fig. 3 for a 2D example). Moreover, all the aforementioned methods assume that there exists an explicit forward mapping z = f (x). Bijections with explicit forward mapping only covers a fraction of the broad class of invertible functions suggested by the first requirement, which may limit the model capacity. In this paper, we propose implicit flows (ImpFlows) to generalize NFs, allowing the transformation to be implicitly defined by an equation F (z, x) = 0. Given x (or z), the other variable can be computed by an implicit root-finding procedure z = RootFind(F (•, x)). An explicit mapping z = f (x) used in prior NFs can viewed as a special case of ImpFlows in the form of F (z, x) = f (x) -z =

