技术摘要：当数据不平衡成为助力：通过快捷方式饱和实现鲁棒泛化

问题陈述

本文研究了在存在**伪相关（spurious correlations）**情况下的鲁棒泛化问题。伪相关是模型失效的核心模式，即模型学会利用在训练集中具有预测性但在测试时不可靠的特征。缓解这一问题的标准做法是数据集平衡——使快捷方式一致（shortcut-consistent）与反快捷方式（anti-shortcut）的样本数量相等，以防止单一特征主导梯度下降。

作者通过研究一个受控的合成设置——最大值-奇偶性-求和-奇偶性（MPSP）任务及其变体，对这一做法提出了挑战。在这些任务中，模型必须预测整数序列之和的奇偶性（真实标签），但可以利用最大值的奇偶性（快捷方式）作为一种伪相关进行预测。本研究改变了伪相关比例（$r$），即训练样本中快捷方式等于真实标签的比例，并考察了其如何与模型容量（特别是 Transformer 层数和注意力头数）相互作用。

研究方法

作者采用受控的实验设置，使用合成数据集和 Transformer 架构来隔离不平衡与容量的影响。

• 任务： 通过将两个真实标签（求和模 2 奇偶性、求和模 3 奇偶性）与两个快捷方式特征（第一个元素、最大元素）进行交叉组合，构建了四种任务。
• 数据划分： 每个实验使用三个固定的划分：
  • 训练集（Train）： 在特定的伪相关比例 $r$ 下包含 16,384 个样本。
  • 验证集（Validation）： 比例 $r=1.0$（全为快捷方式一致样本），用于衡量快捷方式学习。
  • 对抗集（Adversarial, Adv）： 比例 $r=0.0$（全为反快捷方式样本），用于衡量真实规则的泛化能力。
• 泛化指标： 如果在训练过程中的任何时刻实现了 100% 的对抗准确率，则认为该次运行实现了“泛化”。泛化率是指实现该目标的随机种子比例。
• 模型： 单层（1L）和两层（2L）Transformer，具有变化的注意力头数（$n_{heads} \in \{1, 2\}$），参数配置为 $d_{model}=64$，$d_{ff}=128$。使用 AdamW 进行 15,000 个 epoch 的训练。
• 机制分析： 作者利用了以下手段：
  • 梯度冲突动力学（Gradient Conflict Dynamics）： 测量来自快捷方式一致样本和反快捷方式样本之间的梯度余弦相似度。
  • 电路演化（Circuit Evolution）： 通过激活消融（activation ablation）和跨划分补丁（cross-split patching）追踪“鲁棒”电路的形成。
  • QK/OV 电路消融： 分析查询-键（QK）和输出-值（OV）电路的指纹，以区分大小排序（快捷方式）与奇偶编码（鲁棒规则）机制。

核心贡献

本文确立了三个主要贡献：

1. 不平衡 $\times$ 容量的相互作用： 作者展示了一个反直觉的相互作用：在容量充足的模型（2 层 Transformer）中，将伪相关比例 $r$ 从 0.5 增加到 0.9 会提高泛化率；但在容量不足的模型（1 层 Transformer）中，泛化率反而会降低。这识别出了一个容量阈值（位于 1 层和 2 层之间），在此阈值处，不平衡的影响发生了逆转。
2. 机制路径： 通过梯度分析和电路追踪，论文刻画了一条路径，其中**快捷方式饱和（shortcut saturation）**是泛化的先决条件。在容量充足的模型中，高不平衡导致快捷方式电路饱和（在多数样本上损失接近于零），从而导致其梯度消失。这放大了来自反快捷方式少数派的相对梯度信号（放大倍数约为 $r/(1-r)$），进而支持了注意力电路的结构重组。
3. 统一原则： 该效应取决于 $r$ 偏离随机基准（对于二进制为 1/2，对于三进制为 1/3）的程度。无论是快捷方式偏向型（$r > \text{baseline}$）还是反快捷方式偏向型（$r < \text{baseline}$）的比例都能促进泛化，而零偏差比例（随机机会）则会完全抑制泛化。

关键结果

行为结果

• 2 层模型： 在 $r=0.5$（平衡）时，2L 模型在 0% 的种子中实现泛化。当 $r$ 增加到 0.9 时，泛化率上升至 77%。
• 1 层模型： 趋势发生反转。在 $r=0.5$ 时，1L 模型在 33% 的种子中实现泛化。而在 $r=0.9$ 时，这一比例降至 0%。高不平衡现象实际上将这些模型困在了快捷方式上。
• 权重衰减（Weight Decay）： 增加权重衰减（从 0.1 到 0.4）可以使模型发现泛化盆地（generalizing basins），但往往会导致模型无法留在其中，从而产生瞬时泛化。
• 快捷方式类型： 该效应在基于量级（最大元素）和基于位置（第一个元素）的快捷方式，以及二进制和三进制标签任务中均得到了复现。

机制发现

• 梯度冲突： 在 $r$ 较高时的 2L 模型中，梯度冲突（最初为负余弦相似度）最终得到解决（趋于 0），随着模型实现泛化。在 1L 模型中，冲突在整个训练过程中持续存在。
• 梯度放大： 在 $r=0.9$ 时，由于快捷方式电路饱和并停止从多数类中学习，反快捷方式样本产生的梯度信号比快捷方式一致样本高出 9–17 倍。
• 电路重组：
  • 快捷方式阶段： 模型最初开发出强大的快捷方式电路（具有较高的 QK Spearman $\rho$ 量级或较高的位置 0 注意力）。
  • 过渡阶段： 在容量充足的 2L 模型中，当快捷方式电路被取代时，第二个头（或第二层）会开发出鲁棒的奇偶编码电路（具有较高的 OV 奇偶性 Spearman）。
  • 失效模式： 在 1L 模型中，快捷方式电路保持主导地位；模型无法重组其有限的参数来适应真实规则。
• 两阶段动力学： 泛化运行表现出“第二阶段”，即在快捷方式一致准确率保持近乎完美的同时，对抗准确率上升。这一阶段的开启具有种子依赖性，并且在 2L 模型中随头数的增加而加速。

意义与主张

本文声称识别了一个数据不平衡作为鲁棒泛化催化剂而非阻碍的机制。其意义在于挑战了“数据集平衡是减轻伪相关之普遍要求”的标准教条。

作者提出，快捷方式饱和是容量充足模型实现泛化的必要先决条件。当快捷方式被过度代表时，它会迅速饱和，从而使其自身的梯度信号沉默，允许来自少数类的、被放大的对抗性梯度驱动网络内部电路的结构重组。

论文对其研究范围保持了审慎态度，指出：

• 结果源自小型合成 Transformer，在没有进一步验证的情况下不应直接推广到现实世界的任务。
• 效应大小取决于真实计算相对于快捷方式的难度。
• 泛化 epoch 的高方差表明，基于种子条件的分析比基于配置的平均值更具信息量。

最终，这项工作表明，数据集平衡、模型容量与泛化之间的关系比标准假设更为复杂，不平衡可能作为一种机制，迫使容量充足的模型在耗尽简单的快捷方式后，去学习鲁棒的规则。