技术摘要：弱到强泛化的机制

问题陈述

弱到强（W2S）泛化描述了一种现象：一个强大的预训练模型在针对较弱、任务专用的模型输出进行微调后，不仅能在该特定任务上超越弱模型的性能，还能保留其更广泛的能力。尽管实证研究（例如 Burns 等人，2023）表明全量微调（特征学习）优于线性探测（固定表示），但现有的理论分析大多局限于抽象框架或线性/随机特征模型，未能从优化过程本身推导出特征学习。

本文解决的核心开放问题是：在更贴近预训练大语言模型（LLM）的设定下，多步随机梯度下降（SGD）能否在理论上证明在 W2S 特征学习（即学习目标特征方向同时保留多样化的预训练能力）方面取得成功。

方法论与问题设定

作者在奖励模型学习的背景下，利用两层神经网络形式化了 W2S。该设定包括：

1. 加性奖励模型：真实奖励 $r^*(x)$ 被建模为 $K$ 个单索引分量的总和，$r^*(x) = \sum_{k=1}^K \pi_k \sigma^*_k(\theta_k^\top x)$，其中每个任务 $k$ 对应一个特定的特征方向 $\theta_k$。
2. 预训练子空间：与 LLM 将概念编码在低维子空间中的实证发现一致，假设强模型已预训练了任务特定的子空间 $V_k$，其维度 $s \ll d$。强模型知道包含 $\theta_k$ 的子空间 $V_k$，但不知道 $\theta_k$ 本身的确切方向。
3. 弱模型：一个针对目标任务 $\kappa$ 专门化的小型两层神经网络。它被训练以近似奖励分量 $r^*_\kappa$，并具有与 $\theta_\kappa$ 的特征对齐。
4. 强模型：一个更大的两层网络，其神经元被划分为对应 $K$ 个任务的组。它被初始化为神经元定位于各自的子空间。
5. 训练动态：
   • W2S 算法：强模型通过在线 SGD 进行训练，输入来自其自身的生成分布（集中在目标子空间 $V_\kappa$ 上），并由弱模型的输出进行监督。关键在于，弱模型的输出在被用作标签之前会经过非线性变换（$z \mapsto z \exp(z)$）。
   • 标准 SFT（基线）：作为对比，强模型通过在线 SGD 在各项同性输入（$x \sim N(0, I_d)$）上进行训练，标签直接来自真实奖励函数 $r^*_\kappa$。

主要贡献与理论结果

1. 通过多步 SGD 实现特征提取（定理 1）

本文提供了特征学习范式下 W2S 的首个多步 SGD 收敛保证。

• 机制：教师信号的非线性变换降低了学习问题的有效“生成指数”。虽然目标链接函数 $\sigma^*_\kappa$ 可能具有高信息指数 $p > 2$，但该变换使得强模型能够以依赖于生成指数的样本复杂度学习特征方向 $\theta_\kappa$（对于偶次多项式，该指数变为 2）。
• 样本复杂度：强模型以 $\tilde{O}(s^{3/2})$ 的样本复杂度实现与目标特征 $\theta_\kappa$ 的对齐，其中 $s$ 是任务子空间的维度。这显著低于弱模型从头开始学习相同特征所需的 $\tilde{O}(d^{p/2} s^{(p-2)/2})$ 复杂度，也低于标准 SFT 在各项同性采样下的 $\tilde{O}(d^{p-1})$ 复杂度。
• 子空间先验：该结果依赖于强模型的输入分布集中在目标子空间 $V_\kappa$ 上（由 $\lambda_\kappa \approx 1$ 参数化），这有效地将学习维度从 $d$ 降低到了 $s$。

2. 特征保留（命题 2）

一个关键贡献是证明了 W2S 能够保留非目标任务（$k \neq \kappa$）的预训练特征。

• 隐式正则化：强模型的生成分布（集中在 $V_\kappa$ 上）与非线性教师变换的结合充当了隐式正则化器。
• 保留条件：只要非目标任务分配的神经元多于目标任务（或携带较小的奖励权重），非目标特征就能得到保留，这由定位参数 $\chi_k = N_\kappa \pi_k / (N_k \pi_\kappa)$ 量化。
• 结果：非目标神经元与其原始方向 $\theta_k$ 的对齐得以维持，对齐度的下降由与初始化重叠和 SGD 噪声相关的项界定，而非灾难性遗忘。

3. 标准 SFT 中的灾难性遗忘（定理 2）

与 W2S 相反，本文证明了标准监督微调（SFT）会导致非目标特征的灾难性遗忘，特别是当这些特征与目标具有方向性共享时。

• 机制：在各项同性采样下，梯度信号传播到所有神经元。如果非目标神经元与目标方向 $\theta_\kappa$ 之间的初始对齐不可忽略（具体而言，如果 $\theta_k^\top \theta_\kappa \gtrsim s^{-1/2}$），那么非目标神经元收敛到 $\theta_\kappa$ 的速度将快于目标神经元收敛到 $\theta_\kappa$ 的速度。
• 结果：预训练特征 $\theta_k$ 在目标任务被完全学习之前就已丢失，导致模型通用能力的破坏。

数值实验

作者使用输入维度 $d=1024$ 和子空间维度 $s=128$ 的合成数据验证了其理论发现。

• W2S 性能：实验证实，强模型在学习目标特征方向（$\theta_1$）的同时，保持与非目标特征（$\theta_2$）的对齐，这与定理 1 和命题 2 一致。
• SFT 失败：在 SFT 设定下，非目标神经元在目标神经元完全收敛之前就已与目标方向 $\theta_1$ 对齐，证实了灾难性遗忘的理论预测（定理 2）。

意义与主张

本文声称在特征学习范式下为 W2S 泛化奠定了理论基础，超越了以往假设固定表示或线性模型的分析。

• 提取机制：它表明 W2S 不仅仅是“灌输”新能力，而是通过利用预训练子空间和非线性变换来降低样本复杂度，从而有效地“提取”潜在知识。
• 优于 SFT 的优势：该工作从理论上解释了为何 W2S 在保留通用能力方面优于标准 SFT。这归因于算法差异：W2S 在强模型的生成分布（集中在任务子空间上）上查询弱模型，而 SFT 依赖于各项同性采样，这会导致相关特征之间的干扰。
• 局限性：作者指出，他们的分析侧重于第一层特征对齐，并未正式分析第二层学习（输出层的权重更新）。他们还承认，生成指数的降低要求目标链接函数为偶次多项式。

总之，本文认为 W2S 之所以成功，是因为强模型的预训练结构（子空间定位）与特定的训练协议（生成分布 + 非线性变换）相结合，使其能够高效地学习目标特征，同时不破坏编码其更广泛知识的正交或弱相关特征。