Optimizing importance weighting in the presence of sub-population shifts
Authors: Floris Holstege
Bram Wouters
Noud Van Giersbergen
Cees Diks
Venue: ICLR 2025 Poster (2025)
Keywords: Importance weighting
Distribution shift
Group Robustness
Bi-level optimization
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We show that existing importance weighting methods for sub-population shift can be improved using a bi-level optimization procedure over the weights

Abstract: A distribution shift between the training and test data can severely harm performance of machine learning models. Importance weighting addresses this issue by assigning different weights to data points during training. We argue that existing heuristics for determining the weights are suboptimal, as they neglect the increase of the variance of the estimated model due to the limited sample size of the training data. We interpret the optimal weights in terms of a bias-variance trade-off,  and  propose a bi-level optimization procedure in which the weights and model parameters are optimized simultaneously. We apply this framework to existing importance weighting techniques for last-layer retraining of deep neural networks in the presence of sub-population shifts and show empirically that optimizing weights significantly improves generalization performance.

Published as a conference paper at ICLR 2025
OPTIMIZING IMPORTANCE WEIGHTING IN THE PRES- ENCE OF SUB- PUB-POPULATION SHIFTS
Floris Holstege Bram Wouters', Noud van Giersbergen' Cees Diks University of Amsterdam, Department of Quantitative Economics Tinbergen Institute {f. g. holstege, b. wouters, p.p.a.ppp.................. vangiersbergen, c.g. g.g.h w.....di.. di wih.dikksuuuaaaa...... n1
ABSTRACT
A distribution shift between the training and test data can severely harm perfor- mance of machine learning models. Importance weighting addresses this issue by assigning different weights to data points during training. We argue that existing heuristics for determining the weights suboptimal, they neglect the increase of the variance of the estimated model due to the finite sample size of the training data. We interpret the optimal weights in terms of bias-variance trade-off, and propose a bi-level optimization procedure in which the weights and model param- eters are optimized simultaneously. We apply this optimization to existing impor- tance weighting techniques for last-layer retraining of deep neural networks in the presence of sub-population shifts and show empirically that optimizing weights significantly improves generalization performance.
INTRODUCTION
Machine learning (ML) models typically are trained to minimize the average loss on training data, with the implicit aim generalize well to unseen test data. A distribution shift between the training and test data can severely harm generalization performance (Alcorn et al. 2018 Quiñooner-Cannella so-called sub-population shift, where the population is dissected in groups and the distribution shift is entirely due to a change in occurrence frequencies of the different groups (Yang et al., (023). In ML applications, e.g., in the medical field (Jabbour et al.] 2020), such shifts can negatively impact the average accuracy on test. data and/or the accuracy on specific groups. The latter has implications for algorithmic fairness (Shankar et al.) 2017 Buollmwini & Gebru 2018 Wilson et al., (019). Furthermore, sub- population shifts are also relevant regarding spurious correlations (Gururangan et al., 2018 Geirhos iet al. 2020 DeGrave et al. 202I), where the frequency of occurrence of certain combinations of the variable of interest (e.g., a medical condition) and non-causally related attributes (e.g., the sex of the patient) might differ between the training and test distribution.
al. 2022). Distribution shifts come in many forms; a prime example
Importance weighting applied to ML models treats different observations in the training data dif- ferently, for example, by using a weighted loss or weighted sampler (Kahn & Marshall 1953). has long been used to address distribution shifts and other issues, including domain adaptation, covariate shift and class imbalance; see Kimura & Hino (2024) for an overview. In the context of training deep neural networks (DNNs), it was used in a variety of ways to tackle sub-population shifts, e.g., Sagawa al. (2020a); Liu et al. (202T) Idrissi al. (2022); Kirichenko et al. (2023)
This paper addresses the question of how to weigh each observation in an importance weighting scheme that aims mitigate the effects of distribution shift. Typically, the weights are chosen based on heuristic argument. For example, in grouh-wegghted empirical risk minimization (GW- ERM) they are chosen such that the weighted training loss is unbiased for the expected loss with respect to the test distribution. We argue that such choice is usually not optimal: for a limited train- ing sample it does not lead to estimated model parameters that minimize the expected loss with re- spect to the test distribution. Because of the limited size of the sample, heuristically chosen weights tend to put too much weight on certain observations (usually the minority groups). They thereby introduce significant variance, due to strong dependence of the trained model on the particulars of the training data. We illustrate this bias- variance trade-off for importance weighting analytically with example of linear regression model in the presence of sub-population shift.

---

Published as a conference paper at ICLR 2025
Secondly, we introduce an estimation procedure for the optimal weights via a bi-level optimization problem, where the weights and the model parameters are optimized simultaneously. In the context of last-layer retraining of DNNs in the presence of sub- sub-population shifts, we show empirically, for benchmark vision and natural language processing datasets, that existing state-of-the-art importance weighting methods (see Section ( can be improved significantly by optimizing for the weights. We find that both the weighted average and gorst-group accuracy generally improve, and that optimized weights increase robustness against the choice of hyperparameters for training. We show that the effect of optimizing weights is larger when the limited size of the training sample becomes pressing, namely in the case of a small total sample size or when the minority groups are small. We provide an implementation of this procedure as an open-source package (link).
We emphasize that the proposed procedure for optimizing weights should not be positioned as a new importance weighting method. Instead, it should be seen as an addition to existing methods that typically improves their generalization performance.
2 RELATED WORK
In this section we give brief overview of importance weighting techniques, the issue of sub- population shift, and how existing methods aim to address
Importance weighting for distribution shifts in deep learning: importance weighting is clas- sical tool to estimate a model for a test distribution that differs from the distribution of the training sample (Kahn & Marshall 1953
Shimodaira 2000: Koller & Friedman| 2009). When applied to training DNNs, importance weighting only has a noticeable effect in the presence of sufficient reg- ularization (Byrd & Lipton 2018 Xu al. 2021). In applications
focus the use of importance weighting for last-layer retraining, as this has been argued to be sufficient for dealing
with sub-population shifts (Kirichenko et al. 2023 LaBonte al. 2023 Chen et al. 2024).
Sub-population shifts: it has been widely documented that ML models frequently face deteriorating generalization performance due to sub-population shifts (Cai et al., [021] Koh et al., [202]] Yang
et al. 2023). There exist different possible definitions of the groups for which the distribution shift occurs. In the case of spurious correlation, groups are commonly defined by combination of
a label and a non-causally related attribute (Joshi et al. 2022; Makar et al.) 2022). Alternatively, groups can be defined exclusively via an attribute (Martinez al. 2021bb or, the case of class imbalance, exclusively via label (Liu et al. 2019). While our experiments consider spurious correlations, the methodology also applies to alternative group definitions. Generalizing to groups that do not occur in the training data (Santurkar et al. (020)) is beyond the scope of this paper.
Existing approaches of importance weighting for training DNNs in the presence of sub- population shift: a key method is group-weighted empirical risk minimization (GW-ERM), which uses a weighted loss and where the weights are determined by the occurrence frequencies of the groups in the training and test distribution. The weighted sampler analogue subsamples from large groups (SUBG) to create a novel dataset (Sagawa et al. 2020b). Both methods are strong basclines for addressing sub-population shifts (Idrissi et al. (022). An alternative approach is to minimize the worst-group loss via group distributional robust optimization (GDRO, [Sagawa et al.) (2020a))- While these methods require group annotations for all training data, there exist other methods that require limited group annotations in the form of a validation set. An example is Just Train Twice (JTT), where the errors of the model are used to identify data points that need be upweighted (Liu et al. 2021). Other methods in this realm incorporate errors of a secondary model in their estimation procedure (Na al. 2020 Zhhng et al. 2022 Qiu al. 2023) or estimate group attributes (e.g., Nam et al. (2022) Han & Zou (2024); Pezeshki et al. (2024). These methods typically require small validation set for determining hyperparameters. Kirichenko et al. (2023) argue that instead one should use this validation set for the estimation of an ensemble of models via SUBG, method
denoted as deep feature reweighting (DFR).
Estimation of optimal weights: commonly noted downside of importance weighting is that it reduces the effective sample size available for model training (Martino et al., 2017]] Schwartz & Stanovsky 2022). In light of this, Cui et al. (2019) and Wen al. (2020) suggest that the weight given to a particular label should be based on the number of available datapoints. Similar to our approach. Ren et al. (2018) aim to learn weights for neural networks given validation set.
2

---

Published as a conference paper at ICLR 2025
3 PROBLEM SETTING
Consider random variables (y, æ), where y E y is the target variable and x E X' represents the input features. Suppose we dealing with family of models {fo X' + V 0 E } and we aim to find a model that minimizes ((g,a))~PP. [[y. f=)]) where :) is some pre-defined loss function and Pto is a test distribution of interest. Moreover, we have access to a dataset D = (v,, æ; æx);=1 which is an i...d. sample drawn from a training distribution Pte that is different from the test distribution. This distribution shift can be problematic, as the minimum of an empirical loss for the training data does typically not coincide with the minimum expected loss with respect to the test distribution.
One method to address this issue is by estimating the model parameters using a weighted loss,
(w)) arg min Ci (0, w), where Crain (0, w)
we [iyis fx;)).
Og
If there is no unique minimizer, which is often the case for DNNs, the arg min should be interpreted as drawing randomly one element from the set of minimizers. A common choice for the weights (W1, W2,.... WnY Y E R" are the likelihood ratios r Pr., where r; r(y,, x.)) and r(y, x) = (1.().
the (assumed to be known) ratio of the probability density functions (pdfs) associated with the test and training distribution respectively. This is due to key result in importance weighting that, assuming ptr (y, .) when (V: a:) - 0,
Ptr(y.æ]
(v. (rl,, a) fo (x))] ((y.o))
If we choose the weights w to be equal to the likelihood ratios r, then by the law of large numbers train r) converges in the limit of n oc the expectations in Equation 2 As consequence, (r) minimizes for sample size tending to infinity the expected loss with respect to the test distri- bution, ((g.a))-P... [Cy fo (z)].. However, in practice one works with finite amount of training data and the expectation might be poorly approximated by its sample analogue. Therefore, the optimal choice w, for the weights should be defined as the one that leads to a "best" parameter estimate On (w), in the sense that minimizes the expected loss with respect to the test distribution,
" = arg min EDa-Pii T (Fty.e))P. So.. (w) (z))
Note that by marginalizing over D the optimal choice for is not determined by the actual alization of the training data. In practice, with access to only a single dataset D the estimate of the optimal choice will be optimal given that specific training set, although might be able to alleviate this dependence by using a cross-validation or bootstrapping scheme.
The first main claim of this paper is that the optimal weights w typically not equal to the weight choice based on the likelihood ratio, (w;); & r(ys, æ:)). Although the likelihood ratio makes the estimated parameters asymptotically unbiased under certain conditions (Shimodaira) (2000)), neglects the variance that arises due to a limited size of the training set. More concretely, putting more weight on certain training data points could lead overfitting and does not necessarily lead to better generalization, not even on the test distribution for which the weights are supposed to be correction.
The second main contribution of this paper is the introduction of new estimation procedure for the optimal weights w. To address the expectation with respect to Pte in Equation we make use of Equation] 2 and write the optimal weights as
" = arg min EDE-Par [ Ir(y, L(v, C(,, So, (mn) (m)
(4)
w
We propose to estimate the inner expectation using a randomly drawn subset of size Nival from the available training data. We call this subset the validation set. Let n' n nvaal be the remaining number of data points. The result is a bi-level optimization problem. The weights are estimated via
/ n
arg min L ((.((w)). r),
(5)
w
The model parameters are estimated by minimizing this training loss, using the n' data points that are not in the validation set. The creation of the validation set has the advantage of avoiding possible overfitting of w. but comes at the cost of having fewer data points for estimating 0.
where the validation loss is defined analogously to the training loss in Equation I
3

---

Published as a conference paper at ICLR 2025
3.1 OPTIMIZING WEIGHTS FOR SUB-POPULATION SHIFTS
In what follows, we make the previous general claims concrete by focusing on one specific type of distribution shift that has been studied extensively in the machine learning literature, namely sub- population shift. Each observation carries an additional random variable / E 9 which labels a total number of [g] 2, but this is not necessarily the case. For ease of notation, we assume that can take the values 1, 1.2...., G. The distribution shift arises from a difference between the marginal training and test distributions of g, while the conditional pdf p(æ,, V g) remains unchanged:
G subgroups that dissect the population. We note that g can depend on y and
(yy 9. x) = Ptr (p)(y, (l))) (Ve 9, x) = (te(g)p(y, xlg).
(6)
where Ptr (9) * (t) Unless otherwise stated, for the training set defined in Section we presume access to the subgroup labels {g
Since the distribution shift is entirely due to a change in the marginal distribution of g. it is natural to constrain the weights we to be determined by the group labels gi- is convenient to define group weights Py: where g G, such that w; Py: (pr(gg)). Estimating the model parameters can then be seen as a function of the group weights, i.e., (p),), where p = (P1.P2..... -.....-PP .c)) E T Instead of optimal weights w one now aims to find optimal group weights pr, Assuming that the number of
groups is much smaller than the size of the training set, this heavily reduces the dimension of the bi-level optimization problem defined in Equation 5
4 THEORETICAL ANALYSIS OF LINEAR REGRESSION
The optimal weights, as defined in Equation 3 can be analytically derived for the case of linear regression with simple data-generating process (DGP) and sub-population shift. This analysis provides insight into the nature of optimized weights as viri-ner trade-off and the conditions under which optimization of the weights becomes relevant.
We consider a DGP = x 39 + e, where E RR is the target variable, x = (1, T1."T)T represents the intercept and the feature vector to have zero mean, variance o² and to be uncorrelated with æ. The population consists of two groups labeled by g E (0, 1}. Group membership determines the value of the intercept, 3 (Ag. BT¹) for 0, 1. Note that the slope coefficients are common among the two groups. The sub-population shift between training and test data is defined by a change the marginal distribution of the group label, Ptr := P., 1) & ( = 1) Pte.
Rd and E E E2 is the noise. We take /O( Ia) and the noise
Suppose we have a training sample D {ye, 3i g of size n and we use the weighted least squares (WISS) estimator
9 1, 9: = 0.
(.((p) = -X
with W = (iag(w))..... Wr), w.
Here, X E (d+11))) is the design matrix of input features and E IR" contains the target variables. The weights w; are parameterized by p E [o, 1]. we choose the weights to the likelihood ratio between the test and training distribution, the WLS estimator asymptotically minimizes the expected
quadratic loss with respect to the test distribution. This choice of the weights corresponds to = Pte-
For large but finite the approximate expected loss is given by (see Appendix B for details)
EDE-Pie B(y.e))>P.. (v - x (p)) pzl Ps B² (pre, P, a ao 4 (a² P. Ptr, n, d) +o? (8a)
where the result is decomposed in the following bias and variance term, respectively, an)²
B² (prr, 01, uo) [Pra(1 - -p)² +
+1 Pro) Dr=)
(a)
(8b)
p² 21-
d+1
V(O", P. Ptr, n, d) 2²
(8c)
Pir 1
Ptr n
Both terms are quadratic with respect to p. The bias term is minimized for = Pte. provided that ao # 01, whereas the variance minimized for p = Ptr. The minimum of the expected loss is
Pro + P/r
o²
where 1) n(a1 an) Prr P(t)
(9)
1 1+
4

---

Published as a conference paper at ICLR 2025
This is the optimal trade-off between minimizing the bias and variance. In the limit - the variance disappears and the optimal p is indeed equal to the heuristic choice Pte. However, for finite it can differ considerably. This is illustrated in Figure T
Its impact on the expected loss is shown
in Figure in Appendix B
1.0
1.0
0.9
0.9
Pv.
10
0.7
0.8
100
0.8
Py - 0.8
6. 0.7
1000
0.7
0.95
Ple
Ps
0.6
0.6
10F
10
10
10
(a) For Pt- 0.9 and varying d.
(b) For d = 250 and varying
Figure 1: The optimal choice Pi: given by Equation 9 as a function the training dataset size n, for varying feature dimension d (panel (a)) and training distribution Ptr (panel (b)). Other parameters are fixed at 0₃ = 1, ag = 0, o 1. Pte = 0.5. Note that the heuristic choice for p would be at 0.5. Also note that Equation is derived approximately for large (hence the dashed lines for small n).
OPTIMIZING WEIGHTS OF EXISTING IMPORTANCE WEIGHTING TECHNIQUES FOR SUB-POPULATION SHIFTS
Before continuing, we note that the setup described Section 3 when applied to sub-population shifts, largely corresponds to GW-ERM. However, emphasized in Section 1 the methodology proposed in this paper also applies to other importance weighting methods. They differ from GW- ERM in terms of, e.g., the use of weighted sampler instead of weighted loss (SUBG, DFR), the overall objective (GDRO) or the inference of groups (JTT). For all these cases, optimal weights can be defined and subsequently found through a bi-level optimization problem. In what follows, discuss the details of optimizing weights for the different existing importance weighting techniques.
5.1 GROUP- WEIGHTED EMPIRICAL RISK MINIMIZATION
The optimal group weights p are estimated by solving a bi-level optimization problem that for GW-ERM consists of parameter estimation as in Equation [T. with the difference that the estimates are a function of instead w, and group weights optimization
p% = =arg min L (p),r),
(10)
pEPa
where the likelihood ratios are given by (ta Also note that the group weights are restricted Pir
to the G-dimensional probability simplex Pa. This normalization conveniently prevents modifying the effective learning rate of the optimization procedure (Ren et al. 2018).
We solve the bi-level optimization problem iteratively. We find pú in Equation 10 by means of ex-
ponentiated gradient descent (Nemirovsky & Yudin 1983 Kivinen & Warmuth, 1997). where the parameter estimates (p) are kept fixed. After each weight update, we re-calibrate the param- eter estimates of Equation| [ Although convergence to the optimal group weights is not guaran- teed due to potential non- non-convexity, gradient- based optimization of hyperparameters has previously been applied successfully (Pedregosa) 2016). Details of the optimization procedure are provided in Al Algorithm []]
step in this algorithm that requires extra attention is the computation of the gradient of the vali- dation loss with respect to the group weights, commonly referred to as the we use the chain rule
For this,
Sp (p), V Vo,, Cral (Ō. (O. VpQr:
5

---

Published as a conference paper at ICLR 2025
and note that the first factor typically be computed in straightforward manner. For the gradient of the parameters with respect to the weights, we use the implicit function theorem. The idea using the implicit function theorem for gradient-based optimization of hyperparameters was originally introduced by [engio 2000). For the details of calculating the hyper-gradient see Appendix
Algorithm 1: Estimation of the optimal weights pr, for GW-ERM
Input: data {yr, 9s, x4 Ji=1, training set size n' < nt maximum steps T, learning rate 1).
momentum and likelihood ratios
Initialize Po such that Po.9 Pr(())
Peeld) for all g E G. Normalize Po.g Eg Po.g for all g G.
P0.2
Split the data into training set of size n' and validation set. Set Un,g for all g G.
for -= 1, 1,..., T do
Estimate 0n (P:-1) as in Equation
1
Compute "yper-gradienn Ç. V p.-. L (O., (pt.((p)1),)))))) r) with implicit function theorem.
Update weights Pt.g = Pt-1,9 exp(m with U1 = 1, + (1 7)Sr.. for all E g.
Normalize Pt,g Eo'er Puo
Pig for all g 9.
end
Return P. = arg min p. E (po.Pt...... pr) Lall On r).
5.2
SUBSAMPLING LARGE GROUPS
In the case of SUBG Saaawa et al., 2020b), the weighted loss Equation D is replaced by training loss on a subset of the data created by subsampling without replacement. Let s = { lg E 9, 1= 1..,,,. ng), where xa E {0, 1} a random variable denoting whether data point from group g is included in the sample. The weights v, [0, 1] represent the fraction of data points to be subsampled, such that total of [g observations of group g is drawn randomly and included in the training loss
ng
SUBG, train
(0 ΣΣ gEg 1=1
(1I)
where vf of are observations belonging to group 9i m Eyeg Nyngy and (V1, ..,.... v.). This training loss is not differentiable with respect v, which poses a problem for the bi-level optimization procedure. This is resolved by using parameter estimates obtained by averaging over the random subsampling process,
u.))
@UBG (v) == arg min E, dEG
SSLDC (0, arg min ('trinn
Fgg Clv, fo(a)))
(12)
dEG
9E9 1=1
where the approximation [ mated with a weighted validation loss, 6; = arg mill Lval ((SUUB (v), r),, where Vo = {u E [0, | E : fi =1] Note that the weights v are not required to be normalized, but avoid unnecessary loss of training data least
Vg used. The optimized weights are still esti-
of them must be equal to one. Apart from this and some minor differences in the computation of the hyper-gradiens (see Appendix (A). the bi-level optimization procedure for SUBG is similar to the one for GW-ERM as described in Algorithm
A closely related method is DFR Kirichenko et al. 2023)) in which an ensemble of models trained on subsampled groups used to estimate the model parameters by averaging Equation T2 Typically, this is done on validation set. Per model, the hyper- sreer-rrddentt is the same as for SUBG.
5.3
GROUP DISTRIBUTIONAL ROBUST OPTIMIZATION
Instead of minimizing the overall expected loss with respect to a test distribution, GDRO aims to minimize the largest expected loss per group. This means that there is separate test distribution. The associated training procedure minimizes a weighted sum of expected losses per group, where the loss weights q (g1, (21.42... 42..... ... ..)) are updated during training and larger weights are given to groups with a larger expected loss (Oren et al. 2019; (Saawa et al. (2000)). We apply the framework
of optimized importance weighting to the GDRO objective. As in Section 5.1. we work with group

---

Published as a conference paper at ICLR 2025
weights p. The model parameters estimated via weighted loss as in Equation 1 whereas the optimal weights are found by minimizing
P mig E Pip PEPc
sup gEPG
Σ 49 B(g.al))~PPr C(v, So. (o.(() Dus
(13)
where, as before, Pa is the G-dimensional probability simplex. The inner expectation of Equa- tion 13 is estimated with validation set that is withheld from the data used to train the model parameters. The bi-level optimization procedure now consists of updating three quantities per up- date step: the model parameters 0, the group weights p and the loss weights q. See Algorithm 2 in Appendix D for details.
In the standard GDRO setting, the coefficients of the logistic regression trained to minimize worst-group loss. This makes GDRO prone to overfitting (Sagawa et al. (2020a). We, instead, optimize the weights to minimize worst-group loss, whereas the model parameters are trained via Equation [ In typical cases, because of the (much) lower dimension of the group weights, overfitting less likely to be a problem, giving optimized weights an edge over standard GDRO.
5.4 GROUP INFERENCE METHODS
Acquiring group labels often costly. Several methods (see Section 2 resolve this by estimating group membership. Methods in this realm, such as Just Train Twice (Liu et al., 2021), require validation set with the ground-truth group labels to tune hyperparameters. A key hyperparameter of JTT is the extent to which misclassified observations should be upweighted, which is optimized via a grid search on a validation set. We replace this single upweighting factor by multiple group weights, leading to higher flexibility. Furthermore, we optimize them via gradient descent instead of grid search. Like JTT, we use estimated group labels instead of ground-truth group labels, and aim find weights that minimize the worst-group loss.
6 EXPERIMENTS
The goal of this section is to verify empirically that existing of methods for addressing cobpoooulation shifts benefit from using optimized group weights. We use benchmark classification datasets for studying uub-population shifts: two from computer vision (Waterbirds, CelebA) and one from natural language processing (MultiNL.)). Groups are defined via combination of the target label and a spurious attribute. The training and validation set are randomly split. For Waterbirds, this is different from previous implementations, where the validation set is balanced in terms of the groups (Sagawa et al. 2020a). We deviate from this, since it is more realistic that the training and validation sets are drawn from the same distribution. The distribution shift established by measuring the generalization performance details on the datasets, including group sizes, see Appendix E.1
test set for which the groups have equal weights. For
We start by finetuning a pre-trained DNN the respective datasets, ResNet50 model (He al. 2016) for Waterbirds and CelebA, and a BERT model (Devlin et al., 20196) for MultiNLi. For five existing importance weighting methods that address sub-population shifts (GW-ERM, SUBG, DFR, GDRO, JTT), we then retrain the last layer of the DNN, both with the standard choice of weights prescribed by the method and the optimized weights proposed in this paper. We focus on last- layer retraining, which amounts to training a logistic regression on the last- embeddings of the finetuned DNN, because of computational feasibility and because has been shown to be suffi- cient for addressing sub-population shifts (Izmailov et al.) 2022] Kirichenko et al. (003). For each importance weighting method with standard weights choice, we use grid search to optimize the respective hyperparameters, including an L1 penalty for sufficient regularization. To get a fair com- parison, we use the same hyperparameters for optimized weights. This possibly gives an advantage to the standard choice of weights.
We repeat the above procedure, starting from a pre-trained DNN, five times for each dataset, im- portance weighting method and weights choice approach (standard optimized). For each run, we use the test set to compute two metrics for generalization performance: weighted average ac- curacy, which is the equal-weight average of the accuracy per group, and worst-group accuracy.
7

---

Published as a conference paper at ICLR 2025
We compare the difference in generalization performance between the standard weights choice and optimized weights via one-sided t-test of paired samples. See Appendix for details.
Improving existing importance weighting methods: Figure compares the expected generaliza- tion performance the standard weights choice and optimized weights. For a correct interpretation of the results, recall that for GW-ERM, SUBG and DFR, the objective of optimized weights is the weighted average accuracy, whereas for GDRO and JTT this is the worst-group accuracy. We ob- serve that our approach consistently improves the objective for which the weights are optimized, and is at least as good as standard weights with the exception of JTT on CelebA. Out of the 15 cases, we observe 12 times an increase that is statistically significant the 10% level, and times at the 5% level. Regarding the size of the improvements, in cases there is an increase greater than 1%, with a maximum of 5.61 %. To explain the case of JTT on CelebA, we note that even on the validation set optimizing weights did not improve the worst-group accuracy. This indicates that in this case our bi-level optimization procedure was unable find better weights than the grid search of JTT.
Waterbirds
GW-ERM
SUBG
88.2
89.2
90.2
91.3
92.3
93.3
79.0
81.1
83.2
85.2
87.3
89.4
CelebA
GW-ERM
SUBG
GDRO
91.3
MultiLII
GW-FRM
SUBG
DFR
GDRO
ΙTT
83.7
84.0
84.3
84.9
71.8
73.6
75.3
77.1
78.8
80.6
Weighted Avg. Accuracy (%)
Worst-Giroup Accuracy (%)
GW-ERM
Standard
Optimized
90%
Figure 2: Estimated generalization performance (averaged over runs) of standard weights choice (squares) versus optimized weights (triangles). Error bars reflect the paired-sample 90% confidence interval of the difference.
Furthermore, we find that optimizing weights often also improves the metric for which the weights not optimized (worst-group accuracy for GW-ERM, SUBG and DFR; weighted average accu- racy for GDRO and JTT). The only exception SUBG on CclebA. 10 cases, there is an improve- ment by more than 1%, with maximum of 4.03%. This result not necessarily expected, since
8

---

Published as a conference paper at ICLR 2025
weighted average accuracy and group robustness often opposed objectives (Agarwal al. 2018 Martinez et al. 202 See appendix for numerical details of the results.
Role of the dataset size: it can be observed in Figure 2 that the improvement due to optimized weights is less pronounced for CelebA and MultiNLI than for Waterbirds. One possible explanation that the Waterbirds dataset is much smaller, leading (very) small minority groups (see Appendix E.Ip. We have argued in Sections 3 and [ that is in particular this situation where the effect of optimized weights is expected be large.
We test this for GW-ERM on CelebA in Figure 3 where we show the effect of optimized weights using only fractions of the original training and validation set (while keeping the ratio of the sizes the same). The improvement due to optimizing weights peaks when around 10% of the original data used. Note that for smaller fractions of the data (5%) the effect decreases again. This is likely due to the validation set becoming (very) small, increasing the risk of overfitting on this data.
CelebA
1.0
0.5
Fraction
0.25
0.1
86.7
87.6
88.5
89.4
90.3
91.2
71.9
75.1
78.3
8L6
84.5
88.0
Weighted Avg. Accuracy (%)
Wonst-Group Accuracy (%6
Standard
Optimized
90%
Figure 3: Estimated generalization performance (averaged over runs) of standard weights choice (squares) versus optimized weights (triangles) for GW-ERM as function of the sizes of the training and validation set. Error bars reflect the paired-sample 90% confidence interval the difference.
Robustness to other hyperparameters: additional benefit of optimizing weights improved robustness to the choice of the LI penalty. We illustrate this in Figure terbirds dataset. Both generalization metrics vary less with optimized weights. We believe this is because optimized weights improve on the birianar trade-off, thereby compensating for subop- timal regularization choices in the hyperparameters.
4 for GW-ERM and the Wa-
Waterbirds
0.1
3.3
penalty
33.33
H++
Hd-
Hil
300.0
87.9
88.9
90.0
91.0
92.1
93.1
72.5
76.0
79.5
83.1
86.6
90.1
Weighted Avg, Aveurary (15)
Worst-Group Accuracy (%5
Standard
Optimized
90% CT
Figure 4: Estimated generalization performance (averaged over 5 runs) of standard weights choice (squares) versus optimized weights (triangles) for GW-ERM as a function of the LI regularization parameter. Error bars reflect the paired-sample 90% confidence interval of the difference.

---

Published as a conference paper at ICLR 2025
7 CONCLUSION AND DISCUSSION
In this paper we have shown that the standard choice of weights for importance weighting in the presence of sub-population shifts is often not optimal, due to the finite sample size of the training data. We have provided an estimation procedure for the optimal weights, which generally improves existing importance weighting methods in the case of last-layer retraining DNNs.
Although last-layer retraining is known to be sufficient for addressing sub-population shifts (Iz- mailov et al. 2022] Kirichenko et al.) 2023), it might be that retraining the whole DNN is the preferred option. In that case computational costs can become an issue, calculating the hyper- gradient requires refitting the model on the training data and computing the inverted Hessian of the training loss with respect to the model parameters. To alleviate this, one might resort to several existing approaches to reduce computational costs for refitting the model (Maclaurin et al. 2015) or calculating the inverted Hessian (Lorraine et al. 2020).
The optimization properties of our method are another potential topie for future research. Although, with the exception of JTT CclebA in Figure [ the bi-level optimization procedure seems to converge well, formal convergence guarantees or the conditions under which the optimization in Equation TO is convex were not investigated here.
The application to more challenging data- genttaggneratingg processes and models could lead to additional benefits of approach. One example is the case in which the test distribution of interest is un- known (Sugiyama et al. 2012). Another example is concept bottleneck models, which are known suffer from spurious correlations (Margeloiu et al.. 2021 Heidemann al. (023). The existence of many concepts, and thus groups, leads to a severe reduction in the effective sample size. This problematic for existing importance weighting techniques and gives our approach of optimizing weights the potential to lead to big improvements. We leave this for future work.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Alckh Agarwal, Alina Beygelzimer, Miroslav Dudik, John Langford, and Hanna Wallach. A reduc- tions approach to fair classification. Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 60-69. PMLR, 10-15 Jul 2018. URL BtLhttts: /R//tttss//prroeed.gg..o
press/ //////80///// agarwa 118a html
Michael A. Alcom, Qi Li, Zhitao Gong, Chengfei Wang, Long Mai, Wei-Shinn Ku, and Anh Totti Nguyen. Strike (with) a pose: Neural networks are easily fooled by strange poses of familiar objects. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recog- nition (CVPR), pp. 4840-4849, 2018. URL http:pss//pp..sssattccccccbbbbb www.eeennncnnooeeeeeeeeee org/
CorpusID:
Yoshua Bengio. Gradient-based optimization of hyperparameters. Neural Computation, 12(8): 1889-1900, 2000. doi: 10. 10.162288976000000111177..
Y.M. Bishop. R.J. Light, F. Mosteller, S.E. Fienberg, and P.W. Holland. Discrete Multivariate Analysis: Theory and Practice. Springer New York, 2007. ISBN 9780387728056. URL
Joy Buolamwini and Timnit Gebru. Gender shades: Intersectional accuracy disparities in com- mercial gender classification. Sorelle A. Friedler and Christo Wilson (eds.), Proceedings of the 1st Conference Fairness, Accountability and Transparency; volume 81 of Proceed- ings of Machine Learning Research, pp. 77-91. PMLR, 23-24 Feb 2018. URL https: 7 //proceedings.. ml press / v8 ///1//bu amwi n 8a html
Jonathon Byrd and Zachary Chase Lipton. What is the effect of importance weighting in deep learning? In International Conference on Machine Learning, 2018. URL httts://aat
semanticscholad ctt.:aaaaa.aageooooaaa:a:::3330:
Tianle Cai, Ruiqi Gao, Jason Lee, and Qi Lei. A theory of label propagation for subpopulation shift. In Marina Meila and Tong Zhang (eds.). Proceedings of the 38th International Con- ference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp- 1170-1182. PMLR, 18-24 Jul 2021. URL
ngs. ml press/v139/
html
Annie Chen, Yoonho Lee, Amrith Setlur. Sergey Levine, and Chelsea Finn. Project and probe: Samplc-efficient adaptation by interpolating orthogonal features. In The Twelfth International Conference on Learning Representationss ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenRe- net. www.moriioooooo (oorumm1-=ffCCOXXXXV
view.net, 2024. htRLhhttss/
Yin Cui, Menglin Jia, Tsung- Yi Lin, Yang Song, and Serge J. Belongie. Class-balanced loss based on effective number of samples. 2019 IEEE/CVF Conference on Computer Vision and Pat- tern Recognition (CVPR), pp. 9260-9269, 2019. URthtstss:// stct:s/s/ap.ssmaattccchoo.....
Alex J. DeGrave, Joseph D. Janizek, and Su-In Lee. Ai for radiographic covid-19 detection lects shortcuts over signal. Nature Machine Intelligence, (17,,100669, 2021. doi: 10. s42256-021-00338.
CRL/httss://doiorrr httpg//...11///4222200
021 00338-7337
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume I (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019a. Association for Com- putational Linguistics. doi: 10. 10.18853//1/N19-11 14233. URL https: /atps://acllo htt///alanntttt ogy. org/
N19 N19-1423
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of
11

---

Published as a conference paper at ICLR 2025
the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long and Short Papers), pp. 4171-4186 Minneapoliss Minnesota, June 2019b. Association for Com- putational Linguistics. doi: 10.18653/1/N19-14233 URL https:/ https://aclaattoooo.ooo org/
N1 9- 423
Robert Geirhos,
Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A. Wichmann. Shortcut learning in deep neural networks. Na- ture Machine Intelligence, 2(11):665 673, 2020. doi: 10. 020884442226-020000255----- URL
https: ptttps:/ddoi.orr htt///0..rg//....// 1038 0//38/42220 s42256 020 00257-2
Suchin Gururangan, Swabha Swayamdipta, Omer Levy, Roy Schwartz, Samuel Bowman, and Noah A. Smith. Annotation artifacts in natural language inference data. In Marilyn Walker, Heng Ji, and Amanda Stent (eds.). Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Vol- ume 2 (Short Papers), pp. 107-112, New Orleans, Louisiana, June 2018, Association for Com- putational Linguistics. doi: 10. 10.18653/v1/N1--11 2017. URL pttps:/aclanthooogyy org// N18-2017
Yujin Han and Difan Zou. Improving group robustness on spurious correlation requires preciser group inference. In Ruslan Salakhutdinov. Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 17480-17504, PMLR, 21-27 Jul 2024. URL ctttps://proceedingg mwl. mlr.press/ w235/han204 html
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEFE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016.
Lena Heidemann, Maureen Monnet, and Karsten Roscher. Concept correlation and its effects concept-based models. In 2023 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), pp. 4769-4777, 2023. doi: 10.1 1110///C/366688 2WW.AACV6668..202.000044444 2023.00476.
Floris Holstege, Bram Wouters, Noud Van Giersbergen, and Cees Diks. Removing spurious concepts from neural network representations via joint subspace estimation. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.). Proceedings of the 41st International Conference Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 18568-18610. PMLR, 21-27 Jul 2024. URL htttps://proceedie
mwr. wrwss w2w.:/2233////333 2at35/hhlsseeeee/ 4a. html
Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive worst- worst-group-aat group-accuracy In Bernhard Schölkopf, Caroline Uhler, and Kun Zhang (eds.). Proceedings of the First Conference on Causal Learning and Reasoning, volume 177 of Proceedings of Machine Learning Research, pp. 336-351. PMLR, 11-13 Apr 2022. URL htt.tttts://rroeeeeet ngs mlr. www.prrsssvvl... 77/ wtt/////rrrr iitrrisss2222. html
Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew Wilson. On fea- ture learning in the presence of spurious correlations. In S. Koyejo, S. Mo- hamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems, volume 35, pp. 38516-38532. Curran Associates, Inc., cc/ cc/paper files/paper/2022/ Conferenceeeccccc
2022. URL httpps///ppcceednnsss..co ottp://sspccppeeeeeeeeeee///ppe File/ 1wt.1//2eeee31011101333322208880 carerarreee.eaeeeeceeeeeeee... pdf
Sarah Jabbour, David Fouhey, Ella Kazerooni, Michael W. Sjoding, and Jenna Wiens. Deep learning applied to chest x-rays: Exploiting and preventing shortcuts. In Finale Boshi-Velez, Jim Fackler, Ken Jung, David Kale, Rajesh Ranganath, Byron Wallace, and Jenna Wiens (eds.), Proceedings of the 5th Machine Learning for Healthcare Conference, volume 126 of Proceedings of Machine Learning Research, pp. 782. PMLR, 07-08 Aug 2020. URL https: http::////cceeenn....
mlr. mlr. www.press/v126//abbouu
.ttl
12

---

Published as a conference paper at ICLR 2025
Nitish Joshi, Xiang Pan, and He He. Are all spurious features in natural language alike? anal- ysis through a causal lens. In Yoav Goldberg, Zomitsa Kozareva, and Yue Zhang (eds.), Pro- ceedings the 2022 Conference on Empirical Methods in Natural Language Processing, pp- 9804-9817, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 2211665//10022.emmlppy-
URL htt.ttt//ccaaaachoolllo oorg/
2022. emnlp main. 1111.6666 666
H. Kahn and A. W. Marshall. Methods reducing sample size in monte carlo computations. Journal of the Operations Research Society of America, 1(5):263-273 1953. ISSN 00963984. URL
htttp://ww. jstorrssorr.... wtw.oo..rrg/staae//e66777
Masanari Kimura and Hideitsu Hino. A short survey on importance weighting for machine learn- ing. Transactions Machine Learning Research, 2024. ISSN 2835-8856. URL https: / //openeevie.neete for/forunnn Forommmmaa 2tt.f/oorrommm1--11233333aaaaa Survey Certification.
Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training sufficient for robustness to spurious correlations. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023.
Jyrki Kivinen and Manfred K. Warmuth. Exponentiated gradient versus gradient descent for linear predictors. Information and Computation, 132(1):1-63, 1997. ISSN 0890-5401. doi: https:// Noi.org/10..1006/mmmoo 1996 2612. URL httpp::///// /www. www.nceeieeect....cccee.eeeeeeee soinnseeecee/
art art.cle/p11/s0800000 9691999926620 6229221221 6127
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness, Wei Guo, Berton Earnshaw, Imran Haque, Sara M Beery, Jure Leskovec, Anshul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. Wilds: A benchmark of in-the-wild distribution shifts. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 5637-5664. PMLR, 18-24 Jul 2021. URL https://proceett
mlr wttp:///w......eee/////S3
html
Daphne Koller and Nir Friedman. Probabilistic Graphical Models: Principles and Techniques Adaptive Computation and Machine Learning. The MIT Press, Cambridge, Massachusetts, 2009. ISBN 0262013193.
Abhinav Kumar, Chenhao Tan, and Amit Sharma. Probing classifiers are unreli- able for concept removal and detection. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances Neural Information Process- ing Systems, volume 35, pp. 17994-18008. Curran Associates, Inc., 2022. URL
/ttpss//ppoceedings..eeeris http://appe..illeeppppp//////// file/
Tyler LaBonte, Vidya Muthukumar, and Abhishek Kumar. Towards last-layer re- training for group robustness with fewer annotations. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.). Advances in Neural In- formation Processing Systems, volume 36, pp. 11552-11579. Curran Associates, Inc., cc/ http://peeeeeeee///eeeaeaeeeaaaa files//appr/2023/ www.-aaapee---eeeeeeeeeeee pdff
2023. URL pttp:://prooeddngg..e.uu File/ /ttee/25beee2pee waweeee/eeeee6ff//e/ Flle///bbeereeeeedd////363333332
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. In Marina Meila and Tong Zhang (eds.), Proceedings the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 6781-6792 PMLR, 18-24 Jul 2021. URL https: httt::://proceedingg.. ml wwt.mm.pp....111..mmm
1iu21f html
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV) December 2015.
13

---

Published as a conference paper at ICLR 2025
Ziwei Liu, Zhongqi Miao, Xiaohang Zhan, Jiayun Wang. Boqing Gong, and Stella X. Yu. Large-scale long-tailed recognition in an open world. 2019 IEEE/CVF Conference on Com- puter Vision and Pattern Recognition (CVPR), pp. 2532-2541, 2019. URL https://api,
semant i cscho 1 ar. Ctrpg//orrsstt... ID: 11:1133331
Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In Silvia Chiappa and Roberto Calandra (eds.), Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pp. 1540-1552. PMLR, 26-28 Aug 2020. URL https::/proceeddinn m.r.mmmr..m press wtt.aereeeeevvvverrrrrrrar
html
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2017. URL pattps://api... http://pp...eaaniicccc...oaaaaa. /rgg
Corpus ID: 535922292277
Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient- hyperparameter optimiza- tion through reversible learning. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning, volume 37 Proceedings of Machine
Learning Research, pp. 2113-2122 Lille, France, 07-09 Jul 2015. PMLR. URL https: html
//proceedings.. /ttp://ddinsaamipppprreeee////ee
Maggie Makar, Ben Packer, Dan Moldovan, Davis Blalock, Yoni Halpern, and Alexander D' A'mour Causally motivated shortcut removal using auxiliary labels. In Gustau Camps-Valls, Francisco J. R. Ruiz, and Isabel Valera (eds.), Proceedings of The 25th International Conference Arti- ficial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pp- 739-766. PMLR, 28-30 Mar 2022. Eakar22a.htmi
URL
https: http:////occeeeeegseeeee mlr. press//1511
Andrei Margeloiu, Matthew Ashman, Umang Bhatt, Yanzhi Chen, Mateja Jamnik, and Adrian Weller. Do concept bottleneck models learn as intended? arXiv preprint 1442:92 Abs/2105.04289, 2021. URLhtttps:/aaxiv..om http://///aaaa...rror/////../.. 04289
Natalia L Martinez, Martin A Bertran, Afroditi Papadaki, Miguel Rodrigues, and Guillermo Sapiro. Blind pareto fairness and subgroup robustness. In Marina Meila and Tong Zhang (eds.), Pro- ceedings of the 38th International Conference Machine Learning, volume 139 of Proceed- ings of Machine Learning Research, pp. 7492-7501. PMLR, 18-24 Jul 2021a. URL https: http://v13/mmarrriic www.:////mmmaieeeaaaa.mm html
/pproceddings mlr
Natalia Martinez, Martin A Bertran, Afroditi Papadaki, Miguel Rodrigues, and Guillermo Sapiro. Blind pareto faimess and subgroup robustness. In Marina Meila and Tong Zhang (eds.), Pro- ceedings of the 38th International Conference on Machine Learning, volume 139 of Proceed- ings of Machine Learming Research, pp. 7492-7501. PMLR, 18-24 Jul 2021b. URL https: html
www.eng..s.........m.mm... mlr wtt.a/...eeeeee//aaaaaaii.......
Luca Martino, Víctor Elvira, and Francisco Louzada. Effective sample size for importance sampling based on discrepancy measures. Signal Processing, 131:386-401, 2017. ISSN 0165-1684. doi: http://dd.ooorooo...ooo/..iig/// 201..000.......0000...2 URL https:tpp.. /www.wwww....oo.....o..... sciencedirect. com/
Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: training debiased classifier from biased classifier. In Proceedings of the 34th International Con- ference on Neural Information Processing Systems, NIPS '20 Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Junhyun Nam, Jaehyung Kim, Jaeho Lee, and Jinwoo Shin. Spread spurious attribute: Improving worst-group accuracy with spurious attribute estimation. In International Conference on Learning Representations, 2022. URL
net/ /ttp://xummmma__P_ppp000000000
A.S. Nemirovsky and D.B. Yudin. Problem Complexity and Method Efficiency in Optimization. A Wiley-Interseienee publication. Wiley, 1983. ISBN 9780471 1780471103455. URL https: httttps::/bbok...
google.nl/bol http://oooookk..dd..ll.......... 6
14

---

Published as a conference paper at ICLR 2025
Yonatan Oren, Shiori Sagawa, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust language modeling. In Kentaro Inui, Jing Jiang. Vincent Ng. and Xiaojun Wan (eds.), Pro- ceedings of the 2019 Conference Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNIP-IJCNIP), pp- 4227-4233 Hong Kong, China, November 2019. Association for Computational Linguisties. doi: 10.18653/v1/D19911
URL httpppss//ssasaanoooooooooo/////
Fabian Pedregosa. Hyperparameter optimization with approximate gradient. In Maria Florina Balcan and Kilian Q. Weinherger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 737-746, mllr.
New York, New York, USA, 20-22 Jun 2016. PMLR. URL html
press/v48/pedreeoee
Mohammad Pezeshki, Diane Bouchacourt, Mark Ibrahim, Nicolas Ballas, Pascal Vincent, and David Lopez-Paz. Discovering environments with XRM. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Pro- ceedings of Machine Learning Research, pp. 40551 PMLR, 21-27 Jul 2024. URL
Dttts::/pproceeaingg ml www.prrees..3.. http2235/peeesskii2222 http://m.w...aaa////////eeeaaaa2 hat.
Shikai Qiu, Andres Potapezynski, Pavel Izmailov, and Andrew Gordon Wilson. Simple and fast group robustness by automatic feature reweighting. In Andreas Krause, Emma Brun- skill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Pro- ceedings of the 40th International Conference on Machine Learning, volume 202 of Proceed- ings of Machine Learning Research, pp. 28448 28467. PMLR, 23-29 Jul 2023. URL https: //proced htwphh..ww............tptpt.tt
ml www.prresssspp. v/rres//202//iia qiu33c. html
Joaquin Quiñonero- Candela, Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, Massachusetts, 2022.
Mengye Ren, Wenyuan Zeng, Bin Yang, and Raquel Urtasun. Learning to reweight examples for robust deep learning. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th In- ternational Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 4334 4343. PMLR, 10-15 Jul 2018. URL https: http:ss:///occeediggs...
press/ Vress/v80/ae ren 8a htmml
Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. In International Conference on Learning Representations (ICLR), 2020a.
Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Percy Liang. An investigation of why overparameterizatioo exacerbates spurious correlations. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 8346 8356. PMLR, 13-18 Jul 2020b. URL mlr
Shibani Santurkar, Dimitris Tsipras, and Aleksander Madry. Breeds: Benchmarks for subpopulation shift, 2020. URL htt.pp::/axxxv..utt wrt.:bs./////ooee..ooo.o.oo
04859
Roy Schwartz and Gabriel Stanovsky. On the limitations dataset balancing: The lost battle against spurious correlations. In Marine Carpuat, Marie-Catherine de Marneffe, and Ivan Vladimir Meza Ruiz (eds.). Findings of the Association for Computational Linguistics: NAACL 2022, pp. 2182-2194, Seattle, United States, July 2022. Association for Computational Linguistics URL. /attpsss//cclaanthlloy..coo
doi: 10. 2tt25//ww12222..niinnn---- Findings-naacl 168
Shreya Shankar, Yoni Halpem, Eric Breck, James Atwood, Jimbo Wilson, and D. Sculley. No classi- fication without representation: Assessing geodiversity issues in open data sets for the developing world, 2017. UR///ttpss:/araivoooo/////
08536
Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log-likelihood function. Journal of Statistical Planning and Inference, 90(2):227-244 2000.
15

---

Published as a conference paper at ICLR 2025
ISSN 0378-3758. doi: http:://oaaooy///oo0000007777700 URL https: //WW:. sciencedirect
Masashi Sugiyama, Taiji Suzuki, and Takafumi Kanamori. Density Ratio Estimation in Machine Learning. Cambridge University Press, Cambridge, 1st edition, 2012. ISBN 0521190177.
A. W. van der Vaart, Asymptotic Statistics. Cambridge Series in Statistical and Probabilistic Mathe- matics. Cambridge University Press, Cambridge, 1998.
Peter Welinder, Steve Branson, Takeshi Mita, Catherine Wah. Florian Schroff, Serge Belongie, and Pictro Perona. Caltech-ucsd birds 200. Technical Repon CNS-TR-201, Caltech, 2010. URL https:
/ /ww.wwwwwwkmmmmmm vwsiion.i...m....i....iiiiiiii.i ca tech edu/ cub. 200 _00__011]]
Junfeng Wen, Russell Greiner, and Dale Schuurmans. Domain aggregation networks for multi- source domain adaptation. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 10214-10224. PMLR, 13-18 Jul 2020. URL https::/pprceeeeee .ml.. #ress/v119/wente html
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sen- tence understanding through inference. In Proceedings of the 2018 Conference of the North Amer- ican Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pp. 1112 122. Association for Computational Linguistics, 2018. URL http::/accwee org/ antho http://nntooooo//111--1-
Benjamin Wilson, Judy Hoffman, and Jamie Morgenstern. Predictive inequity in object detection. arXiv preprint arXiv:1902.11097 abs/1902.11097, 2019. URL http:://axdv.org//b/
11097
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. Huggingface' 's transformers: State-of-the-art natural language processing. CoRR, abs/1910.03771, 2019. URL
http:///rxixs org/ taabs/1900, 03771
Da Xu, Yuting Ye, and Chuanwei Ruan. Understanding the role of importance weighting for deep learning. In International Conference on Learning Representations, 2021. URL https:// openreview net wtwpe//wwne//ooummmmmm
forum?id Wnwt ieRHxM
Yuzhe Yang, Haoran Zhang, Dina Katabi, and Marzyeh Ghassemi. Change is hard: A closer look at subpopulation shift. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 39584-39622. mlr wwwpess/v222/yaaggae PMLR, 23-29 Jul 2023. URL https::/pprocedd::ii http:://ooooeonn...ooo.errrrrrrr
html
Michael Zhang, Nimit S Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Re. Correct- n-contrast: a contrastive approach for improving robustness to spurious correlations. In Ka- malika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 26484-2000 26516. PMLR, 17-23 Jul 2022. URL 62/ zhwng h
htttps://proceed.nggs
Bolei Zhou, Aditya Khosla, Àgata Lapedriza, Antonio Torralba, and Aude Oliva. Places: An image database for deep scene understanding, arXiv preprint arXiv: abs/1610.02255, 2016. URL
16

---

Published as a conference paper at ICLR 2025
A CALCULATION OF THE SEPER-GRADIENT
The goal of this section is to provide further details on the calculation of the hyper-gradient, which the gradient of the validation loss with respect to our hyperparameter p, introduced in Section 5.1 This section should be read in conjunction with Section Recall that the chain rule gives
Vov (p), VL (p., a0.r 0p
(14)
(p)
Op
Suppose we have a function h(O.p) p) which is twice differentiable, which will play the role of an (expected) training loss. Assume that the gradient of h(0,p) with respect to 0 is zero at the point (n:(p): We can then apply the implicit function theorem to this gradient the point 0 0., (p), and write
a (p) 0p
(h((,p)) p)
Oh p)
(15)
0000T
000pT
(= (p)
(0=
The gradient in Equation 15 depends the form of h, and thus the estimation procedure of the parameters. Here, we discuss two such estimation procedures: (1) a group-weighted loss, and (2) subsampling per group without replacement.
Group-weighted loss: for simplicity of notation, assume that the group label takes values 1,2,...,Ge G. The grightweighted training loss is defined
Fly
h(0,p) p)
Σ
Pg Pur.a
Σ
Cly.. fo(lfo
(16)
n
1=
The partial derivative of h(0,p) with respect to Py becomes a novel (weighted) loss function,
@h(O.p)
Σ Clyl Clv (oll)))
fol.))
(17)
Öp,
Ptr.
(1 Eg' =1 Pur.g"
1=1
where we used that PG = ES Py due to normalization. Then, in order to determine 8 h(0.p) 300p
calculate the derivative with respect to the estimated parameters for the novel loss function in Equation 17 for each Pg-
(p)
Subsampling large groups: here, we are interested in the hyper-gradient for the hyperparameter instead of as defined in Section 5.2 The parameters (w) are not determined via a weighted loss, but rather by minimizing the training loss on a subsample of the data. However, the expected training loss can be rewritten replacement, we say that for each g = I,..., the random variables sy are identically distributed Bern(yg) random variables with constraint
weighted loss. To mimic the procedure of subsampling without
Σ
5 mg. where my [[gmgg]
(18)
Recall from the main text that 0 < Mg < 1, the fraction of samples used from the training data in our subsample. The expected training loss then
G ngg
E,
SUBG," (0,
E.
(19)
m
g=
im1
1
R2
Cly fo (xx)))
(20)
n2
where Eg my is the total number of samples in our subsample. By fixing the value of each my beforehand, m is not stochastic. As explained in Section 5.2 approximate this expected loss
17

---

Published as a conference paper at ICLR 2025
via h0 = Σ Σx Dy Clyl' fo fo()))) The partial derivative of h(0, v) with respect to 1
again becomes a novel loss function
ah(0. v)
G n
[ fa(x)))
(21)
0ug
Note the difference with Equation (00
17 because the v are not normalized. To determine @= (=) we can (again) calculate the derivative with respect to the estimated parameters
800
for the novel loss function in Equation 21 for each Vig. In addition to the change in the calculation of
the syper-gradient the only change to Algorithm iis that we not normalize the weights v.
B THEORETICAL ANALYSIS FOR LINEAR REGRESSION
The goal of this section is derive the approximate expected loss given the data-generating process described in Section 4 It consists of four parts.
In Section B.1 we give description of how EDD-PP ( Iyy - x ((
can be decomposed into a bias and variance term the case of linear regression, and data stemming from G groups.
In Section B.2 we provide the approximations of the OLS coefficients for large but finite n, given the data- gata-generating process described in Section
In Section B.3 combining the results of the previous two sections, we derive the (approxi-
mate) bias-variance decomposition for our data-generating process.
Finally, in Section approximation.
B.4 we provide a numerical comparison between simulations and our
B. THE BIAS-VARIANCE TRADEOFF
For clarity, throughout the next three sections we will the subscript tr indicate random variables are drawn from the training distribution, and the subscript te that they are drawn from the test distribution. We use the mean-squared error as the loss function, and use a linear regression where the coefficients depend on our weights.
EDy- Per Bm (itre -
Dul] EDg Per [e [(ite xx x:BB(w))))\..l))
The following only requires the assumptions stated in Section F regarding the independence of Ete from étc- The expected mean squared error over a test distribution, conditional on the data coming from group 9, be written
Epa~Per Ex [[rr - x
Dul]
EEppPer Earn Wio 2 3(u)) + ((v)) Dtt]
Emm. [wcclg] 21 2. Emeiie (3 lg Dt.
EDR (pp(ow)
(zw))( Puo
18

---

Published as a conference paper at ICLR 2025
For each group 9, the following holds:
E [oclll = Ex.. [<0 + 11) = Eu.,, 3 +o²
E Bpp Bxc..ttr yz 3(w) 1s. D(s
-prr Basassie (1 l 3 Der
Exer 12 9D [Ba. (ew)( Duy)
E Eir Bunnthr
( Dir]]
EDy Var (..T)(w))lg, Dis) + Bmn. ( Dir Dru
We can then write the expected loss conditional on g as:
EDa-Per [mm...nn (1tto - ((1)(()))) Dirl] 2P (ar Bd pop Emi (=TTB(u)))o: Der
Baw
+ EDp Eror (uT)(w))) Des 1D#=
Var (alo((w)) Dtt) Diel Rel+
+ EDDNPer
Suppose we write this expected loss conditional on both are and Then, it can be written as:
Epa-Por
Egg. [(Vte 3T,03 ((rr))2 Paee s]
= -xx x EDEP [(w)) )² + E.E [Var(3(uo)\D..)) Tre + og.
To further specify the expected loss, we will use the data-genen getta-ennrating process specified in Section 4 Most notably, we will use that the Tte does not change with the group g, that there are only two groups, and our definition of the coefficients per group.
By the law total probability, the expected loss conditional on only ahe can written as a sum of the expected losses conditional Øte.
EDa~Per Eyrn.gen 1(yre - x
G
Σ p(gre = gD)BDa~P.
Eyr. (Vire (1hre
2 = Iue.
T -.T ((w))) (1 pre) (1= w -D
+ 2 E.e
T
Var( Ite + 002
(22)
B.2
THE ASYMPTOTIC BIAS AND VARIANCE OF THE COFFFICIENTS
The goal of this subsection to define the expectation and variance of the WLS estimator as defined per equation 7 Because have fixed the amount of observations per group, can state the following
1PPr
n21 Ptr) = (1-ptr:
Pir:
(23)
n
n
71
110
(1 p)
Wi
Ptr
p.
Wo
(1 ptr) =(1 p).
(24)
n
Phr
n
(1 Ptr)
This means we can write
x Wx,
XT X + Wo XT
11,0 Xtro,
X = wi XT 1 Wo XT
t0,00o (25)
19

---

Published as a conference paper at ICLR 2025
and similarly W.W.; wf Xtx,1 + o where Xux,1 Ruixd Xu.o I are matrices where for the observations holds that 9; 1.g; = 0 respectively, and similarly Y1 E , you E "R
In the following analysis we will use that
(w) [wixXTr ur,1 Xtr,1 + wox ir.0 Xamoo)
-
1 Th x tr Xtr.1 + wo 110
(wil
x T Xer.n)
ur.
M1
ni0
x br,
X + (1-p) X
no
(26)
Similarly,
1
H + x
1
(L1 2
"i
X tr, + ww
2 no
x tr. X
n
11
Ti
110
y² 1
x tr,
sr,
(1 (1-p)² 1
x er X.n0 (27)
Per n1
- (1-Ptt) no
Finally, we will use that
XT
11 91 + wo XT
xT.000
= =- =
71
XT
tr. 1/1 + w0 tr 130
7l T
71
1
p 71
tr (1 - p) no 0r
T
(28)
Using Equations 26 and 28,
the WLS estimator can be re-written as
WX.) yor X
X tr
W
yor
1 xxx.wx Xur
-)=
1
p
x tr.
3i + (1 p)
x
ni
n10
The expectation of this estimator, conditional on g. X becomes
Blg(p) gi.. X
W
X 1T
W x yor Mur, X1r
n
72
cr
x
X1r
p
tr.1 3/1
+ (1 111-21
Xx x4.0000
gin. X
1h1
110
1 x X1:
p
XT -
Ely
X + (1 p)
1
X
tr.
B(y0|
T1
n10
p-x 1r,1
X 13 + (1 - pp) no
tr.0 X Xtr.x.of.
Given our DGP we can write the variance of you conditional on gtr, Xte as
Var Var X = -²I.
(29)
We can then define the variance of the WLS estimator conditional g. X: Var(S(p) Ig.. X Var Var((X,)) W x x. ytr Itr: X
= W Xix) x War gt.. X. W.) WX..)
=22(X. WX (X
(30)
20

---

Published as a conference paper at ICLR 2025
Similar to the data-genedn gentratinng process in Section assume here that each follows the same dis- tribution regardless the group. Given this, we can state EI I = E1 x X Xu,ool =
>, c.g. they have the same expectation. Assuming Σ a positive definite matrix with finite ele-
ments, mg
X X1, no X Xo converge in probability it, e.g.
plim - By 11
x,
plim - _ Xo E. Ro 400 10g
(31)
Since we assume each of our observations of a are identically and independently distributed (IID), can write
1
1 x er X E Op(
(32)
See Equation 14.4-7 of Bishop et al. (007) for general version of this result. We can thus write
p- x Xtr.1 + (1 -1) p) X
X..)Xm.))) - -p -1 = E + Og
and also,
x Xxr +
-1
1
p X tr,1 X (1 (1-p)
E Opl
m
no
Vn
=E - -E ',,
Vn E
Σ ---1 Op1
The Op( (n)) indicates 1 X X
n2 1r,1 no
X X bounded in probability (see van der Vaart
1998 Section 2.2). Thus, for large approximate E [X] WX.))
using Σ Σ-1
We use the following approximation
E EDa~PH (3w))
(33)
[x Wx,, -1 ( XT Xor, + (1-p) (1--p)
I1
12
X tr.1
(mm
XT Xrr.oXrrn X.r,n33)) no tr.0
=E -E-1EEE +1 (1 p) E +0p( (
Vn
p3l (1- p),30.
(34)
Similarly, we can define an approximation for the conditional variance. First, we define
plim x ( 2p² n
(1 (1-p)² E.
(35)
Ptr (1 (1ptr))
We can now state
Epy ~Por (Vrr(((w)) = E Eg, [x2 X
p²
+ (1- (1=
2
E
(36)
Per (1 Par)
B.3
THE EXPECTED LOSS
we substitute in our approximations derived in Equations
34 36 in Equation 22
we get
EDg [py...) =TTB((w)))( ~
Σ
(x - x (p3) + (1 - p)))))))2 + +o²
Pre.e
9=1
p2²
+ (1 (1-))
p)²
02
E -
in
Per.1
(1- Pt,1)
21

---

Published as a conference paper at ICLR 2025
Now, will define the expected loss unconditional of Ete. We will start to generally define the squared bias
(x (T( + (1-p) -- - ((x
where 3 is the difference between the coefficients in the data-generating process of group and estimated coefficients. Note that in general
1 (2)
88 By
j=1
( (g) = TT(Ba a T
B| 1),
(Since they are independent)
E 2 T
0 Σ
Now, the challenge to define 83. To do so, we define the difference between the coefficients as Ba (B' ))). We can then define
(p3) + (1 p)33% = p(B 30) + g0,0
( B p(B¹ 80) 30 =(1 (1- () Br),
0 B p(B) Bd) B0 =-p(B' 33
3 0 (1 Ba3
30 00) 3/0 BBBB
Because the difference in the coefficients is only in the intercepts, we can write
B B[ (1 (1-p)= (0) (a)=
B ([ y 2?
Suppose that we define Σ-1 1 write
11 I. Note our DGP as defined in Section 4 uses 7 =
We can now
2 p²
(1 2)²
p² (1 (1- p)
1
1
0
1 Dro]
+
-
Blar
Fia]
n1 Ptr
- Ptr) (1-Prr)
Ptr
(1 Ptr)
71
11
0² p2
(1 p)
d+1
+
Pir
(1 Pt.)
m
Now that we have defined the bias and variance terms, we can combine these with the noise term o² into the expected loss.
B (Pte. P. a1 0o) = Pco - pp) -pp2 + (1
(al- -a))
V(o2: .P,Pe,, n. o2( p² -2p + (1-p) p)2 d+1
Ptr (1 Pr.) n
Epa-Pe B(oraeyP. I( - x T 3( ss B² (pre, P, a1 0n) + V(o² P. PP. n, d) + +
(37)
We then take the derivative of with respect to the bias,
aB? (Pte: P, a, ao) 2(aa 2p))(pp Pre):
Öp
2(a (p Pre) = 0 = p Pre-
22

---

Published as a conference paper at ICLR 2025
In the that 01 # du, the bias is uniquely minimized when Pte. If take the derivative of p with respect to the variance, it is uniquely minimized at p Pt: mppr 1)P + 0.
OV ,po Pre, n. d)
2(d + 1)02(p Ptr)
(38)
Op
n(Per 1)ptr
n(Prr Per # 0, 2(d + 1)02(p Prr)
Per-
(39)
(prr 1)Prr
For the variance, if co, V(o2, Per, P.r, d) + O. Hence, if our sample is large enough, the optimal p will be p for which the bias is minimal. The p. as per Equation is derived by taking the derivative of the approximate expected loss with respect p, and setting it to
2(d + 1) (p Pw) O, (Prr 1)Prr
n(pu- 1)P # 0, 2(a1 2 Pie)
(40)
0²(d 1)
p= Pte Per
where 17 = n(a1 - an) Pte(1 (1 - pir)
(41)
1+1
B.4 NUMERICAL COMPARISON
To illustrate the validity of the approximation used in the previous section, we compare it to the actual mean squared error via simulation. We simulate the mean squared error for the given DGP for the following values: 01 1, ag 0, o² 1, 1, Poz 0.9, Pte 0.5. We compare this to our approximation of the expected loss for several values of n, d. In addition, we plot the bias and variance. For 1,000 simulations, and several values of n, d, the results shown in Figure
We can observe in this figure that as n decreases or d increases, the variance becomes a greater share of the expected loss. This variance can be reduced by selecting a p closer the original However, this subsequently leads to an increase in the bias. The optimal provides trade-off' between these two objectives.
1.0
120
1.5
155
155
145
130
135
04
61-
44-
60-
1.000, 10
(b) 5.000, d 10
5.000, d = 100
Figure 5: Illustraionn of the bias trade-off for Pt: 0.9, Pte 0.5, 1, 00 = 0, 0²
1 = 1, and different values of n and d. The simulated MSE averaged over 1,000 runs. Error bars reflect the 95% Confidence interval.
23

---

Published as a conference paper at ICLR 2025
C
FULL SET OF RESULTS
Method
Weighted Average Accuracy (%)
Worst Group Accuracy (%)
GW-ERM
90.863 (0.351)
82.556 (1.196)
+Optimized weights
91.998 (0.092)
85.435 (0.507)
Difference
1.134 (0.287)**
2.879 (0.878)**
SUBG
90.892 (0.382)
81.908 (1.503)
+Opimized weights
91.709 (0.167)
85,497 (0.543)
Difference
0.817 (0.269)**
3.589 (1 (1.268)**
DFR
90.629 (0.207)
82.333 (0. (0.735)
+Optimized weights
91.039 (0.225)
84.028 (0.978)
Difference
0.410 (0.212)*
1.695 (0.706)**
GDRO
91.523 (0. 188188)
84.306 (0.445)
+Optimized weights
92.595 (0.139)
88.481 (0.551)
Difference
1,072 (0.258)**
4.175 (0.393)**
JTT
88.542 (0.326)
79.319 (0.566)
+Optimized weights
91.408 (0.396)
84.925 (2.202)
Difference
2.866 (0.466)**
5.606 (1 (1.986)**
Table I: Results for Waterbirds dataset comparison between optimized and standard choice of weights: we report the average over 5 runs, with the standard error in parentheses. The "/"* indicate statistical significance 10%/5% significance level respectively for paired sample one-sided test...
Method
Weighted Average Accuracy (%)
Worst Group Accuracy (%)
GW-ERM
90.471 (0.120)
85.667 (0.539)
+Optimized weights
90.700 (0.062)
87.000 (0.136)
Difference
0.229 (0.114)*
1.333 (0.487)**
SUBG
89.141 (0.111)
87.598 (0 (0.284)
+Optimized weights
89.238 (0.162)
83.333 (0.731)
Difference
0.097 (0.185)
-4.665 (0.734)
DFR
88.320 (0. 195)
81.000 (0.920)
+Optimized weights
88.299 (0.176)
81.171 (1.698)
Difference
-0.021 (0.092)
0.171 (1.118)
GDRO
89.113 (0. 173)
85.333 (0.648)
+Optimized weights
90.712 (0.083)
89.667 (0.222)
Difference
1.599 (0.203)**
4.333 (0.727)**
JTT
88.496 (0.121)
79.222 (0.7 16)
+Optimized weights
88.243 (0.165)
75.556 (0.680)
Difference
-0.253 (0.139)
-3.667 (0.372)
Table 2: Results for CelebA dataset of comparison between optimized and standard choice of weights: we report the average over 5 runs, with the standard error in parentheses. The */'" indicate statistical significance at 10%/5% significance level respectively for paired sample one-sided t-test.
24

---

Published as a conference paper at ICLR 2025
Method
Weighted Average Accuracy (%)
Worst Group Accuracy (%)
GW-ERM
83.967 (0.192)
72.087 (0.779)
+Optimized weights
84.315 (0.238)
76.114 (0.660)
Difference
0.349 (0.106)**
4.027 (0.185)"*
SUBG
84.100 (0.233)
73.058 (1.085)
+Optimized weights
84.287 (0.169)
74.163 (0.798)
Difference
0.187 (0.092)*
1,105 (0.621)*
DFR
84.281 (0. 109)
78.207 (0.313)
+Optimized weights
84.654 (0.216)
79.740 (0.130)
Difference
0.373 (0.162)**
1.533 (0.367)**
GDRO
84.452 (0.208)
75.654 (0.585)
+Optimized weights
84.400 (0.190)
76.350 (0,472)
Difference
-0.051 (0.085)
0.696 (0.338)*
JTT
84.071 (0.205)
72.427 (0.680)
+Optimized weights
84.374 (0.257)
76.226 (0.698)
Difference
0.303 (0.081)**
3.799 (0.211)**
Table 3: Results MultiNL. dataset of comparison between optimized and standard choice of weights: we report the average over 5 runs, with the standard error in parentheses. The */** indicate statistical significance at 10%/5% significance level respectively for paired sample one-sided t-test.
D OPTIMIZING THE WEIGHTS FOR THE WORST-GROUP LOSS
Here we present the algorithm for GDRO with optimized weights. It should be read in conjunction with Section 5.3
Algorithm 2: Estimation of the optimal weights Pt, for GDRO
Input: data {y 9s, x, h training set size n' < maximum steps T,, learning rates 7pp- 7g and momentum 7.
Initialize Pa- qo uniformly and normalized 1.
Split the data into a training set of size ní' and validation set.
for { 1,..., do
Estimate 0, (P:-1) as in Equation I
Compute hyper-gradient S. = - P Cral theorem.
9--1) with implicit function Update weights Pt.g = Pt-1,g exp(1pue with u1g = 724-1. + (1 - 7) for all E g.
Normalize Pt.g Eveg Pi.a'
Pr.g for all EG
Update loss weights 9t.9 = 91-1,9 exp (prr))) for all g E G, where
C = Dn, L(y," fo (ax))) and with ngs vf and x. as Section 5.2
for all g G. Normalize Qt.g Eyee end
Return ,; = arg min (po.P.....pr) maxgeg C. (p.)))
25

---

Published as a conference paper at ICLR 2025
E EXPERIMENTAL DETAILS
E. DATASETS
An overview of the datasets, their group definitions and sizes, well as the train/validation split used in the experiments, provided in Table
Target variable
Counts (training set)
Percentage
Dataset size
Water
Land
Train
Waterbirds
Water bird
1,057
56
22.04%
1.17%
4,795
1,199
Land bird
184
3,498
3.84%
72.95%
Female
Male
Train
CelebA
Blonde
22,880
1,387
14.06%
0.85%
12,,770
19,867
Not blonde
71,629
66,874
44.01 %
41.08%
No negation
Negation
Train 50,000
MultiNLI
Entailment Neutral
24,362
638
48.72%
1.28%
20,000
Contradiction
20,930
4,070
41.86%
8.14%
Table 4: Overview of all datasets and their respective group sizes.
this dataset from Sagawa et al. (2020a) a combination of the Places dataset (Zhou et al. 2016) and the CUB dataset (Welinder et al.) 2010). A 'water background' is set by selecting an image from the lake and ocean categories in the places dataset, and the 'land background' is set based on the broadleaf and bamboo forest categories. A waterbird/landbird is then pasted in front of the background. The goal to predict whether or not picture contains a waterbird or landbird.
Waterbirds:
As explained in Section 6 we deviate from (Sagawa et al. (2020a) in that our validation set is not balanced for all groups. This setup of Sagawa et al. (2020a) gives an advantage methods that only fit the validation set, such DFR.
Waterbird water background
Waterbird land background
Landbird water background
Landbird land background
Figure 6: Examples of the different images in the Waterbirds dataset.
CelebA: this dataset contains images of celebrity faces (Liu al. 2015). The goal is to predict whether or not a picture contains a celebrity with blonde or non-blonde hair. This is spuriously correlated with the sex of the celebrity.
L
M
Blende
Non-blorde
Non-blande Maio
Figure 7: Examples of the different images in the CelebA dataset.
26

---

Published as a conference paper at ICLR 2025
the MultiNLI dataset (Williams et al. 2018 contains pairs of hypotheses and premises. The sentences are pasted together with a [SEP] token in between. The original dataset contains three label types: negation, entailment, and neutral. We follow procedure that is similar to Kumar et al. (202)) and Holstege et al. 2024) and create smaller binary version of the dataset (50.000 samples training, 20.000 in validation, 30.000 in test), with only negations (y 0) entailment neutral y 1). This version of the dataset is class balanced, similar to the original MultiNLI dataset. It has been reported that the negation label is spuriously correlated with negation words like nobody, never and nothing (Gururangan et al. 2018). A binary label is created to denote the presence of these words in the hypothesis of given hypothesis sroothesis-rremise pair.
MultiNLI:
E.2 TRAINING OF DNNS
For the Waterbirds and CelebA dataset, we use the Resnet50 architecture implemented in the torchvis si on package: forchvisionn sion. mwdels. resnet 50 (pret ra inneennrrre
Models:
More details on the model can be found in the original paper from He et al. (2016). For Waterbirds, this means using a learning rate of 10-3, 3, a weight decay of 10 a batch size of 32, and for 100 epochs without early stopping. For CelebA, this means using a learning rate of 10-3, a weight
decay of 10-1, a batch size of 128, and for 50 epochs without early stopping. We use stochastic gradient descent (SGD) with a momentum parameter of 0.9. For simplicity, we perform no data augmentation for both datasets.
For the MultiNLI dataset, we use the base BERT model implemented in the transformers pack- age (Wolf al. 2019): BertModel. wttp://ooooeeeerraaeeeeeeeeeeeee pretrained "bert -base-uncased".. The model was pre-trained on BookCorpus, a dataset consisting of 1 11,038 unpublished books, as well as the English Wikipedia (excluding lists, tables and headers). More details on the model can be found in the original paper: Devlin et al. (2019a). For finetuning the BERT model on MultiNLI, we the AdamW optimizer (Loshchilov & Hutter (017) with the standard settings in Pytorch. When finetuning, we use the hyperparameters of lzmailov et al. (2022), training for 10 epochs with a batch size of 16, learning rate of 10-5, and a weight decay 10 and linear learning rate decay.
E.3 IMPLEMENTATION DETAILS OF EXISTING METHODS
For all methods, unless otherwise mentioned, we logistic regression with the sklearn. Logist wwwogiisiiegggeesisioooo cRegression class with the Liblinear solver to optimize the logistic regression, with tolerance of 10-4.
In each case, we demean the training data and scale it such that each column has a variance of 1. We the estimated mean and variances of the training data apply the same transformations the validation and test data.
For each run, the hyperparameters are selected based on the worst-group accuracy on the respective validation set. For the LI penalty, we select it from the following values: 0.1, 3.3, 10.0, 33.33, 100, 300, 500. In the case that multiple L1 penalties lead the same worst- group accuracy, we select the highest one.
GW-ERM and SUBG: we only optimize the LI penalty.
GDRO: We follow the implementation of Sagawa et al. (2020a). We optimize two hyperparameters: the L1 penalty, and the hyperparameter used in adapted worst-group loss from Saaawa et al. (2200)). with the possible values of C' being 0, 1, 2, 3. Instead of the Libl inear solver, we use gradient descent with a learning rate of 10- After observing that the standard options for the LI penalty often led to poor performance, we add the values 0.01, 0.05, 0.2 to the possible L1 penalties.
DFR: We follow the implementation of Kirichenko et al. (0023). We use the validation set of each dataset for fitting the model. For hyperparameter selection, split this validation set in half, and use one half for selecting the Ll penalty. After this, the entire validation set used to fit the model. Parameters estimated by averaging over 10 logistic regression models.
JTT: similar to Liu et al. T) we train a DNN, of which the errors are used to identify the groups For Waterbirds and CelebA, we train Resnet50 model using the hyperparameters chosen by Liu et al. (2021). In the case of Waterbirds, this means a learning rate of 10 and a weight decay of I. In the case of CelebA this means learning rate of 10-5 and weight decay of 0.1. In the case of
27

---

Published as a conference paper at ICLR 2025
MultiNLI, train BERT model with the same hyperparameters that used for finetuning. In all cases, we use early stopping.
After getting the emrors from the respective model, define the groups accordingly and retrain the last layer on the embeddings that are used for all other methods. This is in order to ensure a fair comparison between methods. We optimize over both the L1 penalty and the upweighting factor ASTT- For the latter hyperparameter, we select from the values 1, 5, 10, 25.
E.4 IMPLEMENTATION DETAILS OF OPTIMIZING THE WEIGHTS
Following the notation Algorithm I we use learning rate 7 0.1 for each dataset and method, except in two cases: optimizing the weights for DFR on Waterbirds and MultiNLI, where we use a learning rate of 0.01 and 0.2 respectively. This was done after observing that the optimization procedure did not converge to lower validation loss for the original choice of the learning rate. In the case of Waterbirds, we run Algorithm 1 for 200 steps, and for T 100 in the case of CelebA and MultiNLI.
For each dataset and method, we use a momentum parameter y 0.5, after observing this helped with escaping local minima.
In the case of optimizing the worst-group loss, we use My = 0.1, following Sagawa et al. (2020a), and use Algorithm while using GW-ERM our method for determining the parameters.
In the case of optimizing the weights SUBG and DFR, we set 1n = 1 for the group g with the smallest number of samples. We then optimize the other fractions of groups
28


Reviews:
[{"review_id": "ob4HDQ1f3s", "reviewer": "Reviewer_Ng11", "summary": "The paper addresses the challenge of distribution shifts between training and test data, specifically sub-population shifts. The authors propose a bi-level optimization framework to optimize importance weights while accounting for the finite sample size, thereby improving generalization. The method is validated on benchmark datasets from vision and NLP domains, demonstrating superior performance over traditional importance weighting methods.", "questions": "see weakness section", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The bi-level optimization of importance weights is novel, addressing the bias-variance trade-off more effectively than heuristic-based weighting.\n2. The analysis using a linear regression model provides clear insights into the bias-variance tradeoff, making the proposed method more convincing.\n3. The experiment section is comprehensive, covering multiple datasets.\n4. The proposed method is useful in practice.", "weaknesses": "1. The computational overhead of the proposed methods should be discussed. Since we only need to train on the last layer, I think the computational overhead is not large?\n2. The proposed method need to separate the training dataset into val dataset and remaining train dataset. What's the tradeoff here? What's the best way to split? If the training dataset is limited, will the proposed method work?", "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": "w7PlN8vCub", "reviewer": "Reviewer_xwkt", "summary": "In this work, the authors develop a bi-level optimization algorithm for estimating importance weights (group weights) on top of existing deep learning algorithms, such as Group Weighted ERM and GDRO, to enhance generalization under sub-population shifts. They utilize a standard linear model to illustrate the necessity of this optimization in the presence of limited data and provide mathematical proofs to support their claims. Empirical demonstrations with deep networks using existing training strategies are presented, showing performance improvements on standard datasets.", "questions": "* Can the authors clarify the overall outline of the training procedure? Did they follow the last-layer retraining protocol as described in Izamailov et al. (2022)? If so, additional details about this protocol would be beneficial.\n\n* For the Group Weighted ERM procedure, it is mentioned that the likelihood ratio between the test and training distributions is required to initiate the proposed algorithm. However, in most cases, access to the test distribution is not available. How is this ratio computed for real-world datasets, such as Waterbirds? (I may have missed a key concept here.)\n\n* Can the authors provide qualitative insights on how they believe the performance of their algorithms may vary with an increase in model scale, particularly for architectures like ViTs and SWIN?\n\n* Can the authors include an example with large-scale datasets, such as ImageNet-1K, or a medical imaging benchmark (e.g., ISIC or MedMNIST) to enhance the impact of the proposed algorithms?\n\n## Post-Rebuttal\nI have increased my score and recommend acceptance", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well-written and systematically defines all terms and equations with attention to detail. It is mathematically well-supported and contains relevant experiments on standard benchmarks.", "weaknesses": "Refer to Questions", "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": "CjcJAKr4Cd", "reviewer": "Reviewer_DCrk", "summary": "This work studies importance weighting of sub-population shifts. Conventional naive importance reweighting is $p_{te}(x,y)/ p_{tr}(x,y)$. This work argues that this conventional reweighting is sub-optimal under a limited training data size. Thus this work proposes to uses a learned importance reweighting coefficient. Empirical results show improved generalization performance across benchmarks like image and NLP datasets.", "questions": "- How important are the $p_0$ and $r$ in Algorithm 1 line 278 and line 282? Currently these parameters depend on test dataset directly. And thus weaken the usefulness of the algorithm 1. Is it possible to set $p_0$ and $r$ by assuming a test dataset $p_{te}(g_i) = p_{te}(g_j) , \\forall i,j $? If it is possible, how does it affect the experiments performances?  That is my main concern.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The arguement that \"the conventional importance weighting $p_{te}(x,y)/ p_{tr}(x,y)$ could be sub-optimal\" is interesting. \n\n- This paper estimates model parameters and importance weights iteratively on train and iid-validation datasets. That is interesting, because it reduces the overfitting of training dataset. Different from conventional cross-validation and hyper-parameter search, this approach find the importance weights by optimization.", "weaknesses": "- Experiments shows the good performance of the proposed method. However, Algorithm 1 (line 278 $p_0$ and line 282 $r$) requires the access of test dataset as parameter initialization. This could be problematic. I understand that some explicit reweighting methods, such as JTT, also contains some heuristic & data-dependent importance weights searching space. These heuristic hyperparameter searching space could also depends on test dataset.  However, group distributional robust optimization (GDRO) links to importance weighting in an implicit way [1], and doesn't contain any hints from test dataset. Thus, the comparison between the proposed method and GDRO is prolematic and potentionally unfair. \n\n- minor typos: line 112 [L(y, fθ (x)] --> [L(y, fθ (x))]\n\n \n[1] Zhang, J., Lopez-Paz, D., & Bottou, L. (2022, June). Rich feature construction for the optimization-generalization dilemma. In International Conference on Machine Learning (pp. 26397-26411). PMLR.", "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": "K533DZfx0B", "reviewer": "Reviewer_PJmP", "summary": "The authors present a bi-level optimization strategy to estimate optimal importance weights in the presence of distribution shifts. The authors use a linear regression model to demonstrate how the naïve but often used ratio of likelihoods is not necessarily an optimal choice for importance weights (figure 1.) The authors proceed to introduce their bi-level optimization strategy that estimates optimal weights given a validation set sampled from the training set.\n\nThe authors emphasize that the strategy can be used in conjunction with existing methods and write how they implement their strategy specific to each of five existing methods for fine-tuning models to manage distribution shift. The authors empirically test their claim using a benchmark of three datasets for each of the 5 modified methods and show how in almost all cases the performance on a held-out test dataset is improved using their bi-level optimization strategy.", "questions": "I’ll focus specifically on the empirical experiments. I deem them overall sound, but I have a few questions to be sure. The authors describe fine-tuning a network 5 times to build a sample and perform their t-test. I wanted to confirm how that network was initialized (Appendix suggests you had the same base network that you then fine-tuned 5 different times?). It was then from this set of 5 fine-tuned networks that you then retrained the last layer using each algorithm’s original importance weight scheme and then compared to the paired performance when replacing with your optimal weights? What was controlled in this experiment in terms of random seeds? Were seeds controlling batch ordering varied between runs? Were seeds being varied between the pairs of optimal vs original weights?\n\nThe appendix describes some hyperparameters that had to be found for the optimal weight algorithm (E.4). Can the authors provide some additional detail about how these hyper parameters are found?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "To me the paper is a clear acceptance: it is well written, organized, to the point. You present an intriguing argument that the naïve choice of importance weights used in a variety of prior algorithms can be improved upon, present a practical answer of how to define and find these optimal weights, then show that this definition brings benefit over a controlled baseline. The topic clearly has potential for high impact as distribution shifts are pretty ubiquitous.", "weaknesses": "I think the paper could be improved with some additional discussion about the practical implementation of these strategies. For example, how much additional time is spent on your specific improvements compared to the naïve/heruristic weights? Roughly how does this time scale with the validation dataset fraction (1 – n’/n)? What additional hyperparameter tuning is done to find optimal weights (Appendix E.4)? Would a user need to repeat this hyperparameter tuning for their own model?\n\nThe impact of the paper could be heightened if the authors elected to demonstrate their technique improving upon a real-world challenge; for example, the authors identify in their introduction that importance weighting technique could be used to improve outcomes for algorithmic fairness.", "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)