Optimizing importance weighting in the presence of sub-population shifts

Authors: Floris Holstege, Bram Wouters, Noud Van Giersbergen, Cees Diks
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: Importance weighting, Distribution shift, Group Robustness, Bi-level optimization

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.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730117725848

thread_type:
official_review

<thread_object>
creator: reviewer_pjmp

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

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

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?

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.

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?

The 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730669224798

thread_type:
official_review

<thread_object>
creator: reviewer_dcrk

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

strengths:
- The arguement that "the conventional importance weighting $p_{te}(x,y)/ p_{tr}(x,y)$ could be sub-optimal" is interesting. 

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

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

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

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730683835431

thread_type:
official_review

<thread_object>
creator: reviewer_xwkt

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

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

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.

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

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

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

## Post-Rebuttal
I have increased my score and recommend acceptance
</content>
</thread_object>

</thread>
<thread>
creation date:
1730688121462

thread_type:
official_review

<thread_object>
creator: reviewer_ng11

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

strengths:
1. The bi-level optimization of importance weights is novel, addressing the bias-variance trade-off more effectively than heuristic-based weighting.
2. The analysis using a linear regression model provides clear insights into the bias-variance tradeoff, making the proposed method more convincing.
3. The experiment section is comprehensive, covering multiple datasets.
4. 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?
2. 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?

questions:
see weakness section
</content>
</thread_object>

</thread>
<thread>
creation date:
1731670989448

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General reply to all reviewers

comment:
We are deeply grateful to all reviewers for their careful consideration of the paper. We also thank them for acknowledging the relevance, novelty, and potential usefulness of our procedure for optimizing importance weighting. 
In order to start the discussion phase, we have posted comments to each individual reviewer with detailed replies to specific questions and remarks. 

Based on the initial reviews, we intend to include the results of (at least) two additional experiments in an updated version of the manuscript. 
1. An experiment that investigates the role of the ratio between training/validation split, and how this interacts with the size of the dataset. This in order to address the concern raised by reviewer Ng11 about the impact of having a limited training/validation set. 
2. To address a question from Reviewer PJmP, we have conducted an additional experiment where we vary learning rates ($\eta$) and momentum values ($\gamma$) to assess whether or not our results are dependent on these hyperparameters of the optimization procedure. We present the results in our reply to reviewer PJmP.

We also thank the reviewers PJmp and Xwkt for their suggestion to add a dataset from the medical domain. We are currently searching for a dataset in the medical domain which is relevant for the domain of sub-population shift (e.g. CheXpert [1]), and if time permits will add results for this. 

We look forward to further discussing the paper over the next two weeks. 

[1] 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732559266593

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updated manuscript, including new experiments.

comment:
As a follow-up of our initial replies to all reviewers on the 15th of November, we have updated the manuscript accordingly, and uploaded a revised version. The key updates include:
* An additional experiment investigating the effect of the training/validation split ratio (Appendix D).
* An experiment showing how one can select the hyperparameters of the optimization procedure (Appendix F.4).
* Clarifications in writing, including (but not limited to) the computational complexity of the method,  the experimental setup, and a discussion of the required knowledge of the test distribution. 

We very much appreciate the detailed and constructive feedback from all reviewers that has helped improve our work. We look forward to further discussions and are happy to address any additional questions or concerns the reviewers may have.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733299936583

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General, concluding remarks by the authors at the end of the discussion phase

comment:
To conclude the discussion phase, we would like to reiterate that we are deeply thankful to all reviewers for their attentive consideration of the paper. We also would like to thank the reviewers for acknowledging the relevance and novelty of our work, as well as highlighting it _'is mathematically well-supported’_, that the _‘experiment section is comprehensive’_ and that _‘the topic clearly has potential for high impact’_.

During the discussion phase, we have tried to address the concerns and questions of the reviewers. We are glad to see that several reviewers have expressed their satisfaction with how we handled their initial reviews. Based on the feedback of the reviewers, we have also updated the manuscript. Here is an overview of the key updates:

·        We have added an experiment to illustrate the effect of the ratio between the training and validation set (addressing the concern of reviewer Ng11).

·        We have added an experiment to illustrate how one can select the additional hyperparameters of our method (addressing the concern of reviewer PJmP).

·        We have clarified our experimental procedure in Section 6 (addressing the concern of reviewer xwkt).

·        We have added discussions about the computational resources used for the experiments and the computational complexity of the method (addressing the concern of reviewer Ng11).

·        We have added a paragraph in the section Conclusion & Discussion that discusses the required knowledge of the test distribution (addressing concerns of reviewers xwkt and DCrk). 

 

In response to the suggestions of reviewers PJmP and xwkt, we considered the possibility of applying our procedure to a medical dataset. However, due to time constraints, we were unable to implement this addition. Furthermore, the manuscript already includes applications to three benchmark datasets, encompassing both vision and natural language processing domains, which we believe sufficiently demonstrate the versatility of our approach for now.

Finally, we would also like to thank the reviewers for their engagement during the discussion period. We believe that as a consequence of their feedback and suggestions, the manuscript has greatly improved.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper develops a novel way to assign importance weights to data to deal with sub-population shift.  The authors argue that existing methods to deal with distribution shift between training and test data underperform because they don't take into account the sample size of the training data.  The authors propose a bi-level optimization scheme where they jointly optimize the importance weights and model parameters.  

The reviewers unanimously recommended accept for the paper (8, 8, 6, 6) albeit with relatively low confidence (3, 2, 2, 2).  They found the paper well written, sound, empirically comprehensive and relevant to the community.  The reviewers noted as weaknesses mostly requests for additional details.  One reviewer noted that the method requires access to the test dataset / distribution to estimate importance weights, which is potentially an issue since this is not always available during training.  

Given the unanimous recommendation of the reviewers, and not much cited as weaknesses, there doesn't seem to be any reason to override the reviewers.  Thus the recommendation is to accept.