Robust Learning of a Group DRO Neuron
Guyang Cao, Shuyao Li, Sushrut Karmalkar, Jelena Diakonikolas
TL;DR
This work addresses the challenge of learning a single neuron under Group DRO with arbitrary label noise and groupwise distributional shifts. It introduces a primal–dual algorithm with extrapolated dual updates to tackle the inherent nonconvexity and supports general $f$-divergences (e.g., KL or $\chi^2$) in the group-weights. The authors prove provable constant-factor guarantees relative to the best-fit neuron, with near-optimal sample complexity and iteration counts, and provide a linearization-based analysis to handle nonconvexity and divergence terms. Empirical validation on large-language-model pretraining demonstrates practical benefits of dual-extrapolated reweighting (PD-KL) over prior domain-reweighting methods, achieving higher downstream accuracy with less training time. The framework offers a principled, scalable approach to robust nonconvex learning under distributional shifts and adversarial noise, with potential extensions to broader nonconvex risk settings and applications in LLM pretraining.
Abstract
We study the problem of learning a single neuron under standard squared loss in the presence of arbitrary label noise and group-level distributional shifts, for a broad family of covariate distributions. Our goal is to identify a ''best-fit'' neuron parameterized by $\mathbf{w}_*$ that performs well under the most challenging reweighting of the groups. Specifically, we address a Group Distributionally Robust Optimization problem: given sample access to $K$ distinct distributions $\mathcal p_{[1]},\dots,\mathcal p_{[K]}$, we seek to approximate $\mathbf{w}_*$ that minimizes the worst-case objective over convex combinations of group distributions $\boldsymbolλ \in Δ_K$, where the objective is $\sum_{i \in [K]}λ_{[i]}\,\mathbb E_{(\mathbf x,y)\sim\mathcal p_{[i]}}(σ(\mathbf w\cdot\mathbf x)-y)^2 - νd_f(\boldsymbolλ,\frac{1}{K}\mathbf1)$ and $d_f$ is an $f$-divergence that imposes (optional) penalty on deviations from uniform group weights, scaled by a parameter $ν\geq 0$. We develop a computationally efficient primal-dual algorithm that outputs a vector $\widehat{\mathbf w}$ that is constant-factor competitive with $\mathbf{w}_*$ under the worst-case group weighting. Our analytical framework directly confronts the inherent nonconvexity of the loss function, providing robust learning guarantees in the face of arbitrary label corruptions and group-specific distributional shifts. The implementation of the dual extrapolation update motivated by our algorithmic framework shows promise on LLM pre-training benchmarks.