High-dimensional ridge regression with random features for non-identically distributed data with a variance profile

Random feature ridge regression is often analyzed in the high-dimensional regime under the homogeneous sampling model $x_i={\mathrm{\Sigma }}^{1\mathrm{/}2}{x}_i^{\mathrm{\prime }}$, where the vectors $x_i^{\mathrm{\prime }}$ have iid entries and the same covariance matrix $\mathrm{\Sigma }$ is shared by all samples. In this paper, we move beyond this setting and study non-identically distributed data through a variance-profile model in which the training and test covariates have row-dependent diagonal covariance matrices ${\mathrm{\Sigma }}_i=\text{diag}({\gamma }_{i1}^2,\dots ,{\gamma }_{ip}^2)$ and ${\tilde{\mathrm{\Sigma }}}_i=\text{diag}({\tilde{\gamma }}_{i1}^2,\dots ,{\tilde{\gamma }}_{ip}^2)$. Our main contribution is the derivation of asymptotic equivalents for the training and test risks of ridge regression with random features when $n$, $p$, and $m$ grow proportionally. The first set of equivalents is obtained by combining the linear-plus-chaos approximation with traffic-probability arguments, whereas the second set is deterministic and follows from operator-valued free probability through an amalgamation-over-the-diagonal argument. These equivalents are sharp in numerical experiments. They also reveal how heterogeneous variance profiles, including mixture-type profiles inspired by MNIST, can modify generalization and exhibit double-descent behavior when the ridge parameter is small.