# Robust Learning of a Group DRO Neuron

arXiv: https://arxiv.org/abs/2601.18115 · HTML: https://mlsanity.com/paper/2601.18115/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 34 of its 60 references that are in the map.
- First posted on arXiv 26 Jan 2026.

## Builds on (most cited first, 30 of 34)

- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [A Survey on Deep Transfer Learning](https://mlsanity.com/paper/1808.01974/) (2018)
- [Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning](https://mlsanity.com/paper/2310.06694/) (2023)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Fairness Without Demographics in Repeated Loss Minimization](https://mlsanity.com/paper/1806.08010/) (2018)
- [Certifying Some Distributional Robustness with Principled Adversarial Training](https://mlsanity.com/paper/1710.10571/) (2017)
- [Learning Models with Uniform Performance via Distributionally Robust Optimization](https://mlsanity.com/paper/1810.08750/) (2018)
- [On Gradient Descent Ascent for Nonconvex-Concave Minimax Problems](https://mlsanity.com/paper/1906.00331/) (2019)
- [Distributionally Robust Language Modeling](https://mlsanity.com/paper/1909.02060/) (2019)
- [Weakly-Convex Concave Min-Max Optimization: Provable Algorithms and Applications in Machine Learning](https://mlsanity.com/paper/1810.02060/) (2018)
- [A Single-Loop Smoothed Gradient Descent-Ascent Algorithm for Nonconvex-Concave Min-Max Problems](https://mlsanity.com/paper/2010.15768/) (2020)
- [Near-Optimal SQ Lower Bounds for Agnostically Learning Halfspaces and ReLUs under Gaussian Marginals](https://mlsanity.com/paper/2006.16200/) (2020)
- [Learning ReLUs via Gradient Descent](https://mlsanity.com/paper/1705.04591/) (2017)
- [Distribution-Independent PAC Learning of Halfspaces with Massart Noise](https://mlsanity.com/paper/1906.10075/) (2019)
- [The Optimality of Polynomial Regression for Agnostic Learning under Gaussian Marginals](https://mlsanity.com/paper/2102.04401/) (2021)
- [Learning a Single Neuron with Gradient Methods](https://mlsanity.com/paper/2001.05205/) (2020)
- [Statistical-Query Lower Bounds via Functional Gradients](https://mlsanity.com/paper/2006.15812/) (2020)
- [An Online Method for A Class of Distributionally Robust Optimization with Non-Convex Objectives](https://mlsanity.com/paper/2006.10138/) (2020)
- [On-Demand Sampling: Learning Optimally from Multiple Distributions](https://mlsanity.com/paper/2210.12529/) (2022)
- [Simple and optimal methods for stochastic variational inequalities, I: operator extrapolation](https://mlsanity.com/paper/2011.02987/) (2020)
- [Approximation Schemes for ReLU Regression](https://mlsanity.com/paper/2005.12844/) (2020)
- [Time/Accuracy Tradeoffs for Learning a ReLU with respect to Gaussian Marginals](https://mlsanity.com/paper/1911.01462/) (2019)
- [Minimax Regret Optimization for Robust Machine Learning under Distribution Shift](https://mlsanity.com/paper/2202.05436/) (2022)
- [Near-Optimal Cryptographic Hardness of Agnostically Learning Halfspaces and ReLU Regression under Gaussian Marginals](https://mlsanity.com/paper/2302.06512/) (2023)
- [Class-Weighted Classification: Trade-offs and Robust Approaches](https://mlsanity.com/paper/2005.12914/) (2020)
- [Statistical Analysis of Wasserstein Distributionally Robust Estimators](https://mlsanity.com/paper/2108.02120/) (2021)
- [No-Regret Dynamics in the Fenchel Game: A Unified Framework for Algorithmic Convex Optimization](https://mlsanity.com/paper/2111.11309/) (2021)
- [Learning a Single Neuron with Adversarial Label Noise via Gradient Descent](https://mlsanity.com/paper/2206.08918/) (2022)
- [Nonsmooth Nonconvex-Nonconcave Minimax Optimization: Primal-Dual Balancing and Iteration Complexity Analysis](https://mlsanity.com/paper/2209.10825/) (2022)
- [Robustly Learning a Single Neuron via Sharpness](https://mlsanity.com/paper/2306.07892/) (2023)

## Abstract (arXiv)

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_* 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 w_* that minimizes the worst-case objective over convex combinations of group distributions boldsymbolλ ∈ DeltaK, where the objective is sumi ∈ [K]λ[i],mathbb E_(mathbf x,y)~mathcal p[i](σ(mathbf w·mathbf x)-y)2 - nu df(boldsymbolλ,frac1Kmathbf1) and df is an f-divergence that imposes (optional) penalty on deviations from uniform group weights, scaled by a parameter nu ≥ 0. We develop a computationally efficient primal-dual algorithm that outputs a vector widehatmathbf w that is constant-factor competitive with 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.
