Universal Approximation Using Well-Conditioned Normalizing Flows

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.


Reviews and author responses:

 Thank you for your reply. 

Although the scope of the theoretical result is still limited to log-concave distributions, the problem the authors seek to tackle seems innately non-trivial. The authors have addressed most of my concerns and answered my questions. I believe the general direction and theoretical techniques explored in this paper are beneficial to the research in better understanding the expressive power of invertible architectures, coupling blocks in particular. I have increased my rating to 6 accordingly. 

This paper studies the approximation capability of invertible coupling block. Specifically, they show that log-concave distributions can be approximated by these invertible blocks with good conditioning number (i.e. these transformations have a non-singular Jacobian) with Gaussian padding.    ### Significance and scope

The paper studies the most widely used invertible architecture and its expressivity. However, the class of distributions that this architecture can provably approximate is limited (i.e. log-concave densities). In fact, log-concave distributions are a very small family of distributions; they do not include, for example, multimodal distributions. 

How hard is it to extend the current result to a wider class of distributions? Also, the title of the paper says "Universal Approximation", although traditionally universal approximation refers to the capability to approximate an acceptably big enough class of functions or distributions (such as all C1, L1 functions, or probability measures with smooth enough densities, etc). I'd suggest reconsidering the title of the paper. 

### Relation to prior work

The paper has a fairly thorough review and discussion to prior work. Some distinctions can be made to improve the paper. For example, I would point out the subtlety that in [1], the theorem is indeed for zero padding, but in both [1,2] Gaussian padding is used in practice. [3] proved that "neural ODE" with zero padding can approximate arbitrary homeomorphism, which is stronger than distributional universality. Their result is not directly related to affine coupling; extra steps, such as numerical integration, are needed, so I'm not sure about the discussion between L43-L45. In addition, the singularity limitation of [1] and the possibility to improve the result by choosing a different augmented ODE is pointed out in Section C of [1] (the *Approximate Hamiltonian Flows* paragraph), which I think should be acknowledged. 

I don't see where the relevancy of [4] is discussed, although it's cited. 

[1] Augmented Normalizing Flows: Bridging the Gap Between Generative Flows and Latent Variable Models

[2] VFlow: More Expressive Generative Flows with Variational Data Augmentation

[3] Approximation Capabilities of Neural ODEs and Invertible Residual Networks

[4] Coupling-based invertible neural networks are universal diffeomorphism approximators

### Clarity

I could go over the main paper linearly and roughly see how to put together a proof based on the lemmas. But the organization of the proof steps needs more work, and I would suggest having a bit more hand holding. Concretely, if I understand correctly (correct me if I am wrong), the reason why lemma 3 is needed is because the time-derivative of $x_t$ (resp. $v_t$) of equation (3) depends on $x_t$ itself (resp. $v_t$) in a non-linear way, and we would need an linear relationship so that we can approximate it using an affine coupling block, e.g. equation (7). 

### More detailed suggestions and questions

- Saying "augmented" or "Gaussian-padded" affine-coupling flows can approximate ... e.g. in L104
- Is it possible to directly replace polynomials with a neural network in the theorem? Like would extra care need to be taken care of because neural nets can only approximate a function uniformly well within a compact support?
- In Fact 1, is $\nabla^2$ missing?
- Is Fact 1 the only place where log-concavity is used in the paper, so that a log-Sobolev inequality and convergence of the SDE would be guaranteed?
- Is the derivative in the line above L244 only wrt $x_0$? is $v_0$ involved?
- I don't think (I might be wrong though) a direct Euler discretization of the augmented ODE (i.e. eq (12)) can be expressed as affine-coupling type of invertible map, as claimed in L306-307. Wouldn't mid-point method be needed? The theory developed in this paper depends on quite strict conditions (e.g. log-concavity of the target density), and it's unclear how they can be relaxed without reading the paper in much detail. This limits the scope of the theory; therefore I give a borderline-reject rating. 

 We thank the reviewer for the feedback and kind words! We are happy you found the paper “original, high quality, clearly written, and significant”! 

We respond to the questions/suggestions in the review: 

**Re: “related work section”**

Thank you for the feedback! Space permitting, we will be sure to include a more comprehensive discussion on the existing works on universal approximation (not just for affine-coupling flows, but other models like neural ODEs etc as mentioned) in our related works section in the revised version. Please also see the common response above for more details about how this method extends to other flow models.

**Re: “extend the theory from log-concave”, “approximating polynomials with neural networks”**

Please see the “Where log-concavity is needed in the proof” and “Neural net approximations to polynomial non-linearity” point in the common response.

**Re: “Provide a bound on size of flow model in $O()$”**

A simple bound on the depth of the constructed affine-coupling flow model is $\text{polylog}(\epsilon)/ \epsilon^{O(\kappa)}$. We will include the computation in the revised version.

**Re: “Illustrations”**

Thank you for the suggestions. Space permitting, we will include a table with the conditioning of the Jacobian for existing universal approximation constructions in the appendix in the revised version. Regarding the experiment about Gaussian padding: please see the point “Empirical benefit of Gaussian padding” in the common response above.


 We thank the reviewers for their detailed reading of the paper and their various insightful comments! 

We are encouraged that the reviewers appreciated the originality, significance, and quality of writing in the paper. We believe the main concerns that arose in reviews can be resolved in this rebuttal period and addressed through improvements to the exposition. 

We would like to address a few of the concerns that were common across the reviews here. We have also answered specific reviewer questions in individual replies.

**Where log-concavity is needed in the proof:**

Log-concavity of the target distribution is necessary to ensure that the Jacobian of the transport map from latent to target distribution is well-conditioned. The overall proof strategy of leveraging a connection between symplectic diffeomorphisms and Langevin dynamics seems more general, and it would be interesting to see if our techniques can be extended beyond the log-concave setting.

To be more technically precise, by eq (21), bounding the condition number reduces to bounding $\nabla^2 (\ln p_t - \ln p)$, where $p_t$ is the probability distribution after flowing underdamped Langevin started at $p_0$ (the distribution we wish to approximate) with stationary distribution $p=\mathcal{N}(0,I)$ (standard Gaussian). For this, we require a handle on the Hessian of the log-pdf under convolution (Lemma 7), which we only know to hold for log-concave distributions. We could potentially extend beyond log-concavity, e.g. if we knew that $p_t$ was log-concave after some $t_0$, possibly incurring some exponential dependence in $t_0$.

**Extending to other forms of normalizing flows:**

We only focused on affine couplings in the paper, however our proof implies well-conditioned universal approximation for other related normalizing flow models too. Namely, our Lemma 1 proves that the flow map of underdamped Langevin dynamics is well-conditioned for all $t \in [0,T]$. However, it’s clear by [1] that underdamped Langevin dynamics is a continuous normalizing flow --- so the claim follows. Similarly, the particular affine coupling layers we construct (eq (12)) also form a residual block, so the claim also follows for residual flows [2]. 

We will clarify this in the camera ready version of the paper. 

[1] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. Advances in Neural Information Processing Systems, 2018.

[2] Behrmann, J., Grathwohl, W., Chen, R. T., Duvenaud, D., & Jacobsen, J. H. Invertible residual networks. International Conference on Machine Learning, 2019. 

**Neural net approximations to polynomial non-linearity:**

We assume polynomial nonlinearities for the sake of mathematical convenience. Since our main Theorem 1 regards the restriction to a compact set (arbitrarily large, of mass 1-epsilon), neural net approximation theorems can be applied directly to obtain the statement for neural nets with smooth sigmoidal activation functions (we can choose the neural network to be large enough so as to obtain the final required approximation error). 

**Empirical benefit of Gaussian padding:**

The benefit of Gaussian padding in likelihood-maximization-based training for affine-coupling normalizing flows has already been empirically demonstrated in the work by Koehler et al (2020). In fact, this was one of our motivations to theoretically study the efficacy of Gaussian padding.

 We thank the reviewer for the detailed feedback! We respond to the questions/suggestions in the review:

**Re: “Log-concave distributions are a small family of distributions”, “How hard is it to extend the current result to a wider class of distributions?” and “reconsidering the title of the paper”**

Please see the “Where log-concavity is needed in the proof” point in the common response above for some comments on extending the results to wider classes of distributions. 

We take your point about the title and are very happy to rename our paper in the camera ready version to “Universal Approximation of Log-Concave Distributions Using Well-Conditioned Normalizing Flows”.

**Re: “Relation to prior works”**

Thank you for your suggestions for improving the related works section. We will be sure to incorporate the specific connections and subtleties you pointed out (especially pertaining to paper [1]) in our revised version. We will also discuss the relevance with [4] in slightly more detail. In general, we were limited by the space restriction, and space-permitting, we will expand the related works section.

**Re: “More hand-holding in proof steps” “Question about why lemma 3 is needed”**

You are correct, eqn (10) in Lemma 3 is the augmented Hamiltonian system (which represents the dynamics of underdamped Langevin), but in order to express the transport map of this ODE system as an affine coupling block, we need the ODE system to be of the form in eqn (11). Lemma 3 essentially shows that the transport maps given by both (10) (which is underdamped Langevin) and (11) (which can be discretized to be written as an affine coupling) are close.
Due to space limitations, we tried to be as terse as possible in explaining the proof steps, but we will try to improve the writing so as to be more expository in the camera ready version.

**Re: “Is it possible to directly replace polynomials with a neural network in the theorem? Like would extra care need to be taken care of because neural nets can only approximate a function uniformly well within a compact support?**

Please see the “Neural net approximations to polynomial non-linearity” point in the common response.

**Re: “$\nabla^2$ missing”**

Yes, that is a typo. Thank you for pointing this out!

**Re: “Is Fact 1 the only place where log-concavity is used in the paper, so that a log-Sobolev inequality and convergence of the SDE would be guaranteed?”**

Please see the “Where log-concavity is needed in the proof” point in the common response above. In brief, we also require the assumption of log-concavity to ensure that the Jacobian of the constructed affine-coupling flow is well-conditioned

**Re: “Is the derivative in the line above L244 only wrt x_0? Is v_0 involved?”**

This should be a derivative with respect to both $x_0$ and $v_0$ (so the argument should be $(x_0,v_0)$). We will fix the notation in the revised version.

**Re: “I don't think (I might be wrong though) a direct Euler discretization of the augmented ODE (i.e. eq (12)) can be expressed as affine-coupling type of invertible map, as claimed in L306-307. Wouldn't mid-point method be needed?”**

There is a typo here indeed: the $v_n$ in the definition for $x_{n+1}$ should be $v_{n+1}$ in order to agree with the composition. (See Appendix D.2 for details.) This “alternating Euler” discretization is sufficient to approximate the continuous trajectory within $O(\eta)$ in $C^1$ topology. 


 We thank the reviewer for the feedback and kind words! We respond to the comments in the review:

**Re: “What is B_t?”**

Yes, it is Brownian motion. We will clarify this in the revised version of the paper.

**Re: “liked to see the use of the proposed padded affine-coupling flow in a practical example”**

Please see the “Empirical benefit of Gaussian padding” point in the common response.

**Re: “Perhaps there is a \nable^2 missing”**

You are right, that is a typo, and it should indeed be $\nabla^2 U \ge \lambda I$.

**Re: “Limitations not discussed”**

We will try to make the discussion in our conclusion section more prominent for the camera ready version: we point out there that one clear limitation of our work is that it is restricted to the class of log-concave distributions, and that it would be interesting to try and extend our techniques to other distributions. 


The authors provide a proof for the universal approximation capability of well-conditioned affine coupling network. Their proof requires Gaussian padding, and uses techniques from dynamical systems and stochastic differential equations.  The authors show that well-conditioned affine-coupling normalizing flows have the universal approximation capability of log-concave densities, in Wassertein-1 distance. Their proof requires the target density to be padded by Gaussians, and the latent density to be a Gaussian. Padding by Gaussians avoids having to pad with 0, which have been done in previous works, but leads to ill-conditioned Jacobians, and thus can't be used in practice.
Moreover, their bound on the condition number does not get worse when the precision $\varepsilon \rightarrow 0$.

The contribution of this paper is purely theoretical. The proof technique lies outside of my expertise, and as such I can't comfortably assert that it is correct, nor provide any meaningful feedback for the proof. I didn't find any obvious mistakes, but I will leave a more detailed commentary to the more experienced reviewers.
This being said, I appreciate the novelty and advantages provided by this proof when compared to prior work, and think it is valuable.

Few comments:
- What is $B_t$ in eq (1)? A brownian motion?
- I would have liked to see the use of the proposed padded affine-coupling flow in a practical example. The main motivation for this new theory is that the proposed flow is well-conditioned, and can thus be used in maximum likelihood estimation without a problem.
- In fact 1, you require $U(x) \succeq \lambda I$, but $U$ is a potential? Perhaps there is a $\nabla^2$ missing?


 The contribution is purely theoretical (so no direct societal impact), but the limitations were not clearly discussed. 

The paper presents a universal approximation theorem for affine coupling flows that do not produce ill-conditioned Jacobians. The universal approximation theorem is in the sense of the Wasserstein distance and is for well-conditioned log-concave distributions. The proof is based on the underdamped Langevin dynamics that push the target to the source, and a construction of a sequence of affine coupling layers that approximate the ODE.
  The paper is original, high quality, clearly written, and significant. Normalizing flows are flexible and can offer exact likelihood computation, but these models sacrifice some expressiveness. Therefore, it is very meaningful to prove universal approximation for these models, while there has not been much analysis in the literature. The theoretical background and explanation of the proof are clearly introduced in the main text. 

I find the theoretical framework to be elegant. First, using the fact that the inverse of an affine coupling flow is also an affine coupling flow, the paper switches the source and the target, and so the Langevin dynamics can be naturally used. Then, the paper attempts to approximate the dynamics (in a small step) with affine coupling blocks, which is also natural and simplifies the problem. Combining these analysis, the universal approximation result is obtained.

I do have a few questions and comments for the authors.
- In the related work section, it would be helpful for readers to have a sense on which has been done in this are if the authors can have a more detailed discussion on universal approximation results for different kinds of normalizing flows. For instance, Teshima et al (2020) appears in your references but does not appear in the main text. Some other universal approximation or convergence theory for autoregressive flows, residual flows, and neural ODEs can also be briefly discussed. 
- The class of log-concave distributions are broadly studied in statistics; but, it is questionable how well these distributions can model more complex data, such as real world datasets. One limitation is when there are multiple modes in the target distribution, e.g., mixture of Gaussians as a simple example. Can the authors address how you expect to extend the theory to these scenarios?
- Regarding remark 2 (line 119 ~ line 121), do you mean that you need to approximate these polynomials $s$ and $t$ with neural networks if you want to prove universal approximation of neural network based affine coupling flows? If that is the case, does the latter part of approximation affect the rest part of the proof?
- Regarding remark 5, it would be easier for understanding if you have an explicit bound on depth, even in a very general big O form.
- Although it is acceptable that a theory paper has no experiments, it would be easier for readers to understand your paper if there are some more illustrations: 
    - a table or figure illustrating the ill conditions of Jacobians mentioned in line 47 and 117;
    - some numerical illustration on the Langevin dynamics and your affine coupling approximation to it;
    - and empirical experiments for the Gaussian padding in practice. One limitation, as addressed in the main review, is that the theory is only for log-concave distributions, which have limited capacity. 

Meta-review: The paper introduces a universal approximation proof of normalizing for log-concave distribution when the Jacobians are not ill-conditioned. 
While the scope of the paper remains limited (log-concave distribution), the submission gives a clear exposition of their techniques and could provide a strong foundation for future theoretical studies of flow-based models.
I'm recommending the submission for acceptance as a poster.

Decision: accept