Group-robust Sample Reweighting for Subpopulation Shifts via Influence Functions
Authors: Rui Qiao
Zhaoxuan Wu
Jingtan Wang
Pang Wei Koh
Bryan Kian Hsiang Low
Venue: ICLR 2025 Poster (2025)
Keywords: distribution shift
subpopulation shift
spurious correlation
influence function
sample reweighting
data selection
Primary area: alignment, fairness, safety, privacy, and societal considerations
TL;DR: We introduce Group-robust Sample Reweighting (GSR), which uses group-labeled data to guide the iterative retraining of the model its on group-unlabeled data reweighted using influence functions.

Abstract: Machine learning models often have uneven performance among subpopulations
(a.k.a., groups) in the data distributions. This poses a significant challenge for the
models to generalize when the proportions of the groups shift during deployment.
To improve robustness to such shifts, existing approaches have developed strategies
that train models or perform hyperparameter tuning using the group-labeled data
to minimize the worst-case loss over groups. However, a non-trivial amount of
high-quality labels is often required to obtain noticeable improvements. Given
the costliness of the labels, we propose to adopt a different paradigm to enhance
group label efficiency: utilizing the group-labeled data as a target set to optimize
the weights of other group-unlabeled data. We introduce Group-robust Sample
Reweighting (GSR), a two-stage approach that first learns the representations from
group-unlabeled data, and then tinkers the model by iteratively retraining its last
layer on the reweighted data using influence functions. Our GSR is theoretically
sound, practically lightweight, and effective in improving the robustness to sub-
population shifts. In particular, GSR outperforms the previous state-of-the-art
approaches that require the same amount or even more group labels. Our code is
available at https://github.com/qiaoruiyt/GSR.

Published as a conference paper at ICLR 2025
GROUP-ROBUST SAMPLE REWEIGHTING FOR SUBPOP- ULATION SHIFTS VIA INFLUENCE FUNCTIONS
Rui Qiao! Zhaoxuan 2u? Jingtan Wang"s Pang Wei Koh* Bryan Kian Hsiang Low 12 National University of Singapore "Singapore-Ne Alliance for Research and Technology "Agency for Science, Technology and Research (A*STAR) "University Washington (rui qiao, jingtan. w) @u. nus edu zhaoxuan. wue smart. mit edu pangwei ecs. washington.edu lowkhecomp nus edu.sg
ABSTRACT
Machine learning models often have uneven performance among subpopulations (a.k.a,, groups) in the data distributions. This poses significant challenge for the models generalize when the proportions of the groups shift during deployment. To improve robustness to such shifts, existing approaches have developed strategies that train models or perform hyperparameter tuning using the group-labeled data to minimize the worst-case loss over groups. However, a non-trivial amount of high-quality labels is often required to obtain noticeable improvements. Given the costliness of the labels, we propose to adopt a different paradigm to enhance group label efficiency: utilizing the group-labeled data as a target the weights of other group-unlabeled data. We introduce Group-robust Sample Reweighting (GSR), a two-stage approach that first learns the representations from group- coup-unllbeele data, and then tinkers the model by iteratively retraining its last layer on the reweighted data using influence functions. Our GSR theoretically sound, practically lightweight, and effective in improving the robustness to sub- population shifts. In particular, GSR outperforms the previous state-of-the- approaches that require the same amount even more group labels. Our code is available at /tttttt::/// /ectps://githun comb ctw.///orulyt//ss
optimize
INTRODUCTION
Subpopulation shift refers to the change in the proportion of groups (as defined by the class, attributes. or both) in data distribution between the training and deployment phases (Koh et al. 2021 Yang
al. 2023). Due the nature of subpopulations or the selection bias (Dwork al.) [012]] Zadrozny
2004), the training data is often composed of majority and minority groups that are overrepresented and underrepresented, respectively. Common examples include photos of waterbirds frequently being taken with water background instead of land background (Sagawa et al.) (219)) and people are more likely to be holding a vacuum than fixing it (Nan et al.) (021). Such group imbalance can cause some features from the majority group to be spuriously correlated with the labels, leading to possible undesirable shortcuts in training deep neural networks (Geirhos et al.) 2020) Hermann et al.) 2024. (hah et al. 2020). As result, standard training with empirical risk miniiization (ERM) often overly relies on such shortcuts and fails to attend to the minority groups, leading to worsened model generalization under subpopulation shifts, especially when the proportion of the minority groups increases (Arjovsky et al. 2019 Nagarajan et al.) 2020; Sagawa et al. 2020). Thus, it is important to develop algorithms that are group-robust subpopulation shifts, experiencing minimal performance degradation as measured by the worst-group accuracy.
The group labels can be used to construct balanced groups or minimize the worst-group loss during training (Idriss et al., 2022) Sagawa et al.) 2019). These techniques often lead to improved group robustness compared to ERM. However, group labels can be expensive to obtain, necessitating approaches that maximize the gain in group robustness from limited data. Given labeldl dataset DL. the two mainstream categories are either using DL for model selection (Liu al. 2021
Partial work done while the author was visiting scholar the University of Washington.

---

Published as a conference paper at ICLR 2025
Zhang et al. 2022), or using D to directly train the model parameters (Kirichenko et al. 2022) However, these strategies may not always be the best way to exploit the group information, since they either underutilize the group labels, restrict the most crucial part of the training to only D while neglecting the potential of the more accessible data without group labels. Therefore, we advocate an in-between training paradigm: instead of using DD train the model parameters directly, we use them to iteratively optimize the weights of a group-unlabeled dataset DU which are more realistic to obtain. Then, the model parameters are trained using a weighted ERM objective on these data.
Optimizing the sample weights is an established idea that necessitates solving bilevel optimization problem (Ren et al., 2018). The inner loop of the bilevel optimization performs standard model training on a weighted objective, while the outer loop optimizes the sample weights. The two main challenges hindering its wider adoption include: (1) the unavoidable trade-off between computational cost and accuracy, and (2) the effectiveness when the model is overparameterized For (1), calculating the outer loop gradient requires backpropagating through the entire training process, which compu- tationally prohibitive for deep neural networks. Conventional approaches use a one-step truncated backpropagation that relies only on the last-step gradient of the model training for computational feasibility (Shaban al. 2019 Zhou et al.) (022). However, this approximation can be imprecise as it fails to account for the curvature of the loss landscape and the training dynamics of the model. For (2), training a weighted objective with overparameterized neural networks is likely to converge the same solution as with ERM (Sagawa et al. 2019 ZZai et al.] 2022). Regularization techniques need to be applied such that training with weighted objectives leads to meaningfully different solutions than training with ERM (Byrd & Lipton, 2019 Sagawa et al. 2019; Zhai et al. 2022).
To address the two challenges, we propose a two-stage method named Group-robust Sample Reweighting (GSR), which iteratively reweights individual training samples to improve group robust- ness. We utilize last-layer retraining (LLR), a lightweight method that retrains the last linear layer of neural network (Kang et al.) 2020; Kirichenko et al. 2022). LLR simplifies the inner loop of the bilevel optimization into a convex optimization problem. Importantly, it facilitates our application of the influence function, a technique that can be derived from implicit differentiation (Koh & I Liang] 2017; Krantz& & Parks 2002). to utilize the Hessian to accurately estimate the gradient of sample weight updates without backpropagating through the entire training trajectory. In contrast to methods that rely on one-step truncated backpropagation approximations (Zhou et al. 2022), our approach leverages the fact that LLR is both inexpensive even with Hessian computation and effective for enhancing group robustness. To perform GSR, we split the group- couu-unlabeled dataset D into held-out set DU-h and remaining set DU= The first stage performs unweighted representation learning on DU.r The second stage utilizes the influence function to iteratively optimize the weights of Di-h for the worst-group loss in D achieved through LLR. Our method outperforms the state-of ahe method that uses the same amount of group labels and even outperforms methods that require more group labels.
The specific contributions of this work include the following:
We propose to better leverage the high-quality group-labeled data for group robustness with an alternative paradigm by using them to reweight other samples instead of directly training on them.
We devise an efficient strategy based on implicit differentiation for group-robust sample reweighting, which becomes accurate and computationally feasible via the synergy with last-layer retraining.
We empirically demonstrate the performance advantages of our lightweight approach on improving the group robustness for both vision and natural language tasks.
PROBLEM FORMULATION AND PRELIMINARIES
Let (X,, () denote the space of input, class label, and groups of subpopulation respectively. Denote a dataset by D = (=)...,, ). where each data point 2; == (I;, yı, gi) E X' x y' x G is sampled i.i.d. from data distribution P. For g E G, define Dg = {z+ E Dig = 9) as the subset of data from group g. sampled from the group data distribution 2.p. We assume that P. can be decomposed into P P=Σ R. 2, where the mixing ratio Fg E [0, 1]. I Epec T = As a result, any subpopulation shifts can be represented by adjusting kg. Depending on the group-label availability, datasets can be categorized as DI. DV for labueel group-unlabeled datasets. Meanwhile, depending on the functionality, datasets can be categorized as D DV, Dinr, denoting the training, validation, target sets, respectively. The target set Daa guides the sample weight updates and reserve the notion of
=.
2

---

Published as a conference paper at ICLR 2025
validation set D only for hyperparameter and model selection. Let N {I...., n) be the indices of the training set. We use c for the total derivative and use or V for the partial derivative.
ERM optimizes the model parameters with R(D": 0) = EiEN 1 f(r ys: 0) where is the loss function. Let R" be the weight vector for all training samples. A weighted training objective assigns different importance the training data points, R(D"; we 0) = Lien wel(xi, yı; 0). Since the change of ı is not known at test time, we evaluate the robustness to distribution shifts of a model with parameters using the worst-group risk (Sagawa et al. 2019)
a
R. RwO (D; RW max RDD 959
To minimize the worst-group risk, existing methods (Kirichenko et al. 2022 Sagawa et al.
2019)
have focused on directly training on group-annotated dataset D using the objective R Rwo (D'; 0).
However, they require having a sufficient amount of data for each group. Alternatively, we can assume the availability of some group-unlabeles data DU (e.g., standard training set). By viewing the sample weights of D as another set of parameters and optimizing it jointly with 0 w.r.t. the target set Diar C DL. we arrive at bilevel formulation of the minimax problem (Zhou et al. 2022
min max R(D 0*)
wES gEg 9
s.t. 0+ arg min R(D"; w. w,
where S = {w [w0,, "2,"." -.. wa] E RR/w;; we 2 Ovi E N, and ΣigN w; = 1]. The inner loop performs standard model training on the weighted objective with variants of gradient descent. To optimize the sample weights of the training set in the outer loop, we can similarly perform gradient descent using its total derivative:
dR(DD=;;) = VOR(D (7) d dw
dw
de Since Or is typically obtained via gradient descent, calculating dw involves unrolling the entire training trajectory and backpropagating through it, which is overly expensive compute especially when the model is an overparameterized deep neural network. To efficiently approximate de MAPLE (Zoou et al. 2022) adopts one-step truncated backpropagation (Shaban al. 2019):
de*
d
(D" w. 0r-1)
one step
V =
(2)
dw
a0dw
where 7) is the learning rate. Subsequently, w is updated with projected gradient descent. This approximation essentially relies the last-step gradient w.r.t. each training instance z;, i.e., V Or
-n Ro x Vo Or), Or), derived in Appendix A.2
After dropping 17, we rewrite:
d'R 0*)
MAPLE
-V VOR 0r) V R(z00r)
dw;
To interpret, MAPLE essentially updates sample weights according to the inner product between the gradient of the target loss and the gradient of the unweighted training loss for each sample upon convergence. Performing bilevel optimization with this approximation is reasonable because it tends to upweight training points that share similar gradient directions with the target sets. However, also imprecise, as it neither accounts for the curvature of the loss landscape w.r.t. the current model nor fully captures the training dynamics, where w; could significantly affect the batch gradient and optimization trajectory. Reweighting data points based solely last-step gradients is shortsighted and can lead to suboptimal convergence.
3 GROUP-ROBUST SAMPLE REWEIGHTING VIA IMPLICIT DIFFERENTIATION
Our goal is to develop a technique that updates the sample weights more accurately and efficiently to optimize the bilevel minimax objective in Equation [T] In this section, we first establish the connection between reweighting samples via implicit differentiation and the influence function (Koh & Liang 2017) from the perspective of bilevel optimization. Then, present an effective integration of the
3

---

Published as a conference paper at ICLR 2025
influence function and adaptive aggregation (Sagawa et al. 2019) over the groups to optimize the minimax objective in the outer loop.
Implicit differentiation can be used to calculate the gradient w.r.t. the sample weights without backpropagating through the inner loop of the objective according to the implicit function theorem (IFT) (Krantz& & Parks, 2002). Assuming implies that (1) the gradient w.r.t. the model parameters 8 diminishes to zero upon convergence to the inner loop optimum, i.e., VeR w. 0*) = 0, and (2) the Hessian of the weighted training objective evaluated at 0 0", 11.. - = ViR(D"; w, 0*), is invertible (Boyd & Vandenberghe 2004). This satisfies two of the conditions required to apply the IFT. In addition, further requires the inner objective to be twice continuously differentiable. We formally state the assumptions below.
unconstrained and strictly convex inner objective, it
Assumption 3.1. The inner objective RDD w, 0) is unconstrained, strictly convex, and twice continuously differentiable with respect to 0, Vw E S (Equation
Although appearing as restrictive, these assumptions can be easily satisfied through last-layer retrain- Under Assumption 3.1 as VoR(D"; w. 0*) = O, the outer gradient
ing, as detailed in Section a _00-
dØ" des
db-
E B can be calculated exactly via implicit differentiation
dw
dwss
dun.
d
V V H 6*
w. w. O²)
(4)
dw
H 0+ V Ro, - Vol(z;00).
(5)
dw;
The derivation is in Appendix AA.] To highlight, Equation 5 uses the first-order gradient of the unweighted ERM loss which independent of w;. Hence, the magnitude of w; does not directly affect the scale of its gradient, even if |ua] + O. Instead, w; indirectly affects the gradients through the trained model parameters A and the Hessian H8.... The total derivative of R. w.r.t. wy is thus:
(+)) = V RDD ( HT VoR(zz;0)).
(6)
dw;
In comparison to Equation the key difference lies in the presence of the inverse Hessian term between the two gradients. The Hessian Hå- describes the loss landscape of 0 w.r.t. the training dataset D and their sample weights w. A more detailed discussion on the role of Hessian available in Appendix C We will show later in Section 5 that without the Hessian, the last-step approximated gradient (Equation 3 the outer objective function is inaccurate, often leading to suboptimal results.
The derived form in Equation 6 via implicit differentiation exactly matches the influence function, which estimates the change in model parameters if a training data point 2; is upweighted infinitesimally M
(Cook& & Weisberg 1982; Koh & Liang, 2017). Since the influence function is well-defined for weighted objectives, we slightly generalize the notation by incorporating w into the expressions.
Given a model 0 trained on weighted samples, the influence of upweighting z; E Dir on the loss of model predicting a target instance z' Dtar (or target set Diar)
I(2,, z';0 w) = - VoR(('' H voR(2z;
Thus, we can use the influence function to represent the gradient derived from implicit differentia- tion. For completeness, we define the per-sample influence of a dataset to be ī(D", 2²; 0- w) = I(21, z'; 0-, w) I(22, ple weight update based on the per-sample influence in D on each group g in Dau.
(",,0, w) 0'; I(zn, 2':n=';'00 2';0", w)) E |_xx1, Then, we can calculate the sam-
İ(D" Dir 0- w)
Σ Ī(D" ,''00".
(8)
D
DED
To optimize the minimax objective in the outer loop, the gradient of the worst-group risk w.r.t the sample weights can be calculated via its corresponding influence scores:
d'R WG (D)ar; 0²) I(D"" Dic; 0* w), g" = arg max gEğ
R(D
(9)
dw
4

---

Published as a conference paper at ICLR 2025
However, only optimizing for the worst group in step according to Equation 9 can be inefficient, especially when multiple groups have high error rates. To obtain a more efficient and smoother optimization trajectory, we utilize adaptive aggregation (Sagawa et al. 2019) of the influence scores across groups. The aggregation weights are updated multiplicatively based on group error rates better optimize the sample weights, illustrated in lines 7-14 of Algorithm
EXACT GROUP-ROBUST SAMPLE REWEIGHTING WITH LAST-LAYER RETRAINING
In this section, we introduce the Group-robust Sample Reweighting with last-layer retraining (GSR) algorithm. The algorithm is built on the observation that last-layer retraining (LLR) (Kang et al. 2020 Kirichenko al. 2022) can lead to strongly convex (a subset of strictly convex) inner objective that satisfies the assumptions required for implicit differentiation. Next, we delineate the benefits of LLR and its integration into GSR algorithm.
4.1 SYNERGY WITH LAST-LAYER RETRAINING
LLR is a lightweight method designed to mitigate spurious correlations and improve group robust- ness (Kirichenko et al. 2022 LaBonte et al., 2024 Qiu et al.) 2023). In essence, it performs deep representation learning with ERM. followed by retraining only the last layer a separate group- balanced dataset while freezing the remaining model parameters. Empirical evidence has shown that the initial ERM model has learned sufficient representations for all groups, including minority groups, and over-reliance on spurious correlations can be significantly reduced by simply finc-tuning the last linear classification layer of the model (Izmailov al., 2022, Rosenfeld et al. (022). As a result, despite being lightweight, last-layer retraining while freezing the ERM- representation has become the state-of ofate-of-the-eee of-tthe-rt approach for improving group robustness across various settings.
Notably, LLR synergies well with our proposed sample reweighting via implicit differentiation (Section D. offering three significant benefits. Firstly, employing L.R with cross-entropy loss and positive coefficient À for & regularization creates strongly convex objective that satisfies the Assumption 3.T (proofs are in Appendix (A.4). This allows the exact calculation of the gradients w.r.t. the sample weights via Equation 9 Secondly, calculating the Hessian inverse in Equation 6 is no longer overly expensive, because the only free model parameters are in the last layer instead of the entire deep neural network. Thirdly, LLR acts regularizer by limiting the capacity of the model parameters, which counterintuitively facilitates the training with weighted objectives. This is because weighted objectives do not necessarily yield different predictors compared ERM when training overparameterized models as shown in Zhai et al. (2022). Sufficient regularization is usually needed for the reweighting scheme to be practically meaningful (Byrd & Lipton 2019 Sagawa et al.) 2019). Therefore, regularized LLR ensures that the theoretical soundness of our method translates into practical benefits without violating assumptions or conducting extensive approximations.
4.2 OUR GSR ALGORITHM
comprehensive overview of GSR is presented in Algorithm [ For simplicity, we omit certain details such as regularization and gradient clipping. Specifically, it contains two stages:
(Stage 1) Representation learning. We take out random subset (e.g., 10%) from the training set as a held-out set Dirh Then, train the entire model f on the remaining training set D with ERM until convergence without early stopping. This ensures that the model fits the "remaining" training set D well. Reserving a held-out set Duth that the model has not encountered during training is a crucial step, as it prevents overfitting to the data that will be later used for sample reweighting. Otherwise, changing the sample weights will not make meaningful difference to the last-layer classifier (Zhai ct al. 2022). Note that neither group labels nor sample weights are required at this stage.
(Stage 2) Group-robust sample reweighting with last-layer retraining. Denote the model param- eters as 0 = (ò. e) where & is the feature extractor and 3 is the last-layer linear classifier. In this stage, keep the feature extractor fixed, and jointly train the last layer & and the sample weights w using the held-out set Darh (as a training set) and the target set Diar, Subsequently, we perform model selection with the validation set DV We create Diar and DV with equal size, by splitting the
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Group-robust Sample Reweighting with last-layer retraining (GSR)
Require: Training set D" of size n, validation set DV target set Diar with m groups, held-out set fraction a, outer loop steps T, outer learning rate 3, scaling temperature T.
1: Obtain the held-out set Dirh with size an and the remaining set Dir with size ann
Stage 1:
2: ()), +- arg min)= R(D
Stage
3: Initialize held-out sample weights w()) +- 1 10n
4: Initialize adaptive influence score weights 710) m 1
5: for = 1,..., do
6: 2b(1) f- arg ming R(Drch w((, 10'0,(0'5 (0",,))
> Last-layer retraining
7: for 9 ,,..., do
8:
79 t- 79
(t)
(1-1)
exp (R(Dir: (6', (6',"")))))) ((1))//)
> Update aggregation weights
9:
end for
10:
((() ((() /1t//l//00 (0
11:
<(()) - 0an
12:
for g = I..... m do
13:
((t) g(t) + 19
(1)
I(Dth Dir (Ø", wiit), w)
Adaptive aggregation of the influence
14:
end for
15:
w(++1) < -maaxfw (t) 35(1),000)
Projected gradient descent (element-wise max) D Weight normalization
16:
w(++)) t- -w(++)) (n((+))/|lu((tt /||||
17:
if WERWG (D"; (0) x())) 1 R WC (D); (6", ,))))
then
18:
1/*h ((())
Model selection
19:
end if
20:
end for
21:
return 0 + (*)
initial "standard" validation set (as D in half (see remark below for the reason). The inner loop, which fits linear classifier v on the weighted cross-entropy objective with &2 regularization, is optimized via L-BFGS (Liu & Nocedal 1989) for efficiency. The outer loop performs projected gradient descent with learning rate 3. The gradient is calculated using adaptive aggregation of the influence scores (Equation 9). projected to {w|w 2 Ovi). The Hessian of the inner objective is updated as Horc,, (e)) v RR (Dtr-h, (() (0()) in each iteration. In addition, the gradient of the sample weights can become spiky when high loss values are encountered on the validation set. We employ gradient clipping by norm and weights normalization for better training stability.
Remark. The target set and the validation set cannot be the same. Although the target data is only used to update the weights of the held-out training data and does not directly affect the training of the model parameters, overfitting to the target can still occur as a result of the representer theorem (Schölkopf et al. 2001). Thus, it essential create different splits for the target and validation set to prevent overfitting during sample reweighting, which we illustrate later in Figure 4d
5 EXPERIMENTS
Datasets. We evaluate the effectiveness of algorithms on 4 commonly used datasets. Waterbirds (Wah) al. (011) a binary object recognition dataset for bird types (i.e., waterbird. landbird), which are spuriously correlated with the background (i.c., water, land). CelebA (Liu et al. 2015) is binary object recognition dataset for hair color blondness prediction. There exists spurious correlation between the man gender attribute and con-blondness. MultiNLI (Williams et al. (017)) is multi-class natural language inference dataset. The three classes (i.e., entailment, neutral, contradiction) describe the relationship between pair of sentences. The spurious correlation exists between contradiction and the presence of negation words. CivilComments (WILDS) (Borkan et al. 2019 Koh et al.) 2021) binary text toxicity detection dataset. There are 8 types of identities mentioned in the text, such as male and female. Grouping the samples by the identity and the class results in 16 overlapping groups. Although there obvious spurious correlations, the data is extremely imbalanced among

---

Published as a conference paper at ICLR 2025
Table Performance comparison. We report the worst-group accuracy on four benchmark datasets. For the group label column, X indicates no group label g is used; v indicates is required; wr indicates a proportion of g from the validation set re-purposed beyond hyperparameter tuning, such as training sample weights model parameters. The row sections are ranked in descending order of the total amount group label information required. report the mean standard deviation over 5 random seeds. We use "_" to indicate missing evaluation results from the original paper.
Group label
Worst-group accuracy (%)
Method
Average
Train Val
Waterbirds
CelehA
MultiNLI
CivilComments
Group DRO
vl V
91.4+1.1
88.942.3
77.7+1.4
70.0+2.
82.0
RWG
V1 V
87.6+1.6
84.3+1.s
69.6+1.0
72.0+1.
78.4
XI V
86.7
81.1
72.6
69.3
77.4
CnC
XI V
88 30.3
88.8 88.8
68.942.1
SSA
89.0+0.6
19.8
76.640.7
69.942.0
MAPLE
XI W
91.7
88.0
72.7
64.1
79.1
DFR
XI w
92.9
88.3+1.1
74.7+0.7
70.1
81.5
GSR-HE
X/ w
87.5+0.1
86.3+0.4
74.7+0.0
68.940.2
79.4
XI WV
92.9
87.0+0.4
78.5
71.7
82.5
ERM
XI X
74.9
46.9
65.9
55.640.6
60.8
SELF
X/
93.0=ox
83.940.9
70.7+2.5
that mentions other religions and is classified as non-toxic. Overall, we follow the standard train, validation, and test splits for all datasets. We randomly create target and validation sets by equally splitting the original validation.
these groups, especially
Baselines. We compare our approach against baselines in three categories according to the amount of group labels required. (I) Group labels for both training and validation: Group DRO (Sagawa et al. 2019) optimizes the gorst-group training loss by dynamically adjusting the group weights; RWG (Idrissi et al.) 2022) balances the sampling probability of each group according to their sizes.
(2) Group labels for validation: JTT (Liu et al., 2021) upweights the high-loss training points that
are more likely to be from minority groups. CnC (Zhang et al. 2022) in addition uses contrastive
learning to align the representations of the same-class samples. SSA (Nam et al. [002)) infers the
pseudo group labels of the training set by training a predictor on the validation set. DFR (Kinichenkol
et al. 2022 performs group- group-balance balanced last-layer retraining the validation set. MAPLE (Zhou
2022 is the most similar to ours that uses the validation set to jointly reweight the training et al.
samples and retrain the entire model. (3) No group labels: SELF (LaBonte et al. 2024) construets an approximately group-balanced dataset based on the disagreement from an auxiliary model, then performs class-balanced last-layer retraining. We choose the perttrrrrfomngggg ES disagreement SELF
Setup. For GSR, at stage 1, we use the suggested hyperparameter configuration and data augmentation strategies from DFR all datasets (with minor modifications for MultiNLI and CivilComments). At stage 2, we drop data augmentation and perform a randomized search of hyperparameters using the valuation set D The details are documented in Appendix B.2
5.1
IMPROVEMENT IN GROUP ROBUSTNESS
The main results are in Table Our approach GSR achieves consistent and competitive results compared to the baseline methods terms group robustness, which is measured by the worst- group accuracy. Specifically, GSR has the state-of-the-aa (SoTA) performance for both the MultiNLI and CivilComments datasets, and Moose-to-SoTA worst-group accuracy on the Waterbirds and CelebA datasets. Moreover, GSR demonstrated the highest consistency across all four datasets. It achieves an average improvement of 1.0% in absolute worst-group accuracy compared to DFR (the SoTA method which uses the same amount of group labels as ours), because GSR allows more fine-grained reweighting over the training data. GSR even outperforms Group DRO which requires more group labels, indicating that training the full network with weighted data may not always be necessary. The detailed results including the average accuracy of the methods are in Appendix B.5
We re-implemented MAPLE E. with hetter efficiency and support for MultiNII and CivilComment..
7

---

Published as a conference paper at ICLR 2025
Landa lind wotarceda Inineroo
1-4-
0.35
velarsida wunder
Greug wwaghe
contadiciann
cortrescjat hay
0.0
6.55
antrrc
thinel lemale
danee, mar
6.10
aearat neg Inkoorty
male
0.00
65
100
20
100
15
Steps
Steps
$7830
(a) Waterbirds
(b) CelebA
(c) MNLI
(d) CivilComments
Figure I: The change in the sum of sample weights across different groups throughout the training. The minority groups upweighted and the majority groups are generally downweighted. However, different groups do not have equal sums of weights.
1.00
1.00
133 MOXK, Seeap
Oostaditionn no Eig
7.00
Other 15 grrupe geinortty.
4.00-
Lieetrn water centrrt watartarde or Imaioarty
Toroad Slond, rusil Orinerly)
Frient
2.50
erts krvet, rag (rikwity)
2.00
1.00-
0.90
reutral noinng
3.50
0:47
3:07
2.00
1.00-
0:30
0:50
0:50
6:00
0.00
Wrights
Wi ghtss
3
Wiights
Wrighes
(a) Waterbirds
(b) CelebA
(C) MNLI
(d) CivilComments
Figure The distribution of sample weights for each group that used to train the best models. The minority groups have the weight distribution stretched out towards high values, while the majority- group weights are generally skewed towards O.
Besides, as an ablation study, we include the Hessian-free version of our method, GSR-HF, which uses simplified sample weight update from MAPLE (Equation 3 during the last-layer retraining. As shown in Table [ GSR-HF performs consistently worse than GSR, validating our analysis in Section 2 This highlights the importance of incorporating the Hessian matrix in Equation to ensure the correctness of the gradient when optimizing the sample weights.
5.2 A CLOSER LOOK AT SAMPLE WEIGHTS
Varying weights across groups. Figure D Iillustrates the variation in the sum of sample weights across different groups through the training process of the best-performing model reported in Table [ In general, all the minority-group weights increase, and the majority groups with spurious correlations generally have decreased weights. For example, the weights for minority groups (e.g., landbirds in water in Figure Ta) increase as optimization progresses. In contrast, the weights for certain majority groups with clear spurious correlations (e.g., landbirds on land in Figure Tab decrease over the optimization steps. Note that not all majority groups necessarily experience a decrease in their sum of sample weights. Another key message here is that the better-performing models are not necessarily trained with equally weighted groups In all cases, the worst-group performance is better optimized even though the groups are weighted differently, which effectively mitigates the spurious correlations with still "imbalanced" data. These results obtained from directly optimizing the sample weights align well with the similar findings in Qiu et al. (2023: Sagawa al. (2019).
Varying weights within groups. We visualize the distribution of the sample weights that are used (where the distributions are smoothed using kernel density
to train the best models in Figure 2 estimation). The samples from the majority groups are consistently assigned with lower weights after reweighting, especially for Waterbirds and CelebA in Figures 2a 2b On the other hand, the minority groups are consistently upweighted with their sample weights stretched out to larger values across all datasets. We can also clearly see that our best-performing models are trained on samples with diverse weights, even if they from the same group. This implies that not all samples need to be upweighted or downweighted for group-balame balanced training. Having more fine- -eeeegraindd adjustments over the sample weights offers more fiexibility in improving the robustness subpopulation shifts.
8

---

Published as a conference paper at ICLR 2025
Visualizing high-weight samples. In addition, visualize the held-out samples with the highest weight from Waterbirds in Fi These high-weight samples also tend to be consistent across different runs. All of these images water). As highlighted by the red box, the image has an ambiguous background that blends features of both land and water. Hence, it should be considered as minority. Since the group labels for the held-out sets are not used for retraining, GSR is not affected by the annotation errors for group labels and correctly assigns high weights to these samples, since they are supposedly helpful in improving the worst-group performance on the target set. However, other group- balanced baselines that directly train on these group labels do not offer such a benefit. We will subsequently show that GSR robust to class-label errors in the training data in Section 5.3
from the minority groups (waterbirds on land, landbirds in
Class landbird Closs waterbird waterbird Class: landbird landbird
landaird
Cass. landb andbird
watcrbird
landbird
BCG watee
Iic land
BGG metler
Bic waler
BG:: War
DG: lend
Bic land
BC: Land
BC: water
Bic: lend
(a) The most weighted instances.
(b) The least weighted instances.
Figure 3: We illustrate the selected images with their class label and background label from the held-out split in Waterbirds according to the sample weights. In the top-5 most weighted instances are all from the minority groups with differing bird types and backgrounds. As highlighted in the red box in (a), the background of the waterbird is ambiguous as it blends features of both land and water and hence should be categorized as a minority. GSR correctly identifies it despite the suboptimal annotation. In 3b the top-5 least weighted instances are all from the majority groups where the background spuriously correlated with the bird type.
5.3 ROBUSTNESS TO LABEL NOISE
The real-world data often have annotation errors (Beyer et al., 2020 Gururangan al. 2018).
not addressed properly, the errors can substantially degrade the model 's group robustness when the spurious correlation is strong in the training set (Qiao Low 2024). Assuming efforts are invested into the quality of the group-labeled dataset DI. the other training data without group labels in DU may have unaddressed errors such as mislabeled classes. To evaluate the resilience of GSR on training data with different quality, simulate the variation of data quality by adding class label noise. In particular, we randomly flip up to 40% of the class labels the held-out set to incorrect classes, and leave the target and validation sets unchanged. As shown in Figure 4a| the gorst-group performance of GSR has minimal performance degradation. This is because using high-quality targets for training sample reweighting allows more fine-grained control over the training data. To further understand why, we plot the learned weight distribution of the held-out data in Waterbirds. Compared to that of the clean data in Figure 4b) the weight distribution of the noisy data more skewed towards as shown in Figure 4c In particular, almost all minority instances with noisy labels have close-to-0
00
20-
80
07
Covtul
Muttu
DivlComments
40
60-
parget tarpet
20
20-
og bast large:
0.1
6.2
0.3
0.6
000
0.00
01504
CAW oone
(aoss 1000 0.839 0.6
20
db
60
100
(a) Label noise robustness
(b) Clean data weights
(c) Noisy data weights
(d) Val vs. test acc
Figure 4: In-depth study our algorithm. In 4n the worst-group test accuracy degrades slightly even when up to 40% of the held-out set labels are corrupted. In 4b most of the uncorrupted minority samples received higher weight assignments than non-minority examples. In contrast, in 4c the corrupted minority instances are correctly assigned with close-to-O weights. shows the relationship between validation and test worst-group accuracy on the CelebA dataset. It is important to have separate target and validation sets. Otherwise, overfitting can easily occur.

---

Published as a conference paper at ICLR 2025
weights, equivalent to being removed from the training set. Therefore, the efficacy of GSR can be attributed to the property of automatically cleaning the training data via sample reweighting, achieving similar effect other noisy label detection and removal approaches (Ghorbani & Zou 2019| Wang et al., 2024b). GSR being robust to class label corruption has an important implication: to improve the robustness to distribution shifts under limited annotation budget, more efforts should be dedicated to ensuring the quality of the target set.
6 RELATED WORK
Subpopulation shifts. Group labels are essential for training and evaluating machine learning models under subpopulation shifts. Various approaches have been proposed to handle scenarios with
different levels of availability of group labels: fully known for all data splits (Deng et al. 2024
Idrissi et al. 20022 [Sagawa et al. 2019), partially known (Kirichenko et al. 2022] Liu et al., 2021
2022] Nam et al. 2022.Qui Qiu et al. 20023] Zhang et al. 2022, Zhou et al. 2022]. J. or completely
unknown (Han & Zou| |024, LaBonte et al., (024). In addition. Yang et al. (2023) comprehensively
studied the effectiveness of these approaches on a variety of benchmarks with unknown group labels and demonstrated that optimizing for the worst class can be surprisingly effective for improving the
worst-group performance. Our approach falls into the category of learning with partial group labels (from the validation set) and differs from most of the approaches by using group-labeled data for sample reweighting instead of direct training.
Sample reweighting and implicit differentiation. Importance weighting is classic paradigm for learning under distribution shifts (Chen et al.) 2025 Fang et al.] 2020; 2023). The idea can be further expanded to bilevel optimization of the sample weights and the model parameters (Ren) eet al. 2018 Zhou et al. 2022). These works on bilevel optimization focus on using truncated backpropagation circum vent the challenge of differentiating through the unrolled training trajectory. Implicit differentiation is an alternative solution to the problem and has been applied to areas such as hyperparameter optimization (Foo et al., 2007 Lorraine et al., 2020 Pedregosa 2016), meta- learning (Lee et al.) [2019] Rajeswaran et al. 2019). generative modeling (Domke 2012] Samuel &
Tappen 2009), and fairness (Shui et al.) 2022). )). In addition, [i & Liu (2022); Wang et al. (2024a) have
explored one-step (i.e., not iterative) sample reweighting via the influence function for algorithmic fairness. [Bae et al. (2022) shows an alternative derivation of the influence function via implicit differentiation the training set. In contrast to these works, show the connection between the influence function and bilevel optimization via implicit differentiation and devise an efficient scheme to rigorously apply it iteratively for addressing subpopulation shifts.
7 CONCLUSION AND FUTURE WORK
Motivated from the bilevel minimax objective, GSR utilizes the group-labeled data for fine- ainengrineed reweighting of the training samples via a soft aggregation of the influence functions. It seamlessly integrates with last-layer retraining, resulting in a lightweight and effective strategy for improving group robustness. GSR also has the ability to automatically clean and guide the training of potentially noisy datasets via sample reweighting. Our empirical analysis further supports the alternative training paradigm that emphasizes the use of high-quality, group-labeled instances as the target set.
We identify several future research directions to further improve our method. Despite GSR achieving competitive worst-group performance, results decrease in mean performance. This trade-off is common in efforts to improve group robustness (Chen et al. 2022). ) Additionally, our work is based on the hypothesis that standard ERM learns sufficiently high-quality representations. Although GSR test the limit of last-layer retraining, it does not directly improve the representation learning. To mitigate the trade-off between the worst-group performance and the mean performance, enhancing the quality of representations is still a critical step (Izmailov et al.) 2022). To answer the question of whether sample reweighting can improve the representation learning for the group robustness, is necessary to retrain more components of the model after obtaining the sample weights. Developing more efficient strategies that jointly optimize the sample weights and deep model parameters may offer more insights and advancements in tackling subpopulation shifts and beyond.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This research/project is supported by the National Research Foundation, Singapore under its Al Singapore Programme (AISG Award No: SB5G2-PDD/2021-08-00111
This research is supported by the National Research Foundation Singapore and the Singapore Ministry of Digital Development and Innovation, National AI Group under the AI Visiting Professorship Programme (award number AIVP-2024-001), This research is supported by the National Research Foundation (NRF). Prime Minister's Office, Singapore under its Campus for Research Excellence and Technological Enterprise (CREATE) programme. The Mens. Manus, and Machina (M3S) is an interdisciplinary research group (IRG) of the Singapore MIT Alliance for Research and Technology (SMART) centre. Jingtan Wang is supported by the Institute for Infocomm Research of Agency for Science, Technology and Research (ARS,. We would like to thank the anonymous reviewers and AC for the constructive and helpful feedback.
REFERENCES
Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization for machine learning in linear time. Journal of Machine Learning Research, 18(116);1-40, 2017.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv: arXiv: 1907702897.9.... 2019.
Juhan Bae, Nathan Ng, Alston Lo, Marzyeh Ghassemi, and Roger B Grosse. If influence functions are the answer, then what is the question? Advances in Neurul Information Processing Systems, 35:17953-17967, 2022.
Lucas Beyer, Olivier Hénaff, Alexander Kolesnikov, Xiaohua Zhai, and Aäron van den Oord. Are we done with imagenet? arXiv preprint arXiv:2006.07/55 2020.
Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. In Companion Proceedings of the 2019 World Wide Web Conference, pp. 491-500, 2019.
Stephen Boyd and Lieven Vandenberghe. Convex optimization. Cambridge University Press, 2004.
Jonathon Byrd and Zachary Lipton. What is the effect of importance weighting deep learning? In International Conference Machine Learning, pp. 872 7818 2019.
Yongqiang Chen, Kaiwen Zhou, Yatao Bian, Binghui Xie, Bingzhe Wu, Yonggang Zhang, Kaili Ma, Han Yang, Peilin Zhao, Bo Han, et al. Pareto invariant risk minimization: Towards mitigating the optimization dilemma in out-of-distribution generalization. arXiv preprint arXiv:2206.07700 2022.
Zhiliang Chen, Gregory Kang Ruey Lau, Chuan-Sheng Foo, and Bryan Kian Hsiang Low. Duet: Optimizing training data mixtures via feedback from unseen evaluation tasks. arXiv preprint arXiv:2502.00270. 2025.
Dennis Cook and Sanford Weisberg. Residuals and influence in regression, 1982.
Yihe Deng, Yu Yang, Baharan Mirzasoleiman, and Quanquan Gu. Robust learning with progressive data expansion against spurious correlation. In Advances in Neural Information Processing Systems, volume 36, 2024.
Justin Domke. Generic methods for optimization-bl based modeling. In International Conference on Artificial Intelligence and Statistics, pp. 318-326. PMLR, 2012.
Cynthia Dwork, Moritz Hardt, Toniann Pitassi, Omer Reingold, and Richard Zemel. Faireess through awareness. In Proceedings of the 3rd Innovations in Theoretical Computer Science Conference, pp. 214 2012.
Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Rethinking importance weighting for deep learning under distribution shift. Advances in neural information processing systems, 33: 12997 2020.
11

---

Published as a conference paper at ICLR 2025
Tongtong Fang. Nan Lu, Gang Niu, and Masashi Sugiyama. Generalizing importance weighting to a universal solver for distribution shift problems. Advances in Neural Information Processing Systems, 36:24171-24190 2023.
Chuan-sheng Foo, Chuong B., and Andrew Ng. Efficient multiple hyperparameter learning for log-linear models. In Advances in Neural Information Processing Systems, volume 20, 2007.
Robert Geirhos, Jörn- Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2(1 1):665 673, 2020.
Amirata Ghorbani and James Zou. Data shapley: Equitable valuation of data for machine learning. In International conference on machine learning, pp. 2242-2251. PMLR, 2019.
Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel Bowman, and Noah A Smith. Annotation artifacts in natural language inference data. arXiv preprini 0rX3v:/803.0220 2018.
Yujin Han and Difan Zou. Improving group robustness on spurious correlation requires preciser group inference. arXiv preprint wrXiv:2404.13815 2024.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pp. 770-778, 2016.
Katherine L Hermann, Hossein Mobahi, Thomas Fel, and Michael C Mozer. On the foundations of shortcut learning. In International Conference Learning Representations, 2024.
Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive porst-group-accuracy. In Conference on Causal Learning and Reasoning, pp. 336-351. PMLR, 2022.
Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew G Wilson. On feature learning in the presence of spurious correlations. In Advances in Neural Information Processing Systems, pp. 38516-38532, 2022.
Bingyi Kang, Saining Xie, Marcus Rohrbach, Zhicheng Yan, Albert Gordo, Jiashi Feng, and Yannis Kalantidis. Decoupling representation and classifier for long-tailed recognition. In International Conference on Learning Representations, 2020.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. arXiv preprint arXiv:2204.02937, 2022.
Pang Wei Koh and Percy Liang. Understanding black predictions via influence functions. In International Conference on Machine Learning. pp. 1894-1894. PMLR, 2017.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of wi-the-wild distribution shifts. In International Conference Machine Learning, pp. 5637-5664, 2021.
Steven George Krantz and Harold R Parks. The implicit function theorem: history, theory, and applications. Springer Science Business Media, 2002.
Tyler LaBonte, Vidya Muthukumar, and Abhishek Kumar. Towards last-layer retraining for group robustness with fewer annotations. In Advances in Neural Information Processing Systems, volume 36, 2024.
Kwonjoon Lee, Subhransu Maji, Avinash Ravichandran, and Stefano Soatto. Meta-learning with differentiable convex optimization. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 10657-10665, 2019.
Peizhao Li and Hongfu Liu. Achieving faimess at no utility cost via data reweighing with influence. In International Conference on Machine Learning, pp. 12917-12930. PMLR, 2022.
12

---

Published as a conference paper at ICLR 2025
Weixin Liang and James Zou. Metashift: A dataset of datasets for evaluating contextual distribution shifts and training conflicts. arXiv preprint arXiv:2202.065523, 2022.
Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, 45(1):503-528, 1989.
Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. International Conference on Machine Learning, pp. 6781-6792, 2021.
Sheng Liu, Xu Zhang, Nitesh Sekhar, Yue Wu, Prateek Singhal, and Carlos Srrnandez-Grande Avoiding spurious correlations via logit correction. arXiv preprint arXiv:2212.01433, 2022.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In Proceedings of the IEEE International Conference on Computer Vision, pp. 3730-3738, 2015.
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics, pp. 1540-1552. PMLR, 2020.
Vaishnavh Nagarajan, Anders Andreassen, and Behnam Neyshabur. Understanding the failure modes of out-of-distribution generalization. arXiv preprint 1xXiv:2010.15775: 2020.
Junhyun Nam, Jaehyung Kim, Jaeho Lee, and Jinwoo Shin. Spread spurious attribute: Improving worst-group accuracy with spurious attribute estimation. arXiv preprint arXiv:2204.0200 2022.
Guoshun Nan, Rui Qiao, Yao Xiao, Jun Liu, Sicong Leng, Hao Zhang, and Wei Lu. Interventional video grounding with dual contrastive learning. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 2765-2775, 2021.
Fabian Pedregosa. Hyperparameter optimization with approximate gradient. In International Conference on Machine Learning, pp. 737-746. PMLR, 2016.
Rui Qiao and Bryan Kian Hsiang Low. Understanding domain generalization: A noise robustness perspective. In International Conference Learning Representations, 2024.
Shikai Qiu, Andres Potapczynski, Pavel Izmailov, and Andrew Gordon Wilson. Simple and fast group robustness by automatic feature reweighting. In International Conference on Machine Learning, pp. 28448-28466 PMLR, 2023.
Aravind Rajeswaran, Chelsea Finn, Sham M Kakade, and Sergey Levine. Meta-learning with implicit gradients. In Advances in Neural Information Processing Systems, volume 32, 2019.
Mengye Ren, Wenyuan Zeng, Bin Yang, and Raquel Urtasun. Learning to reweight examples for robust deep learning. In International Conference on Machine Learning, pp. 4334-4343, PMLR, 2018.
Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. Domain-adjusted regression or: ERM may already learn features sufficient for out-of-distribution generalization. arXiv preprint 222.:22002.06866. 2022.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint arXiv: 1911:1991..87.. 2019.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterizatioo exacerbates spurious correlations. arXiv preprint arXiv:2005.044 04345, 2020.
Kegan G. G. Samuel and Marshall F. Tappen. Learning optimized MAP estimates in continuously- valued mrf models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 477-484, 2009.
Bernhard Schölkopf, Ralf Herbrich, and Alex Smola. A generalized representer theorem. In International Conference Computational Learning Theory, pp. 416-426 Springer, 2001.
13

---

Published as a conference paper at ICLR 2025
Amirreza Shaban, Ching-An Cheng, Nathan Hatch, and Byron Boots. Truncated back- propagation for bilevel optimization. In International Conference on Artificial Intelligence and Statistics, pp-
PMLR, 2019.
Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. In Advances in Neural Information Processing Systems, pp. 9573-9585, 2020.
Changjian Shui, Qi Chen, Jiaqi Li, Boyu Wang, and Christian Gagné. Fair representation learning through implicit path alignment. In International Conference on Machine Learning, pp. 20156- 20175. PMLR, 2022.
C. Wah, S.Branson, P. Welinder, P. Perona, and S. Belongie. The caltech-ucsd birds-200-2011 dataset. California Instiuute Technology, 2011.
Haonan Wang. Ziwei Wu, and Jingrui He. Fairif: Boosting fairness in deep learning via influ- ence functions with validation set sensitive attributes, In Proceedings of the ACM International Conference on Web Search and Data Mining, pp. 721-730, 2024a.
Jingtan Wang, Xiaoqiang Lin, Rui Qiao, Chuan-Sheng Foo, and Bryan Kian Hsiang Low. Helpful or harmful data? fine-tuning- free shapley attribution for explaining language model predictions. arXiv preprint arXiv:2406.. 24l.:2400000666, 2024b.
Adina Williams, Nikita Nangia, and Samuel R Bowman, A broad-coverage challenge corpus for sentence understanding through inference. arXiv preprint arXiv: 2017.
Yuzhe Yang, Haoran Zhang, Dina Katabi, and Marzyeh Ghassemi. Change hard: A closer look at subpopulation shift. In International Conference on Machine Learning, 2023.
Bianca Zadrozny. Learning and evaluating classifiers under sample selection bias. In International Conference on Machine Learning, pp. 14, 2004.
Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar. Understanding why generalized reweighting does not improve ERM. arXiv preprint arXiv:2201.12295: 2022.
Michael Zhang, Nimit S Sohoni, Hongyang Zhang, Chelsea Finn, and Christopher Ré. Correct-n- contrast: A contrastive approach for improving robustness to spurious correlations. arXiv preprint arXiv:2203.01517. 2022.
Xiao Zhou, Yong Lin, Renjie Pi, Weizhong Zhang, Renzhe Xu, Peng Cui, and Tong Zhang. Model agnostic sample reweighting for out-of- ditoo--distribttion learning. In International Conference Machine Learning, pp. 272221,. PMLR, 2022.
14

---

Published as a conference paper at ICLR 2025
A DERIVATIONS AND PROOFS
A.1 DERIVATION OF THE JACOBIAN OF THE GRADIENT W.R.T. THE SAMPLE WEIGHTS
The derivative of V Ro (D; w. 0) w.r.t individual sample weights Wi is:
VoR(D"; w 0) Vw
V
wy 12
JEL.....
D
V uz
W/
Vol(ay,
JEI....|Da
VeR(22; 0r)
(10)
A.2 DERIVATION OF LAST-STEP GRADIENT (MAPLE ZZHou ET AL. 2022
According to MAPLE's one-step truncated backpropagation, the gradient descent updates before 0r-1 truncated, that 0r-1 is "constant" and is independent of w. Thus, it is true that
V w. 11- but for w. 0r. Recall that for simplicity, we use the notation RoR(D"; w. 0r)
for VoR(DD;; w, 0)
We show that Equation can be derived as follows:
0=07
V Or
m
V wi
0T-1 17 V RDD w, 0r-1))
(One-step truncation)
V ui
(0r-1
(D": w. 0r)
(0T 0T-1 on convergence, so VoR(D"; w. 0r) VOR w. (r-))) = 0 - VV-7VV. (D 0r)
(V 0w = 0 because of one-step truncation)
(o)(((00)))
(Apply Equation 10)
A.3. DERIVATION OF META-GRADIENT UPDATE
We utilize the technique of implicit differentiation. Upon convergence to an optimal solution, the inner problem has the property of V (o 0*) = 0. By first differentiating w.r.t. w on both sides of the equation, then followed by the chain rule, we have the following derivations:
d0
dV RDD w. 0-)
dw
dw
dV
RD w. 0*))
0
dw
aV 9 R(D"; w. 8*)
0V DVOR(D" w. 8-) d
0
(Chain-rule)
Dw
a
dw
RR Dt w, 0+) d0
0 V ~ VDR(D"; w. 07)
dw
d
(,, V.. VoR w w. 8')
dw
d
H=l g-,
V. (DD ; w. 0-)
dw
For each sample weight wy:
d
= = H11 V.n, VoRiD" w, Ô-) 0-
dw:
= --H R(2 0r)
(Apply Equation I0
15

---

Published as a conference paper at ICLR 2025
A.4 PROOF OF THE STRONG CONVEXITY OF LINEAR CLASSIFICATION WITH CROSS-ENTROPY Loss AND &g- REGULARRRZZAIION
In this section, we add the complete proof for the established result of the strong convexity (a subset of strict convexity) of linear classification with cross-entropy loss and fg-regularization. Before introducing the complete proof, a shortcut intuitive proof for this consists of the following steps:
1. Linear classification with cross-entropy loss is convex, and its Hessian is positive semi-definite, implying !/ ž O.
Adding tg-reuulaiization with a positive regularization strength > produces new Hessian H' H Al = XI, which meets the definition of strong convexity (Definition A.1)
We now present complete proof of the result of the strong convexity by directly showing the H' in Theorem A.6
Definition A.1 (μ-Strong convexity (Boyd & Vandenberghe 2004)) )) For μ > 0, a differentiable function f is p-strongly convex if for some μ U, Vzf(z) = pl.
Definition A.2 (One-hot Vector). For dimension d, u(i) E [0, Σ u(i) is the one-hot vector with i-th dimension u(i); = 1. We omit d in the following context for simplicity.
Definition A.3 (Softmax). Given the input E Rd, define the softmax function or RR - R as:
(xx; E
Its derivative (zx(x)= Rdxd has the form:
V20(z) = diag(o(x)) - (z())(=))
(11)
In addition,
- o()))
(12)
where u(i) E {0, ,0,1) = is the one-hot vector with i-th dimension u(i); =
Proposition A.4. The first-order derivative matrix V2o(2) of the sofimax function is positive semidefinite.
Proof. First (x) defines probability distribution. Hence, has the property that o(x); 2
0 Vi, Σ, (Z)); 1. Then according to the definition of positive semidefiniteness: Vz E Rd, 2 + 0:
z [zr(z>== (diag(o - (x())(x))) diag(( = =o0))))(())
[x(x)12 - C a()))2))²
2 0.
Definition A.5 (Cross-entropy). Given the input z Rd, (1...., .....,KK K} for K-way classification problem, define the cross entropy loss L;
K
Σ 1ly = k) log 24
--
where Il is the indicator function.
Theorem A.6 (Strong convexity of cross-entropy loss with l eg-regularization). For multi-class linear classification with cross-entropy loss, if & regulaaaaiiiii with positive coefficient X is applied, then the objective is X-strongly convex.
16

---

Published as a conference paper at ICLR 2025
Proof. Let E y {1.....K) denote the input, label of K-way classification task. Let E Rd K denote the linear model parameters. Let u(y) denote the one-hot representation of y. The objective function using cross-entropy loss and &g ée-reuularization is:
L(x, y: 2) = CE(O(V 2). + 1
A Σ My k:] log ( x))x) + 2 1
log ( x)ny) (x) + 2 1
The first-order derivative:
V=I( - olv x)v olv F)y (u(y) a( owbTz1) x)) + X
(Equation 12
= u(y)))) + +
The second-order derivative derivaatvv (Hessian):
H = (Z,((z
ax(a(v 2) u(y))
ATaxk
Do(bVx)) + Mlac
Oup
(o(()T)) x² Tv
+ Alaxk
0(v x) Ov
=|a(diag(o (bTTa xx)) - ol () (0)(v))))))) 11T Adax
(Equation II
x (diag(o(()) Tx)) - oow x) xT + Alaxk
Then we show that the Hessian of I, meets the requirement X-strong convexity, i.e., H AI Vz E _xdk
2 + 0:
Txxx (diag(o(V) x)) o( 2)0((" x)= [Waaxk Mlxxk)2 (H 22 = 2 x
= z x (iiag(o(v* (b x)) a( 2) xTz
2 0.
(Proposition A.4
Hence, is strongly convex according to Definition A.I
B EXPERIMENTAL DETAILS
B.
SETUP
For model architectures and initialization, we use ImageNet--retrraiee ResNet-50 (V1) (He al. 2016) for Waterbirds and CelebA datasets, and use BERT (HuggingFace) for MultiNLI and Civil- Comments. The MNLI dataset is preprocessed according the setup in bltttps:: / ////ttttbl com/ The CivilComment dataset is tokenized with Bert-
kehpangwei/ group PROBertToken. Tokenizer with max seq_ lenghh=00. For all experiments, we perform 5 independent runs on seed [1, 2, 4, 5] to obtain the mean and the standard deviations.
B.2 HYPERPARAMETERS
For all experiments, we fix the held-out fraction a as 10% and fix the resulting partition splits (i.e., the held-out set and the remaining set) of the training set D which generated by setting the
17

---

Published as a conference paper at ICLR 2025
numpy random seed 1. Although additionally tuning the held-out fraction and changing the split can most likely improve the worst- worst-group performance by balancing the data for representation learning and sample reweighting, we do not alter these splits for better consistency and only focus the second sample reweighting stage of the training.
For stage 1, better consistency and simplicity by heuristics without any tuning) and record them in Table 2
almost the same hyperparameters used by DFR (some slightly altered for
For stage 2, there two sets of hyperparameters for the inner and the outer loops. We perform a non-exhaustive search by randomly sampling the combination of hyperparameters from the grid. We select the best-performing model checkpoints and hyperparameter configurations based on the worst-group validation accuracy. The selected hyperparameter configurations are documented in Tables and for the outer loop and the inner loop respectively. Besides, we consistently apply weight normalization (to ensure the sum of weights equals one) and a step learning rate scheduler (StepLR), which decays the outer learning rate by a factor of 10 every steps. This facilitates better convergence for the bilevel optimization.
Table Hyperparameters for base model training. For Waterbirds and CelebA, we use almost the same hyperparameters as used by DFR. We slightly change the hyperparameters on MNLI and CivilComments (without any tuiing) so that learning rate scheduler is used. Momentum set as 0.9 for SGD.
Dataset Waterbirds
optimizer
Ir
scheduler
batch size
weight decay
epochs
SGD
le-3
Constant
32
le-3
100
CelebA
SGD
le-3
Constant
128
le-4
50
MultiNLI
AdamW W
2e-5
Constant
32
lc-4
4
CivilComments
AdamW
2e-5
Constant
32
lc-4 4
Table 3: Hyperparameters for sample weight updates (outer loop). For the step learning rate scheduler (StepLR), we simply decay the learning rate by a factor of 10 every 30 steps.
Dataset Waterbirds
optimizer
Ir
scheduler
grad clip
T
steps
GD
StepLR
0.1
100
CelebA
GD
StepLR
le-2
0.01
100
MultiNLI
GD
0.1
StepLR
le-2
50
CivilComments
GD
StepLR
le-1
0.1
50
Table 4: Hyperparameters for last-layer retraining (inner loop)
Dataset
optimizer
Ir
scheduler
batch size
weight decay
line search
Waterbirds
L-BFGS
le-4
Constant
full
le-1
Strong Wolfe
CelebA
L-BFGS
le-4
Constant
full
le-2
Strong Wolfe
MultiNLI
L-BFGS
le-4
Constant
full
le-2
Strong Wolfe
CivilComments
L-BFGS
le-1
Constant
full
le-3
Strong Wolfe
B.3 IMPLEMENTATION DETAILS
Let denote the dimension of 6. Let the size of the held-out training set be here. To efficiently compute the influence scores of the samples for last-layer retraining, we adopt the following strategy:
1. Calculate the Hessian inverse H of the weighted objective w.r.t. 0 = 0". The Hessian has a dimension of d x d.
2. Calculate the per-sample gradient of the unweighted objective w.r.t. 0 0 for the training set. This results a n matrix.
3. Calculate the per proup gradient the unweighted objective w.n.t 0 = 0" for the target set. This results in x d matrix, where is the number of groups
18

---

Published as a conference paper at ICLR 2025
By combining these three matrices, obtain the group- wise influence scores matrix, which can then be aggregated as in Algorithm| and become an n x gradient vector.
B.4 RUNNING TIME COMPARISON
Hardware. For all experiments, (24GB) GPU and AMD EPYC 7543 CPU.
machines with NVIDIA RTX 3080 (10GB) RTX A5000
Running time. We record the detailed running time for both the inner and the outer loop for different datasets in T The results are obtained from running single job on one NVIDIA RTX 3080 GPU with AMD EPYC 7543 CPUs.
Table 5: Running time for GSR
Dataset stats
Running time (s)
Dataset
Held-out size
Target size
of steps
Outer loop
Inner loop
Total
Waterbirds
479
600
100
2.71
0.16
285
CelebA 16277
9934
100
4.03
0.52
451
MultiNLI
20617
46231
50
7.24
0.53
381
CivilComments 26903
22590
50
8.21
0.7
882
Scalability. As used Pytorch for our implementation and the last-layer retraining has very few parameters, the per-sample gradient and the Hessian inverse calculation involves high CPU usage, even though the tensors performance sometimes can be CPU-bound instead of GPU-bound. Thus, the running time is likely increase when the CPU-usage is high.
GPU. When the same processes running in parallel, the
B.5 DETAILED RESULTS
In addition to the worst-group accuracy shown in the main paper, we also report the mean accuracy for these datasets Table for reference. We would like to highlight that the worst-group accuracy gain is usually at the expense of the mean accuracy. This implies that trade-offs generally need to be made for better robustness subpopulation shifts, since the mean accuracy corresponds to an estimation of the model's performance when there is no distribution shift.
C UNDERSTANDING THE ROLE OF THE INVERSE HESSIAN
Koh & Liang 2017) analyzed that the inverse Hessian measures the robustness the resultant model upweighting the training data. As Hessian is symmetric matrix, according the Spectral Theorem, it guarantees the existence of an orthonormal basis the eigenvectors. Hence, the gradient be written linear combination of eigenvectors. As the influence function utilizes the inverse Hessian, which has the reciprocal of the eigenvalues of the Hessian, the influence score has a larger magnitude when the gradient of the training and test points similar and along the direction of the eigenvectors with smaller magnitudes.
Having smaller eigenvalues can be interpreted as the first-order gradient changing relatively slowly. From the optimization perspective (Newton's ' method), this usually means we are allowed to take larger gradient steps when the trajectory has less variation with smoother curvature. On the other hand, the allowed step size becomes smaller when the curvature changes more quickly, indicating more complex loss landscape.
If the Hessian is dropped in between the inner product of the gradient as in Equation 2 (used by MAPLE), then we fail account the variations of step sizes different directions specified by the eigenvectors. Thus, the outer loop gradient becomes incorrect and the optimization will be less reliable.
D VISUALIZATION OF THE SPURIOUS CORRELATIONS
We show illustration of the spurious corrllatioss for the Waterbirds dataset in Figure
19

---

Published as a conference paper at ICLR 2025
Table 6: Performance comparison in detail. We report the worst-group accuracy on four benchmark datasets. For the group label column, X indicates no group label is used; v indicates g is required for training or validation; v indicates g is partially used; w/ indicates proportion of g from the validation set is re-purposed beyond hyperparameter tuning, such training sample weights model parameters. The row sections are ranked in descending order of the total amount of group label information required.
Group Label Train Val
Waterbirds
CelebA
MoltiNII
CivilComments
Method
Worst(S))
Mcan(
Wors((%)
Mcan(())
Worst(S)
Mcan(%)
Worst(S6)
Mcan(S)
87.6+1
81.+1a
90.6
72.0+19
Group DRO
VT
91.4111
00.5
88.9
92.91..
77-7
81.4
70.0
89.9 89.9
JTT
86.7
93.3
81.1
88.0
72.6
78.6
69.3
91.1
CbC
V
88 2..5ao
90.940.1
88 8409
89. 9.na
68 08.942.1
81.740.a
SSA
XI w
89 5
92.2+0.9
89.8+1.
92 820.1
76.6
79.9.0.9
69 942.0
88.2+2.
MAPLE
91.7
92.9
88.0
72.7
89.7
DER
XI 4
92.9100
94.210.4
88.311
91.310.a
74.716-
82.1
70. 70.1
87.210.3
30R-IIP
X7 W
8T 87.520
97.720.0
86.3a 86.3+0
91.5200
71.7400
80.620.0
68. 9..9
89.120.0
GSR
92.9g00
94.920.
87.0go-
90.9
78.5
79.8.00
71.7
85.9.0.4
ERM
74.9
1..1
46.9
35.3
65. 65.9
82.8-0.1
55, 65.6
92.1
FS SELF
X/ X
93.0103
94.011.
83.9109
91.7
70.7
81.2.0.7
background
land
water
landbird
class
3498 samples
184 samples
waterbird
samples
1057 samples
Figure 5: An illustration of the Waterbirds dataset. A spurious correlation exists between the bird class and the background. The majority groups highlighted in orange. The minority groups are in gray.
E WHAT IF WE CAN RETRAIN THE FULL NETWORK WITH GSR?
In this section, a simple setup explore the potential of retraining the full network with GSR, when the compute budget allows.
Setup. We conduct an ablation study on a small binary classification task customized from the ColoredMNIST, where the digits are spuriously correlated with color. We use tiny Convolutional Neural Network (CNN) with 3 layers and about 5000 parameters.
Dataset. ColoredMNIST-S. We modify the original ColoredMNIST dataset (Arjovsky et al., 2019 by first taking only the first digits out of all 10 digits, then assigning y < 3 as 0 and y 2 3 I. The training, validation, and test splits have 10000, 2000, and 27375 instances, respectively. The spurious correlation 7 created by associating most of the instances in class 0 with red color and class with green color. We set y 0.1 for the training split and set 7 0.5 for the validation and test splits. we do not introduce any label corruption. This setting is considered much simpler than the original ColoredMNIST, primarily because we studying with limited network capacity with tiny CNN, which could not fit the original 10-digit classification with more than 90% accuracy.
Baselines. We compare GSR, which is by default trained with last-layer retraining, with the full- parameter retraining variants. GSR-HE discards the Hessian, which is equivalent to MAPLE+LLR. GSR-F-EH retrains the entire network with the influence function and approximation the
20

---

Published as a conference paper at ICLR 2025
Table 7: Performance comparison on ColoredMMNIST-S Dataset with TinyCNN.
ColoredMNNISTSS
Method
Wall-Clock Time
Worst(%)
Mean(%
GSR-F-EH
94.1
96.0
400
GSR-F-AH
80.2
97.6
45
GSR-F-HF
91.9
96.9
42
GSR
91.2
95.4
4+0.8
GSR-HF HF
91.1
96.2
4+0.7
ERM
81.2
97.7
Table Running time and memory consumption of full-pat pall--rattterr retraining on Waterbirds dataset with ResNet-50.
Approach Explicit Hessian
Time
VRAM Consumption
>18h
2PB
Conjugate Gradients
18h
9GB
LiSSA
3.6h
9GB
Hessian. We use Conjugate Gradients calculate the inverse Hessian-vector product (iHVP) instead of explicitly calculating the Hessian to save the memory. GSR-F-AH approximates the Hessian using LiSSA (Agarwal et al. 2017) with d00th=1000 and repeat=1. This approximation calculates the iHVP with 1000 samples from the training set. GSR F HF which discards the Hessian, and this uses the sample gradients as MAPLE.
The results shown in Table 7 we observe that. Full-paramette retraining with Hessian (i.e., GSR-F-EHI is advantageous over all other approaches. However, approximating the Hessian with low quality might hurt the performance, since G underperforms other baselines. When using a small neural network model, last-layer retraining (LLR) can experience slight performance degradation compared retraining the whole model the entire dataset (since we have the capacity to do so), but it is still significantly better than ERM. Overall, these results suggest that we should consider palammeertttr retraining when we have the capacity, and LLR an efficient and effective approximation to it.
E. CHALLENGES OF SCALING UP FULL-PARAMETER RETRAINING WITH GSR
When scaling up full- parameter retraining to more realistic datasets with larger networks, the com- putational challenges for accurately calculating the sample weight gradients are indeed the major obstacle here because of the Hessian. Let n be the number of data points and p be the number of parameters. Directly caleulating the Hessian for the whole network both computationally expensive O(np² + y²) and memory bottlenecked Opp²). In particular, the Hessian of simple neural network with 1M parameters takes 4TB of memory, and ResNet-50 has 23M trainable parameters. Alterna- tively, can calculate the inverse Hessian-vector product (iHVP) to reduce the memory requirements to (p)) and the time complexity to (np²), but is still consumin,,, unless we apply further approximations. Table/8sshww the computational challenges of performing full- parameeeettrr retraining with GSR on the Waterbirds dataset.
It will be extremely costly to perform the experiments that optimize the sample weights for few iterations (e.g., 50 iterations, as done with L.RR). Besides, as shown in the previous ablation study, training the full network with crudely approximated Hessian (using LiSSA) causes inaccurate gradient estimation and may even hurt the performance. We believe these technical challenges require extensive exploration in future works and are beyond the scope of this paper.
F COMPARISON WITH MAPLE
Although the difference seems to be only about the Hessian, GSR and MAPLE are based on two distinct ways solve the bilevel optimization problem.
21

---

Published as a conference paper at ICLR 2025
Table 9: Performance comparison on MetaShift.
MetaShift
Method
(orst(%)
Mean(
GSR
78. 2..0
87 87.9444.1
DFR
71.4_1..
89.7.0.2
ERM
69.2
90.3
Table 10: Performance comparison on modified datasets where the spurious correlations are extremely strong so that the original minority groups no longer exist.
Waterbirds-
MNLI-
GSR
87.8 (-5.1)
74.2 (-4.3)
DFR
89. 40.5 (-33(
66.94002(-7.8
ERM
32.6( 42.3)
9.0( (5699
1. MAPLE directly approximates backpropagating through the training trajectory by one-step truncation, but GSR uses implicit differentiation that results in the Hessian. This calculation is exact when the objective strongly convex.
2. GSR also synergies with last-layer retraining, which was considered by MAPLE.
3. We use adaptive aggregation of the influence scores from the groups the gradient, but MAPLE does not explicitly design group-level aggregation for subpopulation shifts.
4. MAPLE additionally uses coreset selection by training on a subset of all the training data in each step to speed up the training, while GSR utilizes the full training data without approximation.
G ADDITIONAL DATASET FOR SUBPOPULATION SHIFT
To further demonstrate the advantage of GSR, we include an additional experiment with the MetaShift (Liang & Zou 2022) image classification dataset. We use the two-class version im- plemented by Yang et al. (023). where the spurious correlation exists between the class and the background. In particular, cats mostly indoors and dogs are mostly outdoors. We primarily compare with DFR, which uses the same amount of group labels and the same ERM feature extractor ours. The ERM model trained for 100 epochs without early stopping or hyperparameter tuning. GSR outperforms DFR convincingly in this case as shown in Table 9
H ADDITIONAL EMPIRICAL STUDY
H.1 WHAT IF THE ERM-LEARNED REPRESENTATIONS ARE INSUFFICIENT?
To simulate the situation where ERM fails learn sufficiently good representations, we deliberately remove the minority groups from the training data that are used for representation learning, then followed by GSR or DFR with L.R. We experimented with two modified datasets: Waterbirds- and MNLI- where the spurious correlations because other baselines (such as JTT, Group DRO) will not work since they cannot use the minority- group data from training. We report the change in performance in Table T0 (the number in brackets indicates the performance change compared to learning the representations with the minority groups). Clearly, when the representation quality degrades, the performance drops across different methods. In particular, ERM performs much worse than random guesses for the worst group. Nonetheless, some useful representations still learned and can be efficiently utilized by last-layer retraining guided by some minority data.
more severe (s l). We only compare with DFR.
22


Reviews:
[{"review_id": "2FGDHxcHCv", "reviewer": "Reviewer_wcXj", "summary": "The paper considers the case when samples can be grouped (not all training samples are necessarily grouped, but all validation samples and target samples must be grouped), and we consider training the prediction model to minimize the maximum groupwise risk (loss) among all groups, namely the \"worst-group risk\". To train the model under the situation above, the paper considers re-weighting samples during updating the model. Different from an existing method MAPLE, the proposed method considers using the Hessian of the risk when re-weighting. Also, the proposed method introduces the \"last-layer retraining\" to mitigate spurious correlations and improve group robustness and also to reduce the computational cost.", "questions": "-   The proposed method employs the last-layer retraining, but is it known what types of datasets are not sufficiently learned by the last-layer retraining? If so, comparing performances between methods for such datasets may be interesting.\n-   Section 1: It may be better to present an example of spuriously correlated datasets due to groups, e.g., using the \"Waterbirds\" dataset.\n-   Section 3.1: Assumption 3.1 states that, we assume that $\\hat{R}({\\cal D}^\\mathrm{tr}; w,\\theta)$ is strongly convex. However, with respect to what variable(s) should we assume it is strongly convex? All of $(w,\\theta)$? Please specify.\n-   Section 4.1: It states that \"Secondly, calculating the Hessian inverse in Equation 6 is no longer overly expensive, because the only free model parameters are in the last layer instead of the entire deep neural network\", but how the Hessian inverse can be updated in accordance with the update of $\\psi$? It looks a key point of the method, so it is good to be explicitly presented.\n-   Section 5.1: The method GSR-HF (Hessian-free) is examined as an ablation study, but how about the ablation study of using the Hessian but not LLR? To compare with MAPLE, this ablation study is also desired. (Perhaps is it too costly?)", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "The combination of the sample reweighting and the last-layer retraining to tackle with subpopulation shifts and (consequent) spurious correlations. As far as the reviewer's understanding, (i) a baseline method MAPLE (Zhou et al.; 2022) updates sample weights $w$ for group robustness, but it uses not the Hessian $H_{\\hat{\\theta}^*,w}$ but only the last-step gradient. (ii) The use of the Hessian raises the computational cost, however, the proposed method focuses on the fact that the last-layer retraining is both costless even we employ the Hessian, and effective for group robustness.", "weaknesses": "-   Technical novelty looks not well emphasized. Although the advantage above (\"Strengths\" section) seems to be effective, the reviewer could find it by carefully reading the paper. This fact is good to be emphasized in a early part of the paper, e.g., in Section 1.\n-   The experiments are desired to be more extended. One is to increase the number of datasets (perhaps are there not so many appropriate datasets?), and the other is to qualitative discussions. (How the results are well written, but the discussion on the reasons of the results is desired more.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UWJ2xqBHfs", "reviewer": "Reviewer_kehk", "summary": "This paper studies the challenge of improving machine learning model robustness to subpopulation shifts, where models may perform inconsistently across different subgroups within data distributions. Traditional methods typically rely on group-labeled data to tune models and minimize worst-case loss across groups, but these approaches are limited by the high cost of obtaining substantial, high-quality labels. The authors propose an alternative framework, Group-robust Sample Reweighting (GSR), which enhances group label efficiency by utilizing group-labeled data as a target to optimize weights for group-unlabeled data. The two-stage GSR approach first learns representations from group-unlabeled data and then fine-tunes the model’s last layer by iteratively retraining on reweighted data. GSR demonstrates superior performance on benchmark datasets.", "questions": "See Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper highlights that previous methods estimate weight gradients through approximations, but they neither account for the curvature of the loss landscape with respect to the current model nor fully capture the training dynamics. The authors’ proposed method addresses these issues under the assumptions that the inner objective is unconstrained, strongly convex, and twice continuously differentiable.\n\n2. The experiments in this paper are thorough, and the results are promising. Across multiple datasets, the proposed approach outperforms various baseline methods, achieving state-of-the-art performance.", "weaknesses": "1. One limitation of this paper is that the proposed method relies on the inner objective being unconstrained, strongly convex, and twice continuously differentiable—assumptions whose strength is not fully clarified. The authors address these requirements by fixing the feature layers during training and only updating the linear classifier layer. However, further explanation is needed to clarify why training only the classifier layer meets these assumptions. Additionally, limiting training to the classifier layer may risk sacrificing overall model performance, which warrants further investigation.\n\n2. The authors' second motivation is somewhat confusing. The statement that \"regularization techniques must be applied such that training with weighted objectives leads to meaningfully different solutions than using ERM\" may require further clarification. It would be helpful for the authors to elaborate on how this principle is specifically reflected in the design of their method.\n\n3. The authors should conduct additional experiments to examine, without computational constraints, how much performance is lost when training the model with frozen feature layers compared to full-parameter training. This would help determine whether the performance reduction is within an acceptable range.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mjnlF86AIQ", "reviewer": "Reviewer_xdjG", "summary": "This paper introduces a novel approach called Group-robust Sample Reweighting (GSR) aimed at mitigating the issue of spurious correlations, particularly under subpopulation shifts. The paper proposes a two-stage strategy that utilizes group-labeled data more efficiently by leveraging it to reweight group-unlabeled data. The key contribution of the GSR method is its use of last-layer retraining (LLR), which simplifies the bilevel optimization problem by focusing only on the last layer of the model. Through iterative reweighting of the group-unlabeled data based on influence functions and adaptive aggregation of sample gradients, GSR aims to improve the worst-group accuracy.", "questions": "Please see above", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The proposed method is built on a solid theoretical foundation, particularly leveraging bilevel optimization and influence functions, which enhances its robustness and credibility.\n\n2. By focusing on the last layer, the method simplifies the optimization process, making it computationally efficient.\n\n3. The paper provides extensive experimental results across a variety of benchmarks and other understanding experiments.", "weaknesses": "1. In Section 2, it would help to be more explicit about how the proposed approach differs from previous methods. Right now, it feels like Section 3 is mostly about improving an already established objective (MAPLE). Adding the inverse Hessian in the optimization might make things more complicated than before. It’s not clear from the paper whether this extra complexity is worth the performance gains. \n\n2. It’s not obvious how much of the performance boost comes from last-layer retraining (LLR) versus the optimization changes over MAPLE. It would be useful to see a comparison where MAPLE is combined with LLR, so we can better understand where the gains are really coming from.\n\n3. The claim that “over-reliance on spurious correlations mainly occurs in the last layer” feels a bit questionable. Prior research indicates that neural networks can develop biases throughout the entire network, not just in the final layer, e.g, [1]. It would be helpful to either provide more evidence or explain this assumption more clearly.\n\n4. Table 1 is missing some crucial results, especially for MAPLE, which is an important baseline here. The paper mentions that it couldn’t get the results from the original paper, but why not try reproducing them? Was there a specific issue? It feels like an incomplete comparison without those numbers.\n\n[1]What neural networks memorize and why: Discovering the long tail via influence estimation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uiZzC7G76v", "reviewer": "Reviewer_oDPM", "summary": "This paper proposes a method to mitigate spurious correlations and improve group robustness. The task is formulated as a bilevel optimization problem, which the authors solve via implicit differentiation. They employ Last-Layer Retraining (LLR) to ensure that the optimization problem meets the conditions required for implicit differentiation. Experiments on several datasets demonstrate the effectiveness of the proposed method.", "questions": "see weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- This paper utilizes implicit differentiation to calculate the gradient with respect to the sample weights and applies last-layer retraining to ensure that the risk function satisfies the conditions of convexity and continuous differentiability.\n\n- Overall, the paper is written and organized well.\n\n- The authors not only conducted experiments on benchmark datasets but also evaluated the performance of the proposed method in the presence of label noise. The proposed method demonstrated effectiveness in identifying noisy data.", "weaknesses": "- There appears to be an issue with the proof in Appendix A.2. By examining the first and third lines together, we derive $\\nabla_{w_i} \\theta_T \\approx \\nabla_{w_i} \\left( \\theta_T - \\eta \\nabla_\\theta \\hat{\\mathcal{R}}(\\mathcal{D}^{\\text{tr}}; w, \\theta_T) \\right)$ which then leads to $0 \\approx - \\eta \\nabla_\\theta \\hat{\\mathcal{R}}(\\mathcal{D}^{\\text{tr}}; w, \\theta_T)$, which seems unreasonable.\nFurthermore, in the transition from the third to the fourth line, it is unclear why $\\nabla_{w_i} \\theta_T$ was eliminated. Was it because $\\nabla_{w_i} \\theta_T = 0$? If that is the case, the proof in Appendix A.2 could directly conclude that $\\nabla_{w_i} \\hat{\\theta}^* \\approx \\nabla_{w_i} \\theta_T \\approx 0$. It might be worthwhile for the authors to review this part of the proof carefully.\nI do not believe this issue impacts the main contributions of the paper, but the authors may need to consider revising the related explanations in the paper that are based on this proof.\n\n- In Appendix A.3, it is assumed that when $\\nabla_{\\theta} \\hat{\\mathcal{R}}(\\mathcal{D}^{tr}; w, \\hat{\\theta}^*) = 0$, it follows that $\\frac{\\mathrm{d} \\nabla_{\\theta} \\hat{\\mathcal{R}}(\\mathcal{D}^{tr}; w, \\hat{\\theta}^*)}{\\mathrm{d} w} = 0$. Could the authors elaborate on why the gradient with respect to $\\theta$ being zero implies that its derivative with respect to $w$ is also zero in this context?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)