High-dimensional analysis of ridge regression for non-identically distributed data with a variance profile

Technical Summary: High-dimensional analysis of ridge regression for non-identically distributed data with a variance profile

Problem Statement
High-dimensional linear regression has been extensively analyzed under the assumption that predictors are independent and identically distributed (iid). However, in many practical scenarios, data exhibits heteroscedasticity where the variance of predictors varies across samples and features, violating the iid assumption. This paper addresses the statistical properties of ridge regression in a high-dimensional regime where the number of predictors $p$ and the number of observations $n$ grow to infinity at a proportional rate ($p/n \to c > 0$), but the data is independent yet non-identically distributed (non-iid). Specifically, the authors model the predictor matrix $X_n$ as the Hadamard product of a matrix with iid entries and a deterministic variance profile matrix $\Upsilon_n$, such that $X_n = \Upsilon_n \circ Z_n$. The goal is to derive deterministic equivalents for the predictive risk, training risk, and degrees of freedom (DOF) of the ridge estimator under this general variance profile, and to investigate how these profiles influence phenomena such as the "double descent" curve.

Methodology
The analysis relies on tools from Random Matrix Theory (RMT), specifically extending results regarding the spectral distribution of large random matrices with a variance profile.

1. Model Setup: The authors consider a random effect model where the regression coefficients $\beta^*$ are random with zero mean and covariance $\alpha^2 I_p/n$. The predictors are modeled as $X_n = \Upsilon_n \circ Z_n$, where $Z_n$ has iid centered entries with unit variance, and $\Upsilon_n$ contains the deterministic variance scaling factors. The variance profile is governed by $\Gamma_n = (\gamma_{ij}^2)$.
2. Resolvent Analysis: The core of the methodology involves the study of the resolvent matrices $Q_p(z) = (\hat{\Sigma}_n - zI_p)^{-1}$ and $\tilde{Q}_n(z) = (\tilde{\Sigma}_n - zI_n)^{-1}$, where $\hat{\Sigma}_n$ and $\tilde{\Sigma}_n$ are the empirical covariance matrices.
3. Deterministic Equivalents: Building upon the work of Hachem, Loubaton, and Najim [HLN07], the authors establish that the diagonal entries of the resolvent matrices converge to deterministic counterparts. They prove that the diagonal matrix of the resolvent $\Delta[Q_p(z)]$ is equivalent to a deterministic diagonal matrix $T_p(z)$, which is the solution to a system of fixed-point equations (Dyson equations). This extends previous RMT results which typically focused on the normalized trace (Stieltjes transform) rather than the diagonal entries required for risk analysis.
4. Risk Decomposition: Using these deterministic equivalents, the authors derive explicit expressions for the asymptotic predictive risk $\hat{r}^{test}_{\lambda}$ and training risk $\hat{r}^{train}_{\lambda}$, as well as the Generalized Cross-Validation (GCV) criterion.

Key Contributions

1. Deterministic Equivalents for Non-iid Data: The paper provides the first deterministic equivalents for the degrees of freedom and predictive risk of ridge regression in the presence of a general variance profile. Unlike the iid case where these quantities depend on the population covariance, here they depend on the specific structure of the variance profile $\Gamma_n$ via the solution to the Dyson equation.
2. Optimal Regularization Parameter: The authors prove that the optimal regularization parameter $\lambda^*$ that minimizes the predictive risk is universal across all variance profiles and is given by $\lambda^* = \sigma^2 p / (\alpha^2 n)$. This result matches the optimal scaling found in the standard iid setting [DW18], demonstrating the robustness of this shrinkage parameter against heteroscedasticity.
3. Generalization of Double Descent: The study investigates the behavior of the minimum-norm least-squares estimator (the limit as $\lambda \to 0$).
   • For quasi doubly stochastic variance profiles (where row and column sums of the variance profile are constant), the predictive risk recovers the classical double descent phenomenon observed in iid data.
   • For general variance profiles, the paper demonstrates that the shape of the predictive risk can deviate significantly from double descent. Numerical experiments reveal the emergence of "triple descent" and even "quadruple descent" phenomena for specific profiles (e.g., piecewise constant, block, and polynomial profiles).
4. Connection to Mixture Models: The framework is applied to data generated from mixture models (e.g., latent class variables), showing how the variance profile naturally arises when data comes from subpopulations with different covariance structures.

Results

• Risk Formulas: The predictive risk is shown to converge almost surely to a deterministic limit $r^{test}_{\lambda}$, which is a function of the signal strength $\alpha$, noise variance $\sigma^2$, the aspect ratio $c$, and the variance profile $\Gamma_n$ through the deterministic equivalent $T_p(z)$.
• Optimality: Corollary 3.1 confirms that $\lambda^* = \sigma^2 p / (\alpha^2 n)$ minimizes both the empirical and deterministic predictive risks, regardless of the variance profile.
• Risk Landscape: While the optimal $\lambda^*$ is universal, the shape of the risk curve as a function of the aspect ratio $c = p/n$ (when $\lambda=0$) is profile-dependent.
  • In the quasi doubly stochastic case, the risk exhibits the standard double descent: increasing for $c < 1$, peaking at $c=1$, and decreasing for $c > 1$.
  • In non-stochastic cases (e.g., block or piecewise constant profiles), the risk can exhibit multiple peaks (triple or quadruple descent). The authors link this behavior to the spectral properties of the empirical covariance matrix, specifically the behavior of the smallest non-zero eigenvalue near the interpolation threshold.
• Numerical Validation: Extensive numerical experiments using Pareto and Gaussian distributions confirm the accuracy of the derived deterministic equivalents. The theoretical curves match the empirical risks and GCV criteria closely across various variance profiles.

Significance
The paper claims to extend the theoretical understanding of high-dimensional ridge regression beyond the restrictive iid assumption. By providing deterministic equivalents for non-iid data, it offers a "calculus" for predicting performance in heterogeneous settings. The work highlights that while the optimal regularization scaling is robust (universal), the generalization landscape (specifically the double descent phenomenon) is sensitive to the underlying variance structure. This challenges the universality of the double descent phenomenon, showing it is not a guaranteed feature of high-dimensional regression but rather a property that depends on the specific variance profile of the data. The results provide a theoretical basis for analyzing regression in complex settings such as mixture models and offer tools for model selection (via GCV) in the presence of heteroscedasticity.