English

Robust Learning of a Group DRO Neuron

Machine Learning 2026-01-27 v1 Data Structures and Algorithms Optimization and Control

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 w\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 KK distinct distributions p[1],,p[K]\mathcal p_{[1]},\dots,\mathcal p_{[K]}, we seek to approximate w\mathbf{w}_* that minimizes the worst-case objective over convex combinations of group distributions λΔK\boldsymbol{\lambda} \in \Delta_K, where the objective is i[K]λ[i]E(x,y)p[i](σ(wx)y)2νdf(λ,1K1)\sum_{i \in [K]}\lambda_{[i]}\,\mathbb E_{(\mathbf x,y)\sim\mathcal p_{[i]}}(\sigma(\mathbf w\cdot\mathbf x)-y)^2 - \nu d_f(\boldsymbol\lambda,\frac{1}{K}\mathbf1) and dfd_f is an ff-divergence that imposes (optional) penalty on deviations from uniform group weights, scaled by a parameter ν0\nu \geq 0. We develop a computationally efficient primal-dual algorithm that outputs a vector w^\widehat{\mathbf w} that is constant-factor competitive with w\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.

Keywords

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}
}
R2 v1 2026-07-01T09:19:37.792Z