Compositional Risk Minimization

Authors: Divyat Mahajan, Mohammad Pezeshki, Ioannis Mitliagkas, Kartik Ahuja, Pascal Vincent
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: Compositional Generalization, Out of Distribution Generalization, Provable Extrapolation

Abstract:
In this work, we tackle a challenging and extreme form of subpopulation shift, which is termed compositional shift. Under compositional shifts, some combinations of attributes are totally absent from the training distribution but present in the test distribution. We model the data with flexible additive energy distributions, where each energy term represents an attribute, and derive a simple alternative to empirical risk minimization termed compositional risk minimization (CRM). We first train an additive energy classifier to predict the multiple attributes and then adjust this classifier to tackle compositional shifts. We provide an extensive theoretical analysis of CRM, where we show that our proposal extrapolates to special affine hulls of seen attribute combinations. Empirical evaluations on benchmark datasets confirms the improved robustness of CRM compared to other methods from the literature designed to tackle various forms of subpopulation shifts.

Decision: reject
Reviewer scores: 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730098626044

thread_type:
official_review

<thread_object>
creator: reviewer_muzn

<content>
summary:
This paper focus on tackling the challenge of  subpopulation shift through Compositional risk minimization, where the theory of discriminative compositional shifts is established, and a practical algorithm (compositional risk minimization, CRM) is also given. Overall, the theory are technically sound  and the experimental results demonstrates that CRM is . A good paper but require more specifications.

strengths:
1. **Novel Approach**: The paper introduces a novel approach based on flexible additive energy models for compositionality in discriminative tasks, which can provably extrapolate to novel attribute combinations.
2. **Theoretical Grounding**: The approach is theoretically grounded, providing guarantees for generalization under compositional shifts, which is crucial for robust machine learning models.
3. **Comprehensive Evaluation**:  Extensive empirical studies and experiments to demonstrate the impact of different choices and validate the effectiveness of the proposed method.

weaknesses:
1. The Compositional Risk Minimization (CRM) framework lacks explicit constraints on feature representations, which is widely recognized as crucial for generalization under sub-population shifts [1]. While CRM effectively models compositional shifts through additive energy distributions, it does not incorporate mechanisms to ensure the learned features are robust and semantically meaningful across different attribute combinations. Could the author make a further discussion on this problem?
    
    [1] Liu et al., Causal Triplet: An Open Challenge for Intervention-centric Causal Representation Learning, 2023
    
2. For the implementation of Algorithm 1, how to pre-specifies the spurious attribute $a$?

3. The current experimental results are somewhat lacking of conviction. More advanced DG algorithms, e.g. IRM, VREx, Fishr, SelfReg, are recommended to be included in the experiments. 

4. The sample complexity $s$ w.r.t the dimension $d$ should be verified through simulation studies.

questions:
1. I am a little confused of the assumption 1: $p(x|z)=q(x|z)$，i.e.  the data generation mechanism from attributes to the inputs remains invariant, which is the main claim in the paper,  hence the distribution shift is essentially caused by $p(z)\neq q(z)$, which further implies that $p(x)\neq q(x)$, due to  $p(x)= \int p(x|z)p(z)dz$, and $q(x)=\int q(x|z)q(z)dz$. So, the covariate shift is caused by the marginal differences on the attributes. This seems contradict to traditional DG. 

2. For sample complexity $s\geq 8cd\log(d/2)$,  $c$ lacks of specifications?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730432122043

thread_type:
official_review

<thread_object>
creator: reviewer_rmiq

<content>
summary:
This paper addresses a challenging scenario involving compositional shifts, where certain combinations of attributes are not present in the training distribution but appear in the test distribution. To tackle this issue, the authors propose a Compositional Risk Minimization (CRM) framework, which incorporates an additive energy classifier to predict multiple attributes. Furthermore, the framework adapts this classifier to effectively handle compositional shifts. Empirical studies demonstrate the effectiveness of the proposed CRM approach.

strengths:
This paper explores a significant and under-researched area related to compositional shifts, which has been largely overlooked in the existing literature. 

The study is quite rigorous and comprehensive. Notably, it introduces the proposed approach through a running example involving two attributes, which makes the contributions and motivations clear and accessible. Additionally, the formulation of the Compositional Risk Minimization (CRM) framework offers valuable insights for future research.

weaknesses:
One significant concern pertains to the attribute distributions during compositional shifts, which are frequently encountered in real-world applications. For instance, the distribution may shift from uniform to long-tailed, or it may be long-tailed from the outset. In this context, is it still effective to assume the test group prior $q(z)$ to be uniform? I encourage the authors to consider and discuss this issue in their paper.

The efficiency of the proposed CRM approach should also be examined, particularly in the context of multiple attribute abstention cases.

questions:
The authors should address all my concerns presented in the Weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730722192832

thread_type:
official_review

<thread_object>
creator: reviewer_x1t8

<content>
summary:
This paper addresses the compositional shifts, a hard type of sub-population shifts, and proposes compositional risk minimization. The method is well-motivated and some theoretical analyses are provided. Results on the sub-population shift benchmark are shown to support the proposed method.

strengths:
- The compositional risk minimization method is reasonable and well-motivated.
- The formulation of the compositional shift setting provides a foundation for further research.
- Experimental results are good to support the method.

weaknesses:
Based on the theoretical analysis and experimental results, my main concerns are as follows:

1. **Experimental Setting**: The analysis and proposed algorithm are designed to handle multiple attributes, with the theoretical advantages being most relevant for this multi-attribute context. However, the experiments are limited to only two attributes. I suggest that the authors include empirical results with multiple attributes to better align with the theoretical analysis.

2. **Results**: While the CRM results outperform the baselines in this study, I check the original benchmark paper [1] and find that CRM’s performance does not consistently exceed that of the baselines. Furthermore, the authors only report accuracy, which doesn’t fully account for group imbalances. Reporting Macro-F1 scores would provide a more comprehensive evaluation.

3. **“Disjoint” Setting**: I understand that the paper’s chosen disjoint setting adds a level of complexity. However, in real-world scenarios, it is often feasible to obtain a small number of samples for different attribute combinations (particularly with only two attributes, as in these experiments). The proposed method should also be evaluated in traditional settings where all attribute combinations have some representation. This would confirm that the method performs well without requiring group-dropping.

questions:
Please refer to weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732124987500

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank the reviewers for their thoughtful feedback and constructive questions. We are encouraged by their recognition of the significance, novelty, and technical rigor of our work.
- **Reviewer Rmiq** noted that the study is "rigorous and comprehensive" and appreciated the clarity of the contributions through the running example. 
- **Reviewer MuZN** highlighted the theoretical grounding of our approach and its importance for generalization under compositional shifts. 
-  **Reviewer Rmiq** and **Reviewer x1T8** emphasized the significance of exploring compositional shifts, with **Reviewer x1T8** describing our formulation as providing "a foundation for further research".

We address the concerns raised by each reviewer in detail as an individual response to each reviewer. We address the multiple attribute experiment here, which was a shared concern for both **Reviewer x1T8** and **Reviewer Rmiq**.

&nbsp;

## **Multiple Attribute Experiment**

We ran a new experiment using the CelebA dataset, by adding the attribute 'Eyeglasses' as another spurious attribute, leading to a total of *3* attributes. The results are shown below as well as in the newly added Table 13 in Appendix E.1, with more details regarding the setup.

| Method | Average Acc | Balanced Acc | Worst Group Acc | WGA (No Groups Dropped) |
|--------|-------------|--------------|-----------------|-------------------------|
| ERM    | 94.1 (0.1)  | 77.4 (0.3)   | 21.6 (0.6)      | 19.0 (0.0)              |
| G-DRO  | 91.2 (0.5)  | 87.3 (0.3)   | 61.0 (2.1)      | 66.7 (2.3)              |
| LC     | 92.8 (0.0)  | 89.8 (0.1)   | 75.0 (1.3)      | 77.0 (2.0)              |
| sLA    | 93.0 (0.1)  | 89.7 (0.1)   | 76.0 (0.7)      | 77.0 (4.0)              |
| **CRM**    | **92.3 (0.2)**  | **91.6 (0.0)**   | **84.0 (0.3)**      | **86.3 (1.2)**  |

As can be seen from the results, CRM is the best approach w.r.t the worst group accuracy as well as the balanced group accuracy. Note that, as we had mentioned in the paper (line 442), average accuracy is not the best indicator of generalization under spurious correlations, hence we should pay more attention to the other metrics.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732510685895

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Changes to the Draft

comment:
We summarize the new sections added to the paper to address the reviewer's comments. We believe we have thoroughly addressed all the concerns raised by the reviewers. With the rebuttal period nearing its end, we wanted to check if our responses have sufficiently addressed your concerns. If so, we would be grateful if you could update your score accordingly.

- **Appendix E.1**: Multi-Attribute Experiment.  (Reviewer x1T8, Rmiq)

&#8594; Table 13 presents results for benchmarking CRM in the case of multiple attributes, showing that CRM outperforms other approaches in terms of both worst-group accuracy and group-balanced accuracy.


- **Appendix E.2**: Additional Baselines & Macro F1 Score (Reviewer x1T8, MuZN)

&#8594; Table 14 provides results with additional baselines (IRM, VREx) and metrics (Macro F1 score), providing more evidence for  CRM's superior generalization under compositional shifts. To better summarize the results, Table 15 provides the aggregate rank obtained by each method. CRM obtains the best rank w.r.t balanced accuracy, macro F1, and worst group accuracy, and CRM with empirical test prior obtains the best rank w.r.t average accuracy.

- **Appendix E.3:** Sample Complexity Analysis as function of $d$ (Reviewer MuZN)

&#8594; Figure 6 presents the results of our analysis where we evaluate CRM's generalization capabilities as we discard more groups from the training dataset for different values of $d$.

We have also additional details in Appendix D.2 which provides results for the ablation of CRM w.r..t the choice of test prior and extrapolated bias.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733286948438

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Final Author Response

comment:
We thank all the reviewers for their time and engagement during the rebuttal phase! Majority of the reviewer's concerns were about the empirical contributions of our work, which we thoroughly addressed with additional experiments and clarifications.

- We benchmarked CRM against three new baselines (IRM, VREx, Mixup) and additional metrics (Macro F1), and found that CRM outperforms them in its ability to generalize under compositional shifts. (Appendix E.2, E.4)

- We showed the superior generalization of CRM over baselines in the multi-attribute case via experiments on the multi-attribute CelebA dataset comprising of 3 attributes.  (Appendix E.1)

To address the concerns of **Reviewer x1T8** regarding the multi-attribute experiment, we have further augmented the CelebA dataset to include a total of 5 attributes. We present results for the case of 5-attribute CelebA dataset below, where each metric is averaged over the 32 compositional shift scenarios of discarding a group from the training and validation set.


| Dataset         | Method            | Average Acc  | Balanced Acc  | Macro F1  | Worst Group Acc |
|-----------------|-------------------|-------------------|--------------------|----------------|-----------------------|
| CelebAMultiAttr | ERM               | 93.9 (0.0)        | 80.5 (0.4)         | 84.9 (0.3)     | 2.8 (1.3)             |
| | G-DRO             | 86.8 (0.3)        | 90.2 (0.2)         | 94.4 (0.2)     | 53.0 (4.4)            |
| | LC                | 88.1 (0.0)        | 91.1 (0.0)         | 95.0 (0.0)     | 65.5 (1.0)            |
| | sLA               | 88.2 (0.1)        | 91.0 (0.1)         | 94.9 (0.0)     | 68.0 (0.6)            |
| | IRM               | 86.4 (0.6)        | 84.3 (0.1)         | 89.6 (0.1)     | 28.3 (0.6)            |
| | VREx              | 84.5 (0.1)        | 87.5 (0.1)         | 91.4 (0.2)     | 7.6 (2.6)             |
| | Mixup             | 93.7 (0.0)        | 80.4 (0.1)         | 85.2 (0.0)     | 1.2 (0.3)             |
| | **CRM**               | **85.3 (0.0)**        | **88.1 (0.1)**         | **93.6 (0.1)**     | **71.3 (0.5)**            |


We find that CRM still outperforms the baselines for generalizing under compositional shifts, especially w.r.t the worst group accuracy. Hence, the gains with CRM are consistent as we increase the total number of attributes. 

We would also like to restate the significance of our theoretical contribution. Our method is the first to offer provable guarantees for extrapolating under compositional shifts, which also stands the test of experimentation on real benchmark datasets, as CRM systematically outperforms all baselines under such shifts.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This submission is a borderline case. It proposed a problem called "compositional shift" as well as its solution called "compositional risk minimization". After the rebuttal, all three reviewers gave a rating of 6 that is "marginally above the acceptance threshold". During the internal discussions, none of three positive reviewers showed up, even after I cued them, suggesting that they were not truly supportive. In particular, reviewer x1T8 mentioned in his or her post-rebuttal comment that "although I still have concerns about the experimental settings (even if there are three attributes, I don't think they're enough to represent the setting considered in this paper), I would raise my score." My personal feeling is that the motivation of this work may not be strong enough --- as the authors mentioned, the problem under consideration is a **challenging and extreme** form of subpopulation shift, so why is such a hard problem meaningful in practice? Therefore, I don't think we can accept it for publication at such a competitive venue unless the authors can justify their motivation.