Robust Learning of a Group DRO Neuron
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 that performs well under the most challenging reweighting of the groups. Specifically, we address a Group Distributionally Robust Optimization problem: given sample access to distinct distributions , we seek to approximate that minimizes the worst-case objective over convex combinations of group distributions , where the objective is and is an -divergence that imposes (optional) penalty on deviations from uniform group weights, scaled by a parameter . We develop a computationally efficient primal-dual algorithm that outputs a vector that is constant-factor competitive with 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.
Cite
@article{arxiv.2601.18115,
title = {Robust Learning of a Group DRO Neuron},
author = {Guyang Cao and Shuyao Li and Sushrut Karmalkar and Jelena Diakonikolas},
journal= {arXiv preprint arXiv:2601.18115},
year = {2026}
}