Optimizing importance weighting in the presence of sub-population shifts

(continuation of previous discussion)

To show that this approach for selecting the learning rate and momentum can work, we have conducted an additional experiment for the Waterbirds dataset, where we optimize the weights for group-weighted empirical risk minimization (GW-ERM). We keep the setup (i.e., averaging over 5 seeds) and all other hyper-parameters the same (e.g. the $l_1$ penalty strength and $T=200$) and only vary the values of the learning rate and momentum. We show the results in the table below. The bottom row of the table indicates the performance if one selects per seed the best combination of the learning rate and momentum, based on the weighted validation loss.

The table below indicates that the results are relatively robust to the choice of the hyper-parameters. For instance, even the worst combination of learning rate and momentum would result in an improvement (on average) compared to the standard weights. Moreover, this worst combination could be avoided by the user by selecting the learning rate and momentum based on the weighted validation loss. If one were to follow this strategy, this would even lead to an improvement compared to our choice of hyperparameters and the results presented in Figure 2 of the manuscript. We will report this table in Appendix E.4 of an updated version of the manuscript.

With regards to the number of steps $T$, we have been pragmatic. We chose a $T$ large enough such that we observed convergence of the validation loss, but not too large for our computational budget and time available.

The 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.

We welcome this suggestion by the reviewer, and hope to address it by adding results of using our procedure on a dataset in the medical domain, related to algorithmic fairness. In an updated version of the manuscript, we will also add a sentence to the introduction to highlight how importance weighting can play a role in achieving algorithmic fairness.

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...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

For our vision experiments we used Resnet50, with parameters initialized from a model that is pre-trained on Imagenet1K (downloaded from the torchvision package). In the case of the BERT model (used for multiNLI), the parameters are initialized from a model that was pre-trained on BookCorpus, a dataset consisting of 11,038 unpublished books, as well as the English Wikipedia (excluding lists, tables and headers). So, for each of the seeds, the initial parameters of the deep neural network (DNN) are the same.

The seed was set before the finetuning of the DNN started. This means that per seed, the ordering of batches and which samples are contained in them changes. After the DNN was finetuned, we stored the embeddings of the last layer for each seed. Then, a logistic regression is trained on these embeddings - both with standard, and optimized weights. Thus, the seeds are not varied between pairs of standard and optimal weights. In this sense, we performed a paired t-test.

We will clarify both how the DNNs were initialized and which parts of the training run were controlled with random seeds in Appendix E.2 of the updated manuscript.

We are deeply grateful to reviewer PJmP for the detailed feedback and questions. Below, we address these point-by-point.

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)?

We thank the reviewer for raising this question. We have discussed the computational overhead of our optimization procedure in response to reviewer Ng11. Generally speaking, provided we restrict ourselves to last-layer retraining, the computational overhead of our optimization procedure is manageable. To address the specific question of reviewer PJmP, increasing $n'$ relative to $n_{\mathrm{val}}$ (and thus $n$) will increase the computational overhead. This is because $n_{\mathrm{val}}$ plays no role in re-fitting the logistic regression, which is the most time-consuming part of Algorithm 1. Moreover, $n_{\mathrm{val}}$ plays no role in calculating and inverting the Hessian.

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?

Can the authors provide some additional detail about how these hyper parameters are found?

In order to select the the learning rate $\eta$ and the momentum parameter $\gamma$, we observed how the validation loss progressed over the course of the optimization for given iterations $t=1,2,\ldots,T.$ We observed that for some choices of learning rates $\eta$, the validation loss would barely decrease over the course of the optimization. We also observed that if we used no momentum, the optimization sometimes would halt around a particular choice of weights. In all but two cases (optimizing the weights for DFR on Waterbirds and MultiNLI) we observed that for the relatively standard choices of $\eta=0.1$ and $\gamma=0.5$, the validation loss decreased throughout the optimization training procedure, until it plateaued towards the end of the optimization. For the two exceptions, we changed the learning rate to a value (0.01 and 0.2, respectively) for which we observed a similar phenomenon of convergence.

We acknowledge that we did not pursue the most rigorous approach for selecting the additional hyperparameters, although we did follow the standard approach of using the validation set to assess the choices of the hyperparameters. More rigorous would be to conduct a sweep over different values of $\eta$ and $\gamma,$ and observe for which combination of these values the validation loss is lowest. The risk of this approach is that one might 'overfit' on the validation set. This issue can be remedied by existing approaches for preventing overfitting, such as early stopping.

(see next comment for a continuation of this discussion)

The proposed method needs to separate the training dataset into val dataset and remaining train dataset. What's the tradeoff here?

We agree with the reviewer that there is a tradeoff here. The validation loss is used to optimize the weights. A small validation set size $n_{\mathrm{val}}$ might lead to overfitting of the weights. It should be mentioned here that in the case of a sub-population shift with 4 subgroups (which is the case for all experiments in the manuscript) the weight space is effectively 3-dimensional. This makes the risk of overfitting limited. However, there are contexts with many more subgroups (e.g. concept bottleneck models, as mentioned in Section 7), where the risk of overfitting can become an issue. A small $n_{\mathrm{val}}$ can also lead to very small minority subgroups in the validation set. The validation loss might then not be a good estimate of the expected loss, troubling the finding of the optimal weights.

On the other hand, a large validation set size $n_{\mathrm{val}}$, relative to the training set size $n'$, also has disadvantages. The training set is used to optimize the model parameters. A too small $n'$ might lead to overfitting. Since the dimension of the parameter space is typically much larger than the dimension of the weight space, we believe the risk of overfitting the model parameters is more severe. In addition, after optimizing the weights the model gets fitted on the entire dataset of size $n$ (as commonly done after selecting hyperparameters and weights), but the weights are optimized for a training set of size $n'.$ The value of the optimal weights is a function of the training set size (as illustrated in Figure 1). If the difference between $n'$ and $n$ is too big, the weights might not be optimal anymore for the final training of the model on a training set of size $n.$

For our experiments in the paper, we maintained the training/validation split that is standard for each benchmark. Similar to previous work [1] on last-layer retraining, we fitted our models on the training dataset of size $n'$, rather than both the training and validation data. The ratio between training and validation is relatively standard for these datasets (approximately 80/20, 89/11 and 70/30 for Waterbirds, CelebA, and multiNLI, respectively).

We are working on an additional experiment that investigates the role of the ratio between the training and validation set, and how this is impacted by the size of the total dataset. We aim to reproduce Figure 3 of the paper for different ratio's between the training and validation set.

[1] Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023.

We are grateful to reviewer Ng11 for examining the paper and their feedback. Below, we address the feedback point-by-point.

The computational overhead of the proposed methods should be discussed.

We thank the reviewer for pointing out the need to discuss the computational overhead of the optimization procedure for the weights. We will first detail the computational resources used for the paper, as well as the approximate time it takes to optimize the weights. We then analyse the algorithmic complexity of the optimization procedure.

For the results in the manuscript, three types of computations were needed:

1. Each deep neural network (DNN) is finetuned and, after finetuning, the last-layer embeddings are computed and stored.
2. Optimization of the weights for the respective importance weighting procedure.
3. Fitting a logistic regression on the last-layer embeddings, using either standard or optimized weights. Steps 1 and 3 are also needed when using a standard set of weights for importance weighting. For step 1 we used a single GPU. The memory required for storing the embeddings grows as a function of the total datapoints $n$ and the number of features $d$. For step 3 a single CPU of a Macbook Pro laptop (M2) was used. In our experiments, fitting this logistic regression took less than a minute.

The novelty of our paper is in step 2, for which we also used a single CPU of a Macbook Pro laptop (M2). In our experiments, the additional computational resources required for optimizing the weights were relatively manageable. To give an indication, optimizing the weights for group-weighted ERM for the Waterbirds dataset ($n'=4795$) takes approximately 2 minutes for $T=200$ timesteps. For a bigger dataset such as CelebA, where $n'=162,770$, this took approximately 60 minutes for $T=100$ timesteps. For both datasets, the dimension of the embedding space was $d=2048$. The main contributor towards this extended computational time is the re-fitting of the logistic regression at each step (approximately 0.3 seconds for Waterbirds and 30 seconds for CelebA).

Step 2 is described in Algorithm 1 in the manuscript. Its computational overhead scales linearly with the number of iterations $T$. Per iteration, we re-fit the logistic regression using the weights of the previous step, and we determine the hyper-gradient. The re-fitting of the logistic regression is standard and its algorithmic complexity depends on the optimization procedure, but typically will be increasing in $n'$ and $d.$

In the paper we note that the hyper-gradient factorizes into two gradients (see line 269 of the manuscript). Calculating the first factor has an algorithmic complexity of $O(n_{\mathrm{val}}(d+1) )$, since $\hat{\boldsymbol{\theta}}_{n'}$ consists of $d$ coefficients and an intercept. As detailed in Appendix A, in order to calculate the second factor, we need to calculate the Hessian and invert it, which has an algorithmic complexity of $O(n'(d+1)^2 + (d+1)^3).$ As mentioned in Section 7 of the manuscript, this makes it difficult to extend the current form of the optimization procedure beyond last-layer retraining. It is worth pointing out that when using a logistic regression, we have analytical solutions to both parts of the hyper-gradient, allowing for fast computation (see Appendix A).

We will add a paragraph to the updated manuscript that will detail the computational resources, as well as provide details on the algorithmic complexity of the method.

Thanks for the reply. I will maintain my score.

However, group distributional robust optimization (GDRO) links to importance weighting in an implicit way [2], and doesn't contain any hints from test dataset. Thus, the comparison between the proposed method and GDRO is prolematic and potentionally unfair.

This point requires some clarification from our side. The application of our weight optimization procedure to GDRO is described in Section 5.3 and Algorithm 2 in Appendix D of the manuscript. In Algorithm 2 no knowledge of the test distribution is being used. Contrary to Algorithm 1, the likelihood ratio $r$ is not used in Algorithm 2. Instead, and in accordance with the standard GDRO approach, we optimize the weights $\boldsymbol{p}$ for worst-group loss (Equation 13 of the manuscript). For the initial value $p_0$, we simply use the same probability for each group, i.e., $\forall g \in \mathcal{G}, p_{g,0} = \frac{1}{|\mathcal{G}|}.$ This choice can be made regardless of the marginal probabilities of a (hypothetical) test distribution. We thus believe the comparison with GDRO is fair. We agree with the reviewer that in the manuscript this point is not explained with enough clarity and will rectify this in an updated version of the manuscript.

minor typos: line 112 [L(y, fθ (x)] --> [L(y, fθ (x))]

We thank reviewer DCrk for pointing out the typo, and will rectify this in the updated manuscript.

[2] 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.

We are thankful to reviewer DCrk for carefully reading the manuscript and their feedback. Below, we address the feedback point-by-point.

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.

How important are the p0 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 p0 and r by assuming a test dataset pte(gi)=pte(gj),∀i,j? If it is possible, how does it affect the experiments performances? That is my main concern

In the above two paragraphs, the reviewer is raising the issue that the user might not have access to data from the test distribution of interest, and thus might not know the probability at which certain groups occur ($p_{\mathrm{te}}(g)$). We thank the reviewer for pointing this out and we agree with the reviewer that this could be problematic. However, we do believe there are many cases in which it is possible to apply Algorithm 1.

For the experiments in the manuscript we do exactly as the reviewer suggests and assume that the test distribution of interest is one where each group occurs to the same extent - i.e., $\forall g \in\mathcal{G}, p_{\mathrm{te}}(g) = \frac{1}{G}$, where $G = |\mathcal{G}|$. This is in line with previous work [1] that also uses these benchmarks, where performance of the methods is (partially) assessed by an equal-weight average of the accuracy per group. The motivation behind this is that if a model relies on a spurious correlation, e.g. using the water background to classify a waterbird type, it should perform badly on test dataset where is spurious correlation is absent, i.e. when all groups have equal marginal probabilities.

In more practical contexts, the probability at which certain groups occur could be informed by knowledge of the user about the prediction problem. For example, the user knows that male patients are more prevalent than female patients. If a small subset of labels from the test distribution is available, a user could estimate the likelihood ratio (used for $p_0$ and $r$ in Algorithm 1) of each group occuring.

And even if data from the test distribution is not available, a user can still specify the test distribution of interest. The test distribution then reflects the preferences of the user. For example, if a user cares to an equal extent about prediction errors in two groups (e.g., male and female patients), they can specify this via a test distribution in which the groups have an equal marginal probability.

Another alternative would be to optimize the weights for worst-group loss, instead of the weighted average accuracy. This is the GDRO approach with optimized weights, as outlined in Section 5.3 of the paper and Algorithm 2 in Appendix D. In this case, no knowledge of the test distribution is required.

Finally, we believe that the fact that the starting values of $p_0$ in Algorithms 1 and 2 are determined by the test distribution, as pointed out by the reviewer, is a logical choice. As outlined in Sections 3 and 4, this choice is optimal in the case of infinite sample size ($n \rightarrow \infty$). The bi-level optimization procedures (Algorithms 1 and 2) are thus initialized at the heuristic choice of the weights and converge to the optimal weights by deviating from this initial choice. We will clarify the motivation for this choice in an updated version of the manuscript.

[1] Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive worst-group-accuracy. In Bernhard Scholkopf, 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.

Thank authors for the response. The point that GDRO experiments use a uniform p initialization is convincing, even though estimating p from test set is problematic and not solved.

It would be very interesting to initializate p uniformly on all tasks. but...

Overall, I would like to raise my score, because GDRO experiment partically shows the effect of uniform initialization.

We thank reviewer Xwkt for carefully reading the manuscript and their feedback. Below, we address the feedback point-by-point.

Can the authors clarify the overall outline of the training procedure? Did they follow the last-layer retraining protocol as described in Izmailov et al. (2022)?

We indeed perform last-layer retraining in a manner similar to Izmailov et al. [1]. The steps in our training procedure are:

1. We start with a deep neural network (DNN) that is pre-trained on a large dataset (e.g. Imagenet1k).
2. We finetune the DNN on the respective datasets of our experiments.
3. We store the embeddings of the last layer of the DNN.
4. We then train a logistic regression on these embeddings, which is equivalent to retraining the last layer.

With regards to the logistic regression of step 4, we apply all the respective importance weighting techniques (GW-ERM, SUBG, DFR, GDRO, JTT), using both standard and optimized weights. We will clarify this procedure by updating Appendix E.2 of the manuscript.

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?

We thank the author for raising the issue of the choice of the test distribution, and how this related to the benchmark datasets.

For all benchmark datasets, we assume that the test distribution of interest is one where each group has the same probability of occuring - i.e., $\forall g \in\mathcal{G}, p_{\mathrm{te}}(g) = \frac{1}{G}$, where $G = |\mathcal{G}|$. This is in line with previous work [1] that also uses these benchmarks, where performance of the methods is (partially) assessed by an equal-weight average of the accuracy per group. The motivation behind this is that if a model relies on a spurious correlation, e.g. using the water background to classify a waterbird type, it should perform badly on test dataset where is spurious correlation is absent, i.e. when all groups have equal marginal probabilities.

In response to reviewer DCrk we have outlined some general considerations for how a user might select the test distribution of interest.

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?

Because we resort to last-layer retraining, only the phase of finetuning the DNN is affected by the scale of the model. If the last-layer embeddings of the larger model also have a larger dimension $d$, this will increase the computational overhead of our method, as outlined in our reply to reviewer Ng11.

In principle, last-layer retraining can work in combination with (larger) architectures and vision transformers, and even tends to perform better in such cases (see [2]).

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?

We thank the reviewer for raising this suggestion. The datasets used in the manuscript are standard benchmarks for evaluating sub-population shifts (see, e.g., [3]). To the best of our knowledge, a large-scale dataset such as ImageNet-1K is not a standard benchmark in this literature, and it appears not straightforward to evaluate sub-population shifts for this dataset. We are currently searching for a dataset in the medical domain which is relevant for the domain of sub-population shift (e.g. CheXpert [4]), and if time permits will add results for this.

[1] Badr Youbi Idrissi, Martin Arjovsky, Mohammad Pezeshki, and David Lopez-Paz. Simple data balancing achieves competitive worst-group-accuracy. In Bernhard Scholkopf, 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.

[2] Pavel Izmailov, Polina Kirichenko, Nate Gruver, and Andrew G Wilson. On feature learning in the presence of spurious correlations. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 38516–38532. Curran Associates, Inc.,2022

[3] Yuzhe Yang, Haoran Zhang, Dina Katabi, and Marzyeh Ghassemi. Change is hard: a closer look at subpopulation shift. In Proceedings of the 40th International Conference on Machine Learning (ICML'23), Vol. 202. JMLR.org, Article 1652, 39584–39622, 2023

[4] Irvin, J., Rajpurkar, P., Ko, M., Yu, Y., Ciurea-Ilcus, S., Chute, C., Marklund, H., Haghgoo, B., Ball, R., Shpanskaya, K., et al. Chexpert: A large chest radiograph dataset with uncertainty labels and expert comparison. Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 590–597, 2019