Universal Approximation Using Well-Conditioned Normalizing Flows
Authors: Holden Lee
Chirag Pabbaraju
Anish Prasad Sevekari
Andrej Risteski
Venue: NeurIPS 2021 Poster (2021)
Keywords: Generative Models
Normalizing Flows
Universal Approximation
Langevin
Dynamical Systems
Deep Learning Theory

Abstract: Normalizing flows are a widely used class of latent-variable generative models with a tractable likelihood. Affine-coupling models [Dinh et al., 2014, 2016] are a particularly common type of normalizing flows, for which the Jacobian of the latent-to-observable-variable transformation is triangular, allowing the likelihood to be computed in linear time. Despite the widespread usage of affine couplings, the special structure of the architecture makes understanding their representational power challenging. The question of universal approximation was only recently resolved by three parallel papers [Huang et al., 2020, Zhang et al., 2020, Koehler et al., 2020] – who showed reasonably regular distributions can be approximated arbitrarily well using affine couplings – albeit with networks with a nearly-singular Jacobian. As ill-conditioned Jacobians are an obstacle for likelihood-based training, the fundamental question remains: which distributions can be approximated using well-conditioned affine coupling flows? In this paper, we show that any log-concave distribution can be approximated using well-conditioned affine-coupling flows.  In terms of proof techniques, we uncover and leverage deep connections between affine coupling architectures, underdamped Langevin dynamics (a stochastic differential equation often used to sample from Gibbs measures) and Hénon maps (a structured dynamical system that appears in the study of symplectic diffeomorphisms). In terms of informing practice, we approximate a padded version of the input distribution with iid Gaussians – a strategy which Koehler et al. [2020] empirically observed to result in better-conditioned flows, but had hitherto no theoretical grounding. Our proof can thus be seen as providing theoretical evidence for the benefits of Gaussian padding when training normalizing flows.

Universal Approximation For Log-concave Distributions Using Well-conditioned Normalizing Flows
Holden Lee Mathematics Department Duke University Durham, NC 27708 holden. leedduke. edu
Chirag Pabbaraju
Computer Science Department
Stanford University
Stanford, CA 94305
cpabbara@cs stanf ord. edu
Anish Sevekari
Andrej Risteski
Department of Mathematical Sciences Carnegie Mellon University Pittsburgh, PA 15213 asevekar@andreex cmu edu
Machine Learning Department
Carnegie Mellon University
Pittsburgh, PA 15213 aristesk@andrew cmu, edu
Abstract
Normalizing flows are a widely used class of latent-variable generative models with tractable likelihood. Affine-coupling models [Dinh al., 2014, 2016] are a particularly common type of normalizing flows, for which the Jacobian the latent-to-obest wwwettt--bbeevaleevariiee to be computed in linear time. Despite the widespread usage of affine couplings, the special structure of the architecture makes understanding their representational power challenging. The question of universal approximation was only recently resolved by three parallel papers [Huang al., 2020, Zhang al., 2020, Kochler et al., 2020] who showed reasonably regular distributions can be approximated arbitrarily well using affine couplings-albett albeit with networks with nearly-singular Jacobian. As ill-conditioned Jacobians are an obstacle for likelihood- based training, the fundamental question remains: which distributions can be approximated using well-conditioned affine coupling flows?
transformation is triangular, allowing the likelihood
In this paper, we show that any log-concave distribution can be approximated using well-conditioned affine-coupling flows. In terms of proof techniques, we uncover and leverage deep connections between affine coupling architectures, underdamped Langevin dynamics (a stochastic differential equation often used to sample from Gibbs measures) and Hénon maps (a structured dynamical system that appears in the study of symplectic diffeomorphisms) Our results also inform the practice training affine couplings: approximate a padded version of the input distribution with iid Gaussians a strategy which Koehler et al. [2020] empirically observed result in better-conditioned flows, but had hitherto theoretical grounding. Our proof can thus be seen as providing theoretical evidence for the benefits Gaussian padding when training normalizing flows.
Introduction
Normalizing flows [Dinh et al., 2014, Rezende and Mohamed, 2015] are class of generative models parametrizing distribution Rdd the pushfoward simple distribution (e.g. Gaussian) through an invertible map IRd - Rd with trainable parameter 0. The fact that go is invertible allows us to
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

write down an explicit expression for the density of a point through the change-of change-of-variables formula, namely po(x) (2) = (( (de) (gg))deettgaaae (x)). (x)), where & denotes the density of the standard Gaussian. For different choices parametric families for go. one gets different families of normalizing flows, e.g. affine coupling flows [Dinh et al., 2014, 2016, Kingma and Dhariwal, 2018], Gaussianization flows [Meng et al., 2020], sum-of-squares polynomial flows [Jaini al., 2019].
In this paper we focus on affine coupling flows arguably the family that has been most successfully scaled up to high resolution datasets [Kingma and Dhariwal, 2018]. The parametrization of go is chosen to be a composition of so-called affine coupling blocks, which are maps f : IR - R", s.t. ([d)s @ s(rs)) ((xs)). where @ denotes entrywise multiplication and s, I are (typically simple) neural networks. The choice of parametrization is motivated by the fact that the Jacobian of each affine block is triangular, so that the determinant can be calculated in linear time.
f( ()d))s) )==
Despite the empirical success of this architecture, theoretical understanding remains elusive. The most basic questions revolve around the representational power of such models. Even the question of universal approximation was only recently answered by three concurrent papers [Huang et al., ta--tthaansssfattrryy... manner, in light
2020, Zhang et al., 2020, Koehler et al., 2020] though in of how normalizing flows are trained. Namely, Huang et al. [2020], Zhang et al. [2020] show that any (reasonably well-behaved) distribution p. once padded with zeros and treated as a distribution in be
Rd+d', can be arbitrarily closely approximated by affine coupling flow. While such padding operationalized as an algorithm by padding the training image with zeros, is never done in practice, as it results in an ill-conditioned Jacobian. This is expected, as the map that always sends the last d' coordinates to 0 is not injective. Koehler et al. [2020] prove universal approximation without padding; however their construction also gives rise poorly conditioned Jacobian: namely, approximate distribution p to within accuracy in the Wasserstein- I distance, the Jacobian of the network they construct will have smallest singular value on the order of
Importantly, for all these constructions, the condition number of the resulting affine coupling map is poor no matter how nice the underlying distribution it's trying to approximate is. In other words, the source of this phenomenon isn't that the underlying distribution is low-dimensional or otherwise degenerate. Thus the question arises:
Question: Can well-behaved distributions be approximated by an affine coupling flow with a well-conditioned Jacobian?
In this paper, we answer the above question in the affirmative for a broad class of distributions log- log-concave distributions if we pad the input distribution not with zeroes, but with independent Gaussians. This gives theoretical grounding of an empirical observation in Koehler et al. [2020] that Gaussian padding works better than geoo-padiing, well as no padding.
The practical relevance of this question is in providing guidance on the type of distributions we can hope to fit via training using an affine coupling flow. Theoretically, our techniques uncover some deep connections between affine coupling flows and two other (seeming unrelated) areas of mathematies: stochastic differential equations (more precisely underdamped Langevin dynamics, a "momentum" variant of the standard overdamped Langevin dynamics) and dynamical systems (more precisely, a family of dynamical systems called Hénon- like maps).
Overview of results
In order to state our main result, we introduce some notation and definitions.
2.1 Notation
Definition 1. An affine coupling block a map f 1dd + Rd, s.t. F(dl,s) = (zs, Jdyys ? s(rs) + t(xxs) for some set of coordinates S, where @ denotes entrywise multiplication and s, t are trainable (generally non-lineary functions. An affine coupling network is a finite sequence of affine coupling blocks. Note that the partition (S, [d] S), as well as 8, l may be different between blocks. We say that the non-linearities are in a class F (e.g.. neural networks, polynomials, etc.) if s, E F.
The appeal of affine coupling networks comes from the fact that the Jacobian of each affine block triangular, so calculating the determinant is time operation.

---

We will be interested in the conditioning of f-that is, an upper bound on the largest singular value max (Df) and lower bound on the smallest singular value min (D) of the Jacobian f of f. Note that this is a slight abuse of nomenclature most of the time, "condition number" refers to the ratio of the largest and smallest singular value. As training a normalizing flow involves evaluating det D)),, we in fact want to ensure that neither the smallest nor largest singular values are extreme.
The class of distributions we will focus on approximating via affine coupling flows log-concave distributions:
Definition A distribution p IRd + Rt, ,((x) x e-((x) is log-concave v²U(2) =
-V In p(x) ž 0.
Log-concave distributions are typically used to model distributions with Gaussian-like tail behavior. What we will leverage about this class of distributions that special stochastic differential equation (SDE), called underdamped Langevin dynamics, is well-behaved in an analytic sense. Finally, we recall the definitions of positive definite matrices and Wasserstein distance, and introduce notation for truncated distributions.
Definition 3.
We say that a symmetric matrix positive semidefinite (PSD) if all its eigenvalues nen-negative. For symmetric matrices A, B, we write A B if and only if A - B is PSD.
Definition 4. Given two probability measures over metric space (M, d). the Wasserstein-1 distance between them, denoted Wi v/), is defined
Wi(Lee) inf J MMM d( (El((w))
where [' v/) is the set of couplings. i.c measures on M M with marginals p. respectively. For two probability distributions we denote by W1 (p, q) the Wasserstein-1 distance between their associated measures. In this paper, we set M R and d(x; y) = || vl12.
Definition 5. Given distribution g and a compact set we denote by gle the distribution g truncated to the set C. The truncated measure is defined as glc(A)
ACC.
2.2 Main result
Our main result states that we approximate any log-concave distribution in Wasserstein-1 distance by a well- conditioned cofini-coupling flow network. Precisely, we show:
Theorem 1. Let p(x) Rdd - IR+ be of the form p(x) x e-U((), such that:
1. U E C2, ie., (x) exists and is continuous.
2. Inp satisfies 1s 2 In I klg.
Furthermore, let Po == p x (O Ia). Then, for every O, there exists compact set C Rdd and an invertible affine-coupling network f Rddd + R22d with polynomial non-linearities, such that
W (N T2a)lc). po) <.
Furthermore, the map defined by this affine-cuplingg network I well conditioned over C, that is. there are positive constants A(k), B(x) = x. such that for any unit vector
A(k) S [[, S B(k)
for all (x,v) E C, where Dw. is the directional derivative in the direction w. In particular, the condition number of Df(x, v) is bounded by B(x) o11) for all (I, v) E EC
A(x)
We make several remarks regarding the statement of the theorem:
Remark 1. The Gaussian padding (i.e. setting Po p (0, T2)) is essential for our proofs. All the other prior works on the universal approximation properties of normalizing flows (with or without padding) result in ill-conditioned affine coupling networks. This gives theoretical backing of empirical observations on the benefits of Gaussian padding in Kochler et al. [2020].
Remark 2.
The choice of non-linearities 8, / being polynomials is for the sake of convenience in our proofs. Using standard universal approximation results, they can also be chosen be neural networks with smooth activation function.
3

---

The Jacobian Df has both upper-bounde largest singular value, and smallest singular value-whice of course bounds the determinant det(Df). As remarked in Sec- tion 2.1, merely bounding the ratio of the two quantities would not suffice for this. Moreover, the bound we prove only depends on properties of the distribution (i.e., K), and does not worsen as e + 0. in contrast to Koehler et al. [2020].
Remark 3.
Remark 4. The region where the pushforward of the Gaussian through ʃ and po are close is introduced solely for technical reasons- reasons-essentiall essentilly. standard results in analysis for approximating smooth functions by polynomials can only be used the approximation needs to hold on a compact set. Note that C. can be made arbitrarily large by making arbitrarily small.
Remark 5. We do not provide an explicit computation of the number of affine coupling blocks in the constructed network, although a bound of polylog(c)/<0k) can be extracted from our proofs.
Remark 6. Our proof also implies a well-conditioned universal approximation result for other related normalizing flow models. Lemma proves that the flow map of underdamped Langevin dynamics is well conditioned for all t [0. ]. However, as indicated in Chen et al. [2018], underdamped Langevin dynamics is a continuous normalizing flow, thus the claim applies to such flows as well. Similarly, the particular affine coupling layers construct in eq. (13) also form residual block, so the claim also holds for residual flows [Behrmann et al., 2018].
3 Related Work
The landscape of normalizing flow models is rather rich. The inception of the ideas was in Rezende and Mohamed [2015] and Dinh et al. [2014], and in recent years, an immense amount of research has been dedicated to developing different architectures of normalizing flows. The focus of this paper are affine coupling flows, which were introduced in Dinh et al. [2014], introduced the idea of using pushforward maps with triangular Jacobians for computational efficiency. This was further developed in Dinh et al. [2016] and culminated in Kingma and Dhariwal [2018], who introduced Ix1 convolutions in the affine coupling framework to allow for "trainable" choices of partitions. We note, there have been variants of normalizing flows in which the Jacobian is -on-triangular, e.g. [Grathwohl et al., 2018, Dupont et al., 2019, Behrmann et al., 2018], but these models still don't 't scale beyond datasets the size of CIFAR-10.
In terms of theoretical results, the most closely related works are Huang et al. [2020], Zhang et al. [2020], Koehler et al. [2020]. The former two show universal approximation of affine couplings- albeit if the input is padded with zeros. This of course results in maps with singular Jacobians, which is why this strategy isn't used in practice. Koehler et al. [2020] show universal approximation their constructions results in a flow model with condition number 1/e to
without padding-thowe get approximation in the Wasserstein sense, regardless of how well-behaved the distribution to be approximated is. Furthemore, Koehler et al. [2020] provide some empirical evidence that padding with iid Gaussians (as in our paper) is better than both zero padding (as in Huang al. [2020]. Zhang et al. (2020)) and no padding on small-scale data.
Preliminaries
Our techniques leverage tools from stochastic differential equations and dynamical systems. We briefly survey the relevant results.
4.1 Langevin Dynamics
Broadly, Langevin diffusions are families of stochastic differential equations (SDEs) which most frequently used as algorithmic tools for sampling from distributions specified up to a constant of proportionality. They have also recently received lot of attention as tools for designing generative models [Song and Ermon, 2019, Song et al., 2020].
In this paper, we will only make use of underdamped Langevin dynamics, momentum-like analogue of the more familiar overdamped Langevin dynamics, defined below. Our construction will involve simulating underdamped Langevin dynamics using affine coupling blocks.

---

Definition 6 (Underdamped Langevin Dynamics). Underdamped Langevin dynamics with potential U and parameters 5. y is the pair of SDEs
dx
-vpdt
duz = -(vedl (U(x)) + /2 d.
The stationary distribution of the SDEs (limiting distribution + oc) given by p" (az, e) x e-U(x)- [vll?
The variable can be viewed as "velocity" variable and H2 "position" variable in that sense, the above SDE an analogue to momentum methods in optimization.
The convergence of (1) can be bounded when the distribution p(x) & exp(-U(x)) satisfies an analytic condition, namely has bounded log log-Sobolev constant. Though we don't use the log-Soboley constant in any substantive manner in this paper, we include the definition for completeness.
Definition 7. A distribution Rd + IR satisfies a log-Sobolev incquality with constant > 0 Vg Rd + R. s.t. g² .g log g E L (p). we have
Epig² log 9 Eplg") log Eplg²1 2CE, Voll?
(2)
In the context of Markov diffusions (and in particular, designing sampling algorithms using diffusions), the interest in this quantity comes as it governs the convergence rate overdamped Langevin diffusion in the KI. divergence sense. Namely, if pt is the distribution of overdamped Langevin after time t, one can show
(L(p:llp)
Ct "KL(polle
We will only need the following fact about the log- Sob--Solee constant:
Fact 1 (Bakry and Emery [1985], Bakry al. [2013]). Let the distributions p(x) oc exp(-U(x)) be such that Uxr) = AT. Then, has log-Sobolev constant bounded by X.
We will also need the following result characterizing the convergence time underdamped Langevin dynamics in terms of the log-Sobolev constant, as shown in Ma et al. [2019]:
Theorem (Ma et [2019]). Let p" (x) xx exp(-U(x) have a tog-Sobolev constant bounded by p. Furthermore, for a distribution p : Rd 12_ let
Clp] == (L(plpp) + Ep (<
5 (TT((ll))"
(KT(())p") 1
SV
8p
8p
where is a positive definite matrix given by S Laxa (14, which evolve according to (1), we have
[laxa] 2laxd]
Pt is the distribution of
Laxd
Clp.) < - C'p Clpis
(3)
dl
10
whenever p" satisfies a log- Sog-Sobolev inequality with constant p.
We note that the above theorem uses a non-standard Lyapunov function L, which combines KL divergence with an extra term, since the generator of underdamped Langevin self-adjoint-this makes analyzing the drop in K L divergence difficult. As L is clearly an upper bound on K L(pllp"", so it suffices show L decreases rapidly.
We will also need less-wed well-wwell deterministic form of the updates which is equivalent to (1). Precisely, we convert (1) equivalent ODE (with dime-dependent coefficients). The proof of this fact (via a straightforward comparison of the Planckk equation) can be found in Ma et al. [2019].
Theorem 3. Let Pe( v:) be the probability distribution of running (1) for time If started from (zo, vo) Po, the probability distribution of the solution (It, vt) to the ODEs
d
0
la -7la
(Vinpt Vlnp")
(4)
dt
V/
-la
is also pr (2,, v.).
5

---

4.2 Dynamical systems and Henon maps
We also build on work from dynamical systems, more precisely, a family of maps called Hénon-like maps [Hénon, 1976].
Definition 8 ([Turaev, 2002]). A pair of ODEs forms a Hénon-like map if it has the form
dt
(5)
- T VJ(z)
dt
for smooth function J R R.
This special family of ODEs is continuous-time generalization of classical discrete dynamical sys- tem of the same name Hénon [1976]. The property that is useful for that the Euler discretization of this map can be written as a sequence of affine coupling blocks.
proven that these ODEs universal approximators in some sense. Namely, the iterations of this ODE can approximate any symplectic diffeomorphism: a continuous map which preserves volumes (i.e. the Jacobian of the map 1). These kinds of diffeomorphisms have their genesis in Hamiltonian formulations of classical mechanics [Abraham and Marsden, 2008].
In Turaev [2002],
At first blush, symplectic diffcomorphisms and underdamped Langevin seem to have nothing to do with each other. The connection comes through the so-called Hamiltonian representation theorem [Polterovich, 2012], which states that any symplectic diffeomorphism from C S IR2d R2d can be written as the iteration of the following Hamiltonian system of ODEs for some (time-dependent) Hamiltonian (x, v,t):
d HI, do = -5 (I,
(6)
dv
In fact, in our theorem, will techniques inspired by those in Turaev [2002], who shows: Theorem 4 (Turaev [2002]). For any function Hæ, () R2dd R20 - R which is polynomial in (z, v), there exists a polynomial (zz, c, t), s.t. the time-T map of the system
de U(), --
is uniformly O(72)- to the time-2m map of the system
dx d
(8)
dv
--2
TJ ((+,))
Jag
for some integers (94
We will prove a generalization of this theorem that applies underdamped Langevin dynamics.
Proof Sketch of Theorem
5.1 Overview of strategy
We wish to construct an affine coupling network that (approximately) pushes forward Gaussian p" (O, I2a) the distribution we wish model with Gaussian padding, i.e. Po = p X (0, Ia). Because the inverse of an affine coupling network is an affine coupling network, we can invert the problem, and instead attempt to map Po to N(0, I2a).
There a natural map that takes to p" = NO, Iaa). namely, underdamped Langevin dynamics (1). Hence, proof strategy involves understanding and simulating underdamped Langevin dynamics with the initial distribution = x (00 14), and the target distribution p* [O, 12a), and comprises of two important steps.
As an aside, a similar strategy is taken in practice by recent SDE-based generative models (Song et al. [2020]).

---

First, we show that the flow-map for Langevin is well-conditioned (Lemma below). Here, by flow-map, we mean the map which assigns each its evolution over a certain amount of time? according to the equations specified by (1).
Second, we break the simulation of underdamped Langevin dynamics for certain time è into intervals of size т, and show that the inverse flow-map over each r-sized interval of time can be approximated well by composition of affine-coupling maps (Lemma below). To show this, we consider more general system of ODEs than the one in Turaev [2002] (in particular, a non-Hamiltonian system), which can be applied underdamped Langevin dynamics. We then show that the inverse flow-map of this system of ODEs can be approximated by a sequence of aifine-coupling blocks. We note that for this argument, it critical that we use underdamped rather than overdamped Langevin dynamics, as overdamped Langevin dynamics do not have the required form for affine- affine-coupling blocks.
5.2 Underdamped Langevin is well-conditioned
Consider running underdamped Langevin dynamics with stationary distribution p' equal to the standard Gaussian, started log-concave distribution with bounded condition number k. The following lemma says that the flow map is well-conditioned, with condition number depending polynomially on k.
Lemma 1. Consider underdamped Langevin dynamics (1) with = 1, friction coefficient 7 < 2 and starting distribution p which satisfies all the assumptions in Theorem I. Let T₁ denote the flow-map from time 0 to time t induced by (4), Then for any Fo, vo E and unit vector w, the directional derivative T, at To, vo in direction w satisfies
( 1 2 2-7
2 + 1-
-2/7
2 2+7
2/7
< <I
& 1
3-9 (n (N-
4/7
Therefore, the condition number of Te is bounded by (L+ 2+ (n 11) 2-
We sketch the proof below and include complete proof in Section A.
First, using (4) and the chain rule shows that the Jacobian of the flow map at zo, D. = (To( satisfies
d D -Ia
O Ia -7Ia
-
- -lpp'
ie., it bounded by the difference of the Hessians of the log-pdfs of the current distribution and the stationary distribution. We will show that V In pt decays exponentially towards V Inp" Izd-
To accomplish this, consider how 2 Inp, evolves if we replace (1) by its discretization,
3etn Fe + njū Uetn (1 77) Fp & & (O, 2yy]Ta).
Note that because the stationary distribution is Gaussian, VU(z:) 2, in (1), and the above equations take particularly simple form: we apply a linear transformation to
H
and then add Gaussian noise, which corresponds convolving the current distribution by Gaussian. We keep track of upper and lower bounds for 2 In pt, and compute how they evolve under this linear transformation and convolution by a Gaussian. Taking r + we obtain differential equations for the upper and lower bounds for V In pt. which we can solve. A Grönwall argument shows that these bounds decay exponentially towards V Inpp" 1zd- The decay rate can be bounded as a power of +
From (9), we then obtain that the condition number of D, is bounded by the integral of a exponentially decaying function, and hence is bounded independent of t. In particular, we may take large enough so that Pt is e-close to the stationary distribution. Because the decay rate of the exponential is Ro(ay the bound is 011)
Note that we vitally used the fact that the stationary distribution p is a standard Gaussian, as our argument requires that 2 Inp" be constant everywhere.
7

---

5.3 ODE approximation by affine-coupling blocks
Next, we analyze more general version of the Hamiltonian system ODEs considered in Turaev [2002], which we recalled in (7). In particular, the system of ODEs we will be considering is:
dr a H., v, de 8y - H( - H
(10)
de = dz
dt
Note that substituting (z, v, t) = (e(,, -) 1n pn (x. v) above gives us the underdamped Langevin dynamics.
The first step is restrict our considerations H/ being polynomial in v, smooth function. Towards this, we recall the notion of closeness in the C2 topology:
v, rather than general
Definition 9.
Let C I be a compact set. Let f.g + R be two continuously differentiable functions. Then we say that f are uniformly c-close over C in C' topology if
sup(Il[(x) g(>))) (I] Dy(x) Se
=EC
The following lemma (a generalization of the Stone- Weierstrass Theorem) then establishes that suffices to focus on H being polynomial in z; V:
Lemma (Theorem 5, Peet [2007])- Let C C Rd be compact set. For any C2 function H TRd - R. and any 0, there multivariate polynomial P P:Rd - R such that H uniformly c-close over C in C topology.
Focusing the case of polynomials, Lemma below shows that instead of flowing the pair of ODEs given by (10) over an interval of time т, we can instead run a different ODE for time 2n, such that the flow-maps maps corresponding to both these ODEs O(r2)
Lemma 3. Let C Rd2d be compact set. For any function (x,, R2ad R which is polynomial in (x, v), there exist polynomial functions J, F, G, s.t. the time-(tu + T, r,ta) flow map of the system
dx 0 H(, v,t)
dt
3u as
(11)
(x, v,t) - 7 11 v,
dt
uniformly (c7o)s over C in C topology to the time-2x map of the system
dt dx = v TF
(12)
du
dt
Here, & denotes component-wise product, and the constants inside the O(-) depend on C and the coefficients of H.
The complete proof of this lemma is included in Appendix B; we provide brief sketch here. First, we consider the first order (O(r?) approximation of the flow map of standard ODE of the form ý Dy (where is diagonal), and observe that for small т, we can think of (12) as a perturbed version of such ODE with appropriate choice of D. Using standard ODE perturbation techniques, approximately express the time-t evolution of (12) up to first-order in т, in terms of polynomials F, [, J and trigonometric functions.
Then, we compare this map to the first-order approximation of flowing the pair of ODEs (11) for time via Taylor's theorem. Furthermore, this approximation is polynomial in (z, v) since H is polynomial in vx)
The crucial step involves choosing the functional form of (), t), G(z, t) suitably, so that they are polynomials in with coefficients in terms of (i2t)),, cos(ftt. After simplification, both expressions can be expressed in terms of polynomials in æ, where coefficients can be expressed in terms of 12² sin" (ft) cos" ds, which either integrate 0 constant. Thus, to ensure that the two approximations match, we are left with a problem of making two multivariate polynomials in (xr, v)) equal.
This final step can course be written as a linear system of equations. We identify special structure in this system, which helps show that the system is full-rank, and hence has solution. D

---

Finally, consider discretizing the newly constructed ODE (12) into small steps of size 1j by simple Euler schema i.e.
In+11 En + nenn - (e,, 7)1) O In)
(13)
Un+ = Un.j "(( TJy (In,,))) Gynn. (1n:))))
We note that each step above can be written composition of two affine coupling blocks given by (In; Un) + (Ia, Un+1) >> (Izz+1, Namely, the map (In, vn) + (Ing un+1) can be written as
In In
Un+1 = on @ (1 - ()n( mm) 17(22 3 In TJ (z 7m))
This map is an affine coupling block with s(xn) = (1 - 7) @ G(Tns mm)) and = +
In T TJ (z:: mm)). The map (In, Ux+1) + (xn+1, Un-1) can be written
Un+1 Un+1
Fn+1 = In + 771++111 - T FU 171) O an))
which is affine coupling block with s(Un+1) = 1/7 1 ((nn+1, 177) and ((vn41) 1p8+1
The composition of the two maps above yields an affine coupling network (In, Up) + (In+1s Un+1) precisely as given by Equation (13) with non- non-lincarities in each of the blocks given by polynomials. The following lemma bounds the error resulting from this discretization:
Lemma 4 (Euler's discretization method). Let C C 2d be compact set. Consider discretizing the time from 0 to into 1 steps and performing the update given by (13) at each of these steps. Let result of discretizing thus be denoted by Ti and let the original flow map be denoted by Tr. Then Tr and T are uniformly O(1) close over in C topology, and the constants inside the O()) depend on C, and bounds the derivatives of T, over C.
the map obtained
5.4 Simulating by breaking into r-sized intervals
Let T.. denote the time-s, flow-map of (10) from time s to time Since the flow maps are invertible. Ts,: and Tr,, are inverses. We are now ready to state the following lemma which says that the underdamped Langevin flow-map To.0 can be written as a composition of affine-couplings maps:
Lemma 5. Let C C IR2d be a compact set. Suppose that To,0(1, v) the time-(0, 0) flow-map of the ODE's
0 H v,t)
dt
(14)
du 2r
(,, H(r, H(),
8z
where H is C Then for any 61 E Rt, there exists an integer N= N(1,, C) and affine- coupling blocks fi...., SN such that the composition f fN ....o fi is e1-close To,0 in the C topology over C.
The proof of Lemma 5 is in Appendix C. We provide brief sketch here: from Lemma 2, know that it suffices to show the result for a polynomial H. Thereafter, we break the time for which we want flow the ODE given by (14) into small chunks of length r. Lemmas 3 and 4 then show that the flow map over this chunk can be written as an affine coupling network. Composing the affine coupling networks over all the chunks of time gives us the result.
5.5 Putting components together
The previous sections established that for any and any compact set C, there is a affine-coupling network f with polynomial non-linearitie, such that T.o and f are uniformly close over C. We will now pick an appropriate value of and set such that W +f (p" le). S where p* = -O Iza). which is the required result of Theorem 1. First, using Theorem 2, for
-00 1 + log 2 + log L(po]
we have that KL,To.ow (pn), p²) this to a Wasserstein bound.
We the following transportion cost inequality to convert
This result is well known in the O topology, provide analysis for the C bound in Appendix D.I.

---

Theorem 5 (Talagrand [1996]). The standard Gaussian I' satisfies transportation cost inequality: For every distribution q on IR with finite second moment, Wi q)² 2
This gives us that Wi (To, (po), p²) S 61. A simple argument in Lemma 10 (Appendix D.2) gives
then Wi(po: # Wi (To.0 # (To. # (po)), T (p")) - Lip(Te,o)e
(15)
subsequent argument stated as Lemma 11 in Appendix D.2, shows that if f and Ta.0 are uniformly q-closee in Co topology on some C, then their pushforwards through p"le are indeed close, i.e., (16)
Wi (To.0 # (v" le), (p"lc)) EI.
Next, we establish a bound on the Wasserstein distance between the standard Gaussian and its truncation on a compact set, proved in Appendix D.3.
Lemma 6. Let p" (O, I2a)- Then for every d E R., there exists a compact set C B(O, R) such that W1 (p" pp(c) & where B(0, R) denotes the ball of radius centered at the origin.
We now choose compact C' such that Lemma holds for 8 E1. Then Lemma 10 again implies that
Wi (T (p"), To.0 # (p"lc)) - Lip(To.o)e
(17)
Equations (15), (16), (17) and the triangle inequality together imply
W1 (f (p" Po) & (2Lip(To.0) ) +1)c1 S
for small enough 61- We can indeed set 61 small enough so as to satisfy the last inequality above, because of the global bound Lip(Tæ.o) us the statement of Theorem 1. Note that the final value of depends c, K, and C[pu].
(x - 11) 2/7
established in Lemma This gives
6 Conclusion
In this paper, we provide the first guarantees on universal approximation with well-conditioned affine coupling networks. The conditioning of the network is crucial when the networks are trained using gradient-based optimization of the likelihood. Mathematically, we uncover connections between stochastic differential equations, dynamical systems and affine coupling flows. Our construction uses Gaussian padding, which lends support to the empirical observation that this strategy tends to result conditoond flows [Koehler et al., 2020]. We leave it as open problem generalize beyond log-concave distributions.
10

---

References
Ralph Abraham and Jerrold Marsden. Foundations of mechanics. Number 364. American Mathematical Soc., 2008. 6
Uri M Ascher and Chen Greif. A first course numerical methods. SIAM. 2011. 26
Dominique Bakry and Michel Emery. Diffusions hypercontractives. In Séminaire de Probabilités XIX 1983/84, pages 177-206. Springer, 1985. 5
Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators, volume 348. Springer Science & Business Media, 2013. 5
Jens Behrmann, Will Grathwohl, Ricky T. Q. Chen, David Duvenaud, and Jöm-Henrik Jacobsen. Invertible residual networks. nov 2018. URL http:://axttiii org/abs//////eeeoeeo/ooooeee
abs/1811. 01111.009953..... 4
Jens Behrmann, Paul Vicol, Kuan-Chies Wang, Roger Grosse, and Jörn-Henrik Jacobsen. Un- derstanding and mitigating exploding inverses in invertible neural networks. jun 2020. URL 09347v1.
http://arxiv. wwwpeiiiiirrr..oooeeeeee wwwg/bs////e
S.G Bobkov and F Götze. Exponential integrability and transportation related to logarithmic sobolev inequalities. Journal of Functional Analysis, 163(1):1-28, 1999. ISSN 0022-1236. doi: URL httph: httttps::/www... sciencedirect....... wtwpeeedddeeetttttoooooeeeeeeeee
wttpg///...00//faa......om (tticce/pii/S002212665038
Herm Jan Brascamp and Elliott H Lieb. On extensions of the brunn-minkowski and prékopa-leindler theorems, including inequalities for log concave functions, and with an application the diffusion equation. In Inequalities, pages 441-464. Springer. 2002. 14
Jianfei Chen, Cheng Lu, Biqi Chenli, Jun Zhu, and Tian Tian. Vflow: More expressive generative flows with variational data augmentation. feb 2020. URL http:///aart /ttp:////xxvvv/// arggabs//000 09741v2.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. jun 2018. htLthtp///aaii http://aaxxvv..rabbbba/a/ wttp:oo////a//oeeeeoooo htt///bbb///..../// 07366v5. 4
Carmen Chicone. Ordinary differential equations with applications. Springer, New York Berlin, 2006. ISBN 9780387307695. 29
Rob Cornish, Anthony L. Caterini, George Deligiannidis, and Arnaud Doucet. Relaxing bijectivity constraints with continuously indexed normalising flows. sep 2019. URL http://arxiv, org/ abs/1909. 13833v5.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv preprint arXiv:1410.8511 2014. 1, 4
Laurent Dinh, Jascha bock-Diccstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint /rXiv:/605.00803 2016. 2, 4
Emilien Dupont, Araud Doucet, and Yee Whye Teh. Augmented neural odes. In Advances in Neural Information Processing Systems, pages 3134-3144, 2019. 4
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367. 2018. 4
Gilles Hargé. A convex/log-conee eanecccvv concave correlation inequality for Gaussian measure and an application to abstract Wiener spaces. Probability theory and related fields, 130(3):415 440, 2004.
Michel Hénon. two-dimensional mapping with a strange attractor. In The Theory of Chaotic Attractors, pages 94-102. Springer, 1976. 6
Chin- Wei Huang. Laurent Dinh, and Aaron Courville. Augmented normalizing flows: Bridging the gap between generative flows and latent variable models, feb 2020. URL. http::/ttxxiiiiiie org/ 2,

---

Priyank Jaini, Kira A. Selby, and Yaoliang Yu. Sum-of-squares polynomial flow, 2019. 2
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In
Advances in Neural Information Processing Systems, pages 10215-10224, 2018. 2, 4
Frederic Koehler, Viraj Mehta, and Andrej Risteski. Representational aspects of depth and condition-
ing in normalizing flows, 2020. 1, 2, 3, 4, 10
Alexander V Kolesnikov. Mass transportation and contractions. arXiv preprint arXiv:1103.1479, 2011.
Yi-An Ma, Niladri Chatterji, Xiang Cheng, Nicolas Flammarion, Peter Bartlett, and Michael I. Jordan. Is there an analog of nesterov acceleration for meme? feb 2019, URL http: htttp://axxv.. org/ abs// 2902.00996v.. 5
Chenlin Meng. Yang Song, Jiaming Song, and Stefano Ermon. Gaussianization flows, 2020. 2
Felix Otto and Cédric Villani. Generalization of an inequality by talagrand and links with the logarithmic sobolev inequality. Journal of Functional Analysis, 173(2)3661-40,, 2000.
Grigorios Pavliotis. Stochastic processes and applications: diffusion processes, the Fokker-Planck and Langevin equations, volume 60. Springer, 2014.
Matthew M. Peet. Exponentially stable nonlinear systems have polynomial Yyaрunov functions on bounded regions, 2007. 8
Leonid Polterovich. The geometry of the group of symplectic diffeomorphism. Birkhääser, 2012.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. International Conference on Machine Learning, pages 1530-1538. PMLR, 2015. 1,4
Adrien Saumard and Jon A Wellner. Log-concavity and strong concavitt;;: review. Statistics surveys, 8:45, 2014. 13
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. arXiv preprint
2019. 4
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben
Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:201 1. 13456, 2020. 4,
E. G. Tabak and Cristina V. Turner. family of nonparametric density estimation algorithms. 66(2): 145-164, sep 2012. doi: 10. (0.11022/pa.2144) URL https: org/ 10. 1002
21423. 21422.
Esteban G. Tabak and Eric Vanden-Eijnden. Density estimation by dual ascent the log-likelihood. 8(1):217-233, 2010. doi: 1044100cms.2010..8.ll... URL https: :/doii org/10, 4310% 2Fcms 2010.180 v8. mi. a11.
Amirhossein Taghvaei and Prashant Mehta. Accelerated flow for probability distributions. In International Conference on Machine Learning, pages 6076-6085. PMLR, 2019.
Michel Talagrand. Transportation cost for gaussian and other product measures. Geometric & Functional Analysis GAFA, 6(3):587-600, 1996. 10
Takeshi Teshima, Isao Ishikawa, Koichi Tojo, Kenta Oono, Masahiro Ikeda, and Masashi Sugiyama. Coupling-based invertible neural networks universal diffeomorphism approximators, jun 2020.
URL http://arxiv htgg//bs/200066 11469v2.
Dmitry Turaev. Polynomial approximations of symplectic dynamics and richness of chaos in non-hyperbolic area-preserving maps. Nonlinearity, 16:1) nov 2002. doi: 10. 1088/ 7, 17, 20
URL https://d// otte::/000.01100001117777////111
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Han Zhang, Xi Gao, Jacob Unterman, and Tom Arodz. Approximation capabilities of neural odes and invertible residual networks. 2020. 2,
12


Reviews:
[{"review_id": "vB063fnMjE9", "reviewer": "Reviewer_PyGo", "summary": "This paper proves that _well-conditioned_ affine-coupling-based normalizing flows can approximate any _log-concave_ distribution. The proof techinques used in this paper combine underdamped Langevin dynamics and methods from dynamical systems in a novel way. The actual result in this paper involves approximating, to arbitrary precision, a joint distribution comprised of the original target padded with Gaussian noise of the same dimension. Then, universality is implied by the fact that the marginal along the dimensions of the original target distribution must be equal.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I found this paper to be an interesting read and would recommend it for inclusion in the conference. I'll provide further details along the standard review axes below.\n\n### Originality\n\nStudying the approximation properties of $d$-dimensional normalizing flows by embedding in a $2d$-dimensional space is not entirely novel, but the result achieved and the techniques used to do so appear original as far as I am aware. I will note a couple omissions in the related work, however:\n1. The term \"normalizing flow\" was actually coined in _Density estimation by dual ascent of the log-likelihood_  [Tabak & Vanden-Eijnden, 2010], with normalizing flows being further developed in the follow-up _A family of nonparametric density estimation algorithms_ [Tabak & Turner, 2013]. Although Rezende & Mohamed, 2015 and Dinh et al., 2014 may have popularized the notion of flows in the machine learning community, it would be remiss to say that the inception of the ideas were in these papers alone.\n2. Two recent works, _Relaxing bijectivity constraints with continuously-indexed normalizing flows_ [Cornish et al., 2020] and _Understanding and mitigating exploding inverses in invertible neural networks_ [Behrmann et al., 2021] study the theoretical and empirical pitfalls, respectively, of using ill-conditioned normalizing flows. The former in particular notes that the (Bi-)Lipschitz constant of a normalizing flow -- and therefore its condition number -- must necessarily approach infinity if the flow distribution is to approximate a distribution with non-trivial topology to arbitrary precision. The latter provides a complementary empirical view of the matter, demonstrating that normalizing flows quite often display numerical non-invertibility when fully trained. Including these results in the related work may help bolster your motivation even further.\n\n### Quality\n\nAlthough I did not check this closely, the mathematical results in the paper appear to be correct. This paper does a good job of bringing together ideas from different fields to produce a proof. My one criticism here though is that the result is not looked at critically - it is only mentioned that the result is limited to log-concave distributions, without further discussions of the limitations of the paper overall which I've included in the Significance section of the review.\n\n### Clarity\n\nI found the paper to be clearly written, and I very much appreciated the intuitive proof sketches provided throughout.\n\n### Significance\n\nI think that demonstrating how to show the main result is the most important part of the paper, as opposed to the result itself. What I mean is that the combination of techniques provided in this work may allow others to prove stronger results about the approximation properties of certain classes of flows. I do not find the result itself to be particularly shocking, as log-concave distributions do not appear to be a particularly difficult class of distributions to approximate. Furthermore, this work does not really consider the practical implications of using normalizing flows for generic distributions, only noting that they can approximate log-concave distributions without blowing up. I would have liked to see some study or discussion on the practical implications of this result, as for now I cannot really think of any besides this paper possibly inspiring future works with its mathematical techniques.\n\n### Miscellaneous Comments\n\n- How restrictive is the condition $I_d \\leq - \\nabla^2 \\log p(x) \\leq \\kappa I_d$? What happens if $-\\nabla^2 \\log p(x)$ has eigenvalues between $0$ and $1$, and how often might we expect this to occur?\n- I find it quite counterintuitive that $\\mathcal C$ can be made arbitrarily large by making $\\epsilon$ arbitrarily small - can you comment on this further?\n- Is it required, strictly speaking, to pass the proof to a $2d$-dimensional space instead of the original $d$-dimensional space, and why? Is there some smaller dimension where things might still work properly?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think there could have been more discussion about the practical implications of the results beyond just stating that likelihood-based training requires stability in the flow. How can we use the insights here to train flows better? Can we make some recommendations for generic distributions, or can we only say that normaliizing flows can stably approximate log-concave distributions?\n\nI have no concerns about the societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "oYN78WcaSJu", "reviewer": "Reviewer_yz9H", "summary": "This paper proves a universal approximation theorem of well-conditioned affine coupling flows for log-concave probability distributions. This extends previous approximation results with ill-conditioned flows. The paper takes a different proof approach with previous works: it constructs an underdampled Langevin dynamics, which is a well-conditioned pushforward between the Gaussian-augmented data distribution and standard Gaussian. Then, the paper shows that affine coupling blocks can simulate the Langevin dynamic.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The most significant contribution of this paper is it brings in the powerful SDE tool to study affine coupling flows. The proof technique is novel. It's the first time I have seen such an approach. I think the method can be used for analyzing other problems, e.g., extended beyond log-concave distributions. I didn't check the proof in detail, but the idea seems valid. The quality of the work is good. The theorems and background knowledge are nicely organized. The presentation is clear.\n\nIn terms of significance, one limitation of the proposed approach is it seems to be limited to affine coupling flows, since the construction of underdamped Langevin dynamics is specially designed to meet the two-part structure of affine coupling flows. It would be more interesting if the authors can demonstrate the idea for other types of flows, e.g., residual flows.\n\nAnother result I would like to see is the practical implications of the theoretical result. Currently it is not discussed in the paper. For example, can the theoretical result guides us to improve the conditioning of affine coupling flows, or can the Langevin dynamics be used to construct more efficient flows?\n\nThe idea of ODEs and SDEs are also extensively used in Neural ODE, FFJord, residual flow, etc.. It is also interesting to see the connection between the proposed method and the neural ode family of methods. Can we compare their representation capacity / efficiency ?\n\nPost rebuttal: I am satisfied with the authors' feedback regarding to the necessity of log-concavity and applicability to other types of flows. I'm increasing my score to 7.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This work does not have foreseeable negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "jaIsdNlZtUW", "reviewer": "Reviewer_Z2tC", "summary": "In this paper, authors provide the first proof that well conditioned affine coupling flows can approximate any log concave distribution.The authors employ stochastic differential equations and structured dynamical systems to construct their proof, uncovering some unknown connections. ", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Coupling flows have been one of the most widespread used members of the flow family; due to the easy of implementation and computation, they have attracted more attention than other variants. However, as mentioned elsewhere and pointed out by authors, the theoretical understanding of their representation capabilities have been limited. This paper provide the first satisfactory result to that front: well conditioned affine coupling flows can represent any log-concave density. \n\nOverall, the paper is very well written. I applaud the authors for their efforts to make a paper that is basically a long proof to be so readable. Clearly, the paper has been tailored to be as accessible to computer scientists as possible; I commend the authors for their efforts. \n\nI believe the paper has sufficient merit; a valid justification for the use of affine coupling flows is a welcome addition to the literature. Moreover, the proof technique and potentially help future analysis. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I don't think this applies.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4-5 hours"}, "extra_text": {}}, {"review_id": "baPNzIHcc6D", "reviewer": "Reviewer_5zdr", "summary": "The authors study universal approximation with well-conditioned affine coupling networks. They proved that the log-concave distributions can be approximated by an affine coupling flow with a well-conditioned Jacobian. To prove this argument, they study the connections between stochastic differential equations, dynamical systems and affine coupling flows. In their construction, Gaussian padding was adopted, as suggested by the empirical observation from [Koehler et al., 2020]. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The related works are adequately cited. The novelty of this paper is high. The universal approximation results in this paper, as mentioned in the above summary part, will certainly help us have a better understating of affine coupling networks from a theoretical way. I have checked the technique parts and find that the proofs are solid. I think this is a significant contribution to machine learning immunity. My only concern is that, the authors only consider the universal approximation for log-concave distributions, which make the contribution of this paper limited. It will be interesting to explore the universal approximation of affine coupling networks for more general distributions in the future.\n\nQuality: This paper is technically sound.\n\nClarity: This paper is clearly written and well organized. I find it easy to follow.\n\nSignificance: I think the results in this paper are significant, as explained above.   ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors have adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"ethical_concerns": "No", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}, {"review_id": "SiKWVE_M1M", "reviewer": "Reviewer_Jjda", "summary": "The paper gives a proof of universal approximation by well-conditioned (affine-coupling) normalizing flows of arbitrary log-concave distribution. The main difference from previous results on universal approximation by such flows is that transformation is well-conditioned. The proof is constructive and potentially allows to extract the number of flow layers required to achieve given accuracy.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper considers the question of  universal approximation property of affine-coupling normalizing flows. Previous results on this subject established universal approximation for ill-conditioned flows that have their condition number diverge with approximation quality. The authors attack the problem by restricting the target distribution to be log-concave with bounded log-probability Hessian. Using SDE convergence results they define an ODE (equivalent to Langevine SDE), that provides approximate distribution via a well-conditioned transform. Further, they extend a known equivalence (Henon map <-> Simplectic flow) to recast previous approximation into the form of a generalized Henon map. The discrete version of this map has the form of a sequence of (pairs of) affine-coupling flows. This allows the authors to establish the universal approximation property of affine-coupling normalizing flows.\n\n\nThe contribution is original and very relevant for practical applications. The paper gives theoretical grounds for using padding of data distribution with Gaussians for the purpose of approximating it with well-conditioned NF.\nThe paper is well written, giving most definitions and qualitative proofs in the main text, which makes it easy to follow. The authors did a really good job in presenting the complex proofs in a very limited space.\n\n\nRemarks:\n1. in line 174 the equation U(x) >=\t\\lambda I is a bit confusing, because U is scalar and I is a unit matrix\n2.  Even though it may be not physically possible given space constraints, but it would be helpful to add a diagram, representing different pieces that were used in the proof, with e.g. arrows showing connections/equivalences.\n3. I found it hard (not being as expert in the material) to track the connection between \"condition parameters\" k (line 107), lambda (line 174) and rho (eq (2)): may be adding remarks on how they connect to each other around the places where they are used could help with that.\n4. Finally, even though the number of flow layers required to achieve given accuracy can be extracted from the proof, having an explicit formula would be helpful for practitioners.\n\n\n****************************************************\nPost-rebuttal\n\nI'd like to thank the authors for their response and will keep my already high score the same.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "N-brbeFF26", "reviewer": "Reviewer_NQpX", "summary": "Authors answer the question: \"Can well-behaved distributions be approximated by an afﬁne coupling ﬂow with a well-conditioned Jacobian?\" They answer by the affirmative, proving that log-concave distributions belong the the class of distributions that can be approximated by such flows, if the input distribution is padded with independent Gaussians.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\n## Strengths\nThe main strength of this work is its theoretical soundness and the ability to overcome the limitations of previous work, namely the ill-conditioned Jacobian yielded by their assumptions. \nThe bound only depends on properties of the distribution (and does not worsen as epsilon gets small, as opposed to Koehler et al., 2020), which is satisfying.\nAdditionally, the independent Gaussian padding assumption that is made backs previous empirical results that have shown their usefulness in terms of conditioning of the flow.\n\n## Weaknesses\nPerhaps that are two minor limitations in this work that could potentially be alleviated.\n1/ Relation with practical use of affine-coupling flows. This entire approach is based on augmenting the space with an auxiliary variable $v$ (via padding). It has apparently been shown that in practice it does indeed help with the conditioning of the flow (Koehler et al., 2020), yet I don't believe that it is a common practice. It would be nice to remind us what does better conditioning brings for practitioners. Is this work claiming that we should in practice pad the input?\n2/ Discussion of other normalizing flows. I may be missing something, but it seems that From Equation 11 a similar result could be proved for continuous normalizing flows and similarly for Equation 12 and residual flows.\n3/ Relationship with optimal transport map.\nIf I understand correctly, this work is claiming that afﬁne coupling ﬂows are universal approximator of the flow-map from underdamped Langevin dynamics. Why are underdamped Langevin dynamics a natural choice? As opposed to the optimal transport map for instance, which seems quite natural as Theorem 1 is stated in terms of an upper bound on the Wasserstein distance. Is it the same map with some assumptions? Is there any empirical evidence that in practice afﬁne coupling ﬂow approximate this Langevin flow-map?\n\n## Clarity\nThe paper is rigorously written and quite well organized. I believe nonetheless that the flow could still improved to expand the audience of the work. In particular, Theorem 2 could need a bit more background or explanation. Also it is a bit unclear why and where we need Hénon maps. More generally, as Section 5 is a sketch of proof (the full proof being in the appendix), it would be worth appealing perhaps a bit more to the intuition of the reader at times (even though I acknowledge there are already some good effort on that regard).\n\n## Relation to prior work\nIt would be perhaps relevant to mention theoretical work on the representational capacity of other class of normalizing flows (e.g. Huang et al. 2020).\n\n## Additional feedback\n89: Not clear how `underdamped Langevin dynamics' relates to log-concave distribution from this sentence.\n142: \"don't\"\n166: Shouldn't this be stated as a proper result?\n168-169: This intuition is useful. The connection with the Gaussian padding (as a Gaussian prior for the velocity) could be stressed.\n183: It would be useful to briefly explain why L is an upper bound of the KL divergence.\n282: Is there a practical way to compute the value of $\\tau$? That can be useful as it informs us on the number of layers required.\n297: Why are these functions polynomial in $\\cos(\\Omega t)$ and $\\sin(\\Omega t)$?\n306: Would be quite useful to write explicitly the associated non-linearities $s,t$ for both layers. Is appears that $s$ is always of the form $s=(I_d + \\tilde{s})$ due to the Euler discretization. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I wouldn't say that the limitations of this work are clearly discussed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "EOvUO704ODN", "reviewer": "Reviewer_rkxh", "summary": "The paper proves that log-concave distributions can be approximated by normalizing flows defined by affine-coupling models and starting from Gaussian distributions. The special feature of the result, distinguishing it from earlier universal approximation results for this type of flows, is that the constructed map is well-conditioned. The proof first approximates the target distribution using a Langevin dynamics, and then shows that this dynamics can be approximated by affine-coupling models. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This is an interesting and technically complex paper that reveals some interesting connections between normalizing flows, Langevin evolution, and earlier research on Hénon maps as universal approximators for symplectic diffeomorphisms. While there are already several papers on universal approximation with normalizing flows, this paper poses and resolves the question of finding a well-conditioned flow (albeit under relatively restrictive conditions on the target distribution - log-concavity and Gaussian padding). The proposed (rather complex, two-stage) construction combines ideas of the Langevin evolution and earlier research on approximating Hamiltonian flows. The paper is clearly written and includes a sketch of the proof that explains its main points rather well. I think that the paper brings some novel ideas to the study of approximation properties of normalizing flows.     \n\nIf I understand correctly, the well-conditioning in the paper refers only to the full flow map, so that the elementary affine-coupling layers comprising the flow can potentially be ill-conditioned (or ill-conditioning can accumulate because of the large number of these layers)? \n\nI didn't quite understand at which point the Gaussian padding is used in the proof. Is it used in the Langevin convergence result (Lemma 1), because the target velocity distribution needs to be Gaussian?\n\nProbable typos: the minus sign in the first equation in system (1), missing $\\gamma$ in the covariance of $\\xi_t$ in line 250.\n\nPost-rebuttal: I thank the authors for answering my questions.\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "-", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper demonstrates that log-concave distributions can be well approximated by well-conditioned affine-coupling flows with Gaussian base density.

This paper is clearly written and technically sound. It provides novel insights and interesting connections between normalising flows, Langevin diffusions, Hénon maps and their (universal) approximation properties.
Decision: Accept (Poster)