Weak-to-Strong Generalization Even in Random Feature Networks, Provably

We thank the reviewer for their positive feedback. Below we address the main comments:

-The reviewer asks “Why at all introduce the teacher”? Indeed, in the weak-to-strong setup as introduced by Burns, and as we study here, training directly on the target $f^*$ would give better performance than weak-to-strong training. This is what Burns refers to as the “ceiling performance”, denoted by $L_{ST}^{ceil}$ in the definition of the PGR (our equation 9). The weak-to-strong question, as introduced by Burns, asks how training with only the teacher labels, and not the true target, can get close to this ceiling performance. This captures a situation where the target is not available when training the strong student, only the weak teacher. See Burns et al for further discussion. What we are investigating are scenarios in which even if the student is trained only through the teacher, the student is better than the teacher (but not better than the ceiling performance!). Unsurprisingly, this happens when the target is aligned in some way with the inductive bias of the student (and teacher), and the way we capture this in order to obtain crisp possibility results is through the eigenspace assumption (though our general results do not need this specific assumption and are more general).

-Gradient flow: gradient flow can be thought of simply as gradient updates with very small step size (formally, when the step sizes goes to zero), see e.g. “A Continuous-Time View of Early Stopping for Least Squares” by Ali et al 2019 or “Continuous vs. Discrete Optimization of Deep Neural Networks” by Elkabetz et al 2021. Using standard results on the convergence of gradient descent to gradient flow, it is possible to obtain the same results also for early stopped gradient descent, with an additional error term that depends on the stepsize, and can be made as small as we would like. Studying how using large stepsizes effects weak-to-strong generalization, and how small the stepsize needs to be to realize weak-to-strong, is indeed an interesting and difficult problem (this follows the path of many results on the implicit bias of gradient descent, which are also much better understood in the gradient flow, ie small stepsize, limit).

-Regarding the $\kappa_S$: Effective rank is a property of the student’s kernel only, i.e. the eigenvalues of each of the eigendirections of the kernel function. $\kappa_s$ measures, roughly, how much of the teacher error is contained in the $\{e_{\ge S+1}\}$ eigendirections. $\kappa_s$ is a quantity that depends on the projection of the teacher features onto the eigendirections, and is not related to the student at all. So there is no a-priori connection - the eigenvalues of the student kernel do not factor in $\kappa_s$. This can be seen from the fact that $\kappa_S$ and the student eigenvalue factor separately into the final result.

-Regarding the optimal stopping time: we indeed do not have a simple expression for the optimal stopping time. What we show is that with some stopping time we can get W2S generalization. We do however have some insight into the optimal stopping time through the experimental results. We agree that understanding the optimal stopping time is an interesting question. Empirically, optimal stopping can often be achieved through validation.

We thank the reviewer for their positive feedback. Below we address the main comments:

-Regarding $M_{TE}\to\infty$: First, it is important to emphasize that our results are not asymptotic and we show a gap for finite teacher width $M_{TE}$, not only as $M_{TE}\to\infty$. Indeed, we also discuss $M_{TE}\to\infty$, and this is where the gap is largest. Even as $M_{TE}\to\infty$, this is still very much a weak-to-strong setting: the teacher becomes larger, but the student is proportionally even larger. We actually consider this the more interesting and relevant setting, as it captures using an already very large teacher (like GPT2 or 3, which as statistical models go, already has an extremely large number of parameters, well in the “high dimensional” regime) to train an even larger student (eg GPT4). In short, we always consider students which are much larger than their teachers, i.e. $M_{ST} \gg M_{TE}$, and even when we take both widths to be very large with $M_{ST}, M_{TE}\to \infty$ we still have $M_{ST} \gg M_{TE}$.

-Regarding the Gaussian Universality Ansatz: First, we want to emphasize that this is only used in the second part of Theorem 3.1, in order to easily get a lower bound on the teacher, and in a sense is not crucial for the main weak-to-strong effect, which is the improvement in the student as specified in the first part of Theorem 3.1. Regarding the difference between the ReLU setting with the Ansatz and the linear setting: the main difference is that the ReLU network imposes a particular structure on the spectrum of the kernel, which is different from that of the linear model. The difference in exponent is due to the difference in this spectral structure. We can expand on this in the final version of the paper.

-Regarding section 5: As the reviewer pointed out, the goal of section 5 is explain the how weak-to-stron generalization happens in this case. The main takeaway is understanding the implications of how the trajectory of the student predictor behaves in function space. This can be also seen from the closed-form equation of GF, i.e. Equation (13) on page 7. Namely, once we understand the directions along which GF (and approximately GD) shrinks $f_{TE}-f_{ST}$ and the fact that the shrinkage is faster for directions with larger covariance, we can explain the weak-to-strong phenomenon in this case. As it turns out, the relevant directions are eigendirections and they are the coordinate directions for linear networks and spherical harmonics for ReLU networks. How fast the shrinkage happens is determined by the eigenvalue corresponding to the given eigendirection. Therefore, if there is a big eigengap between $e_K$ and $e_{K+1}$, then we can learn the teacher predictor along $e_1,..,e_K$ much quicker than any of the $e_{\ge K+1}$, effectively zeroing out the signal along any of the $e_{\ge K+1}$. We will update section 5 to clarify this in our final version.

We thank the reviewer for their positive feedback.

Full details on the data distributions used in the simulations are provided in the figure captions. These are synthetic data distributions that match Theorems 3.1 and 3.2. More explicitly:

In Figure 2, we use a ReLU networks with the input dimension $d=32$. The student has $M_{ST} = 16384$ hidden units, and different curves correspond to teacher with $M_{TE} \in \{16,...,256\}$ hidden units. The inputs a uniformly distributed on a unit-radius sphere $S^{d-1}$ and the target is given by $f*=\langle \beta,x\rangle$ for a unit norm beta.

In Figure 3, we consider linear networks where the number of hidden units for the student is always $M_{ST}=16384$, and teachers with varying number of hidden units $M_{TE}$ between $16$ and $4096$. The input dimension is $d=M_{TE}^{3/2}$ (depending on the teacher width), as in Theorem 3.1. The input follows a non-spherical Gaussian distribution $x\sim N(0,\Psi)$ where $\Psi = (1,(d-1)^{-2/3},...,(d-1)^{-2/3})$, as in Theorem 3.1, and the target is $f^*=\langle e_1,x\rangle$.

We thank the reviewer for their positive feedback. Below we address the main comments:

-Regarding linear networks achieving optimal exponent: With the linear network model, we choose the covariance directly and so can impose an extreme spectral structure, namely a huge eigengap. This makes our learning algorithm (GF with early stopping) the ideal truncating algorithm — zeroing out all projections onto the eigendirections with low eigenvalues but exactly keeping the projections into the eigendirections with large eigenvalues. This turns out to be sufficient for matching the lower bound. But ReLU networks with spherical Gaussian random weights have a specific spectral structure, which does lead to significant weak-to-strong behavior (as we analyze), but is not as extreme.

-Regarding $M_{TE}\to \infty$: We take $M_{ST}\to\infty$ first, so $M_{TE}/M_{ST}$ should be thought of as small, i.e $o(1)$. Our results do not use the simultaneous limit, i.e. taking $M_{ST}$ and $M_{TE}$ to infinity together.

-Regarding analysis with finite data: It is certainly possible to obtain similar results also when the teacher is trained with a large but finite amount of data. This is possible with a large enough amount of data such that standard generalization guarantees would ensure the teacher is close to its population optimum. There isn’t a significant technical difficulty here, just bookkeeping and additional terms to track and a corresponding complication in the Theorem statements and proofs. We choose to not get into this complication as we felt it does not provide additional insight (the separation and gaps would be the same), would unnecessarily increase complexity of presentation, and that a weak-to-strong effect even if the teacher does has infinite data (and so is in a sense less weak) is in a sense even stronger.

-Ildiz et al also study a two-step stage learning process, first training a linear teacher and then a linear student based on the teacher. But the effect they show is that of “distillation”, i.e. how training the student based on teacher labels is better than training the student based on real labels. This is different from a weak-to-strong effect, where we compare the student trained on teacher labels to the teacher itself. In their Section 3.1 and Proposition 3, Ildiz et al do study a small teacher that, like our weak teacher, has fewer features than the student. But what they show is very different: even in this subsection (which they title “weak to strong”) they show how student-trained-on-teacher-labels can be better than student-trained-on-real-labels (as opposed to weak-to-strong, which is student-trained-on-teacher-labels is better than teacher). Even in their setting of Section 3.1, it is NOT the case that the student-trained-on-teacher-labels is better than the teacher, i.e. there is no “weak to strong” effect in the sense of Burns et al. The type of distillation benefit they show (which we again emphasize is different from the weak-to-strong effect we study) is possible due to the following differences from our work:

--Idliz et al train with either ERM or min-norm-interpolating predictor, not with early stopping. We emphasize that early stopping is crucial for weak-to-strong generalization. What they can show is that the student-trained-on-teacher labels can be better than the student-trained-on-real-labels if the training data is finite, but this does not mean the student can be better than the teacher. In fact, if we increase the number of data to infinity, the student will eventually become the same as the teacher, while in our setting early stopping enables the student to perform better than the teacher.

--The teacher uses a carefully selected subset of features, as opposed to random features.

Similarities include the use of linear models, the two-stage learning process (though without early stopping) and training the teacher model on infinite data.

-Regarding the random features and gradient flow: Our main claim for random features factors into Theorems 6.5 and Theorems 6.9 and is a statement about the random feature features (in the teacher’s model) project onto the top-S eigenbasis. For this, we apply a general Hilbert Space concentration result to a carefully constructed Hilbert space to prove this statement. We are not aware of prior works having a similar approach. For the gradient flow reference, we agree that it is appropriate to include additional works. We will add references in the final version, including “A Continuous-Time View of Early Stopping for Least Squares” by Ali et al 2019 or “Continuous vs. Discrete Optimization of Deep Neural Networks” by Elkabetz et al 2021.