Data Augmentations for Improved (Large) Language Model Generalization

Summary

This paper focuses on counterfactual data augmentation for training robust machine learning models on text data. The concerned applications such as healthcare are safety-critical, highlighting the importance of this work. To generate data, the authors propose to employ a Large Language Model (LLM) to rewrite the documents. The authors also propose utilizing auxiliary data in the generation process to encourage robustness. In the experiments, the proposed method is proved to be both robust and effective.

Strengths

1. The research topic, training robust classifiers on safety-critical applications, is clearly an important task. Therefore, the positive experiment results as reported are expected to bring practical value to the real world. 2. The model performance reported in the experiment section is good. Normally an out-of-distribution generalizable model sacrifices performance in the in-distribution setting. It seems that the proposed method does not suffer from this issue. 3. The paper is overall well-written and easy to follow. The experiment settings in the main paper and Appendix are detailed.

Weaknesses

1. The novelty of the proposed method may be limited. It seems that the main difference between the existing and the proposed method is introducing auxiliary data, which may not be considered as a significant technical advance. 2. The competitors in the experiments are rather basic. In particular, it may be worth considering including IRM (Arjovsky et al., 2019) and GroupDRO (Sagawa et al., 2019) (or Just-Train-Twice (Liu et al., 2021) if the authors are interested in re-weighting) for comparison. 3. As this research targets at safety-critical applications, there are some details regarding model bias and robustness may require more clarification. The questions are in "question to authors".

Questions

1. The data augmentation process relies on an off-the-shelf LLM. Is it possible that the LLM can bring additional and uncontrollable biases to the proposed method? 2. Following the first question, how critical is the selection of LLM in this work? 3. How critical is the quality of the auxiliary data? Also, would it be difficult to identify informative auxiliary data in practice? 4. Regarding the writing, - line 134: Should the term under argmax be c \in [K] instead of y \in [K]? - Would it be more clear if the authors specify that the employed LLM is GPT4 in the main text instead of Appendix?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The authors adequately address the limitation.

Summary

The paper addresses the issue of text classifiers relying on spurious correlations, leading to poor generalization vis-a-vis out of domain, particularly in critical areas like healthcare. The authors propose using counterfactual data augmentation, guided by knowledge of causal data structure, to create text classifiers more robust to distribution shifts. Using an LLM to implement this approach, they demonstrate effectiveness in tasks such as predicting clinical diagnoses from medical narratives and observe improved out-of-distribution accuracy compared to baseline algorithms. They suggest that their method could be beneficial in dealing with various distribution-shift problems in machine learning applications. They demonstrate the utility of using language models to create counterfactuals that could improve model robustness. The paper also goes a step further in formalizing counterfactual data augmentation.

Strengths

- The authors experiment with both real-world clinical diagnoses scenarios and semi-synthetic data, demonstrating the broad applicability and practicality of their proposed methodology. Working with healthcare data isn't always easy either, so credit for that choice as well. - By using the capabilities of large language models for counterfactual generation, this work expands the potential applications of counterfactually augmented data while reducing costs associated with manually constructing counterfactuals. - The paper offers a formalization of counterfactually augmented data, which paired with prior work allows easier understanding and replication of the process for future researchers.

Weaknesses

- The method pre-supposes full knowledge of the causal structure of the data and overlooks complexity involved in real-world datasets where such information might not be readily available or accurately defined. Furthermore, assuming "no unmeasured confounding" isn't always realistic in real-world data sets especially those from healthcare, possibly limiting the model's robustness within certain contexts. Granted it is a common assumption in causal inference, but causal inference methods are also usually restricted to simple worlds that can be represented in a few variables. - The authors acknowledge that generating versions of clinical narratives as if they had been written by different caregivers is difficult to achieve in practice. However, there's not enough discussed on substantive solutions or workarounds they employed (or considered) to meet these challenges. - The paper needs greater discussion of prior work that has attempted to address these issues, including attempts at formalizing counterfactually augmented data and counterfactual invariance (see [1] and [2]). [1] Divyansh Kaushik, Amrith Setlur, Eduard Hovy, and Zachary C. Lipton. "Explaining the efficacy of counterfactually augmented data." ICLR 2021. [2] Victor Veitch, Alexander D'Amour, Steve Yadlowsky, and Jacob Eisenstein. "Counterfactual invariance to spurious correlations: Why and how to pass stress tests." NeurIPS 2021.

Questions

- I think the paper could benefit a discussion of how it differs from [1] and [2]. Could you share how you view the difference (the bounds are clearly new)? - In Section 4 you share using auxiliary data to help with counterfactual generations so the models don't hallucinate. Could you share how often did the models hallucinate (or not) when this auxiliary data was presented vs when it wasn't? - Could you share results of a baseline where you have 2X datapoints from the original distribution (observational data alone), compared against the counterfactually augmented data (X original + X' counterfactuals)? - In line 308, you say "across 5 runs". What is the source of randomness in these 5 runs? - Am I correct in interpreting Figure 4's findings that seemingly more corruption leads to better performance? Could you elaborate?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The limitations are adequately discussed.

Summary

This paper develops and analyzes the counterfactual data augmentation strategies for the setting where the data-generating process is anti-causal. First, they first show the Bayes optimal predictor in the setting they considered and then how can we obtain that predictor using counterfactual data augmentation. They give two different strategies to generate the counterfactual data (1) by prompting LLM and, (2) using the difference in difference method. Then they analyze the sample complexity of counterfactual data augmentation and show that it is better than that of a baseline reweighting method. Finally, they empirically show the effectiveness of their method on synthetic and real-world datasets.

Strengths

**Clarity**: The paper will be well-written and easy to follow. **Sample Complexity Analysis**: The authors show that the sample complexity of counterfactual data augmentation is better than the reweighting baseline introduced in Makar et. al.

Weaknesses

1. Contribution 2 - Counterfactual Data Augmentation (CDA) as a method to deconfound target and spurious attribute: It is not clear how this observation/finding is different from the previous literature on CDA eg. Kaushik et. al. (Explaining The Efficacy of Counterfactually Augmented Data) and Joshi et. al. (An Investigation of the (In)effectiveness of Counterfactually Augmented Data). 2. Lemma 1 is similar to the claim shown in Makar et. al. (Causally motivated shortcut removal using auxiliary labels) 3. Assumption 1 (constant effect): How justified is the assumption that the spurious attribute $c$ change the previous state in an additive manner? 4. The augmentation strategy in this paper is limited to the anti-causal setting which further limits the applicability. I understand this is standard practice in the OOD generalization literature to develop a method under specific data-generating process (causal and anti-causal) but then this paper doesn’t introduce something new than the previous works. Overall this paper seems to re-instantiate the point mentioned in previous work that argues for the effectiveness of counterfactual data augmentation. Though this paper introduces two new ways to perform data augmentation in the context of medical datasets, the overall novelty seems low.

Questions

**Minor Typos/Corrections**: 1. Line 117 (Definition 1): State that $\Delta^{K-1}$ is a simplex

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

NA

Summary

The authors propose to use knowledge from the causal structure of the data to counterfactually simulate interventions on spurious features and to learn more robust classifiers. They focus on text classification tasks and argue that this approach is appropriate in prediction problems for which the label is spuriously correlated with an attribute. Since the authors argue that their approach emulates interventions, one assumes that they referring to the second layer of Pearl's Causal Hierarchy. Here, we assume the data and the Causal Graph (CG). However, the authors left the specification of the causal graph to be addressed as future work. Additionally, following the Pearl's Causal Hierarchy, when leading with counterfactual questions, one must assume the data and the Structural Causal Model (SCM). Given that the authors do not provide neither the CG nor the SCM, I have concerns on how their approach generates counterfactually augmented data.

Strengths

This paper presents an interesting approach to counterfactual-based augmentations focusing on textual data.

Weaknesses

There are some important questions to better understand the proposed approach and the validity of results.

Questions

The authors assume that the writing style of caretaker C affects the vector representation of the clinical note (Fig. 1). Further, they also assume that the label Y causes both X and X^. As far as I understood, the Clinical Condition Prediction consists in, given a patient note (X), predicts the concept(s) associated with it (Y). I want to understand if this is a causal or anti-causal problem. According to Fig. 1, we pick a concept and then generates a clinical note. In Sec. 2, the authors mentioned several approaches to invariant representation learning and counterfactual augmented data generation, however they did not select any of them as baseline(s) in the experiments section. Is it possible to provide a better understand of both (i) why the selected baselines are the most appropriate? and (ii) why the related works are not representative baselines? Does one need to know whether the label is spuriously correlated with an attribute in advance? Is it possible to provide details about the private held-out data used in this paper? All information necessary and sufficient to others reproduce the results using other datasets would be helpful.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

I believe that the authors only mention limitations of their work in the last sentence of the Discussion section. I believe that would be interesting to improve this discussion. For instance, it is very hard to have the correct Causal Graph for most important applications, hence, how this could be an issue, and, in the absence of the causal graph, what is the best approach? Another issue is how can we guarantee that the counterfactual instances generated during augmentation are realistic?

Summary

In this work, the authors develop causally-driven data augmentation methods to improve model robustness.

Strengths

In general, I find the paper to be very, very well written and clear. The authors do a really great job of explicitly stating their assumptions, and acknowledging when certain assumptions are strong. I think this paper makes a very interesting first step in extending OOD generalization to the recent advances in LLMs!

Weaknesses

I would have liked a section that describes some of the limits of using LLMs, and whether certain LLMs would be more appropriate that others. It feels like incorporating LLMs is a big part of this work, so I would have liked more context here. I wasn't completely convinced that the extra information M would be enough to assist the model in achieving more accurate estimates, but the authors do acknowledge that this is a strong assumption.

Questions

see above

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

The authors do acknowledge limitations of their work.

Thanks for the clarification

Thanks for the response. It clarifies many things for me. I am happy to update my score to recommend acceptance assuming the authors would be updating the camera ready to reflect these clarifications as well.

Response to Authors

I thank the authors for taking the time to address my questions. Also apologies for sounding dismissive of your work in the context of overall contribution and I thank you for pointing out the main contribution of the paper i.e. finite-sample complexity analysis and making CAD work in practice using causal estimation techniques. Given this, I resonate with one of the concerns raised by Reviewer 8nR3 around the reliance of this work on the LLMs and the potential bias introduced by them. I would like the authors to include the discussion they had regarding this issue in the final version of the paper. Overall, I am raising my score by 1 and would like the authors to incorporate the changes in the final version.