Technical Summary: The Mechanism of Weak-to-Strong Generalization

Problem Statement

Weak-to-Strong (W2S) generalization describes the phenomenon where a strong, pre-trained model, when fine-tuned on the outputs of a weaker, task-specialized model, can surpass the weak model's performance on that specific task while retaining its broader capabilities. While empirical studies (e.g., Burns et al., 2023) have demonstrated that full fine-tuning (feature learning) outperforms linear probing (fixed representations), existing theoretical analyses have largely been restricted to abstract frameworks or linear/random feature models that do not derive feature learning from the optimization process itself.

The central open question addressed by this paper is whether multi-step Stochastic Gradient Descent (SGD) can provably succeed in W2S feature learning—specifically, learning the target feature direction while preserving diverse pre-trained capabilities—in a setting that more closely resembles pre-trained Large Language Models (LLMs).

Methodology and Problem Setting

The authors formalize W2S in the context of reward-model learning using two-layer neural networks. The setup involves:

1. Additive Reward Model: The true reward $r^*(x)$ is modeled as a sum of $K$ single-index components, $r^*(x) = \sum_{k=1}^K \pi_k \sigma^*_k(\theta_k^\top x)$, where each task $k$ corresponds to a specific feature direction $\theta_k$.
2. Pre-trained Subspaces: Consistent with empirical findings that LLMs encode concepts in low-dimensional subspaces, the strong model is assumed to have pre-trained knowledge of task-specific subspaces $V_k$ of dimension $s \ll d$. The strong model knows the subspace $V_k$ containing $\theta_k$ but not the exact direction $\theta_k$ itself.
3. Weak Model: A small, two-layer neural network specialized to a target task $\kappa$. It is trained to approximate the reward component $r^*_\kappa$ and possesses feature alignment with $\theta_\kappa$.
4. Strong Model: A larger two-layer network with neurons partitioned into groups corresponding to the $K$ tasks. It is initialized with neurons localized to their respective subspaces.
5. Training Dynamics:
   • W2S Algorithm: The strong model is trained via online SGD on inputs drawn from its own generative distribution (concentrated on the target subspace $V_\kappa$) and supervised by the weak model's outputs. Crucially, the weak model's output undergoes a nonlinear transformation ($z \mapsto z \exp(z)$) before being used as a label.
   • Standard SFT (Baseline): For comparison, the strong model is trained via online SGD on isotropic inputs ($x \sim N(0, I_d)$) with direct labels from the true reward function $r^*_\kappa$.

Key Contributions and Theoretical Results

1. Feature Elicitation via Multi-Step SGD (Theorem 1)

The paper provides the first multi-step SGD convergence guarantee for W2S in the feature-learning regime.

• Mechanism: The nonlinear transformation of the teacher signal reduces the effective "generative exponent" of the learning problem. While the target link function $\sigma^*_\kappa$ may have a high information exponent $p > 2$, the transformation allows the strong model to learn the feature direction $\theta_\kappa$ with a sample complexity dependent on the generative exponent (which becomes 2 for even polynomials).
• Sample Complexity: The strong model achieves alignment with the target feature $\theta_\kappa$ with a sample complexity of $\tilde{O}(s^{3/2})$, where $s$ is the dimension of the task subspace. This is significantly lower than the $\tilde{O}(d^{p/2} s^{(p-2)/2})$ complexity required for the weak model to learn the same feature from scratch, and lower than the $\tilde{O}(d^{p-1})$ complexity of standard SFT under isotropic sampling.
• Subspace Prior: The result relies on the strong model's input distribution being concentrated on the target subspace $V_\kappa$ (parameterized by $\lambda_\kappa \approx 1$), which effectively reduces the learning dimension from $d$ to $s$.

2. Feature Preservation (Proposition 2)

A critical contribution is the proof that W2S preserves pre-trained features for off-target tasks ($k \neq \kappa$).

• Implicit Regularization: The combination of the strong model's generative distribution (concentrated on $V_\kappa$) and the nonlinear teacher transformation acts as an implicit regularizer.
• Condition for Preservation: Off-target features are preserved provided the off-target tasks are allocated more neurons than the target task (or carry smaller reward weights), quantified by the localization parameter $\chi_k = N_\kappa \pi_k / (N_k \pi_\kappa)$.
• Result: The alignment of off-target neurons with their original directions $\theta_k$ is maintained, with the drop in alignment bounded by terms related to initialization overlap and SGD noise, rather than catastrophic forgetting.

3. Catastrophic Forgetting in Standard SFT (Theorem 2)

In contrast to W2S, the paper proves that standard Supervised Fine-Tuning (SFT) leads to catastrophic forgetting of off-target features when those features share directionality with the target.

• Mechanism: Under isotropic sampling, the gradient signal propagates to all neurons. If the initial alignment between an off-target neuron and the target direction $\theta_\kappa$ is non-negligible (specifically, if $\theta_k^\top \theta_\kappa \gtrsim s^{-1/2}$), the off-target neurons converge to $\theta_\kappa$ faster than the target neurons converge to $\theta_\kappa$.
• Outcome: The pre-trained feature $\theta_k$ is lost before the target task is fully learned, resulting in the destruction of the model's general capabilities.

Numerical Experiments

The authors validate their theoretical findings using synthetic data with input dimension $d=1024$ and subspace dimension $s=128$.

• W2S Performance: Experiments confirm that the strong model learns the target feature direction ($\theta_1$) while maintaining alignment with the off-target feature ($\theta_2$), consistent with Theorem 1 and Proposition 2.
• SFT Failure: In the SFT setting, the off-target neurons align with the target direction $\theta_1$ before the target neurons fully converge, confirming the theoretical prediction of catastrophic forgetting (Theorem 2).

Significance and Claims

The paper claims to establish a theoretical foundation for W2S generalization in the feature-learning regime, moving beyond previous analyses that assumed fixed representations or linear models.

• Mechanism of Elicitation: It demonstrates that W2S does not merely "instill" new capabilities but effectively "elicits" latent knowledge by leveraging pre-trained subspaces and nonlinear transformations to reduce sample complexity.
• Advantage over SFT: The work provides a theoretical account for why W2S outperforms standard SFT in preserving general capabilities. It attributes this to the algorithmic difference: W2S queries the weak model on the strong model's generative distribution (concentrated on the task subspace), whereas SFT relies on isotropic sampling which causes interference between correlated features.
• Limitations: The authors note that their analysis focuses on first-layer feature alignment and does not formally analyze second-layer learning (weight updates for the output layer). They also acknowledge that the generative exponent reduction requires the target link function to be an even polynomial.

In summary, the paper argues that W2S succeeds because the strong model's pre-trained structure (subspace localization) combined with the specific training protocol (generative distribution + nonlinear transformation) allows it to learn the target feature efficiently without disrupting the orthogonal or weakly correlated features that encode its broader knowledge.