REVIEW 4 major objections 5 minor 23 references
This paper argues that classical Group DRO ignores uncertainty in each group's own data distribution, and shows that adding per-group Wasserstein balls yields a tractable min-max-sup problem with better worst-group accuracy.
Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →
T0 review · deepseek-v4-flash
2026-08-04 19:56 UTC pith:JLSDNQ7I
load-bearing objection Reasonable composition of Group DRO and per-group Wasserstein robustness, but the convergence proof doesn't cover the implemented algorithm and the test-set-level tuning of gamma undercuts the empirical claims. the 4 major comments →
Group Distributionally Robust Machine Learning under Group Level Distributional Uncertainty
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
The paper's central discovery is that the two forms of uncertainty—uncertainty in the mixture of environments and uncertainty in each environment's own distribution—can be handled in a single optimization problem. The proposed objective minimizes over model parameters the worst-case weighted sum of per-group robust losses, where each group's robust loss is the supremum of the expected loss over a Wasserstein ball around the group's empirical distribution. The authors show that this inner supremum can be computed via a Lagrangian relaxation that turns it into an expectation of a penalized loss, making it tractable for nonconvex models. The paper proves that the resulting three-step gradient a
What carries the argument
The central object is the nested objective (5): min over parameters theta of max over group weights q of the sum over groups of q_g times the per-group robust loss, where each robust loss is a supremum over distributions within a 1-Wasserstein ball of radius eps_g around the group's empirical distribution. The enabling identity is Proposition 3.1, a Lagrangian relaxation that converts the inner supremum into an expectation over training samples of a penalized loss L(f_theta; z') - gamma c(z, z'), which can be maximized by gradient ascent. Around this sit a mirror-ascent update for the group weights, a gradient descent update for the parameters, and a Moreau-envelope stationarity argument tha
Load-bearing premise
The proof of convergence requires that each step's inner worst-case distribution be found exactly (or with negligible error), but the implemented algorithm approximates it with a finite, unprojected gradient ascent and never verifies that the penalty parameter gamma exceeds the loss's input-space smoothness constant L_zz.
What would settle it
Estimate L_zz empirically for the two-layer ELU network on the Adult features and check whether the deployed gamma values (including 10^-4) satisfy gamma > L_zz; alternatively, increase T_rob until the inner penalized loss stabilizes and see whether the reported worst-group accuracy and spread change materially. Either test would show whether the convergence theorem covers the method as implemented.
If this is right
- Models trained with this objective are protected against two failure modes at once: shift in the group mixture and drift inside each group's data distribution, both without requiring knowledge of the true group distributions.
- The Lagrangian relaxation makes the inner worst-case distribution search tractable for nonconvex models such as neural networks, where dual formulations of the robust loss are not available.
- The method stays stable across gamma values from 10^-4 to 10^1 in the Adult experiments, so the penalty parameter does not require delicate tuning.
- On the Adult income task with an education covariate shift, the method achieves the highest average accuracy (0.715), highest worst-group accuracy (0.613), and lowest accuracy range (0.193) among the four compared methods.
- Across constructed test environments ranging from 90-10 to 10-90 education splits, the method maintains worst-group accuracy around 0.60-0.63, roughly half the disparity of Group DRO.
Where Pith is reading between the lines
- Because gamma=0 recovers standard DRO and gamma tending to infinity recovers Group DRO, the framework implicitly defines a continuum of robustness levels; the paper leaves unexplored how to choose gamma per group according to estimated shift severity.
- The convergence theorem assumes the inner supremum is computed exactly, but the algorithm uses a fixed number of gradient ascent steps with no projection; an empirical check of the ratio gamma/L_zz on the actual network would show whether the reported results operate in the theorem's regime.
- A natural extension is to test the method under multiple simultaneous covariate shifts (age, occupation, region) rather than a single education shift, which would clarify whether the observed advantage is specific to the constructed split.
- The reported accuracy gains are relative to methods that ignore one of the two uncertainty sources; a more direct comparison would be to a baseline that uses the same inner adversarial training but with a fixed uniform q, isolating the contribution of the mirror ascent over group weights.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes a group DRO framework that adds per-group Wasserstein distributional uncertainty to the standard Group DRO objective. The formulation is the nested problem (5): min over model parameters, max over group weights, and, within each group, a supremum over distributions in a Wasserstein ball around the empirical group distribution. The authors propose Algorithm 2, which alternates gradient ascent on adversarial perturbations (Algorithm 1), mirror ascent on group weights, and gradient descent on the model parameters. They claim a convergence guarantee (Theorem A.1) for the corresponding function P(θ) = max_q Σ_g q_g F_g(θ), where F_g is the exact expected robust group loss. Experiments on the Adult dataset, including a constructed education marginal shift and multiple test environments, report that the method improves worst-group accuracy and reduces accuracy range compared with ERM, DRO, and Group DRO.
Significance. If the convergence guarantee and empirical results were fully supported, the paper would address a genuine gap: existing Group DRO methods typically assume exact knowledge of each group's distribution, while the proposed nested formulation explicitly models within-group distributional uncertainty. The problem setup is well motivated, and the multi-environment Adult evaluation is a sensible stress test. However, the central theoretical claim is not established: Theorem A.1 relies on an unnamed theorem in a prior paper, and the implemented algorithm does not match the theorem's assumptions because the inner supremum is approximated by unprojected gradient ascent without error control. The experimental validation is also weakened by selection of the robustness parameter on the test data and by the absence of a second promised dataset. The contribution is therefore significant in conception but not currently supported in execution.
major comments (4)
- [Appendix A.3, Theorem A.1] The convergence proof is not supplied. The proof of Theorem A.1 states that it follows 'the same steps to combine lemmas A.8, A.9, and A.10 as in Theorem ... in [Sinha et al., 2017]', but no theorem number is given. Lemma A.8 similarly imports a GDA result from [Lin et al., 2020] without a derivation. Since the convergence guarantee is the paper's main theoretical contribution, this is load-bearing and cannot be checked. The authors need to provide a complete proof or a precise, stated theorem from the cited work together with a verification of the conditions.
- [Algorithm 2, lines 4-6; Theorem A.1] Theorem A.1 concerns the exact function F_g(θ) = E sup_z ϕ(f_θ; (x,y), z), i.e., the expected value of the exact supremum. Algorithm 2 instead sets z_i = argmax via Algorithm 1, which performs only T_rob=100 unprojected gradient ascent steps (Section 4.1 reports T_rob=100). No bound is provided for the approximation error between the T_rob-step iterate and the true supremum, nor for its effect on the gradient or on the Moreau-envelope stationarity measure. Consequently the stated convergence result does not apply to the implemented algorithm. The paper must either analyze the inexact inner problem or modify the algorithm so that the exact supremum is used or provably approximated.
- [Lemma A.5 and Section 4.1] The smoothness constant L_f = L_θθ + L_θz L_zθ / [γ - L_zz]_+ is finite only if γ > L_zz. Lemma A.6 also assumes existence of the argmax for γ > L_zz. The authors never estimate L_zz for the two-layer ELU network used in the experiments, and the reported experiments use γ = 10^-4 (Table 1), with a sweep starting at 10^-4. For a neural network on standardized features, the Lipschitz constant of the loss gradient with respect to z is not plausibly below 10^-4. Thus the assumptions of the convergence theorem are at best unverified and very likely violated in the configuration whose results are reported.
- [Section 4, Table 1, Figure 5] Section 4 states that the method is evaluated on 'real-world datasets' (plural), but only the Adult dataset is used; Section 4.2 is a set of constructed test environments on the same Adult data, not a second dataset. More importantly, γ is swept over a wide range and the 'best' value is selected (Section 4.1: 'select the best value via fine-tuning'; Figure 5: 'best-performing γ'), and the final tables report numbers at that selected value. If the selection is made using the test sets that are also used for evaluation, the comparison is fitted to the evaluation data. The authors need a validation-set selection procedure and should report results for a fixed, pre-specified γ or over the whole sweep without test-based selection.
minor comments (5)
- [Section 4.1] The sentence 'we sweep γ ... to select the best value via fine-tuning' is ambiguous. Please specify which data split is used for tuning and how many seeds are used for selection.
- [Algorithm 2, line 8] Line 8 uses the notation 1/N_g and appears to average over the full group dataset. If this is the full-batch gradient, state it explicitly; if minibatches are used, the finite-sample error must be addressed.
- [Table 1] The column header 'Ours (γ=10^-4)' is acceptable, but the header 'Group' is an abbreviation for Group DRO and could be confused with the group index. Please rename for clarity.
- [Section 4 opening] The text says 'real-world datasets from the field of finance', but the Adult dataset is not a finance dataset. Please correct the description.
- [Various] There are several typographical issues, e.g., 'Assimptions' in Lemma A.10 and 'envolope' in the convergence appendix. Typesetting of the many cross-references to equations could also be improved.
Circularity Check
No significant circularity: derivation imports external results and does not reduce to its inputs; the main issues are approximation and validity gaps, not self-reference.
full rationale
The paper's central theoretical chain is: (i) Proposition 3.1 equates the per-group robust loss with an expectation of a Lagrangian-penalized supremum; (ii) Lemma A.5 and Lemma A.6 establish smoothness and Lipschitzness of that robust loss; (iii) Theorem A.1 gives a convergence rate for descent-mirror-ascent on the exact objective P(theta)=max_q sum q_g F_g(theta). Each of these results is imported from external prior work (Sinha et al. 2017, Lin et al. 2020, Davis and Drusvyatskiy 2019), not from the authors' own prior results, and the objective (5) is not defined in terms of the algorithm's approximate quantities. The theorem's statement concerns an idealized version of Algorithm 2 in which the inner supremum is computed exactly (pseudocode line 5 writes z_i = argmax_{z in X x Y} phi(...)), whereas the implemented Algorithm 1 runs only T_rob=100 unprojected gradient ascent steps and supplies no approximation-error bound. This is a validity/applicability gap rather than a circular step, because the theorem's conclusion is not equivalent to the implemented update by construction. The experimental gamma sweep ('we sweep gamma in {10^-4,...,8} to select the best value via fine-tuning') introduces potential selection bias, but the reported worst-group accuracy is an empirical measurement, not a quantity derived from the fitted gamma, so it does not make the derivation circular. The citation [Hong et al., 2023], which includes a co-author, is used only as motivational evidence of disparate subgroup performance and is not load-bearing for the theoretical claims. No uniqueness claim is imported from the authors' own work, no ansatz is smuggled via self-citation, and no known result is merely renamed. The genuine concerns are omitted or sketched proofs (Theorem A.1 says the proof 'follows the same steps ... as in Theorem ... in [Sinha et al., 2017]') and unverified assumptions (gamma > L_zz with L_zz never estimated for the neural network, while experiments use gamma=10^-4). These are correctness risks, not circular steps.
Axiom & Free-Parameter Ledger
free parameters (2)
- gamma =
gamma=1e-4 (ours), gamma=9 (DRO baseline)
- algorithm hyperparameters =
eta_theta=0.1, eta_q=0.1, eta_z=0.05, T_rob=100, T=200
axioms (4)
- domain assumption Loss L is K-Lipschitz and smooth w.r.t. theta and (x,y) (Assumptions A.2-A.3)
- domain assumption Transport cost c is 1-strongly convex and gamma > L_zz (Assumption A.1 and Lemma A.5)
- domain assumption Per-group empirical distribution is an adequate center for a Wasserstein ball
- standard math Standard optimal transport and DRO duality results (Villani 2009, Sinha et al. 2017)
Cite this review
Pith. "Pith review of Group Distributionally Robust Machine Learning under Group Level Distributional Uncertainty." pith.science (2026). https://pith.science/paper/JLSDNQ7I
@misc{pith2026250908942,
author = {Pith},
title = {Pith review of: Group Distributionally Robust Machine Learning under Group Level Distributional Uncertainty},
year = {2026},
howpublished = {\url{https://pith.science/paper/JLSDNQ7I}},
note = {Machine review of arXiv:2509.08942}
}
read the original abstract
The performance of machine learning (ML) models critically depends on the quality and representativeness of the training data. In applications with multiple heterogeneous data generating sources, standard ML methods often learn spurious correlations that perform well on average but degrade performance for atypical or underrepresented groups. Prior work addresses this issue by optimizing the worst-group performance. However, these approaches typically assume that the underlying data distributions for each group can be accurately estimated using the training data, a condition that is frequently violated in noisy, non-stationary, and evolving environments. In this work, we propose a novel framework that relies on Wasserstein-based distributionally robust optimization (DRO) to account for the distributional uncertainty within each group, while simultaneously preserving the objective of improving the worst-group performance. We develop a gradient descent-ascent algorithm to solve the proposed DRO problem and provide convergence results. Finally, we validate the effectiveness of our method on real-world data.
Figures
Reference graph
Works this paper leans on
-
[1]
Becker, B. and Kohavi, R. (1996). Adult . UCI Machine Learning Repository. DOI : https://doi.org/10.24432/C5XW20
doi:10.24432/c5xw20 1996
-
[2]
and Paschalidis, I
Chen, R. and Paschalidis, I. C. (2018). A robust learning approach for regression models based on distributionally robust optimization. Journal of Machine Learning Research , 19(13):1--48
2018
-
[3]
Chen, Y., Guo, Q., Sun, H., Li, Z., Wu, W., and Li, Z. (2018). A distributionally robust optimization model for unit commitment based on kullback--leibler divergence. IEEE Transactions on Power Systems , 33(5):5147--5160
2018
-
[4]
and Drusvyatskiy, D
Davis, D. and Drusvyatskiy, D. (2019). Stochastic model-based minimization of weakly convex functions. SIAM Journal on Optimization , 29(1):207--239
2019
-
[5]
C., Hashimoto, T., and Namkoong, H
Duchi, J. C., Hashimoto, T., and Namkoong, H. (2019). Distributionally robust losses against mixture covariate shifts. Under review , 2(1)
2019
-
[6]
Fuster, A., Goldsmith-Pinkham, P., Ramadorai, T., and Walther, A. (2022). Predictably unequal? the effects of machine learning on credit markets. The Journal of Finance , 77(1):5--47
2022
-
[7]
Ghosal, S. S. and Li, Y. (2023). Distributionally robust optimization with probabilistic group. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 37, pages 11809--11817
2023
-
[8]
and Sim, M
Goh, J. and Sim, M. (2010). Distributionally robust optimization and its tractable approximations. Operations research , 58(4-part-1):902--917
2010
-
[9]
J., Wojdyla, D
Hong, C., Pencina, M. J., Wojdyla, D. M., Hall, J. L., Judd, S. E., Cary, M., Engelhard, M. M., Berchuck, S., Xian, Y., D’Agostino, R., et al. (2023). Predictive accuracy of stroke risk prediction models across black and white race, sex, and age groups. Jama , 329(4):306--317
2023
-
[10]
Hu, W., Niu, G., Sato, I., and Sugiyama, M. (2018). Does distributionally robust supervised learning give robust classifiers? In International Conference on Machine Learning , pages 2029--2037. PMLR
2018
-
[11]
E., Kim, A
Khandani, A. E., Kim, A. J., and Lo, A. W. (2010). Consumer credit-risk models via machine-learning algorithms. Journal of Banking & Finance , 34(11):2767--2787
2010
-
[12]
M., Nguyen, V
Kuhn, D., Esfahani, P. M., Nguyen, V. A., and Shafieezadeh-Abadeh, S. (2019). Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations research & management science in the age of analytics , pages 130--166. Informs
2019
-
[13]
Lin, T., Jin, C., and Jordan, M. (2020). On gradient descent ascent for nonconvex-concave minimax problems. In International conference on machine learning , pages 6083--6093. PMLR
2020
-
[14]
and Duchi, J
Namkoong, H. and Duchi, J. C. (2016). Stochastic gradient methods for distributionally robust optimization with f-divergences. Advances in neural information processing systems , 29
2016
-
[15]
Oren, Y., Sagawa, S., Hashimoto, T. B., and Liang, P. (2019). Distributionally robust language modeling. arXiv preprint arXiv:1909.02060
Pith/arXiv arXiv 2019
-
[16]
Sagawa, S., Koh, P. W., Hashimoto, T. B., and Liang, P. (2019). Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv:1911.08731
Pith/arXiv arXiv 2019
-
[17]
Y., and Ghassemi, M
Seyyed-Kalantari, L., Liu, G., McDermott, M., Chen, I. Y., and Ghassemi, M. (2020). Chexclusion: Fairness gaps in deep chest x-ray classifiers. In BIOCOMPUTING 2021: proceedings of the Pacific symposium , pages 232--243. World Scientific
2020
-
[18]
M., and Kuhn, D
Shafieezadeh Abadeh, S., Mohajerin Esfahani, P. M., and Kuhn, D. (2015). Distributionally robust logistic regression. Advances in neural information processing systems , 28
2015
-
[19]
Sinha, A., Namkoong, H., Volpi, R., and Duchi, J. (2017). Certifying some distributional robustness with principled adversarial training. arXiv preprint arXiv:1710.10571
Pith/arXiv arXiv 2017
-
[20]
Soma, T., Gatmiry, K., and Jegelka, S. (2022). Optimal algorithms for group distributionally robust optimization and beyond. arXiv preprint arXiv:2212.13669
Pith/arXiv arXiv 2022
-
[21]
Villani, C. (2009). Optimal Transport: Old and New , volume 338 of Grundlehren der mathematischen Wissenschaften . Springer
2009
-
[22]
Yu, D., Cai, Y., Jiang, W., and Zhang, L. (2024). Efficient algorithms for empirical group distributionally robust optimization and beyond. arXiv preprint arXiv:2403.03562
Pith/arXiv arXiv 2024
-
[23]
Zhang, L., Zhao, P., Zhuang, Z.-H., Yang, T., and Zhou, Z.-H. (2023). Stochastic approximation approaches to group distributionally robust optimization. Advances in Neural Information Processing Systems , 36:52490--52522
2023
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.