Generalization Error of Generalized Linear Models in High Dimensions

Introduction

A fundamental goal of machine learning is generalization: the ability to draw inferences about unseen data from finite training examples. Methods to quantify the generalization error are therefore critical in assessing the performance of any machine learning approach.

This paper seeks to characterize the generalization error for a class of generalized linear models (GLMs) of the form

We measure the generalization error in a standard manner: we are given training data $(\bm{x}_{i},y_{i})$, $i=1,\ldots,N$ from which we learn some parameter estimate $\widehat{\bm{w}}$ via a regularized empirical risk minimization of the form

where $\mathbf{X}=[\bm{x}_{1}\,\bm{x}_{2}\,\ldots\,\bm{x}_{N}]^{\text{\sf T}}$, is the data matrix, $F_{\rm out}$ is some output loss function, and $F_{\rm in}$ is some regularizer on the weights. We are then given a new test sample, $\bm{x}_{\rm ts}$, for which the true and predicted values are given by

where $d_{\rm ts}$ is the noise in the test sample, and $\phi(\cdot)$ is a postulated inverse link function that may be different from the true function $\phi_{\rm out}(\cdot)$. The generalization error is then defined as the expectation of some expected loss between $y_{\rm ts}$ and $\widehat{y}_{\rm ts}$ of the form

for some test loss function $f_{\rm ts}(\cdot)$ such as squared error or prediction error.

Even for this relatively simple GLM model, the behavior of the generalization error is not fully understood. Recent works (Montanari et al., 2019; Deng et al., 2019; Mei & Montanari, 2019; Salehi et al., 2019) have characterized the generalization error of various linear models for classification and regression in certain large random problem instances. Specifically, the number of samples $N$ and number of features $p$ both grow without bound with their ratio satisfying ${p}/{N}\rightarrow\beta\in(0,\infty)$, and the samples in the training data $\bm{x}_{i}$ are drawn randomly. In this limit, the generalization error can be exactly computed. The analysis can explain the so-called double descent phenomena (Belkin et al., 2019a): in highly under-regularized settings, the test error may initially increase with the number of data samples $N$ before decreasing. See the prior work section below for more details.

Our main result (Theorem 1) provides a procedure for exactly computing the asymptotic value of the generalization error (4) for GLM models in a certain random high-dimensional regime called the Large System Limit (LSL). The procedure enables the generalization error to be related to key problem parameters including the sampling ratio $\beta=p/N$, the regularizer, the output function, and the distributions of the true weights and noise. Importantly, our result holds under very general settings including: (i) arbitrary test metrics $f_{\rm ts}$; (ii) arbitrary training loss functions $F_{\rm out}$ as well as decomposable regularizers $F_{\rm in}$; (iii) arbitrary link functions $\phi_{\rm out}$; (iv) correlated covariates $\bm{x}$; (v) underparameterized ($\beta<1$) and overparameterized regimes ($\beta>1$); and (vi) distributional mismatch in training and test data. Section 4 discusses in detail the general assumptions on the quantities $f_{\rm ts}$, $F_{\rm out}$, $F_{\rm in}$, and $\phi_{\rm out}$ under which Theorem 1 holds.

Prior Work.

Many recent works characterize generalization error of various machine learning models, including special cases of the GLM model considered here. For example, the precise characterization for asymptotics of prediction error for least squares regression has been provided in (Belkin et al., 2019b; Hastie et al., 2019; Muthukumar et al., 2019). The former confirmed the double descent curve of (Belkin et al., 2019a) under a Fourier series model and a noisy Gaussian model for data in the over-parameterized regime. The latter also obtained this scenario under both linear and non-linear feature models for ridge regression and min-norm least squares using random matrix theory. Also, (Advani & Saxe, 2017) studied the same setting for deep linear and shallow non-linear networks.

The analysis of the the generalization for max-margin linear classifiers in the high dimensional regime has been done in (Montanari et al., 2019). The exact expression for asymptotic prediction error is derived and in a specific case for two-layer neural network with random first-layer weights, the double descent curve was obtained. A similar double descent curve for logistic regression as well as linear discriminant analysis has been reported by (Deng et al., 2019). Random feature learning in the same setting has also been studied for ridge regression in (Mei & Montanari, 2019). The authors have, in particular, shown that highly over-parametrized estimators with zero training error are statistically optimal at high signal-to-noise ratio (SNR). The asymptotic performance of regularized logistic regression in high dimensions is studied in (Salehi et al., 2019) using the Convex Gaussian Min-max Theorem in the under-parametrized regime. The results in the current paper can consider all these models as special cases. Bounds on the generalization error of over-parametrized linear models are also given in (Bartlett et al., 2019; Neyshabur et al., 2018).

Although this paper and several other recent works consider only simple linear models and GLMs, much of the motivation is to understand generalization in deep neural networks where classical intuition may not hold (Belkin et al., 2018; Zhang et al., 2016; Neyshabur et al., 2018). In particular, a number of recent papers have shown the connection between neural networks in the over-parametrized regime and kernel methods. The works (Daniely, 2017; Daniely et al., 2016) showed that gradient descent on over-parametrized neural networks learns a function in the RKHS corresponding to the random feature kernel. Training dynamics of overparametrized neural networks has been studied by (Jacot et al., 2018; Du et al., 2018; Arora et al., 2019; Allen-Zhu et al., 2019), and it is shown that the function learned is in an RKHS corresponding to the neural tangent kernel.

Approximate Message Passing.

Our key tool to study the generalization error is approximate message passing (AMP), a class of inference algorithms originally developed in (Donoho et al., 2009, 2010; Bayati & Montanari, 2011) for compressed sensing. We show that the learning problem for the GLM can be formulated as an inference problem on a certain multi-layer network. Multi-layer AMP methods (He et al., 2017; Manoel et al., 2018; Fletcher et al., 2018; Pandit et al., 2019) can then be applied to perform the inference. The specific algorithm we use in this work is the multi-layer vector AMP (ML-VAMP) algorithm of (Fletcher et al., 2018; Pandit et al., 2019) which itself builds on several works (Opper & Winther, 2005; Fletcher et al., 2016; Rangan et al., 2019; Cakmak et al., 2014; Ma & Ping, 2017). The ML-VAMP algorithm is not necessarily the most computationally efficient procedure for the minimization (2). For our purposes, the key property is that ML-VAMP enables exact predictions of its performance in the large system limit. Specifically, the error of the algorithm estimates in each iteration can be predicted by a set of deterministic recursive equations called the state evolution or SE. The fixed points of these equations provide a way of computing the asymptotic performance of the algorithm. In certain cases, the algorithm can be proven to be Bayes optimal (Reeves, 2017; Gabrié et al., 2018; Barbier et al., 2019).

This approach of using AMP methods to characterize the generalization error of GLMs was also explored in (Barbier et al., 2019) for i.i.d. distributions on the data. The explicit formulae for the asymptotic mean squared error for the regularized linear regression with rotationally invarient data matrices is proved in (Gerbelot et al., 2020). The ML-VAMP method in this work enables extensions to correlated features and to mismatch between training and test distributions.

Generalization Error: System Model

where ${\bm{\phi}}_{\rm out}(\mathbf{z},\mathbf{d})$ is the vector-valued function such that $[{\bm{\phi}}_{\rm out}(\mathbf{z},\mathbf{d})]_{n}=\phi_{\rm out}(z_{n},d_{n})$ and $\{d_{n}\}_{n=1}^{N}$ are general noise.

Given the training data $(\mathbf{X},\bm{y})$, we consider estimates of $\bm{w}^{0}$ given by a regularized empirical risk minimization of the form (2). We assume that the loss function $F_{\rm out}$ and regularizer $F_{\rm in}$ are separable functions, i.e., one can write

Large System Limit: We follow the LSL analysis of (Bayati & Montanari, 2011) commonly used for analyzing AMP-based methods. Specifically, we consider a sequence of problems indexed by the number of training samples $N$. For each $N$, we suppose that the number of features $p=p(N)$ grows linearly with $N$, i.e.,

for some constant $\beta\in(0,\infty)$. Note that $\beta>1$ corresponds to the over-parameterized regime and $\beta<1$ corresponds to the under-parameterized regime.

True parameter: We assume the true weight vector $\bm{w}^{0}$ has components whose empirical distribution converges as

Using the covariance (9), we can write the data matrix as

where $S_{\rm mp}\geq 0$ is a non-negative random variable such that $S_{\rm mp}^{2}$ satisfies the Marcencko-Pastur distribution. Details on this distribution are in Appendix H.

Training data output: Given the input data $\mathbf{X}$, we assume that the training outputs $\bm{y}$ are generated from (5), where the noise $\mathbf{d}$ is independent of $\mathbf{X}$ and has an empirical distribution which converges as

Again, the limit (13) will be satisfied if $\{d_{i}\}_{i=1}^{N}$ are i.i.d. draws of random variable $D$ with bounded second moments.

Test data: To measure the generalization error, we assume now that we are given a test point $\bm{x}_{\rm ts}$, and we obtain the true output $y_{\rm ts}$ and predicted output $\widehat{y}_{\rm ts}$ given by (3). We assume that the test data inputs are also Gaussian, i.e.,

In comparison to (9), we see that we are assuming that the eigenvectors of the training and test data are the same, but the eigenvalues may be different. In this way, we can capture distributional mismatch between the training and test data. For example, we will be able to measure the generalization error when the test sample is outside a subspace explored by the training data.

To capture the relation between the training and test distributions, we assume that components of $\mathbf{s}_{\rm tr}$ and $\mathbf{s}_{\rm ts}$ converge as

to some non-negative, bounded random vector $(S_{\rm tr},S_{\rm ts})$. The joint distribution on $(S_{\rm tr},S_{\rm ts})$ captures the relation between the training and test data.

When $S_{\rm tr}=S_{\rm ts}$, our model corresponds to the case when the training and test distribution are matched. Isotropic Gaussian features in both training and test data correspond to covariance matrices ${\bm{\Sigma}}_{\rm tr}=\tfrac{1}{p}{\sigma^{2}_{\rm tr}}\mathbf{I}$, ${\bm{\Sigma}}_{\rm ts}=\tfrac{1}{p}{\sigma^{2}_{\rm ts}}\mathbf{I}$, which can be modeled as $S_{\rm tr}=\sigma_{\rm tr}$, $S_{\rm ts}=\sigma_{\rm ts}$. We also require that the matrix $\mathbf{V}_{0}$ is uniformly distributed on the set of $p\times p$ orthogonal matrices.

Generalization error: From the training data, we obtain an estimate $\widehat{\bm{w}}$ via a regularized empirical risk minimization (2). Given a test sample $\bm{x}_{\rm ts}$ and parameter estimate $\widehat{\bm{w}}$, the true output $y_{\rm ts}$ and predicted output $\widehat{y}_{\rm tr}$ are given by equation (3). We assume the test noise is distributed as $d_{\rm ts}\sim D$, following the same distribution as the training data. The postulated inverse-link function $\phi(\cdot)$ in (3) may be different from the true inverse-link function $\phi_{\rm out}(\cdot)$.

The generalization error is defined as the asymptotic expected loss,

where $f_{\rm ts}(\cdot)$ is some loss function relevant for the test error (which may be different from the training loss). The expectation in (17) is with respect to the randomness in the training as well as test data, and the noise. Our main result provides a formula for the generalization error (17).

Learning GLMs via ML-VAMP

There are many methods for solving the minimization problem (2). We apply the ML-VAMP algorithm of (Fletcher et al., 2018; Pandit et al., 2019). This algorithm is not necessarily the most computationally efficient method. For our purposes, however, the algorithm serves as a constructive proof technique, i.e., it enables exact predictions for generalization error in the LSL as described above. Moreover, in the case when loss function (2) is strictly convex, the problem has a unique global minimum, whereby the generalization error of this minimum is agnostic to the choice of algorithm used to find this minimum. To that end, we start by reformulating (2) in a form that is amicable to the application of ML-VAMP, Algorithm 1.

The first step in applying ML-VAMP to the GLM learning problem is to represent the mapping from the true parameters $\bm{w}^{0}$ to the output $\bm{y}$ as a certain multi-layer network. We combine (5), (10) and (11), so that the mapping $\bm{w}^{0}\mapsto\bm{y}$ can be written as the following sequence of operations (as illustrated in Fig. 1):

The minimization (2) can also be represented using a similar signal flow graph. Given a parameter candidate $\bm{w}$, the mapping $\bm{w}\mapsto\mathbf{X}\bm{w}$ can be written using the sequence of vectors

There are $L=3$ steps in this sequence, and we let

denote the sets of vectors across the steps. The minimization in (2) can then be written in the following equivalent form:

where $\delta_{\mathcal{A}}(\cdot)$ is on the set $\mathcal{A}$, and $+\infty$ on $\mathcal{A}^{c}$.

ML-VAMP for GLM Learning.

Using this multi-layer representation, we can now apply the ML-VAMP algorithm from (Fletcher et al., 2018; Pandit et al., 2019) to solve the optimization (22). The steps are shown in Algorithm 1. These steps are a special case of the “MAP version” of ML-VAMP in (Pandit et al., 2019), but with a slightly different set-up for the GLM problem. We will call these steps the ML-VAMP GLM Learning Algorithm.

For the MAP version of ML-VAMP algorithm in (Pandit et al., 2019), the denoisers are essentially proximal-type operators defined as

An important property of the proximal operator is that for separable functions $F$ of the form (6), we have $[\operatorname{prox}_{F/\gamma}(\bm{u})]_{i}=\operatorname{prox}_{f/\gamma}(\bm{u}_{i})$.

The quantity $\langle{\partial\bm{v}/\partial\bm{u}}\rangle$ on lines 11 and 23 denotes the empirical mean $\tfrac{1}{N}\sum_{n=1}^{N}\partial v_{n}/\partial u_{n}$.

Main Result

The denoisers and link functions satisfy the following continuity conditions:

are uniformly Lipschitz continuous in $\mathbf{r}^{-}_{0}$ and $(\mathbf{r}^{+}_{2},\bm{y})$ over parameters $\gamma^{-}_{0}$ and $\gamma^{+}_{2}$.

The link function $\phi_{\rm out}(p,d)$ is Lipschitz continuous in $(p,d)$. The test error function $f_{\rm ts}(\phi(\widehat{z}),\phi_{\rm out}(z,d))$ is pseduo-Lipschitz continuous in $(\widehat{z},z,d)$ of order 2.

Let $\bm{x}_{k}(N)$ be any finite set of outputs of the ML-VAMP algorithm as above. Then there exist limits

We are now ready to state our main result.

The true parameter $\bm{w}^{0}$ and its estimate $\widehat{\bm{w}}$ empirically converge as

where $W^{0}$ is the random variable from (8) and

with $Q_{0}^{-}={\mathcal{N}}(0,\tau_{0}^{-})$ independent of $W^{0}$.

The asymptotic generalization error (17) with $(y_{\rm ts},\widehat{y}_{\rm ts})$ defined as (3) is given by

where $(Z_{\rm ts},\widehat{Z}_{\rm ts})\sim\mathcal{N}(\mathbf{0}_{2},\mathbf{M})$ and independent of $D$.

Part (b) provides an exact description of the asymptotic statistical relation between the true parameter $\bm{w}^{0}$ and its estimate $\widehat{\bm{w}}$. The parameters $\tau_{0}^{-},\overline{\gamma}_{0}^{+}>0$ and $\mathbf{M}$ can be explicitly computed using a set of recursive equations called the state evolution or SE described in Appendix C in the supplementary material.

We can use the expressions to compute a variety of relevant metrics. For example, the $PL(2)$ convergence shows that the MSE on the parameter estimate is

The expectation on the right hand side of (33) can then be computed via integration over the joint density of $(W^{0},\widehat{W})$ from part (b). In this way, we have a simple and exact method to compute the parameter error. Other metrics such as parameter bias or variance, cosine angle or sparsity detection can also be computed.

Part (c) of Theorem 1 similarly exactly characterizes the asymptotic generalization error. In this case, we would compute the expectation over the three variables $(Z,\widehat{Z},D)$. In this way, we have provided a methodology for exactly predicting the generalization error from the key parameters of the problems such as the sampling ratio $\beta=p/N$, the regularizer, the output function, and the distributions of the true weights and noise. We provide several examples such as linear regression, logistic regression and SVM in the Appendix G. We also recover the result by (Hastie et al., 2019) in Appendix G.

Note that Assumption 1 is satisfied in many practical cases. For example, it can be verified that it is satisfied in the case when $f_{\rm in}(\cdot)$ and $f_{\rm out}(\cdot)$ are convex. Assumption 2 is somewhat more restrictive in that it requires that the ML-VAMP algorithm converges. The convergence properties of ML-VAMP are discussed in (Fletcher et al., 2016). The ML-VAMP algorithm may not always converge, and characterizing conditions under which convergence is possible is an open question. However, experiments in (Rangan et al., 2019) show that the algorithm does indeed often converge, and in these cases, our analysis applies. In any case, we will see below that the predictions from Theorem 1 agree closely with numerical experiments in several relevant cases.

In some special cases equation (32) simplifies to yield quantitative insights for interesting modeling artifacts. We discuss these in Appendix G in the supplementary material.

Experiments

We validate our theoretical results on a number of synthetic data experiments. For all the experiments, the training and test data is generated following the model in Section 2. We generate the training and test eigenvalues as i.i.d. with lognormal distributions,

where $(u_{\rm tr},u_{\rm ts})$ are bivariate zero-mean Gaussian with

In the case when $\sigma^{2}_{u}=0$, we obtain eigenvalues that are equal, corresponding to the i.i.d. case. With $\sigma^{2}_{u}>0$ we can model correlated features. Also, when the correlation coefficient $\rho=1$, $S_{\rm tr}=S_{\rm ts}$, so there is no training and test mismatch. However, we can also select $\rho<1$ to experiment with cases when the training and test distributions differ. In the examples below, we consider the following three cases:

correlated features with matching training and test distributions ($\sigma_{u}=3$ dB, $\rho=1$); and

correlated features with train-test mismatch ($\sigma_{u}=3$ dB, $\rho=0.5$).

For all experiments below, the true model coefficients are generated as i.i.d. Gaussian $w_{j}^{0}\sim{\mathcal{N}}(0,1)$ and we use standard L2-regularization, $f_{\rm in}(w)=\lambda w^{2}/2$ for some $\lambda>0$. Our framework can incorporate arbitrary i.i.d. distributions on $w_{j}$ and regularizers, but we will illustrate just the Gaussian case with L2-regularization here.

Under-regularized linear regression.

Fig. 2 plots the test MSE for the three cases described above for the linear model. In the figure, we take $p=1000$ features and vary the number of samples $n$ from $0.2p$ (over-parametrized) to $3p$ (under-paramertrized). For each value of $n$, we take 100 random instances of the model and compute the ridge regression estimate using the sklearn package and measure the test MSE on the 1000 independent test samples. The simulated values in Fig. 2 are the median test error over the 100 random trials. The test MSE is plotted in a normalized dB scale,

Also plotted is the state evolution (SE) theoretical test MSE from Theorem 1.

In all three cases in Fig. 2, the SE theory exactly matches the simulated values for the test MSE. Note that the case of match training and test distributions for this problem was studied in (Hastie et al., 2019; Mei & Montanari, 2019; Montanari et al., 2019) and we see the double descent phenomenon described in their work. Specifically, with highly under-regularized linear regression, the test MSE actually increases with more samples $n$ in the over-parametrized regime ($n/p<1$) and then decreases again in the under-parametrized regime ($n/p>1$).

Our SE theory can also provide predictions for the correlated feature case. In this particular setting, we see that in the correlated case the test error is slightly lower in the over-parametrized regime since the energy of data is concentrated in a smaller sub-space. Interestingly, there is minimal difference between the correlated and i.i.d. cases for the under-parametrized regime when the training and test data match. When the training and test data are not matched, the test error increases. In all cases, the SE theory can accurately predict these effects.

Logistic Regression.

Nonlinear Regression.

The SE framework can also consider non-convex problems. As an example, we consider a non-linear regression problem where the output function is

The $\tanh(p)$ models saturation in the output. Corresponding to this output, we use a non-linear MSE output loss

For the simulation, the non-convex loss is minimized using Tensorflow where the non-linear model is described as a two-layer model. We use the ADAM optimizer (Kingma & Ba, 2014) with 200 epochs to approach a local minimum of the objective (2). Fig. 4 plots the median test MSE for the estimate along with the SE theoretical test MSE. We again see that the SE theory is able to predict the test MSE in all cases even for this non-convex problem.

Conclusions

In this paper we provide a procedure for exactly computing the asymptotic generalization error of a solution in a generalized linear model (GLM). This procedure is based on scalar quantities which are fixed points of a recursive iteration. The formula holds for a large class of generalization metrics, loss functions, and regularization schemes. Our formula allows analysis of important modeling effects such as (i) overparameterization, (ii) dependence between covariates, and (iii) mismatch between train and test distributions, which play a significant role in the analysis and design of machine learning systems. We experimentally validate our theoretical results for linear as well as non-linear regression and logistic regression, where a strong agreement is seen between our formula and simulated results.

References

Appendix A Empirical Convergence of Vector Sequences

The LSL model in Section 2 and our main result in Section 4 require certain technical definitions.

Observe that for $p=1$, the pseudo-Lipschitz is equivalent to the standard definition of Lipschitz continuity.

In this case, with some abuse of notation, we will write

$PL(p)$ convergence is equivalent to weak convergence plus convergence in $p$ moment (Bayati & Montanari, 2011), and hence $PL(p)$ convergence is also equivalent to convergence in Wasserstein-$p$ metric (See Chapter 6. (Villani, 2008)). We use this fact later in proving Theorem 1.

Appendix B ML-VAMP Denoisers Details

where $\phi_{1}(\cdot)$ is the scalar-valued function,

Finally, the function ${\bm{\phi}}_{3}(\cdot)$ in (20) acts componentwise with

Input denoiser $\mathbf{g}_{0}^{+}(\cdot)$: Since $F_{0}(\mathbf{z}_{0})=F_{\rm in}(\mathbf{z}_{0})$, and $F_{\rm in}(\cdot)$ given in (6), the denoiser (25a) acts componentwise in that,

where $g_{0}^{+}(\cdot)$ is the scalar-valued function,

Thus, the vector optimization in (25a) reduces to a set of scalar optimizations (44) on each component.

Output denoiser $\mathbf{g}_{3}^{-}(\cdot)$: The output penalty $F_{3}(\mathbf{p}_{2},\bm{y})=F_{\rm out}(\mathbf{p}_{2},\bm{y})$ where $F_{\rm out}(\mathbf{p}_{2},\bm{y})$ has the separable form (6). Thus, similar to the case of $\mathbf{g}_{0}(\cdot)$, the denoiser $\mathbf{g}_{3}(\cdot)$ in (25b) also acts componentwise with the function,

and $\widehat{\mathbf{z}}_{1}=\mathbf{S}_{\rm tr}\widehat{\mathbf{p}}_{0}$. This is a simple quadratic minimization and the components of $\widehat{\mathbf{p}}_{0}$ and $\widehat{\mathbf{z}}_{1}$ are given by

Linear denoiser $\mathbf{g}_{2}^{\pm}(\cdot)$: This denoiser is identical to the case $\mathbf{g}_{1}^{\pm}(\cdot)$ in that we need to impose the linear constraint $\mathbf{z}_{2}=\mathbf{S}_{\rm mp}\mathbf{p}_{1}$. However $\mathbf{S}_{\rm{mp}}$ is in general a rectangular matrix and the two resulting cases of $\beta\lessgtr 1$ needs to be treated separately.

with the identical functions $g_{2}^{-}=g_{1}^{-}$ and $g_{2}^{+}=g_{1}^{+}$ as given by (47a) and (47b). Note that in (48a), $n=1,\ldots,p$ and in (48b), $n=1,\ldots,N$.

Appendix C State Evolution Analysis of ML-VAMP

A key property of the ML-VAMP algorithm is that its performance in the LSL can be exactly described by a scalar equivalent system. In the scalar equivalent system, the vector-valued outputs of the algorithm are replaced by scalar random variables representing the typical behavior of the components of the vectors in the large-scale-limit (LSL). Each of the random variables are described by a set of parameters, where the parameters are given by a set of deterministic equations called the state evolution or SE.

The updates in sections labeled “Forward pass” and “Backward pass” in the SE equations in Algorithm 2 parallel those in Algorithm 1. The key quantities in these SE equations are the error variables,

which represent the errors of the estimates to the inputs of the denoisers. We will also be interested in their transforms,

The following Theorem is an adapted version of the main result from (Pandit et al., 2019) to the iterates of Algorithms 1 and 2.

Consider the outputs of the ML-VAMP for GLM Learning Algorithm under the assumptions of Section 2. Assume the denoisers satisfy the continuity conditions in Assumption 1. Also, assume that the outputs of the SE satisfy

A key use of the Theorem is to compute asymptotic empirical limits. Specifically, for a componentwise function $\psi(\cdot)$, let $\langle{\psi(\bm{x})}\rangle$ denotes the average $\tfrac{1}{N}\sum_{n=1}^{N}\psi(x_{n})$ The above theorem then states that for any componentwise pseudo-Lipschitz function $\psi(\cdot)$ of order 2, as $N\rightarrow\infty$, we have the following two properties

That is, we can compute the empirical average over components with the expected value of the random variable limit. This convergence is key to proving Theorem 1.

Appendix D Empirical Convergence of Fixed Points

A consequence of Assumption 2 is that we can take the limit $k\rightarrow\infty$ of the random variables in the SE algorithm. Specifically, let $\bm{x}_{k}=\bm{x}_{k}(N)$ be any set of $d$ outputs from the ML-VAMP for GLM Learning Algorithm under the assumptions of Theorem 2. Under Assumption 2, for each $N$, there exists a vector

representing the limit over $k$. For each $k$, Theorem 2 shows there also exists a random vector limit,

representing the limit over $N$. The following proposition shows that we can take the limits of the random variables $X_{k}$.

The proposition shows that, under the convergence assumption, Assumption 2, we can take the limits as $k\rightarrow\infty$ of the random variables from the SE. To prove the proposition we first need the following simple lemma.

then, there exists a constant $C$ such that,

In particular, the two limits in (60) exist.

For any $\epsilon>0$, the limit (59) implies that there exists a $k_{\epsilon}(\uparrow\infty$ as $\epsilon\downarrow 0)$ such that for all $k>k_{\epsilon}$,

Since this is true for all $\epsilon>0$, it follows that

Similarly, $\limsup_{N\rightarrow\infty}\alpha_{N}\leq\inf_{k>k_{\epsilon}}\beta_{k}+\epsilon$, whereby

Equations (61) and (62) together show that the limits in (60) exists and are equal. $\Box$

Since $\{x_{k,i}(N)\}$ converges $PL(2)$ to $X_{k}$, we have,

where (a) follows from applying the triangle inequality to the definition of $A_{N,k}$ in (65); (b) follows from the definition of pseudo-Lipschitz continuity in Definition 1, $C>0$ is the Lipschitz contant and

and (c) follows from the RMS-AM inequality:

Substituting (67) and (69) into (64) show that $\alpha_{N}$ and $\beta_{k}$ satisfy (59). Therefore, applying Lemma 1 we have that for any pseudo-Lipschitz function $f(\cdot)$, there exists a limit $\Phi(f)$ such that,

Since the numerator and denominator of (71) are $PL(2)$ functions we have that the limit,

Appendix E Proof of Theorem 1

The estimate $\widehat{\bm{w}}$ is the limit,

Also, the true parameter is $\mathbf{z}_{0}^{0}=\bm{w}^{0}$. By Proposition 1, we have that the $PL(2)$ limits of these variables are

From line 15 of the SE Algorithm 2, we have

Since the fixed points are critical points of the constrained optimization (22), $\widehat{\mathbf{p}}_{0}=\mathbf{V}_{0}\widehat{\bm{w}}$. We also have $\mathbf{p}^{0}_{0}=\mathbf{V}_{0}\bm{w}^{0}$. Therefore,

The empirical convergence (73) yields the following limit,

It suffices to show that the distribution of $[z_{\rm ts}^{(N)}\,\widehat{z}_{\rm ts}^{(N)}]$ converges to the distribution of $[Z_{\rm ts}\,\widehat{Z}_{\rm ts}]$ in the Wasserstein-2 metric as $N\rightarrow\infty.$ (See the discussion in Appendix A on the equivalence of convergence in Wasserstein-2 metric and PL(2) convergence.)

Now, Wassestein-2 distance between between two probability measures $\nu_{1}$ and $\nu_{2}$ is defined as

where $\Gamma$ is the set of probability distributions on the product space with marginals consistent with $\nu_{1}$ and $\nu_{2}$. For Gaussian measures $\nu_{1}=\mathcal{N}(\mathbf{0},\Sigma_{1})$ and $\nu_{2}=\mathcal{N}(\mathbf{0},\Sigma_{2})$ we have (Givens et al., 1984)

Therefore, for Gaussian distributions $\nu_{1}^{(N)}=\mathcal{N}(\mathbf{0},\mathbf{M}^{(N)})$, and $\nu_{2}=\mathcal{N}(\mathbf{0},\mathbf{M})$, the convergence (75) implies $W_{2}(\nu_{1}^{(N)},\nu_{2})\rightarrow 0,$ i.e., convergence in Wasserstein-2 distance. Hence,

where $\mathbf{M}$ is the covariance matrix in (75). Hence the convergence holds in the PL(2) sense (see discussion in Appendix A on the equivalence of convergence in $W_{2}$ and PL(2) convergence).

Hence the asymptotic generalization error (17) is

where (a) follows from (3); and step (b) follows from continuity assumption in Assumption 1(b) along with the definition of PL(2) convergence in Def. 3. This proves part (c).

Appendix F Formula for 𝐌𝐌\mathbf{M}

For the special cases in the next Appendix, it is useful to derive expressions for the entries the covariance matrix $\mathbf{M}$ in (75). For the term $m_{11}$,

where $(P_{0}^{0},P_{0}^{+},Q_{0}^{-})$ are independent of $(S_{{\rm tr}},S_{\rm ts})$. Hence,

Appendix G Special Cases

In this section we examine a few special cases of the GLM problem (2). We first consider a linear output with additive Gaussian noise and a squared error training and test loss. Specifically, consider the model,

We consider estimates of $\bm{w}^{0}$ such that:

The factor $\beta$ is added above since the two terms scale with a ratio of $\beta$. It does not change analysis. Consider the ML-VAMP GLM learning algorithm applied to this problem. The following corollary follows from the Main result in Theorem 1.

For linear regression, i.e., $\phi(t)=t,$ $\phi_{\rm out}(t,d)=t+d,$ $f_{\rm ts}(y,\widehat{y})=(y_{\rm ts}-\widehat{y}_{\rm ts})^{2}$, $F_{\rm out}(\mathbf{p}_{2})=\tfrac{1}{N}\left\|\bm{y}-\mathbf{p}_{2}\right\|^{2}$, we have

The quantities $k_{22}$, $\tau_{1}^{-},\overline{\gamma}_{0}^{+},\overline{\gamma}_{1}^{-}$ depend on the choice of regularizer $\lambda$ and the covariance between features.

This follows directly from the following observation:

Substituting equation (81) proves the claim. $\Box$

G.2 Ridge Regression with i.i.d. Covariates

We next the special case when the input features are independent, i.e., (83) where rows of $\mathbf{X}$ corresponding to the training data has i.i.d Gaussian features with covariance $\mathbf{P}_{\rm train}=\frac{\sigma_{\rm tr}^{2}}{p}\mathbf{I}$ and $S_{\rm tr}=\sigma_{\rm tr}$.

Although the solution to (83) exists in closed form $(\mathbf{X}^{\text{\sf T}}\mathbf{X}+\lambda\mathbf{I})^{-1}\mathbf{X}^{\text{\sf T}}\bm{y}$, we can study the effect of the regularization parameter $\lambda$ on the generalization error $\mathcal{E}_{\rm ts}$ as detailed in the result below.

Consider the ridge regression problem (83) with regularization parameter $\lambda>0$. For the squared loss i.e., $f_{\rm ts}(y,\widehat{y})=(y-\widehat{y})^{2}$, i.i.d Gaussian features without train-test mismatch, i.e., $S_{\rm tr}=S_{\rm ts}=\sigma_{\rm tr}$, the generalization error $\mathcal{E}_{{\rm ts}}^{\mathsf{RR}}$ is given by Corollary 1, with constants

We are interested in identifying the following constants appearing in Corollary 1:

In the case of problem (83), the maps $\operatorname{prox}_{f_{\rm in}}$ and $\operatorname{prox}_{f_{\rm out}}$, i.e., $g_{0}^{+}$ and $g_{3}^{-}$ respectively, can be expressed as closed-form formulae. This leads to simplification of the SE equations as explained below.

To begin with, notice that $f_{\rm in}(w)=\frac{\lambda}{2}w^{2}$, and therefore the denoiser $g_{0}^{+}(\cdot)$ in (44) is simply,

Using the random variable $R_{0}^{-}$ and substituting in the expression of the denoiser to get $\widehat{Z}_{0}$, we can now calculate $\overline{\alpha}_{0}^{+}$ using lines 20 and 22,

Similarly, we have $f_{\rm out}(p_{2})={\tfrac{1}{2}}(p_{2}-y)^{2}$, whereby the output denoiser $g_{3}^{-}(\cdot)$ in the last layer for ridge regression is given by,

By substituting this denoiser in line 30 of the algorithm we get $\widehat{P}_{2}^{-}$ and thus, following the lines 35-38 of the algorithm we have

Having identified these constants $\overline{\alpha}_{0}^{+},\overline{\gamma}_{0}^{+},\overline{\alpha}_{2}^{-},\overline{\gamma}_{2}^{-}$, we will now sequentially identify the quantities

in the forward pass, and then the quantities

Notice that from line 23, the pair $(P_{0}^{0},P_{0}^{+})$ is jointly Gaussian with covariance matrix $\mathbf{K}_{0}^{+}$. But the above equation means that $P_{0}^{+}=-P^{0}_{0}$, whereby $R_{0}^{+}=0$ from line 17.

Backward Pass:

Since $Y=\phi_{\rm out}(P_{2}^{0},D)=P_{2}^{0}+D$, line 36 of algorithm on simplification yields $P_{2}^{-}=D$, whereby we can get $\tau_{2}^{-}$,

Next, to calculate the terms $(\overline{\alpha}_{1}^{-},\overline{\gamma}^{-}_{1})$, we use the decoiser $g_{2}^{-}$ defined in (47a) for line 33 of the algorithm to get $\widehat{P}_{1}$.

where we have used $\overline{\gamma}_{2}^{-}=1,$ $R_{1}^{+}=P_{1}^{0}+P_{1}^{+}=0$ due to (90), and $R_{2}^{-}=Z_{2}^{0}+Q_{2}^{-}=S_{\rm{mp}}^{+}P_{1}^{0}+Q_{2}^{-}$ from lines 17, 32 and 4 respectively.

Here, in the overparameterized case $(\beta>1)$, the denoiser $g_{2}^{-}$ outputs $R_{1}^{+}$ with probability $1-\tfrac{1}{\beta}$ and $\tfrac{\lambda}{\sigma_{\rm tr}^{2}\beta}G$ with probability $\tfrac{1}{\beta}$.

Now from line 36 and equation (87) we get,

where $G=G_{\rm{mp}}(-\frac{\lambda}{\sigma_{\rm tr}^{2}\beta})$, with $G_{\rm{mp}}$ given in Appendix H, and $G^{\prime}$ is the derivative of $G_{\rm{mp}}$ calculated at $-\frac{\lambda}{\sigma_{\rm tr}^{2}\beta}$.

Now consider the under-parametrized case ($\beta<1$):

Let $u=-\frac{\lambda}{\sigma_{\rm{tr}}^{2}\beta}$ and $z=G_{\rm{mp}}(u)$. In this case we have

where $R_{\rm{mp}}(.)$ is the R-transform defined in (Tulino et al., 2004) and (a) follows from the relationship between the R- and Stieltjes-transform and (b) follows from the fact that for Marchenko-Pastur distribution we have $R_{\rm{mp}}(z)=\frac{1}{1-z\beta}$. Therefore,

For the over-parametrized case ($\beta>1$) we have:

In this case, as mentioned in Appendix H and following the results from (Tulino et al., 2004), the measure $\mu_{\beta}$ scales with $\beta$ and thus $R_{\rm{mp}}(z)=\frac{\beta}{1-z}$. Therefore, similar to (99a), $z$ satisfies

Now $\tau_{1}^{-}$ can be calculated as follows:

and $z$ is the solution to the fixed points

G.3 Ridgeless Linear Regression

Here we consider the case of Ridge regression (83) when $\lambda\rightarrow 0^{+}$. Note that the solution to the problem (83) is $(\mathbf{X}^{\text{\sf T}}\mathbf{X}+\lambda\mathbf{I})^{-1}\mathbf{X}^{\text{\sf T}}\bm{y}$ remains unique since $\lambda>0$. The following result was stated in (Hastie et al., 2019), and can be recovered using our methodology. Note however, that we calculate the generalization error whereas they have calculated the squared error, whereby we obtain an additional additive factor of $\sigma_{d}^{2}.$ The result explains the double-descent phenomenon for Ridgeless linear regression.

We calculate the parameters $\overline{\gamma}_{0}^{+},\overline{\gamma}_{1}^{-}$, $k_{22}$ and $\tau_{1}^{-}$ when $\lambda\rightarrow 0^{+}$. Before starting off, we note that

as described in Appendix H. Following the derivations in Corollary 2, we have

Now for $\lambda\rightarrow 0^{+},$ we have

Using this in simplifying (95) for $\lambda\rightarrow 0^{+}$, we get

G.4 Train-Test Mismatch

Observe that our formulation allows for analyzing the effect of mismatch in the training and test distribution. One can consider arbitrary joint distributions over $(S_{\rm tr},S_{\rm ts})$ that model the mismatch between training and test features. Here we give a simple example which highlights the effect of this mismatch.

$(S_{\rm ts},S_{\rm tr})$ has a bivariate Bernoulli distribution with

The following result shows that the generalization error increases linearly with the mismatch parameter $\varepsilon.$

Consider the problem of Linear Regression (83) under the conditions of Corollary 1. Additionally suppose we have Bernoulli $\varepsilon$-mismatch between the training and test distributions. Then

where $\gamma^{*}:=\frac{\overline{\gamma}_{0}^{+}}{\overline{\gamma}_{0}^{+}+\overline{\gamma}_{1}^{-}}$. The terms $k_{22},\tau_{1}^{-},\gamma^{*}$ are independent of $\varepsilon$.

This follows directly by calculating the expectations of the terms in Corollary 1, with the joint distribution of $(S_{\rm tr},S_{\rm ts})$ given in Definition 4. $\Box$

The quantities $k_{22}$ and $\tau_{1}^{-}$ in the result above can be calculated similar to the derivation in the proof of Corollary 2 and can in general depend on the regularization parameter $\lambda$ and overparameterization parameter $\beta$.

G.5 Logistic Regression

The precise analysis for the special case of regularized logistic regression estimator with i.i.d Gaussian features is provided in (Salehi et al., 2019). Consider the logistic regression model,

where $\rho(x)=\frac{1}{1+e^{-x}}$ is the standard logistic function.

In this problem we consider estimates of $\bm{w}^{0}$ such that

where $F_{\rm in}$ is the reguralization function. This is a special case of optimization problem (2) where

Similar to the linear regression model, using the ML-VAMP GLM learning algorithm, we can characterize the generalization error for this model with quantities $\mathbf{K}_{0}^{+},\tau_{1}^{-},\overline{\gamma}_{0}^{+},\overline{\gamma}_{1}^{-}$ given by algorithm 2. We note that in this case, the output non-linearity is

where $d\sim\text{Unif}(0,1)$. Also, the denoisers $g_{0}^{+}$, and $g_{3}^{-}$ can be derived as the proximal operators of $F_{\rm in}$, and $F_{\rm out}$ defined in (25).

G.6 Support Vector Machines

The asymptotic generalization error for support vector machine (SVM) is provided in (Deng et al., 2019). Our model can also handle SVMs. Similar to logistic regression, SVM finds a linear classifier using the hinge loss instead of logistic loss. Assuming the class labels are $y=\pm 1$ the hinge loss is

where $\mathbf{X}_{i}$ is the $i^{th}$ row of the data matrix, the ML-VAMP algorithm for GLMs finds the SVM classifier. The algorithm would have proximal map of hinge loss and our theory provides exact predictions for the estimation and prediction error of SVM.

As with all other models considered in this work, the true underlying data generating model could be anything that can be represented by the graphical model in Figure 1, e.g. logistic or probit model, and our theory is able to exactly predict the error when SVM is applied to learn such linear classifiers in the large system limit.

Appendix H Marchenko-Pastur distribution

We describe the random variable $S_{{\rm{mp}}}$ defined in (12) where $S_{\rm{mp}}^{2}$ has a rescaled Marchenko-Pastur distribution. Notice that the positive entries of $\mathbf{s}_{\rm{mp}}$ are the positive eigenvalues of $\mathbf{U}^{\text{\sf T}}\mathbf{U}$ (or $\mathbf{U}\mathbf{U}^{\text{\sf T}}$).

Observe that $U_{ij}\sim N(0,\frac{1}{p})$, whereas, the standard scaling while studying the Marchenko-Pastur distribution is for matrices $\mathbf{H}$ such that $H_{ij}\sim\mathcal{N}(0,\frac{1}{N})$ (for e.g. see equation (1.10) from (Tulino et al., 2004) and the discussion preceding it). Also notice that $\sqrt{\beta}\mathbf{U}$ has the same distribution as $\mathbf{H}$. Thus the results from (Tulino et al., 2004) apply directly to the distributions of eigenvalues of $\beta\mathbf{U}^{\text{\sf T}}\mathbf{U}$ and $\beta\mathbf{U}\mathbf{U}^{\text{\sf T}}$. We state their result below taking into account this disparity in scaling.

The positive eigenvalues of $\beta\mathbf{U}^{\text{\sf T}}\mathbf{U}$ have an empirical distribution which converges to the following density:

where $a_{\beta}=(1-\sqrt{\beta})^{2}$, $b_{\beta}:=(1+\sqrt{\beta})^{2}$. Similarly the positive eigenvalues of $\beta\mathbf{U}\mathbf{U}^{\text{\sf T}}$ have an empirical distribution converging to the density $\beta\mu_{\beta}$. We note the following integral which is useful in our analysis:

More generally, the Stieltjes transform of the density is given by: