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: This paper proved that for a certain modified version of sufficiently-overparametrized univariate normalizing flows where the underlying neural network has only one hidden layer, with high probability it can learn a distribution that is close enough to the target distribution where the distance can be measured in, e.g., KL divergence. The width of the network, number of samples, and the number of quadrature points are required to be at least polynomial in inverse of error rate and complexity measure of the target distribution. The authors also provided theoretical evidence and did experiments on synthetic Gaussian mixture datasets to show that another variation of the normalizing flow model does not benefit from overparametrization under this one-hidden-layer univariate setting.