Learning curves of generic features maps for realistic datasets with a teacher-student model

1 Introduction . Teacher-student models are a popular framework to study the high-dimensional asymptotic performance of learning problems with synthetic data , and have been the subject of intense investigations spanning three decades [ 1–7 ] . In the wake of understanding the limitations of classical statistical learning approaches [ 8–10 ] , this direction is witnessing a renewal of interest [ 10–15 ] . However , this framework is often assuming the input data to be Gaussian i.i.d. , which is arguably too simplistic to be able to capture properties of realistic data . In this paper , we redeem this line of work by defining a Gaussian covariate model where the teacher and student act on different Gaussian correlated spaces with arbitrary covariance . We derive a rigorous asymptotic solution of this model generalizing the formulas found in the above mentioned classical works . We then put forward a theory , supported by universality arguments and numerical experiments , that this model captures learning curves , i.e . the dependence of the training and test errors on the number of samples , for a generic class of feature maps applied to realistic datasets . These maps can be deterministic , random , or even learnt from the data . This analysis thus gives a unified framework to describe the learning curves of , for example , kernel regression and classification , the analysis of 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . feature maps – random projections [ 16 ] , neural tangent kernels [ 17 ] , scattering transforms [ 18 ] – as well as the analysis of transfer learning performance on data generated by generative adversarial networks [ 19 ] . We also discuss limits of applicability of our results , by showing concrete situations where the learning curves of the Gaussian covariate model differ from the actual ones . Model definition — The Gaussian covariate teacher-student model is defined via two vectors u ∈ Rp and v ∈ Rd , with correlation matrices Ψ ∈ Rp×p , Ω ∈ Rd×d and Φ ∈ Rp×d , from which we draw n independent samples : [ uµ vµ ] ∈ Rp+d ∼ i.i.d . N ( 0 , [ Ψ Φ Φ > Ω ] ) , µ = 1 , · · · , n. ( 1 ) The labels yµ are generated by a teacher function that is only using the vectors uµ : yµ = f0 ( 1 √ p θ > 0 u µ ) , ( 2 ) where f0 : R → R is a function that may include randomness such as , for instance , an additive Gaussian noise , and θ0 ∈ Rp is a vector of teacher-weights with finite norm which can be either random or deterministic . Learning is performed by the student with weights w via empirical risk minimization that has access only to the features vµ : ŵ = arg min w∈Rd [ n∑ µ=1 g ( w > vµ√ d , yµ ) + r ( w ) ] , ( 3 ) where r and g are proper , convex , lower-semicontinuous functions ofw ∈ Rd ( e.g . g can be a logistic or a square loss and r a ` p ( p=1 , 2 ) regularization ) . The key quantities we want to compute in this model are the averaged training and generalisation errors for the estimator w , Etrain . ( w ) ≡ 1 n n∑ µ=1 g ( w > vµ√ d , yµ ) and Egen . ( w ) ≡ E [ ĝ ( f̂ ( v > neww√ d ) , f0 ( u > newθ0√ p ) ) ] . ( 4 ) where g is the loss function in eq . ( 3 ) , f̂ is a prediction function ( e.g . f̂ = sign for a classification task ) , ĝ is a performance measure ( e.g . ĝ ( ŷ , y ) = ( ŷ − y ) 2 for regression or ĝ ( ŷ , y ) = P ( ŷ 6= y ) for classification ) and ( unew , vnew ) is a fresh sample from the joint distribution of u and v. Our two main technical contributions are : ( C1 ) In Theorems 1 & 2 , we give a rigorous closed-form characterisation of the properties of the estimator ŵ for the Gaussian covariate model ( 1 ) , and the corresponding training and generalisation errors in the high-dimensional limit . We prove our result using Gaussian comparison inequalities [ 20 ] ; ( C2 ) We show how the same expression can be obtained using the replica method from statistical physics [ 21 ] . This is of additional interest given the wide range of applications of the replica approach in machine learning and computer science [ 22 ] . In particular , this allows to put on a rigorous basis many results previously derived with the replica method . Towards realistic data — In the second part of our paper , we argue that the above Gaussian covariate model ( 1 ) is generic enough to capture the learning behaviour of a broad range of realistic data . Let { xµ } nµ=1 denote a data set with n independent samples on X ⊂ RD . Based on this input , the features u , v are given by ( potentially ) elaborated transformations of x , i.e . u = ϕt ( x ) ∈ Rp and v = ϕs ( x ) ∈ Rd ( 5 ) for given centred feature maps ϕt : X → Rp and ϕs : X → Rd , see Fig . 1 . Uncentered features can be taken into account by shifting the covariances , but we focus on the centred case to lighten notation . The Gaussian covariate model ( 1 ) is exact in the case where x are Gaussian variables and the feature maps ( ϕs , ϕs ) preserve the Gaussianity , for example linear features . In particular , this is the case for u = v = x , which is the widely-studied vanilla teacher-student model [ 24 ] . The interest of the model ( 1 ) is that it also captures a range of cases in which the feature maps ϕt and ϕs are deterministic , or even learnt from the data . The covariance matrices Ψ , Φ , and Ω then represent different aspects of the data-generative process and learning model . The student ( 3 ) then corresponds to the last layer of the learning model . These observation can be distilled into the following conjecture : Conjecture 1 . ( Gaussian equivalent model ) For a wide class of data distributions { xµ } nµ=1 , and features maps u = ϕt ( x ) , v = ϕs ( x ) , the generalisation and training errors of estimator ( 3 ) are asymptotically captured by the equivalent Gaussian model ( 1 ) , where [ u , v ] are jointly Gaussian variables , and thus by the closed-form expressions of Theorem 1 . The second part of our main contributions are : ( C3 ) In Sec . 3.3 we show that the theoretical predictions from ( C1 ) captures the learning curves in non-trivial cases , e.g . when input data are generated using a trained generative adversarial network , while extracting both the feature maps from a neural network trained on real data . ( C4 ) In Sec . 3.4 , we show empirically that for ridge regression the asymptotic formula of Theorem 1 can be applied directly to real data sets , even though the Gaussian hypothesis is not satisfied . This universality-like property is a consequence of Theorem 3 and is illustrated in Fig . 1 ( right ) where the real learning curve of several features maps learning the odd-versus-even digit task on MNIST is compared to the theoretical prediction . Related work — Rigorous results for teacher-student models : The Gaussian covariate model ( 1 ) contains the vanilla teacher-student model as a special case where one takes u and v identical , with unique covariance matrix Ω . This special case has been extensively studied in the statistical physics community using the heuristic replica method [ 1–3 , 24 , 25 ] . Many recent rigorous results for such models can be rederived as a special case of our formula , e.g . refs . [ 10–15 , 26–29 ] . Numerous of these results are based on the same proof technique as we employed here : the Gordon ’ s Gaussian min-max inequalities [ 20 , 30 , 31 ] . The asymptotic analysis of kernel ridge regression [ 32 ] , of margin-based classification [ 33 ] also follow from our theorem . See also Appendix A.6 for the details on these connections . Other examples include models of the double descent phenomenon [ 34 ] . Closer to our work is the recent work of [ 35 ] on the random feature model . For ridge regression , there are also precise predictions thanks to random matrix theory [ 12 , 36–41 ] . A related set of results was obtained in [ 42 ] for orthogonal random matrix models . The main technical novelty of our proof is the handling of a generic loss and regularisation , not only ridge , representing convex empirical risk minimization , for both classification and regression , with the generic correlation structure of the model ( 1 ) . Gaussian equivalence : A similar Gaussian conjecture has been discussed in a series of recent works , and some authors proved partial results in this direction [ 11 , 12 , 28 , 35 , 43–46 ] . Ref . [ 45 ] analyses a special case of the Gaussian model ( corresponding to ϕt = id here ) , and proves a Gaussian equivalence theorem ( GET ) for feature maps ϕs given by single-layer neural networks with fixed weights . They also show that for Gaussian data x ∼ N ( 0 , ID ) , feature maps of the form v = σ ( Wx ) ( with some technical restriction on the weights ) led to the jointly-Gaussian property for the two scalars ( v ·w , u · θ0 ) for almost any vector w. However , their stringent assumptions on random teacher weights limited the scope of applications to unrealistic label models . A related line of work discussed similar universality through the lens of random matrix theory [ 47–49 ] . In particular , Seddik et al . [ 50 ] showed that , in our notations , vectors [ u , v ] obtained from Gaussian inputs x ∼ N ( 0 , ID ) with Lipschitz feature maps satisfy a concentration property . In this case , again , one can expect the two scalars ( v ·w , u · θ0 ) to be jointly Gaussian with high-probability on w. Remarkably , in the case of random feature maps , [ 46 ] could go beyond this central-limit-like behavior and established the universality of the Gaussian covariate model ( 1 ) for the actual learned weights ŵ . 2 Main technical results . Our main technical result is a closed-form expression for the asymptotic training and generalisation errors ( 4 ) of the Gaussian covariate model introduced above . We start by presenting our result in the most relevant setting for the applications of interest in Section 3 , which is the case of the ` 2 regularization . Next , we briefly present our result in larger generality , which includes non-asymptotic results for non-separable losses and regularizations . We start by defining key quantities that we will use to characterize the estimator ŵ . Let Ω = S > diag ( ωi ) S be the spectral decomposition of Ω . Let : ρ ≡ 1 d θ > 0 Ψθ0 ∈ R , θ̄ ≡ SΦ > θ0√ ρ ∈ Rd ( 6 ) and define the joint empirical density µ̂d between ( ωi , θ̄i ) : µ̂d ( ω , θ̄ ) ≡ 1 d d∑ i=1 δ ( ω − ωi ) δ ( θ̄ − θ̄i ) . ( 7 ) Note that Φ > θ0 is the projection of the teacher weights on the student space , and therefore θ̄ is the rotated projection on the basis of the student covariance , rescaled by the teacher variance . Together with the student eigenvalues ωi , these are relevant statistics of the model , encoded here in the joint distribution µ̂d . Assumptions — Consider the high-dimensional limit in which the number of samples n and the dimensions p , d go to infinity with fixed ratios : α ≡ n d , and γ ≡ p d . ( 8 ) Assume that the covariance matrices Ψ , Ω are positive-definite and that the Schur complement of the block covariance in equation ( 1 ) is positive semi-definite . Additionally , the spectral distributions of the matrices Φ , Ψ and Ω converge to distributions such that the limiting joint distribution µ is well-defined , and their maximum singular values are bounded with high probability as n , p , d→∞ . Finally , regularity assumptions are made on the loss and regularization functions mainly to ensure feasibility of the minimization problem . We assume that the cost function r + g is coercive , i.e . lim‖w‖2→+∞ ( r + g ) ( w ) = +∞ and that the following scaling condition holds : for all n , d ∈ N , z ∈ Rn and any constant c > 0 , there exist a finite , positive constant C , such that , for any standard normal random vectors h ∈ Rd and g ∈ Rn : ‖z‖2 6 c √ n =⇒ sup x∈∂g ( z ) ‖x‖2 6 C √ n , 1 d E [ r ( h ) ] < +∞ , 1 n E [ g ( g ) ] < +∞ ( 9 ) The relevance of these assumptions in a supervised machine learning context is discussed in Appendix B.1 . We are now in a position to state our result . Theorem 1 . ( Closed-form asymptotics for ` 2 regularization ) In the asymptotic limit defined above , the training and generalisation errors ( 4 ) of the estimator ŵ ∈ Rd solving the empirical risk minimisation problem in eq . ( 3 ) with ` 2 regularization r ( w ) = λ2 ||w|| 2 2 verify : Etrain . ( ŵ ) P−−−→ d→∞ Es , h∼N ( 0,1 ) g proxV ? g ( . , f0 ( √ρs ) ) m ? √ ρ s+ √ q ? − m ? 2 ρ h , f0 ( √ρs ) Egen . ( ŵ ) P−−−→ d→∞ E ( ν , λ ) [ ĝ ( f̂ ( λ ) , f0 ( ν ) ) ] ( 10 ) where prox stands for the proximal operator defined as proxV g ( . , y ) ( x ) = arg min z { g ( z , y ) + 1 2V ( x− z ) 2 } ( 11 ) and where ( ν , λ ) are jointly Gaussian scalar variables : ( ν , λ ) ∼ N ( 0 , [ ρ m ? m ? q ? ] ) , ( 12 ) and the overlap parameters ( V ? , q ? , m ? ) are prescribed by the unique fixed point of the following set of self-consistent equations : V = E ( ω , θ̄ ) ∼µ [ ω λ+V̂ ω ] m = m̂√γE ( ω , θ̄ ) ∼µ [ θ̄2 λ+V̂ ω ] q = E ( ω , θ̄ ) ∼µ [ m̂2θ̄2ω+q̂ω2 ( λ+V̂ ω ) 2 ] , V̂ = αV ( 1− Es , h∼N ( 0,1 ) [ f ′ g ( V , m , q ) ] ) m̂ = 1√ργ α V Es , h∼N ( 0,1 ) [ sfg ( V , m , q ) − m√ρf ′ g ( V , m , q ) ] q̂ = αV 2Es , h∼N ( 0,1 ) [ ( m√ ρs+ √ q−m2ρ h−fg ( V , m , q ) ) 2 ] ( 13 ) where we defined the scalar random functions fg ( V , m , q ) = proxV g ( . , f0 ( √ρs ) ) ( ρ −1/2ms +√ q − ρ−1m2h ) and f ′g ( V , m , h ) = prox′V g ( . , f0 ( √ρs ) ) ( ρ −1/2ms + √ q − ρ−1m2h ) as the first de- rivative of the proximal operator . Proof : This result is a consequence of Theorem 2 , whose proof can be found in appendix B . The parameters of the model ( θ0 , Ω , Φ , Ψ ) only appear trough ρ , eq . ( 6 ) , and the asymptotic limit µ of the joint distribution eq . ( 7 ) and ( f0 , f̂ , g , λ ) . One can easily iterate the above equations to find their fixed point , and extract ( q∗ , m∗ ) which appear in the expressions for the training and generalisation errors ( E ? train , E ? gen ) , see eq . ( 4 ) . Note that ( q ? , m ? ) have an intuitive interpretation in terms of the estimator ŵ ∈ Rd : q ? ≡ 1 d ŵ > Ωŵ , m ? ≡ 1√ dp θ > 0 Φŵ ( 14 ) Or in words : m ? is the correlation between the estimator projected in the teacher space , while q ? is the reweighted norm of the estimator by the covariance Ω . The parameter V ∗ also has a concrete interpretation : it parametrizes the deformation that must be applied to a Gaussian field specified by the solution of the fixed point equations to obtain the asymptotic behaviour of ẑ . It prescribes the degree of non-linearity given to the linear output by the chosen loss function . This is coherent with the robust regression viewpoint , where one introduces non-square losses to deal with the potential non-linearity of the generative model . V̂ ∗ plays a similar role for the estimator ŵ through the proximal operator of the regularisation , see Theorem 4 and 5 in the Appendix . Two cases are of particular relevance for the experiments that follow . The first is the case of ridge regression , in which f0 ( x ) = f̂ ( x ) and both the loss g and the performance measure ĝ are taken to be the meansquared error mse ( y , ŷ ) = 12 ( y − ŷ ) 2 , and the asymptotic errors are given by the simple closed-form expression : E ? gen = ρ+ q ? − 2m ? , E ? train = E ? gen ( 1 + V ? ) 2 , ( 15 ) The second case of interest is the one of a binary classification task , for which f0 ( x ) = f̂ ( x ) = sign ( x ) , and we choose the performance measure to be the classification error ĝ ( y , ŷ ) = P ( y 6= ŷ ) . In the same notation as before , the asymptotic generalisation error in this case reads : E ? gen = 1 π cos−1 ( m ? √ ρq ? ) , ( 16 ) while the training error E ? train depends on the choice of g - which we will take to be the logistic loss g ( y , x ) = log ( 1 + e−xy ) in all of the binary classification experiments . As mentioned above , this paper includes stronger technical results including finite size corrections and precise characterization of the distribution of the estimator ŵ , for generic , non-separable loss and regularization g and r. This type of distributional statement is encountered for special cases of the model in related works such as [ 28 , 29 , 51 ] . Define V ∈ Rn×d as the matrix of concatenated samples used by the student . Informally , in high-dimension , the estimator ŵ and ẑ = 1√ d Vŵ roughly behave as non-linear transforms of Gaussian random variables centered around the teacher vector θ0 ( or its projection on the covariance spaces ) as follows : w∗ = Ω−1/2 prox 1 V̂ ∗ r ( Ω−1/2 . ) ( 1 V̂ ∗ ( m̂∗t+ √ q̂∗g ) ) , z∗ = prox V ∗g ( . , z ) ( m∗ √ ρ s+ √ q∗ − ( m ∗ ) 2 ρ h ) . where s , h ∼ N ( 0 , In ) and g ∼ N ( 0 , Id ) are random vectors independent of the other quantities , t = Ω−1/2Φ > θ0 , y = f0 ( √ ρs ) , and ( V ∗ , V̂ ∗ , q∗ , q̂∗ , m∗ , m̂∗ ) is the unique solution to the fixed point equations presented in Lemma 12 of appendix B . Those fixed point equations are the generalization of ( 13 ) to generic , non-separable loss function and regularization . The formal concentration of measure result can then be stated in the following way : Theorem 2 . ( Non-asymptotic version , generic loss and regularization ) Under Assumption ( B.1 ) , consider any optimal solution ŵ to 3 . Then , there exist constants C , c , c′ > 0 such that , for any Lipschitz function φ1 : Rd → R , and separable , pseudo-Lipschitz function φ2 : Rn → R and any 0 < < c′ : P ( ∣∣∣∣φ1 ( ŵ√d ) − Eφ1 ( w∗√ d ) ∣∣∣∣ > ) 6 C 2 e−cn 4 , P ( ∣∣∣∣φ2 ( ẑ√n ) − Eφ2 ( z∗√ n ) ∣∣∣∣ > ) 6 C 2 e−cn 4 . Note that in this form , the dimensions n , p , d still appear explicitly , as we are characterizing the convergence of the estimator ’ s distribution for large but finite dimension . The clearer , one-dimensional statements are recovered by taking the n , p , d → ∞ limit with separable functions and an ` 2 regularization . Other simplified formulas can also be obtained from our general result in the case of an ` 1 penalty , but since this breaks rotational invariance , they do look more involved than the ` 2 case . From Theorem 2 , one can deduce the expressions of a number of observables , represented by the test functions φ1 , φ2 , characterizing the performance of ŵ , for instance the training and generalization error . A more detailed statement , along with the proof , is given in appendix B .

Summary: In this papers the author proposed a new Gaussian covariate teacher-student model that is capable of accurately estimating the training and generalization errors for its learned linear coefficients of its student part with an exponentially decaying concentrate bound. Their framework covers a wide range of learning problems that either are either directly in teacher-student style or can be easily adapted into such setup. Their empirical study demonstrated the effectiveness of their framework for a variety of problems even when letting loose some of the theoretical assumptions required.