Learning and Generalization in Univariate Overparameterized Normalizing Flows

1 INTRODUCTION . Neural network models trained using simple first-order iterative algorithms have been very effective in both supervised and unsupervised learning . Theoretical reasoning of this phenomenon requires one to consider simple but quintessential formulations , where this can be demonstrated by mathematical proof , along with experimental evidence for the underlying intuition . First , the minimization of training loss is typically a non-smooth and non-convex optimization over the parameters of neural networks , so it is surprising that neural networks can be trained efficiently by first-order iterative algorithms . Second , even large neural networks whose number parameters are more than the size of training data often generalize well with a small loss on the unseen test data , instead of overfitting the seen training data . Recent work in supervised learning attempts to provide theoretical justification for why overparameterized neural networks can train and generalize efficiently in the above sense . In supervised learning , the empirical risk minimization with quadratic loss is a non-convex optimization problem even for a fully connected neural network with one hidden layer of neurons with ReLU activations . Around 2018 , it was realized that when the hidden layer size is large compared to the dataset size or compared to some measure of complexity of the data , one can provably show efficient training and generalization for these networks , e.g . Jacot et al . ( 2018 ) ; Li & Liang ( 2018 ) ; Du et al . ( 2018 ) ; Allen-Zhu et al . ( 2019 ) ; Arora et al . ( 2019 ) . Of these , Allen-Zhu et al . ( 2019 ) is directly relevant to our paper and will be discussed later . The role of overparameterization , and provable training and generalization guarantees for neural networks are less well understood in unsupervised learning . Generative models or learning a data distribution from given samples is an important problem in unsupervised learning . Popular generative models based on neural networks include Generative Adversarial Networks ( GANs ) ( e.g. , Goodfellow et al . ( 2014 ) ) , Variational AutoEncoders ( VAEs ) ( e.g. , Kingma & Welling ( 2014 ) ) , and Normalizing Flows ( e.g. , Rezende & Mohamed ( 2015 ) ) . GANs and VAEs have shown impressive capability to generate samples of photo-realistic images but they can not give probability density estimates for new data points . Training of GANs and VAEs has various additional challenges such as mode collapse , posterior collapse , vanishing gradients , training instability , etc . as shown in e.g . Bowman et al . ( 2016 ) ; Salimans et al . ( 2016 ) ; Arora et al . ( 2018 ) ; Lucic et al . ( 2018 ) . In contrast to the generative models such as GANs and VAEs , when normalizing flows learn distributions , they can do both sampling and density estimation , leading to wide-ranging applications as mentioned in the surveys by Kobyzev et al . ( 2020 ) and Papamakarios et al . ( 2019 ) . Theoretical understanding of learning and generalization in normalizing flows ( more generally , generative models and unsupervised learning ) is a natural and important open question , and our main technical contribution is to extend known techniques from supervised learning to make progress towards answering this question . In this paper , we study learning and generalization in the case of univariate overparameterized normalizing flows . Restriction to the univariate case is technically non-trivial and interesting in its own right : univariate ReLU networks have been studied in recent supervised learning literature ( e.g. , Savarese et al . ( 2019 ) , Williams et al . ( 2019 ) , Sahs et al . ( 2020 ) and Daubechies et al . ( 2019 ) ) . Multidimensional flows are qualitatively more complex and our 1D analysis sheds some light on them ( see Sec . 4 ) . Before stating our contributions , we briefly introduce normalizing flows ; details appear in Section 2 . Normalizing Flows . We work with one-dimensional probability distributions with continuous density . The general idea behind normalizing flows ( NFs ) , restricted to 1D can be summarized as follows : Let X ∈ R be a random variable denoting the data distribution . We also fix a base distribution with associated random variable Z which is typically standard Gaussian , though in this paper we will work with the exponential distribution as well . Given i.i.d . samples of X , the goal is to learn a continuous strictly monotone increasing map fX : R→ R that transports the distribution of X to the distribution of Z : in other words , the distribution of f−1X ( Z ) is that of X . The learning of fX is done by representing it by a neural network and setting up an appropriate loss function . The monotonicity requirement on f which makes f invertible , while not essential , greatly simplifies the problem and is present in all the works we are aware of . It is not clear how to set up a tractable optimization problem without this requirement . Since the function represented by standard neural networks are not necessarily monotone , the design of the neural net is altered to make it monotone . For our 1D situation , one-hidden layer networks are of the form N ( x ) = ∑m i=1 aiσ ( wix+ bi ) , where m is the size of the hidden layer and the ai , wi , bi are the parameters of the network . We will assume that the activation functions used are monotone . Here we distinguish between two such alterations : ( 1 ) Changing the parametrization of the neural network . This can be done in multiple ways : instead of ai , wi we use a2i , w 2 i ( or other functions , such as the exponential function , of ai , wi that take on only positive values ) ( Huang et al. , 2018 ; Cao et al. , 2019 ) . This approach appears to be the most popular . In this paper , we also suggest another related alteration : we simply restrict the parameters ai , wi to be positive . This is achieved by enforcing this constraint during training . ( 2 ) Instead of using N ( x ) for f ( x ) we use φ ( N ( x ) ) for f ′ ( x ) = dfdx , where φ : R→ R + takes on only positive values . Positivity of f ′ implies monotonicity of f . Note that no restrictions on the parameters are required ; however , because we parametrize f ′ , the function f needs to be reconstructed using numerical quadrature . This approach is used by Wehenkel & Louppe ( 2019 ) . We will refer to the models in the first class as constrained normalizing flows ( CNFs ) and those in the second class as unconstrained normalizing flows ( UNFs ) . Our Contributions . In this paper , we study both constrained and unconstrained univariate NFs theoretically as well as empirically . The existing analyses for overparametrized neural networks in the supervised setting work with a linear approximation of the neural network , termed pseudo network in Allen-Zhu et al . ( 2019 ) . They show that ( 1 ) there is a pseudo network with weights close to the initial ones approximating the target function , ( 2 ) the loss surfaces of the neural network and the pseudo network are close and moreover the latter is convex for convex loss functions . This allows for proof of the convergence of the training of neural network to global optima . One can try to adapt the approach of using a linear approximation of the neural network to analyze training of NFs . However , one immediately encounters some new roadblocks : the loss surface of the pseudo networks is non-convex in both CNFs and UNFs . In both cases , we identify novel variations that make the optimization problem for associated pseudo network convex : For CNFs , instead of using a2i , w 2 i as parameters , we simply impose the constraints ai ≥ and wi ≥ for some small constant . The optimization algorithm now is projected SGD , which in this case incurs essentially no extra cost over SGD due to the simplicity of the positivity constraints . Apart from making the optimization problem convex , in experiments this variation slightly improves the training of NFs compared to the reparametrization approaches , and may be useful in practical settings . Similarly , for UNFs we identify two changes from the model of Wehenkel & Louppe ( 2019 ) that make the associated optimization problem convex , while still retaining empirical effectiveness : ( 1 ) Instead of Clenshaw–Curtis quadrature employed in Wehenkel & Louppe ( 2019 ) which uses positive and negative coefficients , we use the simple rectangle quadrature which uses only positive coefficients . This change makes the model somewhat slow ( it uses twice as many samples and time to get similar performance on the examples we tried ) . ( 2 ) Instead of the standard Gaussian distribution as the base distribution , we use the exponential distribution . In experiments , this does not cause much change . Our results point to a dichotomy between these two classes of NFs : our variant of UNFs can be theoretically analyzed when the networks are overparametrized to prove that the UNF indeed learns the data distribution . To our knowledge , this is the first “ end-to-end ” analysis of an NF model , and a neural generative model using gradient-based algorithms used in practice . This proof , while following the high-level scheme of Allen-Zhu et al . ( 2019 ) proof , has a number of differences , conceptual as well as technical , due to different settings . E.g. , our loss function involves a function and its integral estimated by quadrature . On the other hand , for CNFs , our empirical and theoretical findings provide evidence that overparametrization makes training slower to the extent that models of similar size which learn the data distribution well for UNFs , fail to do so for CNFs . We also analyze CNFs theoretically in the overparametrized setting and point to potential sources of the difficulty . The case of moderatesized networks , where training and generalization do take place empirically , is likely to be difficult to analyze theoretically as presently this setting is open for the simpler supervised learning case . We hope that our results will pave the way for further progress . We make some remarks on the multidimensional case in Sec . 4 . In summary , our contributions include : • To our knowledge , first efficient training and generalization proof for NFs ( in 1D ) . • Identification of architectural variants of UNFs that admit analysis via overparametrization . • Identification of “ barriers ” to the analysis of CNFs . Related Work . Previous work on normalizing flows has studied different variants such as planar and radial flows in Rezende & Mohamed ( 2015 ) , Sylvester flow in van den Berg et al . ( 2018 ) , Householder flow in Tomczak & Welling ( 2016 ) , masked autoregressive flow in Papamakarios et al . ( 2017 ) . Most variants of normalizing flows are specific to certain applications , and the expressive power ( i.e. , which base and data distributions they can map between ) and complexity of normalizing flow models have been studied recently , e.g . Kong & Chaudhuri ( 2020 ) and Teshima et al . ( 2020 ) . Invertible transformations defined by monotonic neural networks can be combined into autoregressive flows that are universal density approximators of continuous probability distributions ; see Masked Autoregressive Flows ( MAF ) Papamakarios et al . ( 2017 ) , UNMM-MAF by Wehenkel & Louppe ( 2019 ) , Neural Autoregressive Flows ( NAF ) by Huang et al . ( 2018 ) , Block Neural Autoregressive Flow ( B-NAF ) by Cao et al . ( 2019 ) . Unconstrained Monotonic Neural Network ( UMNN ) models proposed by Wehenkel & Louppe ( 2019 ) are particularly relevant to the technical part of our paper . Lei et al . ( 2020 ) show that when the generator is a two-layer tanh , sigmoid or leaky ReLU network , Wasserstein GAN trained with stochastic gradient descent-ascent converges to a global solution with polynomial time and sample complexity . Using the moments method and a learning algorithm motivated by tensor decomposition , Li & Dou ( 2020 ) show that GANs can efficiently learn a large class of distributions including those generated by two-layer networks . Nguyen et al . ( 2019b ) show that two-layer autoencoders with ReLU or threshold activations can be trained with normalized gradient descent over the reconstruction loss to provably learn the parameters of any generative bilinear model ( e.g. , mixture of Gaussians , sparse coding model ) . Nguyen et al . ( 2019a ) extend the work of Du et al . ( 2018 ) on supervised learning mentioned earlier to study weakly-trained ( i.e. , only encoder is trained ) and jointly-trained ( i.e. , both encoder and decoder are trained ) two-layer autoencoders , and show joint training requires less overparameterization and converges to a global optimum . The effect of overparameterization in unsupervised learning has also been of recent interest . Buhai et al . ( 2020 ) do an empirical study to show that across a variety of latent variable models and training algorithms , overparameterization can significantly increase the number of recovered ground truth latent variables . Radhakrishnan et al . ( 2020 ) show that overparameterized autoencoders and sequence encoders essentially implement associative memory by storing training samples as attractors in a dynamical system . Outline . A brief outline of our paper is as follows . Section 2 contains preliminaries and an overview of our results about constrained and unconstrained normalizing flows . Appendix B shows the existence of a pseudo network whose loss closely approximates the loss of the target function . Appendix C shows the coupling or closeness of their gradients over random initialization . Appendices D and E contain complete proofs of our optimization and generalization results , respectively . Section 3 and Appendix G contain our empirical studies towards validating our theoretical results .

Summary: The paper studies the role of overparameterization in learning normalizing flow models. More specifically, the authors analyze the optimization and generalization of such a model when the transport map f is parameterized by a two-layer neural network with potentially many hidden units (or highly over-parameterized). Importantly, the focus is on univariate data distributions.