SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chen Xu
School of Industrial and Systems Engineering Georgia Tech
Xiuyuan Cheng
Department of Mathematics  
Duke University
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.
Generative models have wide applications in statistics and machine learning to infer data-generating distributions and to sample from the model distributions learned from the data. In addition to widely used deep generative models such as variational auto-encoders (VAE) Kingma and Welling [2014, 2019] and generative adversarial networks (GAN) Goodfellow et al. [2014], Gulrajani et al. [2017], normalizing flow models Kobyzev et al. [2020] have been popular and with a great potential. The flow model learns the data distribution via an invertible mapping  $F$  between the data density  $p_{X}$  in  $\mathbb{R}^d$  and the target standard multivariate Gaussian density  $p_Z$ ,  $Z \sim \mathcal{N}(0, I_d)$ . While flow models, once trained, can be utilized for efficient data sampling and explicit likelihood evaluation, training of such models is often difficult in practice. To alleviate such difficulties, many prior works [Dinh et al., 2015, 2017, Kingma and Dhariwal, 2018, Behrmann et al., 2019, Ruthotto et al., 2020, Onken et al., 2021], among others, have explored designs of training objectives, network architectures, and
(a) JKO-iFlow
Figure 1: Comparison of JKO-iFlow (proposed) and standard CNF models. In contrast to most existing CNF models, JKO-iFlow learns the unique deterministic transport equation corresponding to the diffusion process by directly performing block-wise training of a neural ODE model.
(b) usual CNF
Among various flow models, continuous normalizing flow (CNF) transports the data density to a target distribution through continuous dynamics, primarily neural ODE [Chen et al., 2018] models. CNF
models have shown promising performance on generative tasks Grathwohl et al. [2019], Kobyzev et al. [2020]. However, a known computational challenge of CNF models is model regularization, primarily due to the non-uniqueness of the flow transport. Without additional regularization, the trained CNF model such as FFJORD Grathwohl et al. [2019] may follow a less regular trajectory in the probability space, see Figure 1(b), which may worsen the generative performance. While regularization of flow models has been developed using different techniques, including using spectral normalization Behrmann et al. [2019] and optimal transport Liutkus et al. [2019], Onken et al. [2021], Finlay et al. [2020], Xu et al. [2022], Huang et al. [2023], only using regularization may not resolve the non-uniqueness of the flow. Besides regularization, several practical difficulties remain when training CNF models, particularly the high computational cost. In many settings, CNF models consist of stacked blocks, and each can be complex. End-to-end training of such deep models often places a high demand on computational resources and memory consumption.
In this work, we propose JKO-iFlow, an invertible normalizing flow network that unfolds the Wasserstein gradient flow via a neural ODE model inspired by the JKO-scheme Jordan et al. [1998]. The JKO scheme, see (5), can be viewed as a proximal step to minimize the Kullback-Leibler (KL) divergence between the current density and the equilibrium density. It recovers the solution of the Fokker-Planck equation (FPE) in the limit of small step size. The proposed JKO-iFlow model thus can be viewed as trained to learn the unique transport map following the FPE which flows from the data distribution toward the normal equilibrium and gives a smooth trajectory of density evolution, see Figure 1(a). Unlike most CNF models, where all the residual blocks are trained end-to-end, each block in the JKO-iFlow network implements one step in the JKO scheme to learn the deterministic transport map by minimizing an objective of that block given the trained previous blocks. The block-wise training significantly reduces memory and computational load. Theoretically, with a small step size, the discrete-time transport map approximates the continuous solution of FPE, which leads to the invertibility of each trained residual block. This leaves the residual blocks to be completely general, such as graph neural network layers and convolutional neural network layers, depending on the structure of data considered in the problem. The theoretical need for small step sizes does not incur a restriction in practice, whereby one can use step size not exceeding a certain maximum value when adopting the neural ODE integration. We further introduce time reparameterization with progressive refinement in computing the flow network, where each block corresponds to a representative point along the density evolution trajectory in the space of probability measures. The algorithm adaptively chooses the number of blocks and step sizes.
The proposed approach is related to the diffusion models [Song and Ermon, 2019, Ho et al., 2020, Block et al., 2020, Song et al., 2021] yet differs fundamentally in that our approach is a type of flow models, which directly computes the data likelihood, and such likelihood estimation is essential for statistical inference. While the diffusion models can also indirectly obtain likelihood estimation, they are more designed as samplers. In terms of implementation, our approach trains a neural ODE model without SDE sampling (injection of noise) nor learning of score matching. It also differs from previous works on progressive training of ResNet generative models Johnson and Zhang [2019], Fan et al. [2022] in that the model trains an invertible flow mapping and avoids inner loops of variational learning. We refer to Section 1.1 for more discussions on related works. Empirically, JKO-iFlow yields competitive performance as other CNF models with significantly less computation. The model is also compatible with general equilibrium density  $p_Z$  having a parametrized potential  $V$ , exemplified by the application to the conditional generation setting where  $p_Z$  is replaced with Gaussian mixtures Xu et al. [2022]. In summary, the contributions of the work include
- We propose an invertible neural ODE model where each residual block corresponds to a JKO step, and the training objective can be computed from pushed data samples through the previous blocks. The residual block has a general form, and the invertibility is ensured due to the regularity and continuity of the approximate solution of the FPE.
- We develop a block-wise procedure to train the JKO-iFlow model, which can adaptively determine the number of blocks. We also propose to adaptively reparameterize the computed trajectory in the probability space with refinement, which improves the model accuracy and the overall computational efficiency.
- We show that JKO-iFlow greatly reduces memory and computational cost when achieving competitive or better generative performance and likelihood estimation compared to existing flow and diffusion models on simulated and real data.
For deep generative models, popular approaches include generative adversarial networks (GAN) [Goodfellow et al., 2014, Gulrajani et al., 2017, Isola et al., 2017] and variational auto-encoder (VAE)[Kingma and Welling, 2014, 2019]. Apart from known training difficulties (e.g., mode collapse [Salimans et al., 2016] and posterior collapse [Lucas et al., 2019]), these models do not provide likelihood or inference of data density. The normalizing flow framework [Kobyzev et al., 2020] has been extensively developed, including continuous flow [Grathwohl et al., 2019], Monge-Ampere flow [Zhang et al., 2018], discrete flow [Chen et al., 2019], and extension to non-Euclidean data [Liu et al., 2019, Mathieu and Nickel, 2020, Xu et al., 2022]. Efforts have been made to develop novel invertible mapping structures [Dinh et al., 2017, Papamakarios et al., 2017] and regularize the flow trajectories by transport cost [Finlay et al., 2020, Onken et al., 2021, Ruthotto et al., 2020, Xu et al., 2022, Huang et al., 2023]. Despite such efforts, the model and computational challenges of normalizing flow models include regularization and the large model size when using a large number of residual blocks, which cannot be determined a priori, and the associated memory and computational load.
In parallel to continuous normalizing flows, which are neural ODE models, neural SDE models become an emerging tool for generative tasks. Diffusion process and Langevin dynamics in deep generative models have been studied in score-based generative models [Song and Ermon, 2019, Ho et al., 2020, Block et al., 2020, Song et al., 2021] under different settings. Specifically, these models estimate the score function (i.e., the gradient of the log probability density with respect to data) of data distribution via neural network parametrization, which may encounter challenges in learning and sampling high dimensional data and call for special techniques [Song and Ermon, 2019]. The recent work of Song et al. [2021] developed reverse-time SDE sampling for score-based generative models and adopted the connection to neural ODE to compute the likelihood; using the same idea of backward SDE, Zhang and Chen [2021] proposed joint training of forward and backward neural SDEs. Theoretically, latent diffusion Tzen and Raginsky [2019b,a] was used to analyze neural SDE models. The current work focuses on a neural ODE model where the deterministic vector field  $\mathbf{f}(x,t)$  is to be learned from data following a JKO scheme of the FPE. Rather than neural SDE, our approach involves no sampling of SDE trajectories nor learning of the score function, and it learns an invertible residual network directly. In contrast, diffusion-based models derive the ODE model from the learned diffusion model to achieve explicit likelihood computation. For example, Song et al. [2021] derived neural ODE model from the learned score function of the diffused data marginal distributions for all  $t$ . We experimentally obtain competitive or improved performance against the diffusion model on simulated two-dimensional and high-dimensional tabular data.
JKO-inspired deep models have been studied in several recent works. [Bunne et al., 2022] reformulated the JKO step for minimizing an energy function over convex functions. JKO scheme has also been used to discretize Wasserstein gradient flow to learn a deep generative model in [Alvarez-Melis et al., 2022, Mokrov et al., 2021], which adopted input convex neural networks (ICNN) [Amos et al., 2017]. ICNN, as a special type of network architecture, may have limited expressiveness [Rout et al., 2022, Korotin et al., 2021]. In addition to using the gradient of ICNN, [Fan et al., 2022] proposed parametrizing the transport in a JKO step by a residual network but identified difficulty in calculating the push-forward distribution. The approach in [Fan et al., 2022] also relies on a variational formulation, which requires training an additional network similar to the discriminator in GAN using inner loops. The idea of progressive additive learning in training generative ResNet, namely training ResNet block-wisely by a variational loss, dates back to Johnson and Zhang [2019] under the GAN framework. Our method trains an invertible neural-ODE flow network that flows from data density to the normal one and backward, which enables the computation of model likelihood as in other neural-ODE approaches. The objective in each JKO step to minimize KL divergence can also be computed directly without any inner-loop training, see Section 4.
Compared to score-based neural SDE methods, our approach is closer to the more recent flow-based models related to diffusion models [Lipman et al., 2023, Albergo and Vanden-Eijnden, 2023, Boffi and Vanden-Eijnden, 2023]. These works proposed to learn the transport equation (a deterministic ODE) corresponding to the SDE process. Specifically, [Lipman et al., 2023, Albergo and Vanden-Eijnden,
2023] matched the velocity field from interpolated distributions between initial and terminal ones; [Boffi and Vanden-Eijnden, 2023] proposed to learn the score  $\nabla \log \rho_t(x)$  (where  $\rho_{t}$  solves the FPE and is unknown a priori) to solve high-dimensional FPE. Using Stein's identity (which is equivalent to the derivation in Section 3.2), the step-wise training objective in [Boffi and Vanden-Eijnden, 2023] optimizes to learn the score without the need to simulate the entire SDE. The idea of approximating the solution of FPE by a deterministic transport equation dates back to the 90s Degond and Mustieles [1990], Degond and Mas-Gallic [1989], and has been used in kernel-based solver of FPE in [Maoutsa et al., 2020] and studied via a self-consistency equation in [Shen et al., 2022]. While our approach learns an equivalent velocity field at the infinitesimal time step, see section 3.2, the formulation in JKO-iFlow is at a finite time-step motivated by the JKO scheme. We also mention that an independent concurrent work [Vidal et al., 2023] proposed a similar block-wise training algorithm using the JKO scheme under the framework of [Onken et al., 2021] and demonstrated benefits in avoiding tuning the penalty hyperparameter associated with the KL-divergence objective. Our model is applied to generative tasks of high-dimensional data, including image data, and we also develop additional techniques for computing the flow probability trajectory.
For the expressiveness of deep generative models, approximation properties of deep neural networks for representing probability distributions have been developed in several works. Lee et al. [2017] established approximation by composition of Barron functions [Barron, 1993]; Bailey and Telgarsky [2018] developed space-filling approach, which was generalized in Perekrestenko et al. [2020, 2021]; Lu and Lu [2020] constructed a deep ReLU network with guaranteed approximation under integral probability metrics, using techniques of empirical measures and optimal transport. These results show that deep neural networks can provably transport one source distribution to a target one with sufficient model capacity under certain regularity conditions of the pair of densities. Our JKO-iFlow model potentially leads to a constructive approximation analysis of the neural ODE flow model.
Normalizing flow. A normalizing flow can be mathematically expressed via a density evolution equation of  $\rho(x,t)$  such that  $\rho(x,0) = p_X$  and as  $t$  increases  $\rho(x,t)$  approaches  $p_Z \sim \mathcal{N}(0, I_d)$  [Tabak and Vanden-Eijnden, 2010]. Given an initial distribution  $\rho(x,0)$ , such a flow typically is not unique. We consider when the flow is induced by an ODE of  $x(t)$  in  $\mathbb{R}^d$
$$
\dot {x} (t) = \mathbf {f} (x (t), t), \tag {1}
$$
where  $x(0) \sim p_X$ . The marginal density of  $x(t)$  is denoted as  $p(x,t)$ , and it evolves according to the continuity equation (Liouville equation) of (1) written as
$$
\partial_ {t} p + \nabla \cdot (p \mathbf {f}) = 0, \quad p (x, 0) = p _ {X} (x). \tag {2}
$$
Ornstein-Uhlenbeck (OU) process. Consider a Langevin dynamic denoted by the SDE  $dX_{t} = -\nabla V(X_{t})dt + \sqrt{2} dW_{t}$ , where  $V$  is the potential of the equilibrium density  $p_Z$ . We focus on the case of normal equilibrium, that is,  $V(x) = |x|^2 / 2$  and then  $p_Z \propto e^{-V}$ . In this case, the process is known as the (multivariate) OU process. Suppose  $X_0 \sim p_X$ , and let the density of  $X_{t}$  be  $\rho(x,t)$  also denoted as  $\rho_t(\cdot)$ . The Fokker-Planck equation (FPE) describes the evolution of  $\rho_{t}$  towards the equilibrium  $p_Z$  as follows, where  $V(x) \coloneqq |x|^2 / 2$ ,
$$
\partial_ {t} \rho = \nabla \cdot (\rho \nabla V + \nabla \rho), \quad \rho (x, 0) = p _ {X} (x). \tag {3}
$$
Under generic conditions,  $\rho_{t}$  converges to  $p_Z$  exponentially fast. For Wasserstein-2 distance and the standard normal  $p_Z$ , classical argument gives that (take  $C = 1$  in Eqn (6) of Bolley et al. [2012])
$$
W _ {2} \left(\rho_ {t}, p _ {Z}\right) \leq e ^ {- t} W _ {2} \left(\rho_ {0}, p _ {Z}\right), \quad t > 0. \tag {4}
$$
JKO scheme. The seminal work Jordan et al. [1998] established a time discretization scheme of the solution to (3) by the gradient flow to minimize  $\mathrm{KL}(\rho ||p_Z)$  under the Wasserstein-2 metric in probability space. Denote by  $\mathcal{P}$  the space of all probability densities on  $\mathbb{R}^d$  with a finite second moment. The JKO scheme computes a sequence of distributions  $p_k$ ,  $k = 0,1,\dots$ , starting from  $p_0 = \rho_0\in \mathcal{P}$ . With step size  $h > 0$ , the scheme at the  $k$ -th step is written as
$$
p _ {k + 1} = \arg \min  _ {\rho \in \mathcal {P}} F [ \rho ] + \frac {1}{2 h} W _ {2} ^ {2} \left(p _ {k}, \rho\right), \tag {5}
$$
where  $F[\rho] \coloneqq \mathrm{KL}(\rho || p_Z)$ . It was proved in Jordan et al. [1998] that as  $h \to 0$ , the solution  $p_k$  converges to the solution  $\rho(\cdot, kh)$  of (3) for all  $k$ , and the convergence  $\rho_{(h)}(\cdot, t) \to \rho(\cdot, t)$  is strongly in  $L^1(\mathbb{R}^d, (0, T))$  for finite  $T$  where  $\rho_{(h)}$  is piece-wise constant interpolated from  $p_k$ .
3 JKO scheme by neural ODE
Given i.i.d. observed data samples  $X_{i} \in \mathbb{R}^{d}$ ,  $i = 1, \dots, N$ , drawn from some unknown density  $p_{X}$ , the goal is to train an invertible neural network to transport the density  $p_{X}$  to an a priori specified density  $p_{Z}$  in  $\mathbb{R}^{d}$ , where each data sample  $X_{i}$  is mapped to a code  $Z_{i}$ . A prototypical choice of  $p_{Z}$  is the standard multivariate Gaussian  $\mathcal{N}(0, I_{d})$ . In this work, we leave the potential of  $p_{Z}$  abstract and denote by  $V$ , that is,  $p_{Z} \propto e^{-V}$  and  $V(x) = |x|^{2}/2$  for normal  $p_{Z}$ . By a slight abuse of notation, we denote by  $p_{X}$  and  $p_{Z}$  both the distributions and the density functions of data  $X$  and code  $Z$  respectively.
3.1 Objective of JKO step
We are to specify  $\mathbf{f}(x,t)$  in the ODE (1), to be parametrized and learned by a neural ODE, such that the induced density evolution of  $p(x,t)$  converges to  $p_Z$  as  $t$  increases. We start by dividing the time horizon  $[0,T]$  into finite subintervals with step size  $h$ , let  $t_k = kh$  and  $I_{k + 1}\coloneqq [t_k,t_{k + 1})$ . Define  $p_k(x)\coloneqq p(x,kh)$ , namely the density of  $x(t)$  at  $t = kh$ . The solution of (1) determined by the vector-field  $\mathbf{f}(x,t)$  on  $t\in I_{k + 1}$  (assuming the ODE is well-posed [Sideris, 2013]) gives a one-to-one mapping  $T_{k + 1}$  on  $\mathbb{R}^d$ , s.t.  $T_{k + 1}(x(t_k)) = x(t_{k + 1})$  and  $T_{k + 1}$  transports  $p_k$  into  $p_{k + 1}$ , i.e.,  $(T_{k})_{\#}p_{k - 1} = p_{k}$ , where we denote by  $T_{\#}p$  the push-forward of distribution  $p$  by  $T$ , such that  $(T_{\#}p)(A) = p(T^{-1}(A))$  for a measurable set  $A$ . In other words, the mapping  $T_{k + 1}$  is the solution map of the ODE from time  $t_k$  to  $t_{k + 1}$ .
Suppose we can find  $\mathbf{f}(\cdot, t)$  on  $I_{k+1}$  such that the corresponding  $T_{k+1}$  solves the JKO scheme (5), then with small  $h$ ,  $p_k$  approximates the solution to the Fokker-Planck equation 3, which then flows towards  $p_Z$ . By the Monge formulation of the Wasserstein-2 distance between  $p$  and  $q$  as  $W_2^2(p, q) = \min_{T: T_\# p = q} \mathbb{E}_{x \sim p} \| x - T(x) \|^2$ , solving for the transported density  $p_k$  by (5) is equivalent to solving for the transport  $T_{k+1}$  by
$$
T _ {k + 1} = \arg \min  _ {T: \mathbb {R} ^ {d} \rightarrow \mathbb {R} ^ {d}} F [ T ] + \frac {1}{2 h} \mathbb {E} _ {x \sim p _ {k}} \| x - T (x) \| ^ {2}, \tag {6}
$$
where  $F[T] = \mathrm{KL}(T_{\#}p_k||p_Z)$ . The equivalence between (5) and (6) is proved in Lemma A.1.
Furthermore, the following proposition gives that, once  $p_k$  is determined by  $\mathbf{f}(x,t)$  for  $t \leq t_k$ , the value of  $F[T]$  can be computed from  $\mathbf{f}(x,t)$  on  $t \in I_{k+1}$  only. The counterpart for convex function-based parametrization of  $T_k$  was given in Theorem 1 of [Mokrov et al., 2021], where the computation using the change-of-variable differs as we adopt an invertible neural ODE approach here. The proof is left to Appendix A.
Proposition 3.1. Given  $p_k$ , up to a constant  $c$  independent from  $\mathbf{f}(x, t)$  on  $t \in I_{k+1}$ ,
$$
\operatorname {K L} \left(T _ {\#} p _ {k} \| p _ {Z}\right) = \mathbb {E} _ {x \left(t _ {k}\right) \sim p _ {k}} \left(V \left(x \left(t _ {k + 1}\right)\right) - \int_ {t _ {k}} ^ {t _ {k + 1}} \nabla \cdot \mathbf {f} (x (s), s) d s\right) + c. \tag {7}
$$
By Proposition 3.1, the minimization (6) is equivalent to
$$
\min  _ {\left\{\mathbf {f} (x, t) \right\} _ {t \in I _ {k + 1}}} \mathbb {E} _ {x (t _ {k}) \sim p _ {k}} \left(V (x (t _ {k + 1})) - \int_ {t _ {k}} ^ {t _ {k + 1}} \nabla \cdot \mathbf {f} (x (s), s) d s + \frac {1}{2 h} \| x (t _ {k + 1}) - x (t _ {k}) \| ^ {2}\right), \tag {8}
$$
where  $x(t_{k + 1}) = x(t_k) + \int_{t_k}^{t_{k + 1}}\mathbf{f}(x(s),s)ds$ . Taking a neural ODE approach, we parametrize  $\{\mathbf{f}(x,t)\}_{t\in I_{k + 1}}$  as a residual block with parameter  $\theta_{k + 1}$ , and then (8) is reduced to minimizing over  $\theta_{k + 1}$ . This leads to a block-wise learning algorithm to be introduced in Section 4, where we further allow the step-size  $h$  to vary for different  $k$  as well.
3.2 Infinitesimal optimal  $\mathbf{f}(x,t)$
In each JKO step of (8), let  $p = p_k$  denote the current density,  $q = p_Z$  be the target equilibrium density. In this subsection, we show that the optimal  $\mathbf{f}$  in (8) with small  $h$  reveals the difference between score functions between target and current densities. Thus, minimizing the objective (8) searches for a neural network parametrization of the score function  $\nabla \log \rho_t$  implicitly, in contrast to diffusion-based models which learn the score function explicitly [Ho et al., 2020, Song et al.,
2021], e.g., via denoising score matching. At infinitesimal  $h$ , this is equivalent to solving the FPE by learning a deterministic transport equation as in [Boffi and Vanden-Eijnden, 2023, Shen et al., 2022].
Consider general equilibrium distribution  $q$  with a differentiable potential  $V$ . To analyze the optimal pushforward mapping in the small  $h$  limit, we shift the time interval  $[kh, (k + 1)h]$  to be  $[0, h]$  to simplify the notation. Then (8) is reduced to
$$
\min  _ {\left\{\mathbf {f} (x, t) \right\} _ {t \in [ 0, h)}} \mathbb {E} _ {x (0) \sim p} \left(V (x (h)) - \int_ {0} ^ {h} \nabla \cdot \mathbf {f} (x (s), s) d s + \frac {1}{2 h} \| x (h) - x (0) \| ^ {2}\right), \tag {9}
$$