SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yatin Dandi, Ludovic Stephan, Florent Krzakala Idephics, EPFL, Switzerland
Bruno Loureiro
DI, Ecole Normale Superieure, Paris, France
Lenka Zdeborova
SPOC, EPFL, Switzerland
A recent line of work in high-dimensional statistics working under the Gaussian mixture hypothesis has led to a number of results in the context of empirical risk minimization, Bayesian uncertainty quantification, separation of kernel methods and neural networks, ensembling and fluctuation of random features. We provide rigorous proofs for the applicability of these results to a general class of datasets  $(\boldsymbol{x}_i, y_i)_{i=1,\dots,n}$  containing independent samples from mixture distribution  $\sum_{c \in \mathcal{C}} \rho_c P_c^{\boldsymbol{x}}$ . Specifically, we consider the hypothesis class of generalized linear models  $\hat{y} = F(\boldsymbol{\Theta}^\top \boldsymbol{x})$  and investigate the asymptotic joint statistics of a family of generalized linear estimators  $(\boldsymbol{\Theta}^{(1)}, \dots, \boldsymbol{\Theta}^{(M)})$ , obtained either from (a) minimizing an empirical risk  $\hat{R}_n^{(m)}(\boldsymbol{\Theta}^{(m)}; \boldsymbol{X}, \boldsymbol{y})$  or (b) sampling from the associated Gibbs measure  $\exp(-\beta n \hat{R}_n^{(m)}(\boldsymbol{\Theta}^{(m)}; \boldsymbol{X}, \boldsymbol{y}))$ . Our main contribution is to characterize under which conditions the asymptotic joint statistics of this family depend (on a weak sense) only on the means and covariances of the class conditional feature distribution  $P_c^{\boldsymbol{x}}$ . This allows us to prove the universality of different quantities of interest, including training and generalization errors, as well as the geometrical properties and correlations of the estimators.
A recurrent topic in high-dimensional statistics is the investigation of the typical properties of signal processing and machine learning methods on synthetic, i.i.d. Gaussian data, a scenario often known under the umbrella of Gaussian design [1, 2, 3, 4, 5]. A less restrictive assumption arises when considering that many machine learning tasks deal with data partitioned into a fixed number of classes. In these cases, the data distribution is naturally described by a mixture model, where each sample is generated conditionally on the class. In other words: data is generated by first sampling the class assignment and then generating the input conditioned on the class. Arguably the simplest example of such distributions is that of a Gaussian mixture, which shall be our focus in this work.
Gaussian mixtures are a popular model in high-dimensional statistics since, besides being an universal approximator, they often lead to mathematically tractable problems. Indeed, a recent line of work has analyzed the asymptotic performance of a large class of machine learning problems in the proportional high-dimensional limit under the Gaussian mixture data assumption, see e.g. [6, 7, 8, 9, 10, 11, 12]. The goal of the present work is to show that this assumption, and the conclusions derived therein, are far more general than previously anticipated.
A recent line of work [13, 14], initiated by the work of [15] for Kernel matrices, posits that for generalized linear estimation on non-linear feature maps satisfying certain regularity conditions and a "one-dimensional CLT", the data distribution can be replaced by equivalent Gaussian data without affecting the training and generalization errors. This was recently proven by [16] for empirical
risk minimization under the setup of strongly convex objectives, and extended to a larger class of objectives by [17].
However, there is strong empirical evidence that Gaussian universality holds in a more general sense [18]. First, existing results rely on the assumption of a target function depending on linear projections in the latent or feature space. This excludes the rich class of classification on mixture distributions, where the target function is given by the label. For such distributions, a more appropriate equivalent distribution is given by a mixture of Gaussians. Such a "Gaussian mixture equivalence" has been conjectured and used in existing works, such as [11, 12] and was found to closely approximately classification on real datasets.
Furthermore, equivalence with mixtures of Gaussians has been observed to hold not only for training, generalization errors but other quantities of the estimators such as overlaps, variance, etc. For instance, [19, 20] empirically observed that the equivalence holds even while considering the joint distribution of multiple uncertainty estimators or ensembles of multiple random feature mappings. This suggests the equivalence of the distributions of the minimizers themselves.
Our results fill these gaps and provide rigorous justification for the universality in all the aforementioned works. Namely, we show that the joint statistics of multiple generalized estimators obtained either from ERM or sampling on a mixture model asymptotically agrees (in a weak sense) with the statistics of estimators from the same class trained on a Gaussian mixture model with matching first and second order moments. Our main contributions are as follows:
- Through a generalization of recent developments in the Gaussian equivalence principle [14, 17, 16], we prove the universality of empirical risk minimization and sampling for a generic mixture distribution and an equivalent mixture of Gaussians. In particular, we show that a Gaussian mixture observed through a random feature map is also a Gaussian mixture in the high-dimensional limit, a fact used for instance (without rigorous justification) in [11, 12, 21, 22].
- A consequence of our results is that, with conditions on the matrix weights, data generated by conditional Generative Adversarial Networks (cGAN) behave as a Gaussian mixture when observed through the prism of generalized linear models (kernels, feature maps, etc...), as illustrated in Figs 1 and 2. This further generalizes the work of [23] that only considered the universality of Gram matrices for GAN generated data through the prism of random matrix theory.
- We construct a unified framework involving multiple sets of parameters arising from simultaneous minimization of different objectives as well as sampling from Gibbs distributions defined by the empirical risk. Through the design of suitable reductions and a convexity-based argument, we establish conditions for the asymptotic universality of arbitrary functions of the set of minimizers or samples from different Gibbs distributions (Theorem 4). For instance, it includes ensembling [20], and the uncertainty quantification and Bayesian setting assumed (without proof) in [19, 24].
- Finally, we show that for multi-class classification, the conditions leading to universality hold for a large class of functions of the minimizers, such as overlaps and sparsity measures, leading to the equivalence between their distributions of themselves, and provide a theorem for their weak convergence (Theorem 5).
- As a technical contribution of independent interest, our proof of Theorem 5 demonstrates a principled approach for leveraging existing results on the exact asymptotics for simple data distributions (such as for Gaussian mixture models in [12]) to prove the weak convergence and universality of the joint-empirical measure of the estimators and parameters (means, covariances) of the data-distribution.
Related works — Universality is an important topic in applied mathematics, as it motivates the scope of tractable mathematical models. It has been extensively studied in the context of random matrix theory [25, 26], signal processing problems [1, 2, 3, 27, 28, 29, 30, 31] and kernel methods [15, 32, 33]. Closer to us is the recent stream of works that investigated the Gaussian universality of the asymptotic error of generalized linear models trained on non-linear features, starting from single-layer random feature maps [34, 35, 16, 36, 37] and later extended to single-layer NTK [17] and deep random features [38]. These results, together with numerical observations that Gaussian universality holds for more general classes of features, led to the formulation of different Gaussian equivalence conjectures [13, 14, 18]. Our results build on the line of works on the proofs of these conjectures, through the use of the one-dimensional CLT (Central Limit Theorem), stated formally in [14] who proved it for random features of Gaussian data. We generalize this principle to a Gaussian equivalence conditioned on the cluster assignment in a mixture-model with a corresponding conditional 1d-CLT (Assumption 10).
A complementary line of research has investigated cases in which the distribution of the features is multi-modal, suggesting a Gaussian mixture universality class instead [39, 23, 40]. A bridge between these two lines of work has been recently investigated with random labels and teachers in [21, 22]. Our results provide rigorous extensions of Gaussian universality to the setups of mixture models as well as uncertainty quantificatio and ensembling.
1 Setting and models
Consider a supervised learning problem where the training data  $(\pmb{x}_i, y_i) \in \mathbb{R}^p \times \mathcal{V}$ ,  $i \in [n] := \{1, \dots, n\}$  is drawn i.i.d. from a mixture distribution:
$$
\boldsymbol {x} _ {i} \sim \sum_ {c \in \mathcal {C}} \rho_ {c} P _ {c} ^ {\boldsymbol {x}}, \quad \mathbb {P} (c _ {i} = c) = \rho_ {c}, \tag {1}
$$
with  $c_{i}$  a categorical random variable denoting the cluster assignment for the  $i_{th}$  example  $\mathbf{x}_i$ . Let  $\pmb{\mu}_{c},\pmb{\Sigma}_{c}$  denote the mean and covariance of  $P_{c}^{\pmb{x}}$ , and  $k = |\mathcal{C}|$ . Further, assume that the labels  $y_{i}$  are generated from the following target function:
$$
y _ {i} (\boldsymbol {X}) = \eta \left(\boldsymbol {\Theta} _ {\star} ^ {\top} \boldsymbol {x} _ {i}, \varepsilon_ {i}, c _ {i}\right), \tag {2}
$$
where  $\eta : \mathbb{R}^3 \to \mathbb{R}$  is a general label generating function,  $\Theta_{\star} \in \mathbb{R}^{k \times p}$  and  $\varepsilon_i$  is an i.i.d source of randomness. It is important to stress that the class labels (2) are themselves not constrained to arise from a simple function of the inputs  $x_i$ . For instance, the functional form in (2) includes the case where the labels are exclusively given by a function of the mixture index  $y_i = \eta(c_i)$ . This will allow us to handle complex targets, such as data generated using conditional Generative Adversarial Networks (cGANs).
In this manuscript, we will be interested in hypothesis classes defined by parametric predictors of the form  $y_{\Theta}(\boldsymbol{x}) = F(\boldsymbol{\Theta}^{\top}\boldsymbol{x})$ , where  $\boldsymbol{\Theta} \in \mathbb{R}^{k \times p}$  are the parameters and  $F: \mathbb{R}^k \to \mathcal{V}$  a possibly nonlinear function. For a given loss function  $\ell: \mathbb{R}^k \times \mathcal{V} \to \mathbb{R}_+$  and regularization term  $r: \mathbb{R}^{k \times p} \to \mathbb{R}_+$ , define the (regularized) empirical risk over the training data:
$$
\widehat {\mathcal {R}} _ {n} (\boldsymbol {\Theta}; \boldsymbol {X}, \boldsymbol {y}) := \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(\boldsymbol {\Theta} ^ {\top} \boldsymbol {x} _ {i}, y _ {i}\right) + r (\boldsymbol {\Theta}), \tag {3}
$$
where we have defined the feature matrix  $\mathbf{X} \in \mathbb{R}^{p \times n}$  by stacking the features  $\mathbf{x}_i$  column-wise and the labels  $y_i$  in a vector  $\mathbf{y} \in \mathcal{V}^n$ . In what follows, we will be interested in the following two tasks:
(i) Minimization: in a minimization task, the statistician's goal is to find a good predictor by minimizing the empirical risk (3), possibly over a constraint set  $S_{p}$ :
$$
\hat {\Theta} _ {\operatorname {e r m}} (\boldsymbol {X}, \boldsymbol {y}) \in \underset {\boldsymbol {\Theta} \in \mathcal {S} _ {p}} {\arg \min } \widehat {\mathcal {R}} _ {n} (\boldsymbol {\Theta}; \boldsymbol {X}, \boldsymbol {y}), \tag {4}
$$
This encompasses diverse settings such as generalized linear models with noise, two-layer networks with a constant number of neurons and fixed second layer, mixture classification, but also the random label setting (with  $\eta (\Theta_{\star}^{\top}\pmb{x}_i,\varepsilon_i,c_i) = \varepsilon_i$ ). In the following, we denote  $\widehat{\mathcal{R}}_n^\star (X,y)\coloneqq \min_\Theta \widehat{\mathcal{R}}_n(\Theta ;X,y)$
(ii) Sampling: here, instead of minimizing the empirical risk (3), the statistician's goal is to sample from a Gibbs distribution that weights different hypothesis according to their empirical error:
$$
\Theta_ {\text {B a y e s}} (\boldsymbol {X}, \boldsymbol {y}) \sim P _ {\text {B a y e s}} (\boldsymbol {\Theta}) \propto \exp \left(- \beta n \widehat {\mathcal {R}} _ {n} (\boldsymbol {\Theta}; \boldsymbol {X}, \boldsymbol {y})\right) \mathrm {d} \mu (\boldsymbol {\Theta}) \tag {5}
$$
where  $\mu$  is reference prior measure and  $\beta > 0$  is a parameter known as the inverse temperature. Note that minimization can be seen as a particular example of sampling when  $\beta \to \infty$ , since in this limit the above measure peaks on the global minima of (4).
Applications of interest— So far, the setting defined above is quite generic, and the motivation to study this problem might not appear evident to the reader. Therefore, we briefly discuss a few scenarios of interest which are covered by this model.
Figure 1: Illustration of Corollary 2: high-dimensional data generated by generative neural networks starting from a mixture of Gaussian in latent space  $(\mathbf{z} \in \mathbb{R}^H)$  are (with conditions on the weights matrices) equivalent, in high-dimension and for generalized linear models, to data sampled from a Gaussian mixture. A concrete example is shown in Fig. 2.
(i) Conditional GANs (cGANs): These were introduced by [41] as a generative model to learn mixture distributions. Once trained in samples from the target distribution, they define a function  $\Psi$  that maps Gaussian mixtures (defining the latent space) to samples from the target mixture that preserve the label structure. In other words, conditioned on the label:
$$
\forall c \in \mathcal {C}, \quad \boldsymbol {z} \sim \mathcal {N} \left(\boldsymbol {\mu} _ {c}, \boldsymbol {\Sigma} _ {c}\right) \mapsto \boldsymbol {x} _ {c} = \Psi (\boldsymbol {z}, c) \sim P _ {c} ^ {\boldsymbol {x}} \tag {6}
$$
The connection to model (1) is immediate. This scenario was extensively studied by [39, 23, 40], and is illustrated in Fig. 1. In Fig. 2 we report on a concrete experiment with a cGAN trained on the fashion-MNIST dataset.
(ii) Multiple objectives: Our framework also allows to characterize the joint statistics of estimators  $(\Theta_{1},\dots ,\Theta_{M})$  obtained from empirical risk minimization and/or sampling from different objective functions  $\hat{R}_n^m$  defined on the same training data  $(X,y)$ . This can be of interest in different scenarios. For instance, [19, 24] has characterized the correlation in the calibration of different uncertainty measures of interest, e.g. last-layer scores and Bayesian training of last-layer weights. This crucially depends on the correlation matrix  $\hat{\Theta}_{\mathrm{erm}}\Theta_{\mathrm{Bayes}}^{\top}\in \mathbb{R}^{k\times k}$  which fits our framework.
(iii) Ensemble of features: Another example covered by the multi-objective framework above is that of ensembling. Let  $(z_{i},y_{i})\in \mathbb{R}^{d}\times \mathcal{V}$  denote some training data from a mixture model akin to (1). A popular ensembling scheme often employed in the context of deep learning [42] is to take a family of  $M$  feature maps  $\pmb {z}_i\mapsto \pmb {x}_i^{(m)} = \varphi_m(\pmb {z}_i)$  (e.g. neural network features trained from different random initialization) and train  $M$  independent learners:
$$
\hat {\boldsymbol {\Theta}} _ {\mathrm {e r m}} ^ {(m)} \in \underset {\boldsymbol {\Theta} \in \mathcal {S} _ {p}} {\arg \min } \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(\boldsymbol {\Theta} ^ {\top} \boldsymbol {x} _ {i} ^ {(m)}, y _ {i}\right) + r (\boldsymbol {\Theta}) \tag {7}
$$
Prediction on a new sample  $\mathbf{z}$  is then made by ensembling the independent learners, e.g. by taking their average  $\hat{\mathbf{y}} = 1 / M\sum_{m = 1}^{M}\hat{\Theta}_{\mathrm{erm}}^{(m)\top}\varphi_{m}(\mathbf{z})$ . A closely related model was studied in [43, 44, 20].
Note that in all the applications above, having the labels depending on the features  $\mathbf{X}$  would not be natural, since they are either generated from a latent space, as in (i), or chosen by the statistician, as in (ii), (iii). Indeed, in these cases the most natural label model is given by the mixture index  $y = c$  itself, which is a particular case of (2). This highlights the flexibility of our target model with respect to prior work [17]. Instead, [16] assumes that the target is a function of a latent variable, which would correspond to a mismatched setting. The discussion here can be generalized also to this case, but require an additional assumption discussed in Appendix B.
Universality — Given these tasks, the goal of the statistician is to characterize different statistical properties of these predictors. These can be, for instance, point performance metrics such as the empirical and population risks, or uncertainty metrics such as the calibration of the predictor or moments of the posterior distribution (5). These examples, as well as many different other quantities of interest, are functions of the joint statistics of the pre-activations  $(\Theta_{\star}^{\top}\boldsymbol{x},\Theta^{\top}\boldsymbol{x})$ , for  $\boldsymbol{x}$  either a test or
Figure 2: Illustration of the universality scenario described in Fig.1. Logistic (left) & ridge (right) regression test (up) and training (bottom) errors are shown versus the sample complexity  $\alpha = n / d$  for an odd vs. even binary classification task on two data models: Blue dots data generated from a conditional GAN [41] trained on the fashion-MNIST dataset [45] and pre-processed with a random features map  $x \mapsto \tanh(Wx)$  with Gaussian weights  $W \in \mathbb{R}^{1176 \times 784}$ ; Red dots are the 10- clusters Gaussian mixture model with means and covariances matching each fashion-MNIST cluster conditioned on labels ( $\ell_2$  regularization is  $\lambda = 10^{-4}$ ). Details on the simulations are discussed in Appendix D.
training sample from (1). For instance, in a Gaussian mixture model, where  $\pmb{x} \sim \sum_{c \in \mathcal{C}} \rho_c \mathcal{N}(\pmb{\mu}_c, \pmb{\Sigma}_c)$ , the sufficient statistics are simply given by the first two moments of these pre-activations. However, for a general mixture model (1), the sufficient statistics will generically depend on all moments of these pre-activations. Surprisingly, our key result in this work is to show that in the high-dimensional limit this is not the case. In other words, under some conditions which are made precise in Section 2, we show that expectations with respect to (1) can be exchanged by expectations over a Gaussian mixture with matching moments. This can be formalized as follows. Define an equivalent Gaussian data set  $(\pmb{g}_i, y_i)_{i=1}^n \in \mathbb{R}^p \times \mathcal{V}$  with samples drawn i.i.d. from the equivalent Gaussian mixture model:
$$
\boldsymbol {g} _ {i} \sim \sum_ {c \in \mathcal {C}} \rho_ {c} \mathcal {N} \left(\boldsymbol {\mu} _ {c}, \boldsymbol {\Sigma} _ {c}\right), \quad y _ {i} (\boldsymbol {G}) = \eta \left(\boldsymbol {\Theta} _ {\star} ^ {\top} \boldsymbol {g} _ {i}, \varepsilon_ {i}, c _ {i}\right). \tag {8}
$$
We recall that  $\mu_c, \Sigma_c$  denotes the mean and covariance of  $P_c^x$  from (1). Consider a family of estimators  $(\Theta_1, \dots, \Theta_M)$  defined by minimization (3) and/or sampling (5) over the training data  $(X, y)$  from the mixture model (1). Let  $h$  be a statistical metric of interest. Then, in the proportional high-dimensional limit where  $n, p \to \infty$  at a fixed  $\alpha = n / d > 0$ , and where  $\langle \cdot \rangle$  denote the expectation with respect to the Gibbs distribution (5), we define universality as:
$$
\mathbb {E} _ {\boldsymbol {X}} \left[ \langle h \left(\boldsymbol {\Theta} _ {1}, \dots , \boldsymbol {\Theta} _ {M}\right) \rangle_ {\boldsymbol {X}} \right] \underset {n \rightarrow \infty} {\simeq} \mathbb {E} _ {\boldsymbol {G}} \left[ \langle h \left(\boldsymbol {\Theta} _ {1}, \dots , \boldsymbol {\Theta} _ {M}\right) \rangle_ {\boldsymbol {G}} \right] \tag {9}
$$
The goal of the next section is to make this statement precise.
We now present the main theoretical contributions of the present work and discuss its consequences. Our proofs for Theorems 4 and 6 build upon existing results on the universality of empirical risk minimization for uni-model distributions [16, 17] and therefore rely on similar technical regularity and concentration assumptions. Concretely, our work relies on the following assumptions:
Assumption 1 (Loss and regularization). The loss function  $\ell : \mathbb{R}^{k+1} \to \mathbb{R}$  is nonnegative and Lipschitz, and the regularization function  $r: \mathbb{R}^{p \times k} \to \mathbb{R}$  is locally Lipschitz, with constants independent from  $p$ .
Assumption 2 (Boundedness and concentration). The constraint set  $\mathcal{S}_p$  is a compact subset of  $\mathbb{R}^{k \times p}$ . Further, there exists a constant  $M > 0$  such that for any  $c \geq 0$ ,
$$
\sup  _ {\boldsymbol {\theta} \in \mathcal {K} _ {p}, \| \boldsymbol {\theta} \| _ {2} \leq 1} \| \boldsymbol {\theta} ^ {\top} \boldsymbol {x} \| _ {\psi_ {2}} \leq M, \quad \sup  _ {\boldsymbol {\theta} \in \mathcal {K} _ {p}, \| \boldsymbol {\theta} \| _ {2} \leq 1} \| \boldsymbol {\Sigma} _ {c} ^ {1 / 2} \boldsymbol {\theta} \| _ {2} \leq M, \quad a n d \quad \| \boldsymbol {\mu} _ {c} \| _ {2} \leq M \tag {10}
$$
where  $\| \cdot \|_{\psi_2}$  is the sub-gaussian norm, and  $\mathcal{K}_p \subseteq \mathbb{R}^p$  is such that  $\mathcal{S}_p \subseteq \mathcal{K}_p^k$ .
Assumption 3 (Labels). The labeling function  $\eta$  is Lipschitz, the teacher vector  $\Theta$  belongs to  $S_{p}$  and the noise variables  $\varepsilon_{i}$  are i.i.d sub-gaussian with  $\| \varepsilon_i\|_{\psi_2}\leq M$  for some constant  $M > 0$ .
Those three assumptions are fairly technical, and it is possible that the universality properties proven in this article hold irrespective of these conditions. The crucial assumption in our theorems is that of a conditional one-dimensional CLT:
Assumption 4. For any Lipschitz function  $\varphi :\mathbb{R}\to \mathbb{R}$
$$
\lim  _ {n, p \rightarrow \infty} \sup  _ {\boldsymbol {\theta} \in \mathcal {K} _ {p}} \left| \mathbb {E} \left[ \varphi \left(\boldsymbol {\theta} ^ {\top} \boldsymbol {x}\right) \mid c _ {\boldsymbol {x}} = c \right] - \mathbb {E} \left[ \varphi \left(\boldsymbol {\theta} ^ {\top} \boldsymbol {g}\right) \mid c _ {\boldsymbol {g}} = c \right]\right| = 0, \quad \forall c \in \mathcal {C} \tag {11}
$$
where  $\pmb{x}$  and  $\pmb{g}$  denote samples from the given mixture distribution and the equivalent gaussian mixture distribution in equations (1) and (8) respectively.
The above assumption is a generalization of the "one-dimensional CLT" underlying the line of work based on the Gaussian equivalence (GE) Principle [13, 14, 17, 16]. The above assumption splits the proof of universality for a general mixture distribution into two parts. First, one shows that asymptotic universality of an observable  $h$  can be reduced to the proof of a one-dimensional CLT. Second, one proves this CLT holds for the particular class of features of interest. This proof scheme streamlines universality proofs. Our work provides a general proof of the first step in Theorem 4, conditioned on the second, later showing that Assumption 4 holds for some natural feature maps of interest, i.e. random features applied to a Gaussian mixture (Theorem 6). However, Assumption 4 is expected to hold for a large class of features, as supported by our empirical observations in Figure 2 and arguments in Appendix C.
2.1 Universality of Mixture Models