Generalization Error of Generalized Linear Models in High Dimensions

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.

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.

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.

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):

where ${\bm{\xi}}_{\ell}$ are the following vectors:

and the functions ${\bm{\phi}}_{\ell}(\cdot)$ are given by

We see from Fig. 1 that the mapping of true parameters $\bm{w}^{0}=\mathbf{z}^{0}_{0}$ to the observed response vector $\bm{y}=\mathbf{z}^{0}_{3}$ is described by a multi-layer network of alternating orthogonal operators $\mathbf{V}_{\ell}$ and non-linear functions ${\bm{\phi}}_{\ell}(\cdot)$. Let $L=3$ denote the number of layers in this multi-layer network.

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 the penalty functions $F_{\ell}$ are defined as

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

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.

The algorithm operates in a set of iterations indexed by $k$. In each iteration, a “forward pass” through the layers generates estimates $\widehat{\mathbf{z}}_{k\ell}$ for the hidden variables $\mathbf{z}^{0}_{\ell}$, while a “backward pass” generates estimates $\widehat{\mathbf{p}}_{k\ell}$ for the variables $\mathbf{p}^{0}_{\ell}$. In each step, the estimates $\widehat{\mathbf{z}}_{k\ell}$ and $\widehat{\mathbf{p}}_{k\ell}$ are produced by functions $\mathbf{g}^{+}_{\ell}(\cdot)$ and $\mathbf{g}^{-}_{\ell}(\cdot)$ called estimators or denoisers.

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})$.

In the case of the GLM model, for $\ell=0$ and $L$, on lines 7 and 19, the denoisers are proximal operators given by

Note that in (25b), there is a dependence on $\bm{y}$ through the term $F_{\rm out}(\bm{y},\cdot)$. For the middle terms, $\ell=1,2$, i.e., lines 9 and 21, the denoisers are given by

where $(\widehat{\mathbf{p}}_{\ell\!-\!1},\widehat{\mathbf{z}}_{\ell})$ are the solutions to the minimization

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}$.

Thus, the ML-VAMP algorithm in Algorithm 1 reduces the joint constrained minimization (22) over variables $(\mathbf{z}_{0},\mathbf{z}_{1},\mathbf{z}_{2})$ and $(\mathbf{p}_{0},\mathbf{p}_{1},\mathbf{p}_{2})$ to a set of proximal operations on pairs of variables $(\mathbf{p}_{\ell\!-\!1},\mathbf{z}_{\ell})$. As discussed in (Pandit et al., 2019), this type of minimization is similar to ADMM with adaptive step-sizes. Details of the denoisers $\mathbf{g}_{\ell}^{\pm}$ and other aspects of the algorithm are given in Appendix B.

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.

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:

1. (1)

   i.i.d. features ($\sigma_{u}=0$);

2. (2)

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

3. (3)

   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.

We first consider the case of under-regularized linear regression where the output channel is $\phi_{\rm out}(p,d)=p+d$ with $d\sim{\mathcal{N}}(0,\sigma^{2}_{d})$. The noise variance $\sigma^{2}_{d}$ is set for an SNR level of 10 dB. We use a standard mean-square error (MSE) output loss, $f_{\rm out}(y,p)=(y-p)^{2}/(2\sigma^{2}_{d})$. Since we are using the L2-regularizer, $f_{\rm in}(w)=\lambda w^{2}/2$, the minimization (2) is standard ridge regression. Moreover, if we were to select $\lambda=1/\mathbb{E}(w_{j}^{0})^{2}$, then the ridge regression estimate would correspond to the minimum mean-squared error (MMSE) estimate of the coefficients $\bm{w}^{0}$. However, to study the under-regularized regime, we take $\lambda=(10)^{-4}/\mathbb{E}(w_{j}^{0})^{2}$.

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.

Fig. 3 shows a similar plot as Fig. 2 for a logistic model. Specifically, we use a logistic output $P(y=1)=1/(1+e^{-p})$, a binary cross entropy output loss $f_{\rm out}(y,p)$, and $\ell_{2}$-regularization level $\lambda$ so that the output corresponds to the MAP estimate (we do not perform ridgeless regression in this case). The mean of the training and test eigenvalues $\mathbb{E}S_{\rm tr}^{2}=\mathbb{E}S_{\rm ts}^{2}$ are selected such that, if the true coefficients $\bm{w}^{0}$ were known, we could obtain a 5% prediction error. As in the linear case, we generate random instances of the model, use the sklearn package to perform the logistic regression, and evaluate the estimates on 1000 new test samples. We compute the median error rate ($1-$ accuracy) and compare the simulated values with the SE theoretical estimates. The i.i.d. case was considered in (Salehi et al., 2019). Fig. 3 shows that our SE theory is able to predict the test error rate exactly in i.i.d. cases along with a correlated case and a case with training and test mismatch.

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

This output loss is non-convex. We scale the data matrix so that the input $\mathbb{E}(p^{2})=9$ so that the $\tanh(p)$ is driven well into the non-linear regime. We also take $\sigma^{2}_{d}=0.01$.

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.