Learning the Robustness Mechanism with Bilevel Optimization

• ERM can fail under subpopulation shift because groups defined by spurious attributes may have different proportions in training and test data.The paper notes that classifiers may inadvertently depend on those attributes in applications including healthcare, fairness, robotics, and autonomous driving.Source: §1 · p. 1
• Group DRO addresses inter-group shifts, but fixed group distributions overlook intra-group variation that can make test samples underrepresented within a group.The appropriate robustness level may differ across groups and requires extensive tuning; worst-group membership predictors also require tuning when group labels are unavailable.Source: §1 · p. 2
• The proposed framework learns robustness parameters from validation data using an upper-level minimax problem, alongside lower-level minimax robust training.Two instantiations cover settings with and without group labels, and the bilevel problem is tractably solved with a first-order method.Source: §1 · p. 2
• The authors derive sample-complexity generalization guarantees for the framework and report that continuous hyperparameter optimization avoids grid-search discretization error.The guarantees are described as advantageous relative to grid search.Source: §1 · p. 2
• The framework improves worst-group prediction under simultaneous intra-group and inter-group test shifts, both with and without group labels during training.The paper presents this as validation of an enhanced evaluation setting.Source: §1 · p. 2

2.1 GROUP DISTRIBUTIONALLY ROBUST OPTIMIZATION

• GDRO weights groups according to their average losses, and its robustness parameter η controls how far group weights can deviate from uniform weighting.Selecting η using validation loss over a grid requires training a model for each candidate and does not directly extend to missing group labels.Source: §2.1 · p. 3

2.2 WARM-UP: BILEVEL GROUP DRO (BI-GDRO)

• Bi-GDRO makes η a continuous upper-level decision variable, learning it and representation parameters θ from validation performance while optimizing classifier weights W in the lower level.The lower-level problem is convex in W and concave in q.Source: §2.2 · p. 3

2.3 BILEVEL-HIERARCHICAL DRO (BI-HDRO)

• HDRO models inter-group shifts with group weights and intra-group shifts with Wasserstein ambiguity sets around each group distribution; setting ϵ_g to zero reduces it to GDRO.The section notes that directly optimizing over the distributions in those sets is generally intractable.Source: §2.3 · p. 4
• For HDRO, the proposed modified loss adds variable u; it is jointly convex in W and u for binary hinge or logistic loss, while the multiclass cross-entropy formulation provides an upper bound.The formulation permits exact gradient evaluation without solving the inner maximization, and projection onto its constraint set has a closed form.Source: §2.3 · p. 4

2.4 BILEVEL-PROBABILISTIC GROUP DRO (BI-PG-DRO)

• Bi-PG-DRO integrates training the predictor and soft group predictor with tuning the robustness parameter in one bilevel model.Its soft group predictor is optimized for validation performance rather than separately trained for group-label prediction accuracy; the KL regularizer makes optimization more lightweight than the alternative described.Source: §2.4 · p. 5

2.5 BILEVEL MINIMAX ALGORITHM

• The bilevel minimax formulations are solved using a first-order method that penalizes the lower-level primal-dual gap and applies an inexact proximal-point method.This reduces the problem to approximately solved strongly-convex-strongly-concave subproblems, handled with the stochastic accelerated primal-dual algorithm.Source: §2.5 · p. 6

2.6 GENERALIZATION THEORY

• The generalization analysis assumes standard conditions including bounded continuous hyperparameter spaces and strongly convex lower-level regularization, which yield a bounded, Lipschitz-continuous hypothesis space.The stated theory also captures a bias-variance trade-off through approximation bias and a stability gap.Source: §2.6 · pp. 6–7
• Continuous bilevel tuning achieves a worst-group risk bound with statistical gaps and a continuous-tuning penalty, avoiding discretization error while incurring a penalty no worse than grid search.The framework’s tuning penalty scales logarithmically with the algorithmic Lipschitz constant; exhaustive grid search has exponential computational complexity in high-dimensional spaces.Source: §2.6 · p. 7

3 RELATED WORKS

• HDRO addresses both inter-group and intra-group uncertainty, while its fixed perturbation radii are sensitive and require extensive tuning.GDRO targets the highest-loss group, but minimizing worst-group training loss alone does not ensure test robustness; DFR and AFR retrain the last-layer classifier on group-balanced data.Source: §3 · p. 7
• SSA and PG-DRO use predicted groups for robust training, while AGRO jointly changes group assignments to increase group-wise loss.Other approaches identify spurious attributes through contrastive learning, concept banks, reference data, sample removal, or twin classifiers before robust training.Source: §3 · pp. 7–8
• Bilevel optimization has been applied to hyperparameter tuning, meta-learning, reinforcement learning, and neural architecture search.Prior hyperparameter-tuning work optimizes validation performance in the upper level; this paper treats the robustness mechanism, potentially including non-convex neural networks, as hyperparameters.Source: §3 · p. 8

4.1 DATASETS WITH TEST DISTRIBUTION SHIFT

• The shifted benchmarks rotate red minority digits in CMNIST, shift minority CelebA test images to glasses variants, and alter the minority toxic/black group’s “white” attribute in CivilComments.The setup combines intra-group and inter-group shifts; the authors do not investigate inter-group-only performance because baselines performed similarly and further tuning offered no improvement.Source: §4.1 · p. 8

4.2 RESULTS ON BI-HDRO

• On CelebA, HDRO improves from 59.2 of GDRO to 72.4, while added intra-group shift decimates strong baselines.The authors report tangible benefits from ambiguity sets across all datasets.Source: §4.2 · p. 8, Table 1 · p. 9
• Bi-HDRO outperforms lower-level HDRO alone and variants with individually fixed, manually tuned ϵg or η.The comparison is reported in the bottom half of Table 1.Source: §4.2 · p. 9

4.3 RESULTS ON BI-PG-DRO

• Without training-time group labels, the experiments use small labeled validation splits and separate them for attribute-predictor training and model tuning.The labeled fractions are 5% for CMNIST, 15% for CelebA, and 3% for CivilComments.Source: §4.3 · p. 9
• Bi-PG-DRO with KL-based implicit attribute-predictor tuning and η tuning performs best and outperforms all baseline methods.Implicit KL-based predictor tuning works just as well as explicit cross-entropy training; the complete bilevel method combines KL with η tuning.Source: §4.3 · p. 9

5 CONCLUSION

• The current framework is somewhat restrictive because its solver requires convex-concave lower-level training; no algorithm is available for the broader non-convex-concave case.The authors identify machine unlearning and language-model alignment as possible extensions.Source: §5 · p. 9

AI USE STATEMENT

• The authors report using generative AI for implementation, dataset cleanup and reformatting, proof-writing assistance, and writing, and say AI-assisted work was reviewed.They state that LLM-generated code was tested and proof steps were reviewed and revised by the authors.Source: AI USE STATEMENT · p. 10

A A FIRST-ORDER METHOD FOR BILEVEL MINIMAX PROBLEMS

• The first-order method constructs a proximal-penalized objective, solves its strongly-convex-strongly-concave form with SAPD, and returns a uniformly sampled iterate.Algorithmic hyperparameters are tuned once for the datasets and proposed methods; validation accuracy, rather than objective convergence, is used for model selection.Source: A A FIRST-ORDER METHOD FOR BILEVEL… · p. 16

B CLOSED FORM OF LOWER LEVEL ADVERSARY

• Computing the closed-form lower-level adversary requires losses for all groups, making it prohibitively expensive during stochastic training.Source: B CLOSED FORM OF LOWER LEVEL ADVERSARY · p. 16

C.1 CLOSED-FORM PERTURBATION FOR BINARY CLASSIFICATION

• For binary classification, the appendix derives closed-form robust hinge and logistic losses for a linear final layer, smoothing the norm term to obtain a jointly convex formulation.The broader binary-classification treatment uses projected descent for a second-order-cone constraint; the experiments use BCE, while hinge loss is also formulated.Source: §C.1 · pp. 16–17

C.2 PERTURBATION FOR MULTI-CLASS CLASSIFICATION

• For multiclass hinge loss, latent perturbation yields a closed-form robust loss convex in final-layer parameters; for softmax cross-entropy, the appendix instead derives an efficient upper-bound relaxation.The relaxation is jointly convex in the final-layer parameters and auxiliary variables for fixed features, and its gradient can be evaluated without solving the inner maximization.Source: §C.2 · pp. 18–19

D.1 SHIFTED DATASETS

• CMNIST rotations are applied to both validation and test samples of the minority group to reduce tuning variance and stabilize prediction performance.Samples are not moved; in the original HDRO setup, rotation was applied only to test data, producing 8%-10% variance across random seeds.Source: §D.1 · p. 19, Table 3 · p. 19
• CelebA and CivilComments shifts reallocate minority-group samples between train, validation, and test, including a CivilComments train/test split by white annotation.For CelebA, 164 blond males without glasses move from test to train, 90 with glasses from train to test, and 10 with glasses from validation to test. For CivilComments, 1,278 toxic, black samples with white=1 move from train to test, while 905 with white=0 move from test to train.Source: §D.1 · p. 20, Table 5 · p. 20, Table 4 · p. 20
• Waterbirds is omitted because its known mislabeled attributes, small training set, and saturated performance make comparisons less informative.The training set has 4,795 examples, including 56 waterbirds with land backgrounds in the minority group.Source: §D.1 · p. 20

D.2 SETTINGS

• Hyperparameters are described as relatively dataset-insensitive, so they are tuned on CMNIST for one method and reused for the other two datasets.The perturbation parameter ϵg is clipped to [0,1] to follow the HDRO preset range.Source: §D.2 · p. 20, Table 6 · p. 21

D.3.1 ABLATION PROCEDURE

• The ablation compares Bi-HDRO with learnable η and ϵ against Bi-GDRO, fixed-ϵ grid search, and joint grid search over ϵ and η.Bi-GDRO fixes ϵ to zero; the grid-search schedules test ϵ values {60,72,84,96,108}/255 and η values {0.01, 0.1, 1.0, 10.0}.Source: §D.3.1 · p. 21

D.3.2 BASELINES

• DFRTr trains ERM on a larger data partition before retraining the last layer on a smaller group-balanced subset; the comparison uses DFRTr for fairness of group-information access.DFRVal uses all minority-group data in validation and samples an equal amount from other groups.Source: §D.3.2 · p. 21
• PDE warms up the full model on a group-balanced subset, then progressively adds training data while continuing to train the model.The procedure tunes warm-up epochs in {10,20} and added samples in {50,100,500}.Source: §D.3.2 · p. 21
• HDRO is a single-level min-max-max baseline whose perturbation radius and C are tuned across all three datasets.The search uses ϵ in {60,72,84,96}/255 and C in {0, 1, 2, 3}.Source: §D.3.2 · p. 21

D.4.1 ABLATION PROCEDURE

• The ablation varies BCE loss and KL prior regularization for the attribute predictor over a grid of {1,10,30}.Source: §D.4.1 · p. 21

D.4.2 BASELINES

• GICCy trains an attribute predictor to maximize the KL difference between predicted and true spurious labels before GDRO; its γ is tuned on the shifted datasets.The γ search is {2,5,10}.Source: §D.4.2 · p. 21
• XRM identifies minority samples through confident held-out misclassification after flipping training labels, then trains GDRO with pseudo-group labels.Three hyperparameter combinations are sampled, the one with the highest flip rates is selected, and GDRO is trained with three seeds.Source: §D.4.2 · pp. 21–22
• AGRO jointly trains an attribute predictor and robust model with soft memberships, which the paper says make GDRO difficult; SSA and PG-DRO instead generate hard or probabilistic pseudo-group labels before GDRO.AGRO’s α is tuned in {0.2,0.3,0.4}; SSA and PG-DRO tune C in {0, 1, 2, 3}.Source: §D.4.2 · p. 22

E.1 EFFICIENCY

• Bi-HDRO reaches a validation level in one run that took HDRO about 6,000 optimization steps to reach, while HDRO also requires extensive grid search.Accounting for grid search would make HDRO's runtime longer; Table 7 reports wall-clock time per run.Source: §E.1 · p. 22, Table 7 · p. 22

E.2 DOES PERTURBATION HELP SOFT GROUP ASSIGNMENTS (BI-PG-DRO)?

• For Bi-PG-DRO, perturbation converges to zero across tested initializations, and active perturbation worsens validation performance.The tested cases include initial values of 0 and 96/255, plus an unclipped initialization of 0.Source: §E.2 · p. 23, Figure 2 · p. 23

E.3 SHOULD MEMBERSHIP INFERENCE BE PRECISE FOR WORST-GROUP PERFORMANCE?

• Worst-group validation performance need not depend on precise predicted group labels when bilevel optimization focuses on worst-group validation loss.Without upper-level regularization, post-hoc attribute accuracy is poor, while KL regularization gives closer attribute accuracy and a small worst-group validation boost.Source: §E.3 · p. 23, Figure 3 · p. 24

F GENERALIZATION THEORY FOR CONTINUOUS BILEVEL HYPERPARAMETER TUNING

• The proof roadmap uses a Lipschitz-continuous lower-level hypothesis mapping and Rademacher complexity to bound upper-level validation error across continuous hyperparameters.It also applies the machinery to regularized ERM before deriving the Group DRO result.Source: F GENERALIZATION THEORY FOR CONTINUOUS… · p. 24
• The theory establishes a Group DRO oracle inequality and extends continuous generalization guarantees to Bi-HDRO with multidimensional perturbation radii.The analysis describes a continuous oracle inequality as navigating the trade-off between structural complexity and worst-case group robustness.Source: F GENERALIZATION THEORY FOR CONTINUOUS… · p. 24

F.1 SETUP AND ASSUMPTIONS

• The setup evaluates predictors induced by a continuous hyperparameter mapping on an independent validation set under boundedness, convexity, and Lipschitz assumptions.The stated assumptions include a bounded continuous hyperparameter domain and a strongly convex lower-level regularizer.Source: §F.1 · p. 25

F.2 ALGORITHMIC LIPSCHITZ CONTINUITY EXAMPLES

• For regularized ERM, strong convexity and a bounded optimization domain make the predictor mapping Lipschitz-continuous in the regularization parameter.The stated Lipschitz constant is (B + ∥a∥)/λmin.Source: §F.2 · pp. 25–26
• For Group DRO, a fixed strongly convex regularizer makes the lower-level classifier Lipschitz-continuous in robustness parameter η when η is bounded away from zero.The continuous choice of η tunes the trade-off between average-case empirical risk and worst-group robustness on unseen data.Source: §F.2 · pp. 26–27

F.3 UNIFORM CONVERGENCE OVER MULTI-DIMENSIONAL CONTINUOUS HYPOTHESIS SPACES

• The validation-loss class over a bounded k-dimensional continuous hyperparameter space has Rademacher complexity bounded using the lower-level mapping's Lipschitz constant.The bound depends on the hyperparameter-space dimension and diameter, and on the loss and mapping Lipschitz constants.Source: §F.3 · pp. 27–28
• Combining the complexity bound with concentration yields a high-probability uniform deviation bound between population and validation loss across all continuous hyperparameters.The stated probability is 1 − δ.Source: §F.3 · pp. 27–28

F.4 WARM-UP: CONTINUOUS ORACLE INEQUALITY FOR REGULARIZED ERM

• The ERM warm-up tunes a continuous L2 regularization coefficient using training and independent validation data.The resulting predictors form a one-dimensional continuous hypothesis class.Source: §F.4 · p. 28
• With probability at least 1 − δ, the validation-selected regularized ERM predictor satisfies the continuous oracle inequality for any fixed reference predictor.The theorem assumes bounded Lipschitz loss, strong convexity, a bounded tuning interval, and a Lipschitz lower-level mapping.Source: §F.4 · p. 29
• A finite grid for regularization can introduce discretization error, while making the grid denser worsens the stated generalization guarantee.The grid-based approach selects a model using validation data and uses a union bound over its discrete options.Source: §F.4 · p. 29
• The inequality adapts to the reference predictor's regularization trade-off without discretization error, with a logarithmic uniform-convergence penalty and continuous optimization instead of repeated grid training.The bound combines validation uniform convergence with lower-level algorithmic stability.Source: §F.4 · p. 29

F.5 CONTINUOUS ORACLE INEQUALITY FOR GROUP DRO

• Under bounded, Lipschitz group losses, the lower-level Group DRO algorithm is uniformly stable with a modified constant governed by the smallest training group.The robustness penalty η is tuned while λ is held fixed to satisfy strong convexity.Source: §F.5 · p. 30
• The analysis obtains a final continuous oracle inequality for the true worst-group risk, with total failure probability exactly δ from uniform-convergence, stability-gap, and fixed-reference Hoeffding bounds.Source: §F.5 · pp. 31, 34
• Corollary F.6 extends the oracle guarantee to jointly tuned λ and η by treating them as a two-dimensional hyperparameter vector with a jointly Lipschitz algorithmic mapping.The resulting bound trades the validation uniform-deviation dimension factor against the local Group DRO stability gap.Source: §F.5 · pp. 34–35
• With deep encoder parameters tuned jointly, the stability constant remains unchanged, while the bound’s dimension factor becomes 2 + dθ and its covering radius uses diam(Ψ × Θ).The corollary assumes the encoder output is Lipschitz in θ and the classification-loss gradient is Lipschitz in the encoder output.Source: §F.5 · pp. 35–36
• The oracle bound controls predictor deviation through regularization, group-weight, and encoder-representation perturbations, with the combined bound divided by λ ≥ λ_min.The encoder contribution is bounded using loss smoothness and encoder Lipschitzness.Source: §F.5 · p. 36
• For a fixed hyperparameter configuration, uniform stability is unchanged from the linear case because θ is fixed during lower-level optimization and stability depends on loss properties and λ.The stated condition is that ℓ(W, z) is L-Lipschitz over the bounded representation space.Source: §F.5 · p. 36
• Specialized neural-network generalization theory, such as the Neural Tangent Kernel (Jacot et al., 2018), could further tighten the sample-complexity bounds.The passage presents this as a possible refinement to the generic continuous treatment.Source: §F.5 · p. 36

F.6 EXTENSION TO HIERARCHICAL DRO (BI-HDRO)

• Bi-HDRO tunes multidimensional perturbation radii ϵ alongside λ and η; with a smoothed robust loss, its lower-level mapping is jointly Lipschitz under the stated assumptions.The assumptions include bounded gradient shifts in ϵ and λ-strong convexity of the lower-level objective.Source: §F.6 · pp. 36–37
• For binary classification, the bounded-gradient-shift assumption holds for exact closed-form perturbations, yielding robust margin y(W ⊤z+b)−ϵ∥W∥∗.For L2 perturbations, the dual norm is ∥W∥2; the appendix also extends joint Lipschitz tuning to deep encoder parameters under additional assumptions.Source: §F.6 · pp. 37–38
• Bi-HDRO retains the same lower-level uniform stability constant as standard group DRO, even though tuning ϵ expands the algorithmic Lipschitz constant.The stability claim applies for a fixed configuration and uses global bounds on the base loss across perturbed inputs.Source: §F.6 · p. 38

F.7 BACKGROUND: UNIFORM STABILITY

• Uniform stability bounds the expected generalization gap, and the appendix states a high-probability risk bound for uniformly stable algorithms with bounded loss.For λ-strongly convex regularized loss minimization, the stability analysis bounds the parameter displacement after replacing one training example.Source: §F.7 · p. 39

F.8 HELPFUL LEMMAS

• The auxiliary lemmas bound the expected maximum of sub-Gaussian variables and show that taking a maximum is non-expansive under elementwise perturbations.These are stated as general inequalities for finite collections.Source: §F.8 · p. 41