Learning Gaussian Mixtures with Generalized Linear Models: Precise Asymptotics in High-dimensions

1 Introduction . A recurring observation in modern deep learning practice is that neural networks often defy the standard wisdom of classical statistical theory . For instance , deep neural networks typically achieve good generalisation performances at a regime in which it interpolates the data , a fact at odds with the intuitive bias-variance trade-off picture stemming from classical theory [ 1–3 ] . Surprisingly , many of the “ exotic ” behaviours encountered in deep neural networks have recently been shown to be shared by models as simple as overparametrised linear classifiers [ 4 , 5 ] , e.g. , the aforementioned benign over-fitting [ 6 ] . Therefore , understanding the generalisation properties of simple models in high-dimensions has proven to be a fertile ground for elucidating some of the challenging statistical questions posed by modern machine learning practice [ 7–16 ] . In this manuscript , we pursue this enterprise in the context of a commonly used model for highdimensional classification problems : the Gaussian mixture . Indeed , it has been recently argued that the features learned by deep neural networks trained on the cross-entropy loss “ collapse ” in a mixture of well-separated clusters , with the last layer acting as a simple linear classifier [ 17 ] . Another observation put forward in [ 18 ] is that data obtained using generative adversarial networks behave as Gaussian mixtures . Here , we derive an exact asymptotic formula characterising the performance 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . of generalised linear classifiers trained on K Gaussian clusters with generic covariances and means . Our formula is valid for any convex loss and penalty , encompassing popular tasks in the machine learning literature such as ridge regression , basis pursuit , cross-entropy minimisation and max-margin estimation . This allow us to answer relevant questions for statistical learning , such as : what is the separability threshold for K-clustered data ? How does regularisation affects estimation ? Can different penalties help when the means are sparse ? We also extend the observation of [ 18 ] showing that the learning curves of binary classification tasks on real data are indeed well captured by our asymptotic analysis . Model definition — We consider learning from a d-dimensional mixture of K Gaussian clusters Ck∈ [ K ] . The data set is obtained by sampling n pairs ( xν , yν ) ν∈ [ n ] ∈ Rd+K identically and independently . We adopt the one-hot encoded representation of the labels , i.e. , if xν ∈ Ck , then yν = ek , kth basis vector of RK . We will denote the matrix of concatenated samplesX ∈ Rd×n . The mixture density then reads : P ( x , y ) = K∑ k=1 ykρkN ( x |µk , Σk ) , ( 1 ) where N ( x|µ , Σ ) is the multivariate normal distribution with mean µ and covariance matrix Σ . The matrix of concatenated means is denoted M ∈ Rd×K . In Eq . ( 1 ) , ∀k , ρk = P ( y = ek ) ≥ 0 , µk ∈ Rd and Σk ∈ Rd×d is positive-definite . We will consider the estimator obtained by minimising the following empirical risk : R ( W , b ) ≡ n∑ ν=1 ` ( yν , Wxν√ d + b ) + λr ( W ) , ( 2 ) ( W ? , b ? ) ≡ argmin W∈RK×d , b∈RK R ( W , b ) , ( 3 ) whereW ∈ RK×d and b ∈ RK are the weights and bias to be learned , ` is a convex loss function , and r is a regularisation function whose strength is tuned by the parameter λ ≥ 0 . For example the loss function ` can represent the composition of a cross-entropy loss with a softmax thresholding on the linear part of Eq . ( 2 ) . We will characterise the distribution of the estimator ( W ? , b ? ) , and we will evaluate the average training loss defined as ` = 1 n n∑ ν=1 ` ( yν , W ? xν√ d + b ? ) , ( 4 ) as well as the average training error t and generalisation error g , defined as the misclassification rates : t = 1 n n∑ ν=1 I [ yν 6= ŷ ( W ? xν√ d + b ? ) ] , g = E ( xnew , ynew ) [ I [ ynew 6= ŷ ( W ? xnew√ d + b ? ) ] ] , where ( xnew , ynew ) is a new unseen data point sampled from the distribution in Eq . ( 1 ) . In the previous equations , we have used the function ŷ : RK → RK , so that ŷk ( x ) : = I ( maxκ xκ = xk ) . The main contributions in this manuscript are the following : ( C1 ) In Sec . 2 and Appendix A we prove closed-form equations characterizing the asymptotic distribution of the matrix of weightsW ? ∈ RK×d , enabling the exact computation of key quantities such as the training and generalisation error . Our proof method solves shortcomings of previous approaches by introducing a novel approximate message-passing sequence , building on recent advances in this framework , that is of independent interest . ( C2 ) In Sec . 3.1 we study the problem of classifying an anisotropic mixture with sparse means , where the strong or weak directions in the data are correlated with the non-zero components of the mean as in [ 19 ] . We study how learning the sparsity with an ` 1 penalty improves the classification performance . ( C3 ) In Sec . 3.2 we study the performance of the cross-entropy estimator in the limit of vanishing regularisation λ → 0+ for K Gaussian clusters as a function of the sample complexity α = n/d ; we show that a phase transition takes place at a certain value α ? K between a regime of complete separability of the data and a regime in which the correct classification of almost all points in the data set is not possible . We also investigate the effect of λ > 0 regularisation on the generalisation error , comparing the K > 2 case with the results given in the literature for K = 2 [ 14 , 20 ] . ( C4 ) In Sec . 3.3 we investigate the applicability of our formula beyond the Gaussian assumption by applying it to classification tasks on real data . We show that for different tasks and losses , it closely captures the real learning curves , even when data is mapped through a non-linear feature map . This further shows that Gaussian mixtures are a good surrogate model for investigating real classification tasks , as put forward in [ 18 ] . Relation to previous work — The analysis of Gaussian mixture models in the high-dimensional regime has been the subject of many recent works . Exact asymptotics has been derived for the binary classification case with diagonal covariances in [ 21–23 ] for the logistic loss and in [ 24 , 25 ] for the square loss , both with ` 2 penalty . A similar analysis has been performed in [ 26 ] for the hard-margin SVM . These works were generalised to generic convex losses and ` 2 penalty in [ 14 ] , where it has been also shown that the regularisation term can play an important role in reaching Bayes-optimal performances . Hinge regression with ` 1 penalty and diagonal covariance was treated in [ 13 ] . Recently , these asymptotic results were generalised to the case in which both clusters share the same covariance in [ 27 ] , and finite rate bounds were given in [ 28 , 29 ] in the case of sub-Gaussian mixtures . Asymptotic results for the multiclass problem with diagonal covariance were derived in [ 20 ] for the restricted case of the square loss with ` 2 penalty . Our result unifies all the aforementioned asymptotic formulas , and extends them to the general case of a multiclass problem with generic covariances and arbitrary convex losses and penalties . From a technical standpoint , in [ 13 , 14 , 20 , 21 , 25 , 27 , 30 ] the authors use convex Gaussian comparison inequalities , see e.g . [ 31 , 32 ] , to prove their result . In particular , the proof given in [ 20 ] for the multiclass problem harnesses the geometry of least-squares , and it is then stressed that this method breaks down for multiclass problems in which the risk does not factorise over the K clusters ( as for the cross-entropy , for example ) . We solve this problem using an innovative proof technique which has an interest in its own . Our approach is to capture the effect of non-linearity and generic covariances via the rigorous study of an approximate message-passing ( AMP ) sequence , a family of iterations that admit closed-form asymptotics at each step called state evolution equations [ 33 ] . Our proof relies on several refinements of AMP methods to handle the full complexity of the problem , notably spatial coupling with matrix valued variables [ 34–36 ] and non-separable update functions [ 37 ] , via a multi-layer approach to AMP [ 38 ] . The sparse Gaussian mixture model analysed in Section 3.1 is closely related to the rare/weak features model introduced in [ 19 ] and widely studied in the context of sparse linear discriminant analysis [ 39–42 ] . It was recently revisited in [ 28 , 29 ] in the context of ERM with max-margin classifiers . Here , we consider a correlated variation of the model and study the benefit of using a sparsity inducing ` 1 penalty . The separability transition is a classical topic [ 43 , 44 ] that has recently witnessed a renewal of interest thanks to its connection to overparametrization . It was studied in [ 16 ] in the context of uncorrelated Gaussian data , in [ 8 ] in the random features model and in [ 14 , 21 ] for binary Gaussian mixtures . Recently , [ 12 , 45 , 46 ] showed that the performance of different regression tasks on real data are well-captured by a teacher-student Gaussian model in high-dimensions for ridge regression , but this turned not to be true for non-linear problems such as logistic classification [ 12 ] . Authors of [ 18 ] showed instead that data from generative adversarial networks behave like Gaussian mixtures , motivating the modeling of such mixture for real-data in the present paper . 2 Technical results . Our main technical result is an exact asymptotic characterization of the distribution of the estimator W ? . Informally , the estimator W ? and the quantity W ? X/ √ d behave asymptotically as non-linear transforms of multivariate Gaussian distributions . These transforms are directly linked to the proximal operators [ 47 , 48 ] associated to the loss and regularisation functions , summarizing the effect of the cost function landscape on the estimator . The parameters of these Gaussian distributions and proximals can then be computed from the fixed point of a self-contained set of equations . We start by presenting the most generic form of our result in a concentration of measure-like statement in Theorem 1 , and discuss an intuitive interpretation of the different quantities involved . Theorem 2 then states how the training and generalisation errors can be computed . All results presented in the experiments section can be obtained from Theorem 1 . In Corollary 3 we discuss a particular case where explicit simplifications can be obtained . But first , let ’ s summarise the required assumptions for our result to hold . ( A1 ) The functions ` ( as a function of its second argument ) and r are proper , closed , lower semicontinuous convex functions . We assume additionally that either the cost function ` ( y , •X ) + r ( • ) is strictly convex , or that ` ( y , • ) is strictly convex in its second argument and r is the ` 1 norm , and that the optimization problem 2 is feasible . ( A2 ) The covariance matrices are positive definite and their spectral norms are bounded . ( A3 ) The mean vectors µk are distributed according to some density Pµ ( M ) such that the following quantity is finite ∀d E [ ∥∥M > M∥∥ F ] < +∞ , ( 5 ) where ‖ • ‖F denotes the Frobenius norm . ( A4 ) The number of samples n and dimension d both go to infinity with fixed ratio α = n/d , called hereafter the sample complexity . The number of clusters K is finite . ( A5 ) The fixed point of the set of self-consistent equations Eq . ( 8 ) exists and is unique . As specified by assumption ( A1 ) , our proof does not apply to any convex problem . We discuss this assumption further in Appendix A.5 . We also comment on the existence and uniqueness of the solution to the set of self consistent equations Eq . ( 8 ) in Appendix A.6 . Before proceeding further , let us specify a useful notation . Suppose that the matrix G = ( Gki ) ki ∈ RK×d is given , alongside the four-index tensor A = ( Aki k′i′ ) ki k′i′ ∈ RK×d ⊗ RK×d . We will use the notation G A = ∑ kiGkiAki k′i′ ∈ RK×d . Similarly , given a four-index tensor A , we will define √ A as the tensor such that A = √ A √ A . We are now in a position to state our main result . Theorem 1 ( Concentration properties of the estimator ) . Let ξk∈ [ K ] ∼ N ( 0 , IK ) be collection of K-dimensional standard normal vectors independent of other quantities . Let also be { Ξk } a set of K matrices , Ξk ∈ RK×d , with i.i.d . standard normal entries , independent of other quantities . Under the set of assumptions ( A1–A5 ) , for any pseudo-Lispchitz functions of finite order φ1 : RK×d → R , φ2 : RK×n → R , the estimatorW ? and the matrix Z ? = 1√dW ? X verify : φ1 ( W ? ) P−−−−−−→ n , d→+∞ EΞ [ φ1 ( G ) ] , φ2 ( Z ? ) P−−−−−−→ n , d→+∞ Eξ [ φ2 ( H ) ] , ( 6 ) where we have introduced the proximal for the loss : hk = V 1/2 k Prox ` ( ek , V 1/2k • ) ( V −1/2 k ωk ) ∈ R K , ωk ≡mk + b+Q1/2k ξk , ( 7 ) and H ∈ RK×n is obtained by concatenating each hk , ρkn times . We have also introduced the matrix proximalG ∈ RK×d : G = A 1 2 Prox r ( A 1 2 • ) ( A 1 2 B ) , A−1 ≡ ∑ k V̂k⊗Σk , B≡ ∑ k ( µkm̂ > k +Ξk √ Q̂k⊗Σk ) . The collection of parameters ( Qk , mk , Vk , Q̂k , m̂k , V̂k ) k∈ [ K ] is given by the fixed point of the following self-consistent equations : Qk= 1 dEΞ [ GΣkG > ] mk= 1√ d EΞ [ Gµk ] Vk= 1 dEΞ [ ( G ( Q̂k ⊗Σk ) − 12 ( IK ⊗Σk ) ) Ξ > k ] Q̂k= αρkEξ [ fkf > k ] V̂k= −αρkQ − 12 k Eξ [ fkξ > ] m̂k= αρkEξ [ fk ] ( 8 ) where fk ≡ V −1k ( hk − ωk ) , and the vector b ? is such that ∑ k ρkEξ [ Vkfk ] = 0 holds . The purpose of this statement is to have an asymptotically exact description of the distribution of the estimator , where the dimensions going to infinity are effectively summarized as averages over simple , independent distributions . Those distributions are parametrised by the set of finite-size parameters ( Qk , mk , Vk , Q̂k , m̂k , V̂k ) k∈ [ K ] that can be exactly evaluated and have a clear interpretation . Indeed , the parameters ( mk , m̂k ) and ( Qk , Q̂k ) respectively represent means and covariances of multivariate Gaussians ( combined with the original µk , Σk ) , and the ( Vk , V̂k ) parametrise the deformations that should be applied to these Gaussians to obtain the distribution of W ? , Z ? . The distribution is characterized in a weak sense with concentration of pseudo-Lipschitz ( i.e. , sufficiently regular ) functions , whose definition is reminded in the Appendix A . From this result one can work out a number of properties of the weightsW ? , e.g. , training and generalisation error , but also hypothesis tests as done in [ 49 ] for the LASSO . Due to the generality of the statement , no direct simplification is possible . However , we will see that in certain specific cases all quantities can be greatly simplified . This is notably the case for diagonal covariance matrices and separable estimators and observables φ1 , φ2 , where the sums over high-dimensional Gaussians concentrate explicitly to one-dimensional expectations . For instance the results of [ 14 , 20 ] can be recovered as special cases of the present work . Theorem 1 then allows to obtain the asymptotic values of the generalisation error , of the training loss and of the training error . Their explicit expression is given in the following Theorem . Theorem 2 ( generalisation error and training loss ) . In the hypotheses of Theorem 1 , the training loss , the training error and the generalisation error are given by ` = K∑ k=1 ρkEξ [ ` ( ek , hk ) ] , t = 1− K∑ k=1 ρkEξ [ ŷk ( hk ) ] , g = 1− K∑ k=1 ρkEξ [ ŷk ( ωk ) ] . ( 9 ) The case of ridge regularisation and diagonal Σk — The general formulas given above can be remarkably simplified under some assumptions about the choice of the regularisation and about the structure of the covariance matrices Σk . This is the case for instance for the ridge regularisation r ( W ) = ‖W ‖2F/2 and jointly diagonalizable covariances . In this case , Theorem 1 simplifies as follows . Corollary 3 . Under the hypotheses of Theorem 1 , let us further assume that a ridge regularisation is adopted , r ( W ) = ‖W ‖2F/2 , and that the covariance matrices Σk have a common set of orthonormal eigenvectors { vi } di=1 , so that , for each Σk = ∑d i=1 σ k i viv > i . Let us also introduce , in the d→ +∞ limit , the joint distribution for the K-dimensional vectors σ = ( σ1 , . . . , σK ) and µ = ( µ1 , . . . , µK ) , 1 d d∑ i=1 K∏ k=1 δ ( σk − σki ) δ ( µk − √ dµ > k vi ) d→+∞−−−−−→ p ( σ , µ ) , ( 10 ) Then , the first three saddle point equations in eqs . ( 8 ) take the form Qk = Eσ , µ [ σk ( λIK + ∑K κ=1 σ κV̂k ) −2 ( ∑ κκ′ µ κµκ ′ m̂κm̂ > κ′ + ∑K κ=1 σ κQ̂k ) ] , mk = Eσ , µ [ µk ( λIK + ∑K κ=1 σ κV̂k ) −1∑K κ=1 µ κm̂κ ] , Vk = Eσ , µ [ σk ( λIK + ∑K κ=1 σ κV̂k ) −1 ] . ( 11 ) Narrative of the proof — The proof is detailed in Appendix A . It overcomes problems that existing methods , notably convex Gaussian comparison inequalities [ 20 ] , have yet to be adapted to . The first main technical difficulty resides in the estimator of interest being a matrix learned with non-linear functions . This makes it impossible to decompose the problem on each row of the estimator , which must be characterized in a probabilistic sense directly as a matrix . The second main difficulty is brought by the mixture of arbitrary covariances . Intuitively , the covariances correlate the estimator with the individual clusters , and therefore the correlation function can not be represented by a single quantity . In our proof , these points are handled using the AMP and related state-evolution techniques [ 33 , 50–52 ] . The main idea of the proof is to express the estimatorW ? as the limit of a convergent sequence whose structure enables the decomposition of all correlations and distributions in closed form . AMP iterations can handle matrix valued variables [ 36 , 53 ] , correlations in block-structure [ 36 ] , non-separable functions [ 37 , 38 ] and compositions of the previous three , leaving a large choice of possibilities in their design . We thus reformulate the problem in a way that makes the interaction between the estimator and each cluster explicit , effectively introducing a block structure to the problem , and isolate the overlaps with the means { µk } . We then design a matrix-valued sequence that obeys the update rule of an AMP sequence , in order to benefit from its exact asymptotics , and whose fixed point condition matches the optimality condition of the ERM problem , Eq . ( 2 ) . Our proof builds on the spatial coupling framework in the AMP literature [ 36 , 54 ] , which shows that the effect of random matrices defined with non-identically distributed blocks can be embedded in an AMP iteration while explicitly keeping the effect of each block . The non-linearities are then obtained by a block decomposition of the proximal operators defined on sets of matrices , acting on different variables of the AMP sequence and representing the effect of each cluster . The convergence analysis is made possible by the convexity of the problem : the sequence is defined with proximal operators of convex functions which are roughly contractions , and results in converging sequences when combined with the high-dimensional properties of the iteration . It is also interesting to note that the replica method , although heuristic , yet again gives the correct prediction without any hindering from the aforementioned main difficulties , as detailed in Appendix B. Universality — AMP-type proofs are amenable to both finite sample size analysis and universality proofs . For instance , in [ 55 ] it is shown that simpler instances of AMP for the LASSO exhibit exponential concentration in the system size , and the i.i.d . Gaussian assumption can be relaxed to independently sampled sub-Gaussian distributions , as shown in [ 56 , 57 ] . Although these results do not formally encompass our case , their proof method contains most of the required technicalities , and it should be possible to prove similar results in the present setting . Indeed , recent results in [ 18 ] suggest that the formula of Theorem 1 and 2 should be universal for all mixtures of concentrated distribution in high-dimension , not only Gaussian ones . As we discuss Sec . 3.3 , even real data learning curves are empirically found to follow the behavior of the mixture of Gaussians .

Summary: This paper provides a general asymptotic characterization of the solution that comes from learning a mixture of K Gaussians under a supervised learning setting with arbitrary convex loss and regularization. The technique for proving the characterization is novel. Specifically, it constructs an approximate message passing sequence that admits state evolution equations, whose fixed point satisfies the optimality condition of the original problem. Equipped with the theoretical tools, the authors study both synthetic data (sparse Gaussian mixture and general K Gaussian mixtures) and real datasets to further understand (1) $\ell_{1}$ and $\ell_{2}$ regularization for sparse mixture models; (2) separability transition; and (3) quality of Gaussian mixtures as surrogate model. Interestingly, the experiments exhibit a high agreement between the theoretical and numerical values.