Robust Learning of a Group DRO Neuron

arXiv:2601.18115v2 Announce Type: replace 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]},ots,mathcal p_{[K]}, we seek to approximate mathbf{w}* that minimizes the worst-case objective over convex combinations of group distributions oldsymbol{lambda} in Delta_K, where the objective is sum{i in [K]}lambda_{[i]},mathbb E_{(mathbf x,y)simmathcal p_{[i]}}(sigma(mathbf wdotmathbf x)-y)^2 - nu d_f(oldsymbollambda,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 nu 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.

Source: arXiv cs.LG | 2026-06-02