Group-robust Sample Reweighting for Subpopulation Shifts via Influence Functions

Authors: Rui Qiao, Zhaoxuan Wu, Jingtan Wang, Pang Wei Koh, Bryan Kian Hsiang Low
Primary area: alignment, fairness, safety, privacy, and societal considerations
Keywords: distribution shift, subpopulation shift, spurious correlation, influence function, sample reweighting, data selection

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

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

Reviews and discussion:
<thread>
creation date:
1730278833838

thread_type:
official_review

<thread_object>
creator: reviewer_odpm

<content>
summary:
This paper proposes a method to mitigate spurious correlations and improve group robustness. The task is formulated as a bilevel optimization problem, which the authors solve via implicit differentiation. They employ Last-Layer Retraining (LLR) to ensure that the optimization problem meets the conditions required for implicit differentiation. Experiments on several datasets demonstrate the effectiveness of the proposed method.

strengths:
- This paper utilizes implicit differentiation to calculate the gradient with respect to the sample weights and applies last-layer retraining to ensure that the risk function satisfies the conditions of convexity and continuous differentiability.

- Overall, the paper is written and organized well.

- The authors not only conducted experiments on benchmark datasets but also evaluated the performance of the proposed method in the presence of label noise. The proposed method demonstrated effectiveness in identifying noisy data.

weaknesses:
- There appears to be an issue with the proof in Appendix A.2. By examining the first and third lines together, we derive $\nabla_{w_i} \theta_T \approx \nabla_{w_i} \left( \theta_T - \eta \nabla_\theta \hat{\mathcal{R}}(\mathcal{D}^{\text{tr}}; w, \theta_T) \right)$ which then leads to $0 \approx - \eta \nabla_\theta \hat{\mathcal{R}}(\mathcal{D}^{\text{tr}}; w, \theta_T)$, which seems unreasonable.
Furthermore, in the transition from the third to the fourth line, it is unclear why $\nabla_{w_i} \theta_T$ was eliminated. Was it because $\nabla_{w_i} \theta_T = 0$? If that is the case, the proof in Appendix A.2 could directly conclude that $\nabla_{w_i} \hat{\theta}^* \approx \nabla_{w_i} \theta_T \approx 0$. It might be worthwhile for the authors to review this part of the proof carefully.
I do not believe this issue impacts the main contributions of the paper, but the authors may need to consider revising the related explanations in the paper that are based on this proof.

- In Appendix A.3, it is assumed that when $\nabla_{\theta} \hat{\mathcal{R}}(\mathcal{D}^{tr}; w, \hat{\theta}^*) = 0$, it follows that $\frac{\mathrm{d} \nabla_{\theta} \hat{\mathcal{R}}(\mathcal{D}^{tr}; w, \hat{\theta}^*)}{\mathrm{d} w} = 0$. Could the authors elaborate on why the gradient with respect to $\theta$ being zero implies that its derivative with respect to $w$ is also zero in this context?

questions:
see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730288484571

thread_type:
official_review

<thread_object>
creator: reviewer_xdjg

<content>
summary:
This paper introduces a novel approach called Group-robust Sample Reweighting (GSR) aimed at mitigating the issue of spurious correlations, particularly under subpopulation shifts. The paper proposes a two-stage strategy that utilizes group-labeled data more efficiently by leveraging it to reweight group-unlabeled data. The key contribution of the GSR method is its use of last-layer retraining (LLR), which simplifies the bilevel optimization problem by focusing only on the last layer of the model. Through iterative reweighting of the group-unlabeled data based on influence functions and adaptive aggregation of sample gradients, GSR aims to improve the worst-group accuracy.

strengths:
1. The proposed method is built on a solid theoretical foundation, particularly leveraging bilevel optimization and influence functions, which enhances its robustness and credibility.

2. By focusing on the last layer, the method simplifies the optimization process, making it computationally efficient.

3. The paper provides extensive experimental results across a variety of benchmarks and other understanding experiments.

weaknesses:
1. In Section 2, it would help to be more explicit about how the proposed approach differs from previous methods. Right now, it feels like Section 3 is mostly about improving an already established objective (MAPLE). Adding the inverse Hessian in the optimization might make things more complicated than before. It’s not clear from the paper whether this extra complexity is worth the performance gains. 

2. It’s not obvious how much of the performance boost comes from last-layer retraining (LLR) versus the optimization changes over MAPLE. It would be useful to see a comparison where MAPLE is combined with LLR, so we can better understand where the gains are really coming from.

3. The claim that “over-reliance on spurious correlations mainly occurs in the last layer” feels a bit questionable. Prior research indicates that neural networks can develop biases throughout the entire network, not just in the final layer, e.g, [1]. It would be helpful to either provide more evidence or explain this assumption more clearly.

4. Table 1 is missing some crucial results, especially for MAPLE, which is an important baseline here. The paper mentions that it couldn’t get the results from the original paper, but why not try reproducing them? Was there a specific issue? It feels like an incomplete comparison without those numbers.

[1]What neural networks memorize and why: Discovering the long tail via influence estimation

questions:
Please see above
</content>
</thread_object>

</thread>
<thread>
creation date:
1730364314604

thread_type:
official_review

<thread_object>
creator: reviewer_kehk

<content>
summary:
This paper studies the challenge of improving machine learning model robustness to subpopulation shifts, where models may perform inconsistently across different subgroups within data distributions. Traditional methods typically rely on group-labeled data to tune models and minimize worst-case loss across groups, but these approaches are limited by the high cost of obtaining substantial, high-quality labels. The authors propose an alternative framework, Group-robust Sample Reweighting (GSR), which enhances group label efficiency by utilizing group-labeled data as a target to optimize weights for group-unlabeled data. The two-stage GSR approach first learns representations from group-unlabeled data and then fine-tunes the model’s last layer by iteratively retraining on reweighted data. GSR demonstrates superior performance on benchmark datasets.

strengths:
1. The paper highlights that previous methods estimate weight gradients through approximations, but they neither account for the curvature of the loss landscape with respect to the current model nor fully capture the training dynamics. The authors’ proposed method addresses these issues under the assumptions that the inner objective is unconstrained, strongly convex, and twice continuously differentiable.

2. The experiments in this paper are thorough, and the results are promising. Across multiple datasets, the proposed approach outperforms various baseline methods, achieving state-of-the-art performance.

weaknesses:
1. One limitation of this paper is that the proposed method relies on the inner objective being unconstrained, strongly convex, and twice continuously differentiable—assumptions whose strength is not fully clarified. The authors address these requirements by fixing the feature layers during training and only updating the linear classifier layer. However, further explanation is needed to clarify why training only the classifier layer meets these assumptions. Additionally, limiting training to the classifier layer may risk sacrificing overall model performance, which warrants further investigation.

2. The authors' second motivation is somewhat confusing. The statement that "regularization techniques must be applied such that training with weighted objectives leads to meaningfully different solutions than using ERM" may require further clarification. It would be helpful for the authors to elaborate on how this principle is specifically reflected in the design of their method.

3. The authors should conduct additional experiments to examine, without computational constraints, how much performance is lost when training the model with frozen feature layers compared to full-parameter training. This would help determine whether the performance reduction is within an acceptable range.

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730461937805

thread_type:
official_review

<thread_object>
creator: reviewer_wcxj

<content>
summary:
The paper considers the case when samples can be grouped (not all training samples are necessarily grouped, but all validation samples and target samples must be grouped), and we consider training the prediction model to minimize the maximum groupwise risk (loss) among all groups, namely the "worst-group risk". To train the model under the situation above, the paper considers re-weighting samples during updating the model. Different from an existing method MAPLE, the proposed method considers using the Hessian of the risk when re-weighting. Also, the proposed method introduces the "last-layer retraining" to mitigate spurious correlations and improve group robustness and also to reduce the computational cost.

strengths:
The combination of the sample reweighting and the last-layer retraining to tackle with subpopulation shifts and (consequent) spurious correlations. As far as the reviewer's understanding, (i) a baseline method MAPLE (Zhou et al.; 2022) updates sample weights $w$ for group robustness, but it uses not the Hessian $H_{\hat{\theta}^*,w}$ but only the last-step gradient. (ii) The use of the Hessian raises the computational cost, however, the proposed method focuses on the fact that the last-layer retraining is both costless even we employ the Hessian, and effective for group robustness.

weaknesses:
-   Technical novelty looks not well emphasized. Although the advantage above ("Strengths" section) seems to be effective, the reviewer could find it by carefully reading the paper. This fact is good to be emphasized in a early part of the paper, e.g., in Section 1.
-   The experiments are desired to be more extended. One is to increase the number of datasets (perhaps are there not so many appropriate datasets?), and the other is to qualitative discussions. (How the results are well written, but the discussion on the reasons of the results is desired more.)

questions:
-   The proposed method employs the last-layer retraining, but is it known what types of datasets are not sufficiently learned by the last-layer retraining? If so, comparing performances between methods for such datasets may be interesting.
-   Section 1: It may be better to present an example of spuriously correlated datasets due to groups, e.g., using the "Waterbirds" dataset.
-   Section 3.1: Assumption 3.1 states that, we assume that $\hat{R}({\cal D}^\mathrm{tr}; w,\theta)$ is strongly convex. However, with respect to what variable(s) should we assume it is strongly convex? All of $(w,\theta)$? Please specify.
-   Section 4.1: It states that "Secondly, calculating the Hessian inverse in Equation 6 is no longer overly expensive, because the only free model parameters are in the last layer instead of the entire deep neural network", but how the Hessian inverse can be updated in accordance with the update of $\psi$? It looks a key point of the method, so it is good to be explicitly presented.
-   Section 5.1: The method GSR-HF (Hessian-free) is examined as an ablation study, but how about the ablation study of using the Hessian but not LLR? To compare with MAPLE, this ablation study is also desired. (Perhaps is it too costly?)
</content>
</thread_object>

</thread>
<thread>
creation date:
1732174179326

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank all reviewers for their valuable feedbacks. We have revised the main paper accordingly, with changes highlighted in **blue**. 

As suggested by Reviewer wcXj and kehk, we first include an ablation study on **full-parameter retraining** vs. **last-layer retraining** with GSR at a **manageable** scale, then demonstrate the challenges of performing GSR with full-parameter retraining at a **larger** scale such as for the subpopulation shift benchmarks. 

### **Comparing full-parameter retraining (FR) vs. last-layer retraining (LLR) with GSR**

Due to the computational constraints, we design a small binary classification task *ColoredMNIST-S*, customized from the ColoredMNIST dataset [1] with 10000 training examples, where the digits are spuriously correlated with color (90% of all 1s are green and 90% of all 0s are red), creating 4 groups. We use a tiny Convolutional Neural Network with 3 layers and about 5000 parameters. A more detailed description of the setup and dataset is available in the Appendix E. We compare the results of training with:

1. GSR-FR-EH: Full-parameter retraining (FR) on GSR with “exact” Hessian, calculated with Conjugate Gradients (CG) method. 
2. GSR-FR-AH: FR on GSR  approximated Hessian (calculated with LiSSA [2], which approximates the Hessian with 10% of the training data), 
3. GSR-FR-HF: FR with gradient inner product (i.e., MAPLE)
4. Our two-stage GSR (with LLR) and GSR-HF (i.e., MAPLE+LLR). 

| **Approach** | Worst-group Accuracy | Total Wall-clock Time (min) |
| --- | --- | --- |
| GSR-FR-EH (Full-param retraining with “exact” Hessian (CG)) | **94.1** | 400 |
| GSR-FR-AH (Full-param retraining with approximated Hessian (LiSSA [2])) | 80.2 | 45 |
| GSR-FR-HF (Hessian-free Full-param retraining, i.e., MAPLE) | 91.9 | 42 |
| GSR (LLR with Hessian) | 91.2 | 4+0.8 |
| GSR-HF (LLR without Hessian, i.e., MAPLE+LLR) | 91.1 | 4+0.7 |
| ERM | 81.2 | 4 |

From the table, we can observe that at a smaller scale:

1. Full-parameter retraining with Hessian is advantageous over all other approaches, but at a much higher cost. 
2. Approximating the Hessian with low quality might hurt the performance.
3. When using a small neural network model, LLR can cause slight performance degradation when we have the capability to retrain the whole model on the entire dataset, but it is still significantly better than ERM. 

This result suggests that we should consider FR when we have the capacity, and LLR is an effective approximation to FR.

---

### **Challenges of scaling up the full-parameter retraining with GSR**

However, when scaling up full-parameter retraining to more realistic datasets with larger networks, **the computational challenges for accurately calculating the sample weight gradients are indeed the major obstacle here because of the Hessian.** Let $n$ be the number of data points and $p$ be the number of parameters. Directly calculating the Hessian for the whole network is both computationally expensive $O(np^2+p^3)$ and memory bottlenecked $O(p^2)$. In particular, the Hessian of a simple neural network with 1M parameters takes 4TB of memory, and ResNet-50 has 23M trainable parameters. Alternatively, we can calculate the inverse Hessian-vector product (iHVP) to reduce the memory requirements to $O(p)$ and the time complexity to $O(np^2)$, but it is still time-consuming, unless we apply further approximations. The table below shows the computational challenges of performing full-parameter tuning with GSR on the Waterbirds dataset:

| **Approach** | **Time (One-iteration)** | **VRAM Consumption** |
| --- | --- | --- |
| Explicit Hessian | >18h | 2PB |
| Conjugate Gradients | 18h | 9GB |
| LiSSA | 3.6h | 9GB |

It will be extremely costly to perform the experiments that optimizes the sample weights for 50 iterations as done with LLR. Besides, as shown in the previous ablation study, training the full network with crudely approximated Hessian (using LiSSA) causes inaccurate gradient estimation and may even hurt the performance. We believe these technical challenges require extensive exploration in future works and are beyond the scope of this paper, as originally discussed in Section 7 (lines 536-539),

[1] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019.

[2] Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization for machine learning in linear time. Journal of Machine Learning Research, 18(116):1–40, 2017.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733297711618

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Rebuttal Summary

comment:
We thank all reviewers for their precious time and feedback during the rebuttal. Below, we summarize the reviews and discussions: 

- **Reviewer wcXj** found our approach effective and computationally efficient in tackling subpopulation shifts, but raised the following concerns:
    - *Presentation:* We have carefully revised the paper and incorporated additional diagrams.
    - *Empirical studies for additional datasets and cases where last-layer retraining might be insufficient:* we added experiments with additional datasets and demonstrated that our approach remains advantageous.
    - *Empirical study involving full-parameter retraining with Hessian:* we introduced an additional ablation study that exhaustively compares the baselines and variants at a manageable scale.
- **Reviewer kehk** recognized the technical contribution of our approach and the thoroughness of our experiments, but was concerned about:
    - *The relationship between last-layer retraining and other factors, such as our assumption and regularization:* We have revised the paper to include additional clarifications and detailed explanations.
    - *Comparison with full-parameter retraining:* We exhaustively compared baselines involving full-parameter and last-layer retraining at a manageable scale, highlighting the advantages of using the Hessian.
- **Reviewer xdjG** recognized our theoretical solidity and the comprehensiveness of the experiments, but raised the following concerns:
    - *Comparisons with MAPLE:* we have updated the main paper to clearly state our differences, amended the empirical results for MAPLE, and clarified that the GSR-HF baseline corresponds to MAPLE+LLR, as requested by the reviewer.
    - *Validity of a claim in Section 4:* we have revised the confusing claim and provided supporting evidence including references to relevant literature and additional experiments.
- **Reviewer oDPM** acknowledged the validity, effectiveness, and noise robustness of our approach, but was concerned about the *theoretical derivations*. We have addressed these concerns by incorporating further clarifications and explanations in the revision.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This submission is a borderline case. It worked on subpopulation shift (a special case of distribution shift) with spurious correlations, and proposed group-robust sample reweighting (GSR) utilizing group-labeled data to optimize the weights of other group-unlabeled data, to effectively reduce the required amount of group labels. After the rebuttal, all four reviewers gave a rating of 6 that is "marginally above the acceptance threshold". During the internal discussions, none of four positive reviewers showed up, even after I cued them, suggesting that they were not truly supportive.

I quickly checked this paper and found that I shared exactly the same concerns with Reviewer xdjG. Subsequently, I found that the authors did a particularly good job of addressing his or her concerns (and thus mine). I like very much the first point in the rebuttal about the conceptual or philosophical difference between the proposed GSR and the baseline MAPLE, which emphasized the novelty of this work. The third point is also interesting about the motivation or justification of last-layer retraining in GSR. Since the motivation is strong and the contributions are solid (at least to me), I think we should accept this paper for publication even though our reviewers are only marginally positive.

I feel the authors considered themselves doing distribution shift research (they mentioned so in the rebuttal letters). There are some related papers you might be interested in about distribution shift and importance weighting for deep learning: "Rethinking importance weighting for deep learning under distribution shift", NeurIPS 2020 spotlight; "Generalizing importance weighting to a universal solver for distribution shift problems", NeurIPS 2023 spotlight. For certain hard distribution shift problems, last-layer retraining is not enough, and the papers tried to improve "the representations from a poor ERM-trained model" by iteratively updating all layers of the model and updating the weights on the data.