Robust Learning of a Group DRO Neuron

技术摘要：鲁棒学习一个 Group DRO 神经元

问题陈述
本文研究了在平方损失 $\mathbb{E}[(\sigma(w^\top x) - y)^2]$ 下学习单个神经元（由向量 $w$ 参数化）的问题。该设定具有两个同时存在的挑战：

1. 对抗性标签噪声： 标签 $y$ 可能包含任意形式的损坏，使问题进入不可实现（non-realizable）的对抗性学习（agnostic learning）范畴。
2. 组分布偏移（Group Distributional Shifts）： 训练数据源自 $K$ 个不同的组，其分布分别为 $p^{[1]}, \dots, p^{[K]}$。目标是找到一个在最坏情况下的组重加权下表现良好的模型，这被形式化为一个组分布鲁棒优化（Group DRO）问题。

目标函数是最小化正则化的最坏情况损失：
$$\min_{w \in B(W)} \max_{\lambda \in \Delta^K} \left( \sum_{i=1}^K \lambda^{[i]} \mathbb{E}_{p^{[i]}}[(\sigma(w^\top x) - y)^2] - \nu d_f(\lambda, \tfrac{1}{K}\mathbf{1}) \right)$$
其中 $\nu \ge 0$ 是正则化参数，$d_f$ 是强凸的 $f$-散度（例如 KL 或 $\chi^2$），目标是输出一个估计值 $\hat{w}$，使其在最坏情况组权重下的表现与最优参数 $w^*$ 具有常数因子竞争性。

方法论
作者开发了一种高效的原对偶算法（Primal-Dual Algorithm，算法 1）来解决这个非凸-凹鞍点问题。该方法依赖于以下几个关键技术组件：

• 结构性假设： 分析假设激活函数 $\sigma$ 是凸的且满足 $(\alpha, \beta)$-无界性（例如 ReLU）。协变量分布满足亚指数尾部界限（假设 1.2）以及一致边际条件（Uniform Margin Condition，假设 1.3），该条件确保了在边际区域内的协方差矩阵在所有组和权重向量下都是一致良置的。
• 代理梯度（Surrogate Gradients）： 由于 $\sigma$ 可能不可微，算法使用代理梯度 $v(w; x, y) = 2\beta(\sigma(w^\top x) - y)x$。
• 对偶侧外推（Dual-Side Extrapolation）： 该算法的一个显著特征是在对偶变量 $\lambda \in \Delta^K$ 上而非原变量 $w$ 上使用外推（动量），这是为了内存效率的考虑，因为通常组的数量 $K$ 远小于维度 $d$。
• 线性化引理（Linearization Lemma）： 为了处理损失函数的非凸性和外推步骤（外推可能会产生带有负项的对偶向量），作者证明了一个“线性化”引理（引理 3.5）。该结果利用代理梯度从下方界定了均方损失，从而避免了依赖于对偶迭代之间散度的复杂高阶误差项。这使得该框架能够处理 $\nu \to 0$ 以及像 KL 这样的替代散度，而不像之前的工作那样受限于需要大 $\nu$ 的 $\chi^2$ 散度。
• 收敛性分析： 证明通过相对于混合参考点 $(w^*, \lambda^*)$ 的间隙函数（gap function）来建立收敛性。通过结合来自经验锐度（empirical sharpness）的下界和来自算法更新规则的上界，作者表明迭代过程收敛到一个解，其中与 $w^*$ 的平方距离被最坏组损失的一个常数因子加上一个误差项 $\epsilon$ 所界定。

核心贡献

1. 首个具有保证的高效算法： 本文提供了第一个在存在对抗性标签噪声和组分布偏移的情况下，学习单个神经元并具有可证明常数因子近似保证的高效算法。
2. 放宽的假设： 与近期工作 [LKDD24] 相比，本方法：
   • 适用于更广泛的 $f$-散度类（包括 KL），而不仅仅是 $\chi^2$。
   • 消除了对大正则化参数 $\nu$ 的要求，能够处理无正则化情形（$\nu=0$）。
   • 消除了对损失函数第四阶矩的限制性假设。
   • 使用对偶侧外推，这比原侧外推更具内存效率，且与标准求解器更兼容。
3. 样本复杂度： 该算法实现了 $\tilde{O}(Kd/\epsilon^2)$ 的样本复杂度，在忽略对数因子的情况下，达到了已知最优凸无正则化 Group DRO 问题的速率。

结果

• 理论保证： 定理 3.1 指出，在足够的样本下，算法输出的 $\hat{w}$ 满足 $\|\hat{w} - w^*\|_2^2 \le C \cdot \text{OPT}_m + \epsilon$，其中 $\text{OPT}_m$ 是最优参数在其表现最差组上的损失，$C > 1$ 是一个普适常数。
• 实证评估： 作者进行了初步实验，将他们的原对偶法结合 KL 散度正则化（PD-KL）集成到 Sheared-LLaMA-1.3B 模型的预训练中。
  • 与 DoReMi 基线（动态批次加载）相比，PD-KL 算法在 11 个下游任务的无权重平均准确率上表现出持续的提升。
  • 在 33.6M token 时，PD-KL 模型达到了 47.87% 的准确率，而基线达到相同水平需要多出 1.5 倍的训练时间。
  • 作者指出这些结果仅具说明性，并非全面的验证，但表明理论更新规则为训练稳定性提供了实际益处。

意义与主张
本文声称填补了非凸 Group DRO 在理论理解上的关键空白，这种设定在深度学习中普遍存在，但在面对任意噪声时此前缺乏严格保证。通过证明在这些严苛条件下，学习单个神经元的常数因子近似是可行的，这项工作为异构数据环境下的鲁棒学习提供了理论基础。作者谦逊地将其实验结果描述为“初步探索”，暗示其正则化对偶外推重加权的核心思想值得在更大规模语言模型预训练及其他复杂非凸设置中进一步研究。他们明确指出，未来的工作可以将这些更新扩展到更广泛的应用，并强化相对于全局最优 $\text{OPT}$ 而非最坏组最优 $\text{OPT}_m$ 的误差保证。