Data Augmentations for Improved (Large) Language Model Generalization
Authors: Amir Feder
Yoav Wald
Claudia Shi
Suchi Saria
David Blei
Venue: NeurIPS 2023 poster (2023)
Keywords: Counterfactually Augmented Data
Invariant Learning
Out-of-distribution Generalization
Clinical NLP
TL;DR: We propose counterfactual data augmentation methods, guided by knowledge on the causal structure of the data, to simulate interventions on spurious features.

Abstract: The reliance of text classifiers on spurious correlations can lead to poor generalization at deployment, raising concerns about their use in safety-critical domains such as healthcare. In this work, we propose to use counterfactual data augmentation, guided by knowledge of the causal structure of the data, to simulate interventions on spurious features and to learn more robust text classifiers. We show that this strategy is appropriate in prediction problems where the label is spuriously correlated with an attribute. Under the assumptions of such problems, we discuss the favorable sample complexity of counterfactual data augmentation, compared to importance re-weighting. Pragmatically, we match examples using auxiliary data, based on diff-in-diff methodology, and use a large language model (LLM) to represent a conditional probability of text. Through extensive experimentation on learning caregiver-invariant predictors of clinical diagnoses from medical narratives and on semi-synthetic data, we demonstrate that our method for simulating interventions improves out-of-distribution (OOD) accuracy compared to baseline invariant learning algorithms.

Data Augmentations for Improved (Large) Language Model Generalization
Amir Feder *1.2, Yoav Wald *3 Claudia Shi 1, Suchi Saria and David Blei
CColumiaa University, Google Research, Johns Hopkins University
Abstract
The reliance of text classifiers spurious correlations lead to poor general- ization at deployment, raising concerns about their use in safety-critical domains such as healthcare. In this work, we propose to counterfactual data augmenta- tion, guided by knowledge of the causal structure of the data, to simulate interven- tions spurious features and to learn more robust text classifiers. We show that this strategy is appropriate in prediction problems where the label is spuriously correlated with attribute. Under the assumptions of such problems, we discuss the favorable sample complexity of counterfactual data augmentation, compared importance re-weighting. Pragmatically, we match examples using auxiliary data, based on difff-n-diff methodology, and use a large language model (LLM) to represent conditional probability of text. Through extensive experimentation on learning pnegiver-innvaian predictors of clinical diagnoses from medical nar- ratives and on semi comi-synthetic data, we demonstrate that our method for simulating interventions improves out-of-distribution (OOD) accuracy compared to baseline invariant learning algorithms.
Introduction
The reliance on spurious correlations a significant challenge for Machine Learning (ML) safety as it can lead to performance degradation of deployed models. Spurious correlations prevalent in various applications such as medical imaging [1. 2], text classification [3]. and risk prediction sys- tems [4]. Failures due spurious correlations occur under distribution shift [5-7], which may result from differences data recording protocols, shifts in the underlying population being monitored, or the way the ML tool is being used. In this paper, focus on text classification and explore how ing language models in a domain domain-informan informed way can help us avoid reliance on spurious correlations.
Consider scenario where we want to make robust predictions about patients' conditions, proba- bility of readmission, etc., using clinical narratives written hospitals [8-10]. In this setting, common issue arises due to clinical practice, where patients with certain conditions directed to specific caregivers in the hospital. When we train a predictor from a single dataset that exhibits some correlation between caregiver-specific style and clinical outcomes, the predictor may unintention- ally rely on the style to make predictions. This leads to poor generalization on unseen hospitals, i.e. failure to generalize out of (astrbution(OOOD) due to changes in clinical practice [7]. However, col- lecting a dataset that is large enough to avoid such spurious associations is infeasible due to various reasons such as rare conditions, privacy concerns, etc. To tackle this problem, propose leverag- ing available auxiliary data (e.g., time, document type, demographics) and incorporating knowledge about the causal structure of the problem build more robust classifier. For example, in the note classification task, we can use our knowledge that some auxiliary data, such as the patient's 's current state, can affect doctor assignment, improve the classifier's robustness.
Equal Contribution. Correspondence to amir www..fffee@@@lmbbaa.. edu
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Causal inference often makes use of such auxiliary data and has now been used in variety of ways to improve OOD generalization [6, 11-14]. Data augmentation methods have demonstrated impressive performance in these tasks well [15-17], and with recent improvements in generative models, forming additional principles to incorporate domain knowledge into data augmentations seems like promising path forward.
In this work we pursue this and develop cansally-drive data augmentation methods, that leverage auxiliary data and domain knowledge. Intuitively, generating versions of clinical narratives as if they had been written by different caregivers, de-correlates the writing style from the patient condition we wish to predict. However, such data generation can be difficult to achieve in practice and problem- specific traits must be taken into account [18]. Observing that data augmentation can be treated as counterfactual outcome estimation under a causal formalism, motivates the use of causal inference methods that are commonly used for such tasks across the sciences. While our approach can be applied to many modalities of data, in this work we focus on text classification and harness the recent advances in LLMMs towards counterfactual estimation. Our contributions are:
1. Through extensive experiments, we show how the of language models in manner that is informed by causal knowledge improves model robustness challenging safety-critica tasks in healthcare. Furthermore, our findings reinforced by experiments that incorporate semi- synthetic scenarios, and simulations where there are ground-truth counterfactuals.
2. We formalize counterfactual data augmentation in prediction setting as a method to decon- found the target and spuriously correlated attribute. We show how deconfounding improves OOD generalization. In setting where sample complexities for alternative methods (re- weighting and invariance penalties) can be derived, we show favorable generalization bounds for accurately performed data-augmentation
3. Our autta-augmetatiion methods rely on common assumptions in the causal inference literature such as no unmeasured confounding and parallel trends in diff-in-diff [19], applied with LLMs. We believe that leveraging auxiliary data and assumptions about causal structure, along with the use of L.MSs and other generative models, can be fruitful framework for addressing many out-of-distribution generalization problems.
Next, we provide a brief survey of relevant work ($2). We then present a formal setting motivating counterfactual augmentation for OOD generalization ($3), our methods for counterfactual estimation and reason formally about the preferable sample complexity of our approach ($4). Finally, we present our main experimental results ($5) and discuss limitations and future directions ($6).
2 Related Work
Invariant and Shift-stable Learning. This paper contributes to the growing literature on invariant and shift-stable learning, which tackles the problem of learning models that generalizes across dif- ferent distributions or settings. Invariant learning through feature pruning was pioneered by Peters et al. [11], and has since been developed for variable selection [12, 20] and representation learn- ing [13, 21-26]. These methods have been applied in range of domains, including natural science [11. 12, 20], causal estimation [27, 28], computer vision [13, 23], and NLP [29-32]. However, recent studies have highlighted limitations in many invariant learning approaches, particularly in achieving conditional independence [33-36]. Others have investigated learning of stable models by leveraging causal methods through techniques like graph-surgery [6, 14], that come with generalization guar- antees. Yet others have explored the advantages of data augmentation [37, 38]. In this work, we combine the latter two approaches to improve OOD generalization for text based classification.
Counterfactually Augmented Data. To learn invariant predictors, popular and straightforward approach data augmentation. When data augmentation involves actions that go beyond simple manipulations (e.g. image rotations, crops etc.), it is often referred to as counterfactual data aug- mentation [37]. Constructing counterfactual instances that involve perturbations factors [39], or to the label [37, 38, 40], and incorporating them into the training data, breaks up correlations that we do not wish our model to exploit towards prediction. Most work on counterfac- tual data augmentation in text involve manual editing by humans, heuristic keyword replacement, or automated text rewriting [37, 39, 41-50]. Manual editing accurate and effective [38, 51] but expensive, hence our goal is to make counterfactual data augmentation scalable, demanding smaller
confounding

---

human effort. Keyword- based methods can be limited in coverage and difficult to generalize across languages [52]. Generative approaches offer balance of fluency and coverage [53], but generat- ing meaningful counterfactuals is challenging [54]. Our work departs from previous techniques by using causal auxiliary data structure and L. to alleviate this challenge and generate plausible
counterfactual data augmentations.
Clinical Notes. Clinical notes are the backbone of electronic health records, often containing vital information not observed in other structured data Kreimeyer et al. [55]. Clinical NLP involves iden- tifying this information, and standardized datasets and competitions exist for this purpose [56-60]. Best performing approaches have leveraged transformer architectures both for token-level classi- fication tasks [61-64], and for using complete clinical records [65, 66]. Recently, large language models (LLMs), similar to those we use to generate counterfactual notes, were shown to have clear potential for improving clinical NI.P systems [67, 68]. In our experiments, we follow recent papers in clinical NLP addressing challenges of degraded performance across different hospitals [69-71].
3 Problem Setting
To formally analyze how counterfactual data augmentation helps OOD generalization, we con- sider a setting where the label is spuriously correlated with known attribute. This set- ting has been used previously to study learning with "shortcuts" [25] and spurious correla- tions [29]. We note that our approach applicable and valid under additional settings and causal graphs (e.g. "purely spurious" problems defined in Wang and Veitch [72]) and we elab- orate on this at ??. The data generating process used here motivates counterfactual data aug- mentation in principled manner, as it describes the main problem we study and it is pos- sible to analytically compare sample complexity with an alternative solution (see section 4.3).
Consider a classification problem with L classes, where the label Y corre- lated with certain attribute C' in the training data and this correlation may change arbitrarily test time (denoted by red edge C+Y in fig. 1). our medical notes example, C' is the caregiver writing the note and Y is the un- derlying condition we wish to diagnose. We denote the number of caregivers
C
Y
X*
in our training data by [K]. For given loss function f: RR x [L] - R and distribution P. we denote the expected loss of a hypothesis h X + R by (h) and its expected accuracy by R for (h). The data geata-generating process is depicted by the causal model in fig. for our motivating example of clinical notes classification X is vector representation of the clinical note and X** an unobserved sufficient statistic, representing all the relevant information about Y in the note that is unaffected by the writing style of the caregiver. Let us formally define this setting.
Figure 1: Predic- tion problem with a spuriously corre- lated attribute.
We denote the set of distributions induced by interventions on a causal model with the structure in fig. by
Definition 1.
P (P(XIXXXX))) (Y)())P((IT) : (CClY=y)EAS [L]]
where all distributions other than P(C YY are fixed. In prediction problem with spuriously correlated attribute, the learner is provided with a set {(x,, Wi, G)) sampled i.i.d from Pimin P. We assume that X eXX almost surely for some e: IRd Rd".
In this problem, once X* is recovered no additional information from X is needed to predict Y. We can also see from the graph that interventions on P(C Y) do not change the con- ditional distribution P(Y X*). Therefore an optimal solution that does not rely on is h"" (x) arg m P(Y y ). In clinical note classification, X* represents all the informa- tion the note about the patient conditions, unsullied by the writing style of caretaker C. To obtain h² (x) we will rely on risk minimization w.r.t distribution where Y and C' are uncorrelated.
3

---

3.1 Learning Robust Classifiers when Counterfactuals are Available
Consider the unconfounded distribution P1 P that is given by intervening on C, setting it inde- pendent of Y and uniformly distributed, P(C (Y) (C). An optimal classifier under P1 has the following min-max optimality guarantee. 2
Lemma For the prediction problem in definition the Bayes optimal classifier under the unconfounded distribution P1 E P where C is uniformly distributed and independent of Y is h" (x) arg g Pı(Y = X = (. It is a minimizer of min mBxppp paxpep Rp (h)
and R.p (h²) R.p (h*) for all P.
Hence we would like to minimize risk w.r.t P1 and we cannot do that directly by via ERM since our training data sampled from Pinin # P1. Instead we consider risk minimization over an augmented dataset that contains counterfactual instantiations of our training data under different values of C.
Minimizing R.p. via Counterfactual Data Augmentation. Returning our motivating example, assume that we could generate clinical notes for all alternative scenarios. That is, obtain the clinical notes that would have been written each patient had been seen by all possible caregivers c é [K] and each caregiver had written their own version of the note (c). Given these counterfactual clinical notes, we seek hypothesis that minimizes the average loss over all such possible scenarios, denoted by Re
Definition 2. Consider prediction problem with spuriousty-correlated attribute (see Defini- tion 1). For a given example x,, we denote its counterfactual with attribute value [K] as derived from the corresponding causal model, by x.(c). For estimates of the counterfactuals (&,(c)]se[N))e [ee and a hypothesis h H. the counterfactually augmented empirical risk
1
R Ran (h)
Σ (()())()). NK se[[],c[[K]
We use approximate counterfactuals (c) in our definition to highlight that in practice we cannot obtain a precise estimate of x, (c). the ideal case where xx())) (c). the expected loss R (h) where N + 0o, satisfies Ring (h) = R2p (h). This follows by simple derivation and it is part of claim give later in Lemma Hence obtaining this dataset is useful for goal of minimizing
risk under P1. Our main challenge is then to derive effective approximations for counterfactuals such clinical notes under alternative writing styles.
4 Assumptions and Algorithms for Estimating Counterfactuals
Perfectly capturing writing style a strong assumption. Even if we could perfectly model writing styles, we only observe limited set of variables the actual notes I, outcomes y, and assigned caregivers c. We do not observe all factors that could influence what each caregiver would write. To alleviate this problem, we make use of auxiliary data M that is available during training, but might not be available in deployment.
As an example, consider two caregivers e and ट where a note x; was written by G = C. We want to estimate what xc).., the note caregiver would have written, might look like. To this end will build a model T-(-) that takes data and generates a note in caregiver c's style. Now suppose caregiver usually sees patients with high blood pressure and always includes blood pressure values in notes, while ē rarely does. A naive model estimating (c) Tx(x;) based only on c's notes may fill in false blood pressure information, conflating that with c's style. Including vitals data like blood pressure, typically recorded in a patient's 's health record, can provide additional context for our model. This extra information can assist the model in reasoning about external/background variables. leading to more accurate estimates.
"This claim is shown in Makar et al. [25], appendix A includes proof for completeness. We set the distribution over C in P. uniform for simplicity, the derivation for non-um non-uniform distributions analogous.

---

4.1 Identification of the Counterfactual Distributions
To make effective use of this data, we suggest that the input to the model Te X M X will include a baseline text to be edited and auxiliary data m. Intuitively, accounting for confounding between the identity of the caregiver C and the text X, with auxiliary data M should result in improved augmentation.
We formalize this intuition using an assumption from causal inference. To identify the counterfactual text distributions using the observed distribution, we assume strong ignorability [73-75] Assumption (Strong ignorability). For all P e P it holds that (()) C M, and for all values of m E M, P(m)> O.
Under this assumption, we can rewrite the counterfactual distribution with the observed distribution,
P(X(c)) f P(X(O) M P(M- m)dm f
P(X C c. c, M P) - 1m)mmm)dm m)dm.
However, in practice, we do not observe many samples from P(X C = C. m). making it poor approximation for the counterfactual distribution. We address this by using counterfactual data augmentation [37]. Formally, we assume that for all possible counterfactual distributions e [K], there exist a function Te that maps from the observed distribution P(X) M m) to the target counterfactual distribution ((X(c) M=m).
=
We approximate the loss under the counterfactual distributions through the empirical loss produced by data augmentation. That is, for a hypothesis h E H
(p(x(en)(((((().)))
N Σ Tx(xx, s[[N]
Note that whenever the text in the training set is already written by caregiver c, i.e. G; will simply keep the original text x;
Evaluation of Augmented Distribution. The right hand-side of the above equation is Monte-Carlo estimator of the distribution of augmented notes, which averages the distributions Te,c( X, M)) over all caregivers E [K]. The distribution Prrann PX,,an(,, M)) is aimed to follow the style of caregiver c. While the observed samples from one counterfactual distribution may not be sufficient to approximate the whole distribution, they can be used to assess the quality of the counterfactual augmentation algorithm To.
High-quality counterfactual estimation, as measured by small distributional divergence between our estimator and the target distribution, will help lowering the upper bound on the risk R.p, (h) (see lemma 2 in section 4.3). Then to estimate divergences between these two distributions, we may use validation sets from our training data. A sample from T. Te (Pirain (X, M)) is obtained simply by running training data through Te, while a sample from P(X(c)) can be obtained either by adjusting for M, or we can obtain sample from P(X C c, M m) for each value of m and compare that to sample obtained by augmenting validation data where M = m. In both cases two-sample tests can be applied and obtain estimates divergences between the two distributions. That of course as long as positivity holds, i.e. the second part of the assumption, as otherwise we will not be able to obtain samples of P(X C = c, M m) for certain values of m and c.
We now describe the estimation methods that obtain Te, The methods are based on classical causal inference methods, applied to our high-dimensional setting, and relying on the auxiliary data M.
4.2 Methods for Estimation of Counterfactuals
Counterfactual estimation is an established problem in causal effect estimation [74, 76, 77]. Here we adapt identification strategies and estimation procedures in the causal literature to estimate xc)... Our framework for estimating counterfactuals CATO (Causal-struccta Text OOD Generalization) involves the use of an LLM to model the conditional probability distri- bution of text. Counterfactuals are formed by matching similar auxiliary data examples or manipu- lating texts' vector representations, described below.
Driven Augmentations for
Prompting with matched examples. Our first estimation method in Algorithm 1(B) draws insights from matching [76]. We construct prompt for an LLM, that given an original text x and set of
5

---

Algorithm/ CATO
Pre-process CATO (A)
Input:
Training set {(xi, Me. Cj m;)
Assume:
m includes the label y and pre- treatment attribute Opre among other aux-
Hypothesis class
Version (B)}
iliary data. We are given (x), ) j=1'
Optional tre-treatment data { (xpre,i) i
1: Set p(ej, mý) Xj Xipe for j [N].
Output: hypothesis Maug (x)
2: return To (x, m) Xpre p(e,
1: if Version (A) then
2:
(.mm m. x) with preprocess (A) Get To
Pre-process CATO (B)
3:
Get (c) = To (Xi.pre: m,) Vie [N]
4: else
Assume:
m includes the label among other auxiliary data.
5:
Get To (r.mm))) with preprocess (B)
6:
Get (c) = To (x, m.) Vi [N]
1: return prompt 7x(x, m) that rewrites in the style of matching examples with at- tribute c, i... {xj (m,, (m,c)).
7:
end if
return Nung H that minimizes R.
8:
context notes, asks the LLM to rewrite in their style. Now given text x with auxiliary data that we wish to estimate with counterfactual value e (i.e. writing style), Ta(x, m) runs this prompt with context notes whose auxiliary data is similar
and their attribute value equals the desired
Diff-in-diff estimation. The procedure we use for medical note generation relies on additional structure involving panel data (i.e. data collected over time intervals across several individuals). In our case of clinical narratives, a narrative is usually consisted of several notes taken over the course of patient's visit and each may be written by different caregiver. Prediction is made using the release note from the hospital whose embedding consists our features x. For simplicity let us [19, 78, 79] estimation of causal
consider single note Xre taken prior to snfference-in-difffer effect is based on the parallel-trends, or constant effect assumption that two units i, with similar pre-treatment conditions would have seen the same effect had they been given the same treatment. In our case, the treatment is an assignment to a certain caregiver. Hence we assume our auxiliary data m includes Cpre. the caregiver assigned pre-treatnnnn treatment.
Assumption (constant effect). Let Xi.gre be the pre-treatment features for unit i, and assume m; includes the pre-treatment attribute Ci.pre- There exists a function p [K] M X such that (c) pie ((c, m;).
Under this assumption, calculate (c) unit for which m, = my and has Cy to estimate p(c,m.) = -x--Xpe,ji The resulting estimation procedure is given in algorithm ((B and illustrated in section 4.2.
any
T - (Progress)
T (Discharge)
Palient
Caretaker
Note
Caretaker
Note
i
X,
X
Before empirically evaluating alternatives for learning robust classifiers in our setting, and how their properties fair compared to counterfactual augmentation.
methods,
discuss
j
Xyone
X
4.3 Why Bother with Counterfactual Data Augmentation?
Panel Maiching using suxllary dataa
Reasoning about counterfactuals with problem-specific domain knowledge is considerable challenge, and it interesting to whether this has any advantage in learning robust classifiers compared to methods that rely on less stringent assumptions. A simple alternative to approximating counterfactuals involves re-weighting the loss function e.g. Makar et al. [25], Shimodaira [80]).
(dc)
X
+
X
X
81e)
X
+
X
X
Panal B: Generating counterfactual discharge summariss
Generating counterfactual clinical notes for patients using auxil- iary data with Algorithm 1(A).
Figure 2:
Reweighting baseline. Intuitively, re-weighting sam- ples from the uncorrelated distribution P(Y P(Y) by setting for each example weight w; Prain y1) Punin (C Prain) ys, C c:) and

---

minimizing the weighted empirical risk:
1 Σ wil (h(x:). m =c[m]
RSd (h)
It can be proved that at the limit of infinite data the method learns a min-max optimal hypothesis, as it also effectively minimizes R/p, (see [25]). While augmentations may not seem advantageous for identifying the correct hypothesis, reweighting can demand a larger sample to identify the correct hypothesis, particularly when Y and C are highly correlated."
Comparing sample complexities. To make this statement precise, we can apply the bounds from Cortes et al. [81] and compare them with upper bound that will derive for our method in Lemma 2. To this end, let us consider the exponent of the Rényi divergence as a measure of dependence between Y and C in the training data. The divergence is given by da,uain (Y, C)
[E Lnain (Y= = y, C c) (Y = y) Pram (C = e)] and we may derive the follow- ing bound for hypothesis e H and any ò E [0:
R2p,()) (h) Re, (h)+
V
2d2 main (Y, 11/(1/0)
dea train (Y,C) 1og(1/0)
(2)
N
N
A complementary lower bound on Rww can also be derived based on results in Cortes al. [81]. To compare this with counterfactual augmentations, denote our augmentation model by T X' x M X* which is some measurable function whose output's c-th coordinate is the counterfactual estimate w.r.t caregiver c, i.e. x(c) Tx m). The following statement quantifies the relation between the accuracy of 7()) in approximating counterfactuals and the classification accuracy of a model learned from the augmented data, via minimization of Ri (h) in eq. (1).
Lemma 2. Consider a prediction problem with a spuriousty-correlated attribute (definition 1), measurable function T X x M + x and let di (P, Q) denote the total variation distance be- tween two distributions P, Q. Further let h" Wanug denote the optimal hypotheses w.rl Rep 801 R
respectively and let Xang- [1eTT (lmmg) R2p (h")]. For any hypothesis h E H. and any 8 E (0,1) it holds that with probability at least 1 5 over the draw of the training set,
log(1/8) N
for R/o (h)
E01 (h) Rang
+ K
Σ di (Tep (Prain (x, Xmug. [[[]]
The divergence di (T,, (Pirain (X, M)), (c))) is distance between the true distribution over counterfactual instances (cX( and augmented data Te,, (Parin (X. M)). Divergences other than varlatarittion can be used, resulting in tighter bounds, e.g. see Ben-David et al. [82]. As generate better counterfactuals this divergence decreases, and it can also be shown that h and hing coincide. Hence Amug vanishes and the bound scales with N 3 resulting in gain of factor d2 draii (Y, C) over the upper bound on Ria (h) in Equation (2). We discuss the details in the ap- pendix, and Section 5 we show this empirically through simulations.
Takeaways and additional baselines. We emphasize that that the counterfactual datapoints should be interpreted as "more data" in the sense of i.id training examples, they rather embody knowl- edge about how the causal mechanism that generates features X acts under interventions on the attribute C. (as formalized in e.g. [74, 83]). This translates into an improved sample complexity towards risk minimization on PI. Counterfactuals are not the only type of causal knowledge that may be leveraged for learning more stable models. Many data dependent penalty terms have been proposed to impose conditional independence constraints drawn from the causal structure of the problem. Theory on these methods usually shows improved OOD performance under infinite data [13, 22, 24, 29). Our baselines include a method based on the Maximum-Mee Discrepency (MMD) from Makar et al. [25] who show improved sample complexity under a linear hypothesis class.
We remark that other works discuss the potential benefits of data augmentation for identification in other problem settings, e.g. [72, Thm. and [17].
'The notation Te,x (-) denotes the pushforward measure. We note that in our implementation Te is data dependent and ignore this dependence to enable simple analysis.
7

---

Experiments
We empirically study the following questions: (1) Can CATO enhance OOD performance of down- stream classifiers? (2) Does it surpass the combination of reweighting and invariance penalties? (3) Is it more effective than alternative augmentation techniques, thus demonstrating the usefulness of the causal graph? (4) How sensitive is CATO to quality counterfactuals?
These questions seek to establish muusall-motivateed augmentations as practical approach for im- proving OOD performance. We address Q#1,#2 and #3 through our theoretical foundation and across all empirical studies, while Q#4 is explored in the synthetic experiments. Further details about the experimental setup, including data statistics, model hyperparameters, and data splits, can be found in Appendix B. Table 1 provides overview of the tasks we experiment with.
Input (xz)
Label (v)
ID Data
OOD Data
Spurious Feature
auxiliary data (m)
Condition Prediction Note Segmentation
i2b2-2010
Medications,
Clinical Narratives
MIMIC-111
partner data
Caregiver ID
Lab Results,
Demographic Traits
i2b2-2006
Vitals
Restaurant Reviews
Restaurant Rating
CEBaB
CeBAB- Spurious
Food-mention
Service, Noise, Ambiance, Food
Synthetic Data
[0.1)
Gaussians
(0. 7
Table 1: Description of all our tasks and their corresponding experimental setup.
Baselines. We compare CATO to several baselines:
Observational Baseline model trained on the original data. PubMED BERT [84] for clinical
narratives, logistic regression for the restaurant reviews and synthetic experiments. 5
Reweighting Baseline model with sample reweighting as in Makar al. [25].
MMD Baseline model with an MMD penalty as in Makar et al. [25]. Veitch et al. [29].
IRM Baseline model with the IRMv1 penalty as in Arjovsky et al. [13].
GroupDRO- - Baseline model trained with the GroupDRO objective as in Sagawa et al. [85].
Naive Augmentations Baseline model on a dataset that also includes augmentations, generated by prompting an LLM to create more examples (without matching or diff-in-diff).
Conditional Augmentations Augmentations are generated by matching on auxiliary data and prompting an LLM to create one example in the the style of the other.
The reweighting and MMD approaches are discussed and contrasted to counterfactual augmentation in Section 4. IRM and GroupDRO are the most well-known principled methods for OOD generaliza- tion that are used in the literature. The augmentation approaches are compared here demonstrate the importance of using the causal structure of the data.
5.1 Clinical Narratives
Data. We consider three representative clinical NLP tasks, clinical condition prediction, note seg- mentation and demographic traits identification". for which we have both ID and OOD data. We utilize several electronic health records (EHR) datasets. We train on MIMIC-III [86], widely-used medical dataset containing over 2 million notes from 38, 597 adult patients, 49, 785 hospital admis- sions, and 3, 500 healthcare professionals between 2001 and 2012. MIMIC-III is commonly used in NLP research for clinically-reltted tasks and for pre-training language models for the medical i2b2 2006 and 2010 competitions our held-out hospital dataset. In the note segmentation task, we use private held-out data.
domain [87]. When available,
Generating notes from counterfactual caregivers. To generate augmentations, we select care- givers with multiple patients and notes for more than one patient. For each caregiver-paticn pair where both their last progress note and discharge summary were written by that caregiver', we match them to similar patients having the same initial caregiver but a different one for their discharge sum- mary. In matching, we select patients with similar medications and lab results (denoted patient's
Appendix B includes results where the Baseline model is also BioBERT, SentenceBERT or GPT3. See Appendix for results on the demographic traiss identification task. 'During patient's stay, progress notes capture its current state. When leaving the hospital, discharge summary is written.
8

---

auxiliary data m in Table 1). We then generate counterfactual discharge summaries for matched pa- tients using Algorithm ((A) and train the model using original data and generated counterfactuals.
Figure presents results for CATO (A) using language model representations generated using these matched examples. See Appendix B for training details and results for CATO (A) with LLM prompts, and Appendix C for synthetic note examples and the prompts used.
(A) Clinical Conditions
(B) Note Segmentation
90
90
80
80
F1
70
70
60
60
ID (MIMIC-II)
OOD (12h2-2010)
ID (MIMIC-II III)
OOD (Privare Held-Out)
Observational I+ Reweighting D1++ MMD IRM lI GRo I N Aug. 1CAAOO (A)
Figure 3: Results (F1 averaged across runs) for predicting clinical conditions (A) and for clinical note segmentation (B) from the text narratives. CATO (A) outperforms all baselines on OOD data.
Clinical Condition Prediction. Clinical condition prediction is concept extraction task focused on medical concepts in patient reports [88]. Here we trained PubMED BERT models on a subset of MIMIC-III labelled using the same annotation guidelines as in i2b2-2010, the OOD dataset the models are tested on. As can be seen in the Figure 3(A), in the ID setting only the naive augmen- tations improve performance slightly. In the OOD setting, all OOD methods help (reweighting, MMD, IRM, GroupDRO, CATO (A)), but our matsally augmentation approach is substan- tially better than the alternatives. On average (across 5 runs), CATO (A) improves precision above the baseline by more than 7% (absolute), and recall by more than 8%. The naive augmentation ap- proach improves over the vanilla PubMED BERT model, but outperformed by all OOD methods.
Note Segmentation. In this task, models need to recognize sections in free-form clinical notes [89]. Given that section headers vary between hospitals, the models must discern sections based solely on the note content, excluding headers. As can be seen in Figure 3(B), similarly to clinical condi- tion prediction, the diff-in-diff approach to augmentations (CATO (A)) substantially improved OOD) performance, and as expected does not help ID. The naive augmentations are the best performing method ID, but is again outperformed by all other methods OOD.
5.2 Restaurant Reviews
Data. We use the CEBaB dataset [49], which consists of short restaurant reviews and ratings from OpenTable, including evaluations for food, service, noise, ambiance, and an overall rating. We used the train-exclusive split of the dataset, which contains 1, 755 examples. We construct two experimental settings: the original CeBAB dataset, and a modified version, denoted Ce- BAB-Spurious, where there's a spurious correla- tion between training and deployment.
Method
CeBAB
CeBAB-Spur.
Observational Reweighting Naive Aug.
0.85
0.64
0.84
0.68
0.80
0.62
Conditional Aug.
0.84
0.70
CATO (B)
0.84
0.75
Accuracy on CeBAB and CeBAB- Spurious. CATO (B) outperforms all baselines when we introduce a spurious correlation.
Table 2:
To construct CeBAB-Spurious, we leverage the availability of both the original and perceived rat- ings for each review in CeBAB. The original rating represents the reviewer's initial thoughts when writing the review, while the perceived rating indicates whether the review contains information

---

about various restaurant attributes (e.g., food, service, noise, ambiance) and their associated senti- ment. We utilize this unique data structure to capture reviewers' writing styles. Some reviewers are concise and provide limited descriptions, while others are more descriptive and include more infor- mation. To incorporate this variability, we introduce a new attribute called food-mention to signify the presence of food-related information in review. If the perceived food rating is either negative or positive, we assign a value to the food-mention attribute; otherwise, it is set to O. We subsample the data such that there is correlation of 0.72 between food-mention and the outcome.
Generating reviews with counterfactual food mentions. Following Algorithm 1, we generate counterfactual restaurant reviews conditional on food and overall ratings. We find matched examples for each review, select those with different food-mentions and prompt an LLM rewrite them, reflecting how the reviews would appear if the reviewer was more/less concise.
Results. As shown in Table 2, adding counterfactual augmentations leads better OOD general- ization, while naive data augmentation hurts model performance In line with the sample complexity argument in Section conditional augmentation effectively doesn't 't add new data and therefore doesn't improve model performance.
5.3 Synthetic Data
To test sensitivity of CATO to quality of counterfactuals (Q#4), generate synthetic data for binary classification problem where K 8 (cardinality of C). We sample P(C Y) to simulate varying degrees of spurious correlations. Then we draw x = [x" X from a Gaussian distribution,
0 o
X;
x Xspu,à
NV[
ply Mc.
Idl
spa
In this case (c) is obtained by adding Me- Pee Xspo.i- To corrupt augmentation, in- stead add & (pe Pe,) where & is drawn from truncated Gaussian centered A (0, 1). We train models with fixed sample size (in the appendix we also examine varying sample sizes and additional types of corruption) and evaluate the trained models' accuracy on P1 to examine the interplay between spurious corre- lation strength (measured by mutual informa- tion T(Y; C)), and counterfactual augmenta- tion quality. As can be seen Figure 4, corrup- tions degrade performance under stronger spu- rious correlations, though a strong corruption is required for reweighting to become preferable.
0.9
0.8
BRM
0.7
Reweighting
Aug. A= 0.2
A
0.6
A= 0.4
0
0.2
0.4
0.6
0.8
I(Y:C)
Figure 4: OOD accuracy (1 - R/p, (h)) and Y,C correlation strength (I(Y:))). Lower values of À correspond to stronger corruptions of the aug- mentations. Even with substantial corruption (A = 0.2) and strong correlation, augmentations outper- form baselines.
6 Discussion
In this work, we have presented data augmen- tation approach based on the causal structure of auxiliary data for improving OOD generaliza- tion, specifically focusing text classification tasks. However, our approach is not without limitations. The validity of our assumptions, the specification of the causal graph and the qual- ity of the counterfactual approximation all present challenges to address in future work. Further, our results suggest that performing data augmentation in an unprincipled manner can also hurt model performance. Utilizing additional techniques for OOD generalization. learning the causal structure directly from the data, and improving quality and reliability of the counterfactual approximation process can help mitigate these concerns. Overall, we believe that motivvate data aug- mentation methods like ours can help address challenges developing robust and reliable machine learning systems, particularly in safety-critical applications.
10

---

References
[1] John R Zech, Marcus A Badgeley, Manway Liu, Anthony B Costa, Joseph J Titano, and Eric Karl Oermann. Variable generalization performance of deep learning model detect pneumonia in chest radiographs: a cross-sectional study. PLoS medicine, 2018.
[2] Alex DeGrave, Joseph D Janizek, and Su-In Lee. Ai for radiographic covid-19 detection selects shortcuts over signal. Nature Machine Intelligence, 3(7):610-619)) 2021.
[3] Thomas McCoy, Ellie Pavlick, and Tal Linzen. Right for the wrong reasons: Diagnosing syntactic heuristics in natural language inference. arXiv preprint arXiv: 1002000000000077.. 2019.
[4] Rich Caruana, Yin Lou, Johannes Gehrke, Paul Koch, Marc Sturm, and Noemie Elhadad. Intelligible models for healthcare: Predicting pneumonia risk and hospital 30-day readmission. In Proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining, pages 1721-1730 2015.
[5] Joaquin Hainonero-Candele Masashi Sugiyama, Anton Schwaighofer, and Neil D Lawrence. Dataset shift in machine learning. Mit Press, 2008.
[6] Adarsh Subbaswamy, Peter Schulam, and Suchi Saria. Preventing failures due to dataset shift: Learning predictive models that transport. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 31 5118-3127. PMLR, 2019.
[7] Samuel G Finlayson, Adarsh Subbaswamy, Karandeep Singh, John Bowers, Annabel Kupke, Jonathan Zittrain, Isaac S Kohane, and Suchi Saria. The clinician and dataset shift in artificial intelligence. New England Journal of Medicine, 388(3))233-286, 2021.
[8] Peter Spyns. Natural language processing in medicine: an overview. Methods of information in medicine, 35(04/05):85-300. 1996.
[9] Li Zhou and George Hripesak. Temporal reasoning with medical data review with emphasis on medical natural language processing. Journal of biomedical informatics, 2007.
[10] Stephen WU, Kirk Roberts, Surabhi Datta, Jingcheng Du Zongcheng Ji, Yugi Si, Sarvesh Soni, Qiong Wang, Qiang Wei, Yang Xiang, et al. Deep learning clinical natural language processing: a methodical review. Journal of the American Medical Informatics Association, 27(3):457-470, 2020.
[11] Jonas Peters, Peter Bühlmann, and Nicolai Meinshausen. Causal inference by using invariant prediction: identification and confidence intervals. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 2016.
[12] Sara Magliacane, Thijs Ommen, Tom Claassen, Stephan Bongers, Philip Versteeg, and Joris M Mooij. Domain adaptation by using causal inference to predict invariant conditional distributions. In Neural Information Processing Systems (NeurIPS), pages 10869-10879, 2018.
[13] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint arXiv: /907.02893, 2019.
[14] Adarsh Subbaswamy, Bryant Chen, and Suchi Saria. A unifying causal framework for analyz- ing dataset shift-stable learning algorithms. Journal of Causal Inference, 10(1):64-8,, 2022.
[15] Alexander Robey, George J. Pappas, and Hamed Hassani. Model-based domain generaliza- tion. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Neural Information Processing Systems (NeurIPS), 2021. URL htttpp://openreviee..ee.. id=JOx89h40A-1.
[16] Huaxiu Yao, Yu Wang, Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Im- proving out-s--distiibuttoon robustness via selective augmentation. In International Conference on Machine Learning, pages 25407-25437. PMLR, 2022.
11

---

[17] Irena Gao, Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto. and Percy Liang. Out-of- domain robustness via targeted augmentations. arXiv preprint arXiv:2302. 18861, 2023.
[18] Murat Kocaoglu, Christopher Snyder, Alexandros G. Dimakis, and Sriram Vishwanath. CausalGAN: Learning causal implicit generative models with adversarial training. In Inter- national Conference on Learning Representations, 2018. URL htwph:t/t///// /ttttp:://ppereviee...
forumf2dd-BJR-4x00U
[19] Alberto Abadie. Semiparametric differeecee-n-diiffeeeeee estimators. The review of economic studies, 72(1):1-19, 2005.
[20] Christina Heinze-Demi. Jonas Peters, and Nicolai Meinshausen. Invariant causal prediction for nonlinear models. Journal of Causal Inference, 6(2), 2018.
[21] Ya Li, Xinmei Tian, Mingming Gong, Yajing Liu, Tongliang Liu, Kun Zhang, and Dacheng Tao. Deep domain generalization via conditional invariant adversarial networks. In Proceed- ings of the European conference on computer vision (ECCV), pages 624-639, 2018.
[22] Yoav Wald, Amir Feder, Daniel Greenfeld, and Uri Shalit. On calibration and out-of-domain generalization. Neural Information Processing Systems (NeurIPS), 34:2215-223 2021.
[23] David Krueger, Ethan Caballero, foern-Henrie Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distribution generalization via risk extrapolation (rex). In International Conference on Machine Learning, pages 5815-5826. PMLR, 2021.
[24] Aahlad Manas Puli, Lily H Zhang, Eric Karl Oermann, and Rajesh Ranganath. Out-of- distribution generalization in the presence of muisance-induced spurious correlations. In In- ternational Conference on Learning Representations, 2022. URL https://openrevierr net/forumiid 122o2222221...2..2.
[25] Maggie Makar, Ben Packer, Dan Moldovan, Davis Blalock, Yoni Halpem, and Alexander D' Amour. Causally motivated shortcut removal using auxiliary labels. In International Con- ference on Artificial Intelligence and Statistics, pages 739-766. PMLR, 2022.
[26] Yibo Jiang and Victor Veitch. Invariant and transportable representations for anti-causal do- main shifts. arXiv preprint arXiv:2207.00 2000000,7,,,,,, 2022.
[27] Claudia Shi, Victor Veitch, and David M Blei. Invariant representation learning for treatment effect estimation. In Uncertainty in Artificial Intelligence, pages 1546-1555. PMLR, 2021.
[28] Mingzhang Yin, Yixin Wang, and David M Blei. Optimization-basel causal estimation from heterogenous environments. arXiv preprini arXiv:2109. 11900, 2021.
[29] Victor Veitch, Alexander D' Amour, Steve Yadlowsky, and Jacob Eisenstein. Counterfactual in- variance to spurious correlations in text classification. Neural Information Processing Systems (NeurIPS), 14:16166-16208 2021.
[30] Yana Dranker, He He, and Yonatan Belinkov. Irm-when it works and when it doesn't: A test case of natural language inference. Advances in Neural Information Processing Systems, 34: 19212-18224, 2021.
[31] Amir Feder, Katherine A Keith, Emaad Manzoor, Reid Pryzant, Dhanya Sridhar, Zach Wood- Doughty, Jacob Eisenstein, Justin Grimmer, Roi Reichart, Margaret E Roberts, et al. Causal inference in natural language processing: Estimation, prediction. interpretation and beyond. Transactions of the Association for Computational Linguistics, 10: 138-1158, 2022.
[32] Amir Feder, Guy Horowitz, Yoav Wald, Roi Reichart, and Nir Rosenfeld. In the eye of the beholder: Robust prediction with causal user modeling. In Neural Information Processing Systems (NeurIPS), 2022.
[33] Pritish Kamath, Akilesh Tangella, Danica Sutherland, and Nathan Srebro. Does invariant risk minimization capture invariance? International Conference on Artificial Intelligence and Statistics, pages 4069-4077. PMLR, 2021.
12

---

[34] Elan Rosenfeld, Pradeep Ravikumar, and Andrej Risteski. The risks of invariant risk mini- mization, arXiv preprint arXiv:2010.05761, 2020.
[35] Ruocheng Guo, Pengchuan Zhang, Hao Liu, and Emre Kiciman. Out-of-distribution pre- diction with invariant risk minimization: The limitation and an effective fix. arXiv preprint arXiv:2101.07732 2021.
[36] Yoav Wald, Gal Yona, Uri Shalit, and Yair Carmon. Malign overlitting: Interpolation can provably preclude invariance. arXiv preprint arXiv:2211.15724 2022.
[37] Divyansh Kaushik, Eduard Hovy, and Zachary C Lipton. Learning the difference that makes a difference with counterfactuully-auuggee
data. arXiv preprint (9ti:::999..243339 2019.
[38] Divyansh Kaushik, Amrith Setlur, Eduard Hovy, and Zachary C.iptton, Explaining the efficacy of wwwperrrfactaalllaggeeeeeee data. arXiv preprint arXiv:2010.02/ 14, 2020.
[39] Sahaj Garg, Vincent Perot, Nicole Limtiaco, Ankur Taly, Ed H Chi, and Alex Beutel. Counter- factual faireess text classification through robustness. In Proceedings of the 2019 AAAI/ACM 2019.
Conference on AI, Ethics, and Society, pages
[40] Rohan Jha, Charles Lovering, and Ellie Pavlick. Does data augmentation improve generaliza- tion in nip? arXiv preprint arXiv:2004. 15012, 2020.
[41] Matt Gardner, Yoav Artzi, Victoria Basmov, Jonathan Berant, Ben Bogin, Sihao Chen, Pradeep Dasigi, Dheeru Dua, Yanai Elazar, Ananth Gottumukkala, Nitish Gupta, Hannaneh Hajishirzi, Gabriel Ilharco. Daniel Khashabi, Kevin Lin, Jiangming Liu, Nelson F. Liu, Phoebe Mulcaire, Qiang Ning. Sameer Singh, Noah A. Smith, Sanjay Subramanian, Reut Tsarfaty, Eric Wallace, Ally Zhang, and Ben Zhou. Evaluating models' local decision boundaries via contrast sets, In Findings of the Association for Computational Linguistics: EMNLP 2020, pages 1307-1323, Online, November 2020. Association for Computational Linguistics. doi: 10. (0.18653/v/20000. findings-emmlp.111 17. URL https::::: http//cacnncootoccoco/oo httpooooogoorg//200000 findings--eiinn emmlp.....pp..... 117.
[42] Ravi Shekhar, Sandro Pezzelle, Yauhen Klimovich, Aurélic Herbelot, Moin Nabi, Enver Sangineto, and Raffaella Bernardi. FOIL it! find one mismatch between image and lan- guage caption. In Proceedings of the 55th Annual Meeting of the Association for Com- putational Linguistics (Volume 1: Long Popers), pages 255-265, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10. 10.18653/11/P11-10044 URL https: // //aclanthology. wwtp://ooooy...o/111-//1 org/P17-1022
[43] Amir Feder, Nadav Oved, Uri Shalit, and Roi Reichart. Causalm: Causal model explanation through counterfactual language models. Computational Linguistics, 47(2))333-332 -3633666366633 2021.
[44] Ran Zmigrod, Sabrina J. Mielke, Hanna Wallach, and Ryan Cotterell. Counterfactual data augmentation for mitigating gender stereotypes in languages with rich morphology. In Pro- ceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 1651-1661, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/199-1100 URL https: h/alaanthol//l hwwphhooooyooog/911-111 11.//199111..
[45] Parker Riley, Noah Constant, Mandy Guo, Girish Kumar, David Uthus, and Zarana Parekh. Textsettr: Label-free text style extraction and tunable targeted restyling. arXiv preprint arXiv:2010.03802, 2020.
[46] Tongshuang Wu, Marco Tulio Ribeiro, Jeffrey Heer, and Daniel Weld. Polyjuice: Automated, gererae-purpose counterfactual generation. arXiv preprint wtX.i:2101.002250 2021.
[47] Chengzhi Mao, Augustine Cha, Amogh Gupta, Hao Wang, Junfeng Yang, and Carl Vondrick. Generative interventions for causal learning. In Proceedings of the IEEF/VV Conference on Computer Vision and Pattern Recognition, pages 3947-3956. 2021.
[48] Daniel Rosenberg, Itai Gat, Amir Feder, and Roi Reichart. Are vqa systems rad? measuring robustness to augmented data with focused interventions. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume 2: Short Papers), pages 61-70, 2021.
13

---

[49] Eldar D Abraham, Karel D' D'Oosterlinck, Amir Feder, Yair Gat, Atticus Geiger, Christopher Potts, Roi Reichart, and Zhengxuan Wu. CEBaB: concepts on NLP model behavior. Neural Information Processing Systems (NeurIPS), 35: 17582-17596, 2022.
Estimating the causal effects of real-world
[50] Zhengxuan Wu, Karel D'Oosterlinck, Atticus Geiger, Amir Zur, and Christopher Potts. Causal proxy models for concept-based model explanations. In International Conference on Machine Learning, pages 37313-37334. PMLR, 2023.
[51] Nitish Joshi and He He. An investigation of the (in) effectiveness of counterfactually aug- mented data. In 60th Annual Meeting of the Association for Computational Linguistics, ACL 2022, pages 3668-3681. Association for Computational Linguisties (ACL), 2022.
[52] Maria Antoniak and David Mimno. Bad seeds: Evaluating lexical methods for bias measure- ment. In Proceedings of the 59th Annual Meeting of the Association for Computational Lin- guistics and the 1 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 1889-1904, Online, August 2021. Association for Computational Lin- guistics. doi: 10. 10.18653/v1/202 Lacl-llng. 148. URL https: //ap:::///aaaaaaa hattp:://allnthoooogy acl-long. 148.
anthology org/2021
[53] Xiaoling Zhou and Ou Wu. Implicit counterfactual data augmentation for deep neural net- works. arXiv preprint TDt:2304.133311 2023.
[54] Nitay Calderon, Eyal Ben-David, Amir Feder, and Roi Reichart. DoCoGin: Domain Counter- factual Generation for Low Resource Domain Adaptation. In Proceedings of the 60th Annual Meeting of the Association of Computational Linguistics (ACL., 2022.
[55] Kory Kreimeyer, Matthew Foster, Abhishek Pandey, Nina Arya, Gwendolyn Halford, Sandra F Jones, Richard Forshee, Mark Walderhaug, and Taxiarchis Botsis. Natural language processing systems for capturing and standardizing unstructured clinical information: a systematic review. Journal of biomedical informatics, 73:14-29, 2017.
[56] Özlem Uzuner. Recognizing obesity and comorbidities in sparse data. Journal of the American Medical Informatics Association, 16(4):561-570, 2009.
[57] Guergana K Savova, James J Masanz, Philip V Ogren, Jiaping Zheng, Sunghwan Sohn, Karin C Kipper-Schuler, and Christopher G Chute. Mayo clinical text analysis and knowl- edge extraction system (ctakes): architecture, component evaluation and applications. Journal of the American Medical Informatics Association, 17(5):507-513, 2010.
[58] Peter B Jensen, Lars J Jensen, and Søren Brunak. Mining electronic health records: towards better research applications and clinical care. Nature Reviews Genetics, 13(6):395-400) 2012.
[59] Elizabeth Ford, John A Carroll, Helen E Smith, Donia Scott, and Jackie A Cassell. Extracting information from the text of electronic medical records improve case detection: systematic review. Journal of the American Medical Informatics Association, 23(5): 2016.
[60] Henghui Zhu, Ioannis Ch Paschalidis, and Amir Tahmasebi. Clinical concept extraction with contextual word embedding. arXiv preprint arXiv:1810.1055 2018.
[61] Yifan Peng, Shankai Yan, and Zhiyong Lu. Transfer learning biomedical natural language processing: an evaluation of bert and elmo on ten benchmarking datasets. arXiv preprint arXiv:1906.054477 2019.
[62] Vikas Yadav and Steven Bethard. survey on recent advances in named entity recognition from deep learning models. arXiv preprint arXiv: 1470, 2019.
[63] Yugi Si, Jingqi Wang, Hua Xu, and Kirk Roberts. Enhancing clinical concept extraction with contextual embeddings. Journal of the American Medical Informatics Association, 26(11): 1297-1304, 2019.
[64] Jinhyuk Lee, Wonjin Yoon, Sungdong Kim, Donghyeon Kim, Sunkyu Kim, Chan Ho So, and Jaewoo Kang. BioBERT: a pre-trained biomedical language representation model for biomedical text mining. Bioinformatics, 16(4):1234-1240, 2020.
14

---

[65] Dmitri Roussinov, Andrew Conkie, Andrew Patterson, and Christopher Sainsbury. Predicting clinical events based on raw text: from bag-of-words attention-baun -baiid transformers. Frontiers in Digital Health, 3:214, 2022.
[66] Tom M Seinen, Egill A Fridgeirsson, Solomon Ioannou, Daniel Jeannetot, Luis H John, Jan A Kors, Aniek F Markus, Victor Pera, Alexandros Rekkas, Ross D Williams, et al. Use of unstructured text in prognostic clinical prediction models: a systematic review. Journal of the American Medical Informatics Association, 29(7):1292-1300 2022.
[67] Karan Singhal, Shekoofeh Azizi, Tao Tu, S Sara Mahdavi, Jason Wei, Hyung Won Chung, Nathan Scales, Ajay Tanwani, Heather Cole-Lewis, Stephen Pfohl, et al. Large language models encode clinical knowledge. arXiv preprint arXiv:2212.131538 2022.
[68] John W Ayers, Adam Poliak, Mark Dredze, Eric Leas, Zechariah Zhu, Jessica B Kelley, Dennis J Faix, Aaron M Goodman, Christopher A Longhurst, Michael Hogarth, et al. Com- paring physician and artificial intelligence chatbot responses to patient questions posted public social media forum. JAMA Internal Medicine, 2023.
[69] Amir Feder, Itay Laish, Shashank Agarwal, Uri Lerner, Avel Atias, Cathy Cheung, Peter Clardy, Alon Peled-Cohen, Rachana Fellinger, Hengrui Liu, al. Building a clinically-focused problem list from medical notes. In Proceedings of the 13th International Workshop on Health Text Mining and Information Analysis (LOUHH), pages 60-68, 2022.
[70] Fan Zhang, Itay Laish, Ayelet Benjamini, and Amir Feder. Section classification in clinical notes with multi-task transformers. In Proceedings of the 13th International Workshop on Health Text Mining and Information Analysis (LOUHI), pages 54-59, 2022.
[71] Amir Feder, Danny Vainstein, Roni Rosenfeld, Tzvika Hartman, Avinatan Hassidim, and Yossi Matias. Active deep learning to detect demographic traits in free-form clinical notes. Journal of Biomedical Informatics, 107:103436, 2020.
[72] Zihao Wang and Victor Veitch. A unified causal view of domain invariant representation learn- ing. arXiv preprint arXiv:2208.00 06987, 2022.
[73] Guido W Imbens and Jeffrey M Wooldridge. Recent developments in the econometrics of program evaluation. Journal of economic literature, 47(1):5-86, 2009.
[74] Judea Pearl. Causality. Cambridge university press, 2009.
[75] Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. International conference on machine learning, pages 9076-3085.1 PMLR, 2017.
[76] Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observa- tional studies for causal effects. Biometrika, 70(1):41-55, 1983.
[77] Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sciences. Cambridge University Press, 2015.
[78] David Card and Alan B Krueger. Minimum wages and employment: A case study of the fast food industry in new jersey and pennsylvania, 1993.
[79] Joshua D Angrist and
Pischke. Mostly harmless econometries: An empiricist's
companion. Princeton university press, 2009.
[80] Hidetoshi Shimodaira. Improving predictive inference under covariate shift by weighting the log- ligeliidl function. Journal of statistical planning and inference, 90(2):227-24: 2000.
[81] Corinna Cortes, Yishay Mansour, and Mehryar Mohri. Learning bounds for importance weighting. In J. Lafferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta, editors, Neural Information Processing Systems (NeurIPS), volume 23, Curran Associates, Inc., 2010.
[82] Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jen- nifer Wortman Vaughan. theory of learning from different domains. Machine learning, 79:151-175, 2010.
15

---

[83] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foum- dations and learning algorithms. The MIT Press, 2017.
[84] Yu Gu, Robert Tinn, Hao Cheng, Michael Lucas, Naoto Usuyama, Xiaodong Liu, Tristan Naumann, Jianfeng Gao, and Hoifung Poon. Domain-specific language model pretraining for biomedical natural language processing. ACM Transactions Computing for Healthcare (HEALTH), 3(1):1 23, 2021.
[85] Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2019.
[86] Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G Mark. Mimic- iii, freely accessible critical care database. Scientific data, 3(1): 2016.
[87] Emily Alsentzer, John R Murphy, Willie Boag, Wei-Hung Weng, Di Jin, Tristan Naumann, and Matthew McDermott. Publicly available clinical bert embeddings. arXiv preprint arXiv:1904.035 033233, 2019.
[88] Özlem Uzumer, Brett R South, Shuying Shen, and Scott L DuVall. 2010 i2b2/va challenge on concepts, assertions, and relations in clinical text. Journal of the American Medical Informatics Association, 1865)5522-556, 2011.
[89] Alexandra Pomares- Saaires-uummbaaaa..m Markus Kreuzthaler, and Stefan Schulz. Current approaches to identify sections within clinical narratives from electronic health records: a systematic review. BMC medical research methodology, 19:1-20, 2019.
[90] Terence Tao. An introduction to measure theory, volume 126. American Mathematical Soc., 2011.
[91] Mehryar Mohri, Afshin Rostamizadeh, and Ameet Talwalkar. Foundations of machine learn- ing. MIT press, 2018.
[92] Koby Crammer, Michael Kearns, and Jennifer Wortman. Learning from multiple sources. Journal of Machine Learning Research, 9(8), 2008
[93] Aahlad Puli, Nitish Joshi, He He, and Rajesh Ranganath. Nuisances via negativa: Adjusting for spurious correlations via data augmentation. arXiv preprint arXiv:22 2022.
[94] Aahlad Puli, Lily Zhang, Yoav Wald, and Rajesh Ranganath. Don't blame dataset shift! short- cut learning due to gradients and cross entropy. arXiv preprint 2rXiv:2308.1255.. 2023.
[95] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 180.04805, 2018.
[96] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gre- gory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Al- ban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetioo Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024 8035. Curran Associates, Inc., 2019. URL http://papers. nww..ps. cc/paper/ 2015-pytorch-an an- pdf.
neurips.cc/ppper
Learning- -ibngry.
[97] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's transform- ers: State-of-the-art natural language processing. arXiv preprint arXiv: 191.00/71000771 2019.
[98] Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, et al. Scikit- learn: Machine learning in python. Journal of machine learning research, 12((ct):2825-22 2830, 2011.
[99] OpenAl. Gpt-4 technical report, 2023,
16


Reviews:
[{"review_id": "3eGZYFwCfP", "reviewer": "Reviewer_MXTx", "summary": "In this work, the authors develop causally-driven data augmentation methods to improve model robustness.", "questions": "see above", "limitations": "The authors do acknowledge limitations of their work.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "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.\n\nI 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful and positive review. We are happy to see that you think the paper is a very interesting first step in utilizing recent advances in LLMs for OOD generalization, which is exactly our goal in this work.\n\nThe points for improvement that you lay out are very helpful, and we comment on them below.\n\n**Limitations of LLMs**: following this comment, we will add to the appendix a discussion on the possible limitations of LLMs for our approach. Since reviewer 8nR3 also asked about possible limitations of LLMs, we included a summary of this discussion in the general rebuttal and we repeat it below for convenience.\n* *LLM generation quality*: We acknowledge that LLMs vary in their ability to generate realistic text. It is possible that LLMs introduce biases into our problem, inherited from their own training data. This requires further study, however from our manual examination we found their quality satisfactory (see Appendix C for generation examples) and that OOD generalization also improved for models trained on the augmented data they generate. We'll include this analysis and results with different LLMs, namely Bio-BERT, Sentence-BERT and GPT4, in the revised paper.\n\n* *Counterfactual approximation*: Other than generation quality, the additional challenge in using LLMs for CDA is our ability to elicit a good approximation to the counterfactual text. Our methods rely on principles from causal inference to advance disciplined approaches for this task. While further studies are required (e.g. systematically comparing small sets of manual re-writes of texts to the elicited LLM output), we view our work as a promising first step in this direction, which we expect to be significantly extended and improved in future work.\n\n* *Effect of biases on OOD generalization*: Since we focus on OOD generalization, the limitations and possible biases mentioned above must be weighed within this context. Namely, we should bear in mind that even though generation may be biased, this bias is only harmful when it affects the generalization of a downstream classifier, and this is what we evaluate. Further, in OOD generalization we consider cases where the training data is biased in the first place, and training a standard predictive model also results in a biased solution. Hence we must weigh risks and limitations of alternative solutions vs. those of LLMs.\n\n\n**Auxiliary information may not be enough to assist a model**: As you correctly point out, we cannot be sure that the auxiliary information is enough to account for all the unobserved features relating to the writer and the text. One step we can take to further validate our method in real world settings is to obtain a small dataset where caretakers rewrite other caretakers’ existing notes in their own style (similarly to the counterfactual data augmentation suggested in Kaushik et al. 2020) and compare these against the synthetic estimates generated by the LLM. While such an evaluation would typically be of small scale, it offers another check for the validity of the approach.\n\nAnother point worth mentioning on this topic is that our end goal is to obtain a classifier that generalizes well in an OOD setting. As suggested by our generalization bound in lemma 2 of the paper, this requires distributional similarity between the distribution of counterfactuals and the distribution of our estimates. So even if our estimation methods are imperfect, they may still suffice to achieve better generalization w.r.t the baselines. Understandably, this is why we view the accuracy on both ID and OOD test data as the main evaluation for our method and the baselines.\n\nThank you again for your effort and help in improving our paper.\n\n**References**:  \n1. Kaushik et al. 2020 Learning the Difference that Makes a Difference with Counterfactually-Augmented Data"}, {"review_id": "mRrXaFh5K8", "reviewer": "Reviewer_YqAy", "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.\n\nSince 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. \n\nAdditionally, 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.", "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.\n\nIn 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?\n\nDoes one need to know whether the label is spuriously correlated with an attribute in advance?\n\nIs 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.", "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?\n\nAnother issue is how can we guarantee that the counterfactual instances generated during augmentation are realistic?", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your review and questions, and are eager to clarify all points regarding our data, baselines, and assumptions.\n\nIn response to questions:  \n**Causal vs. anti-causal prediction problem**: As you correctly suggest in the review, we solve an anti-causal prediction problem. A true condition in the world, $Y$, along with the identity of the caretaker, $C$, determines the text we observe $X$. Recovering $Y$ is thus an anti-causal prediction problem.  \nSome additional points regarding this setting: The causal model we use has several roles.\n* To provide a coarse description of the real world problem motivating this work, namely anti-causal prediction on clinical notes. We say the description is coarse since we omit (on purpose) some variables from the graph, like the auxiliary data $M$, since one can consider multiple possible graphical structures that include it where our method is adequate (we will elaborate on this in the appendix of our revision).  \n* Another role of this setting is to motivate Counterfactual Data Augmentation (CDA), showing that given infinite training data it retrieves the min-max optimal solution for the problem. Note that Wang & Veitch 23 formalize a family of graphical structures they call “purely spurious”, and show CDA retrieves a min-max optimal predictor in such problems. Hence our methods are well motivated under any of these structures (which also include some causal prediction problems), and we choose the one most adequate for our problem of interest.  \n* Finally, the setting allows for a formal comparison in terms of sample complexity, between approximate data augmentation and a reweighting baseline that also provably solves the problem with enough data.\n\n**Baselines**:\n1. Following our model for the problem in Fig. 1, we choose baselines that are appropriate for this task since they (a) are proposed in a paper that studies this anti-causal prediction problem, and (b) achieve state-of-the-art performance on it. \n2. We have experimented with other baselines on the medical notes tasks and following your comment we will include results on GroupDRO and IRMv1 in the paper. These are arguably the two most popular approaches in the literature, and the results can be found in the PDF attached to the rebuttal. IRMv1 underperforms the baseline from our original submission and GroupDRO is comparable, while CATO (A) is still the best performing method.\n\nWe refrained from including these baselines in the original submission, since previous work already documented their limited success on most large scale problems (e.g. [Rosenfeld et al. 20, Kamath et al. 21, Gulrajani and Lopez-Paz 20]) and we did not want to make that the focus of our work. The medical notes task is far larger in scale than datasets where these methods have shown significant gains, and devising methods that show gains in large scale problems was among our main motivations for this work. Following your comment, we see why it is important to include these results, and it led us to revise our decision. Thank you for helping us improve this aspect of the paper.\n\n**Knowing whether at attribute is spurious**: Our setup is suitable for problems where we have some domain knowledge, and we know that the correlation between the attribute and label is spurious (whether or not a large correlation exists in training data is something we can estimate numerically). This is a standard assumption in many works on spurious correlations (Makar et al. 22, Puli et al. 22, to name a couple). Nonetheless, we think that discovery of spurious correlations is an important direction for future work.\n\n**Private held-out data**: In all clinical data experiments, we use notes from different hospitals as the held-out data. Whenever possible, we use publicly available datasets (i2b2 competitions) that have the same structure as MIMIC-III. Whenever such data is unavailable we use data from real hospitals. While we can’t publish this data for privacy reasons, it also has the same structure as MIMIC-III and i2b2.\n\nIn response to limitations:  \n**Knowledge of causal graphs**: We agree that to justify our methods some strong assumptions are required, as is often the case when a method involves causal estimation. However, we note that the causal graph assumed in figure 1 is rather simple to reason about, and as explained earlier CDA is also suitable for several other graphs (see Wang & Veitch 23). The strongest assumptions we make in our view are the constant effect assumption for CATO (A), and that auxiliary data M accounts for all unobserved factors for CATO (B). Following this comment, we will add a section to the appendix that describes these assumptions much more extensively and discusses the possible limitations they pose.\n\n**Guaranteeing realistic counterfactual**: While we cannot guarantee realistic counterfactuals with total confidence, we can take some measures such as letting caretakers rewrite a handful of medical notes taken by other caretakers (in the spirit of what’s proposed in Kaushik et al. 2020) and compare these to the ones generated by the LLM. The scale of this validation may be too small to certify realistic counterfactuals, but it could be valuable. In addition, it would be interesting to complement this in future work with sensitivity analysis and uncertainty estimation to gain a better sense of the possible inaccuracies of the counterfactuals.\n\n**References**:\n1. Makar et al. 22, Causally motivated Shortcut Removal Using Auxiliary Labels\n2. Puli et al. 22, Out-of-distribution Generalization in the Presence of Nuisance-Induced Spurious Correlations\n3. Rosenfeld et al. 20, The Risks of Invariant Risk Minimization\n4. Kamath et al. 21, Does Invariant Risk Minimization Capture Invariance?\n5. Gulrajani & Lopez-Paz 20, In Search of Lost Domain Generalization\n6. Wang & Veitch 23, The Causal Structure of Domain Invariant Supervised Representation Learning"}, {"review_id": "CQD4ABchRw", "reviewer": "Reviewer_J1UG", "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. \n", "questions": "**Minor Typos/Corrections**: \n1. Line 117 (Definition 1): State that $\\Delta^{K-1}$ is a simplex", "limitations": "NA", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**Clarity**: The paper will be well-written and easy to follow.\n\n**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.\n", "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). \n2. Lemma 1 is similar to the claim shown in Makar et. al. (Causally motivated shortcut removal using auxiliary labels)\n3. Assumption 1 (constant effect): How justified is the assumption that the spurious attribute $c$ change the previous state in an additive manner?\n4. 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.\n\nOverall 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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper. We greatly value your feedback and would like to address each comment you’ve provided while highlighting the essential aspects crucial for evaluating our work.\n\n**Distinction from previous work**:\nOur contribution stands out first and foremost by the study of **counterfactual estimation methods**, and also the treatment of **finite sample effects**. These distinguish our work from those of Kaushik et al. 21 and Joshi & He 22.\n\nIn reference to these papers, we would like to also add Wang & Veitch 23 to our discussion. It is important to emphasize that *all the aforementioned papers do not discuss estimation methods from observed data* nor address sample complexity. The first point is extremely crucial since the type of Counterfactual Data Augmentation (CDA) proposed by Kaushik et al. 20 is infeasible in our task. \nTo put this in perspective, consider letting 3000 caretakers rewrite each other’s notes enough times to obtain an adequately large training set. To perform CDA at this scale, estimation methods are necessary and we are unaware of papers that propose causal estimation methods for purposes of data augmentation.\n\nAs for the finite sample analysis, your review points to this as a strength of our paper. We emphasize that this departs from previous explanations on the effectiveness of CDA. Kaushik et al. and Joshi & He discuss the effect of CDA with an infinite sample and a very specific data-generating process. Thm. 9 in Wang & Veitch gives a more general structure where CDA provably helps generalization in what they call “purely spurious” problems. In the language of Wang & Veitch, we treat is an instance of these “purely spurious” problems. In our case, there are other methods besides CDA that learn the optimal model (both methods require infinite data). Hence a finite sample analysis illuminates CDA's favorable performance w.r.t baselines (provided our CDA is of sufficient quality) instead of simply reiterating known observations.  \nThank you for this comment, we see that the point on finite samples should be better emphasized in contribution 2 instead of the other points it makes. We will rephrase it accordingly.\n\n**Lemma 1 and Makar et al.**: Footnote 2 of our paper explicitly states that this claim was proved in Makar et al. We included it in the paper using our notation since the lemma is important for motivating our method. We do not see how this is a weakness of the paper.\n\n**Constant effect assumption**: Before discussing the validity of the assumption, we clarify that its focus is on the effect being constant, rather than being additive. By definition, the quantity $x_i(c) - x_{i, pre}$ exists and can be defined as an additive effect. The strong assumption is that once we fix $c$ and the auxiliary data $m_i$, this effect is constant (i.e. does not depend on $i$). The assumption is more reasonable as we control for more auxiliary features. In experiments, we control for many factors such as prescribed medications, vitals etc. While assuming a purely constant effect is unrealistic, our generalization bound includes a distributional distance between counterfactuals and their estimations. Hence even if the assumption is imprecise, the method may still suffice to improve generalization. Weaker versions of the assumption can thus be phrased and reasoned on, yet in this work, we opt for a simple version that is easy to communicate and fit in the scope of the paper.\n\n**Augmentation limited to anti-causal case**: There is no inherent limitation to applying the estimation methods in causal prediction problems. We use the specific anti-causal problem to motivate CDA in a principled manner, as it describes the main problem we study and it is possible to analytically compare sample complexity with an alternative solution (i.e. reweighting). CDA is also appropriate in the “purely spurious” problems mentioned in response to article 1 in your review. These are more general and include some causal prediction problems. Then to obtain the counterfactual estimates, our techniques are still viable when the required assumptions hold (constant effect, or auxiliary data accounts for all unobserved factors).\nNonetheless, we agree that like all methods for OOD generalization, ours are not effective for all possible problems and must be used with caution. Thank you for turning our attention to this, our revised version will mention the usefulness of CDA in the more general setting of Wang & Veitch 23.\n\n**The paper re-instantiates points on CDA**:  \nWe believe that dismissing our work as “re-instantiating the point made in previous work” on CDA is not justified for several reasons. As explained in our rebuttal, estimation methods are crucial in practice and they are applicable beyond the medical datasets we work with. Our main interest is OOD generalization, and given that there are hundreds of papers on methods for this problem, we believe that a method that significantly outperforms the baselines on a large-scale real-world task is of interest to the community. At least from our familiarity with the literature, the large majority of proposed methods do not demonstrate results of this type. We are also unaware of other works that utilize causal estimation methods for data augmentation, and our finite sample bounds are different from analyses of previous work on CDA. Overall, we think the focus of our work has a significant but small overlap with those mentioned in the review.\n\n**References**:  \n1. Wang & Veitch 23 The Causal Structure of Domain Invariant Supervised Representation Learning  \n2. Makar et al. 22 Causally motivated Shortcut Removal Using Auxiliary Labels  \n3. Kaushik et al. 20 Learning the Difference that Makes a Difference with Counterfactually-Augmented Data  \n4. Kaushik et al. 21 Explaining the Efficacy of Counterfactually Augmented Data  \n5. Joshi & He 22 An Investigation of the (In)effectiveness of Counterfactually Augmented Data"}, {"review_id": "yEUHysGwuo", "reviewer": "Reviewer_Erys", "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.", "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)?\n- 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?\n- 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)?\n- In line 308, you say \"across 5 runs\". What is the source of randomness in these 5 runs?\n- Am I correct in interpreting Figure 4's findings that seemingly more corruption leads to better performance? Could you elaborate?", "limitations": "The limitations are adequately discussed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n- 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.\n- 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. \n- 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.\n- 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]).\n\n[1] Divyansh Kaushik, Amrith Setlur, Eduard Hovy, and Zachary C. Lipton. \"Explaining the efficacy of counterfactually augmented data.\" ICLR 2021.\n\n[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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the in-depth review of the paper and for raising important questions that help us improve it. Please find our responses to the concerns and questions below.\n\n**Knowledge of Causal Structure**:\nWhile our method assumes some knowledge of causal structure (i.e. Figure 1, an anti-causal prediction problem), we clarify that it does not require a fully specified structure that includes each and every feature in the problem (e.g. every component of the auxiliary data $M$). There are multiple graphs that comply with a “no unobserved confounding” assumption given $M$, and also with the constant effect assumption (which is a functional assumption, rather than a graphical one).\nNonetheless we agree that these are strong assumptions that require domain knowledge. In this context, it is important to note that unlike standard causal inference problems, we are interested in downstream OOD prediction abilities instead of the estimation of a specific counterfactual text (i.e. the estimation accuracy). Our generalization bound suggests that the distribution of augmented data should resemble the distribution of true counterfactuals, this can still hold when specific estimations suffer from some errors.  \nEventually, the evaluation that matters to us most is prediction accuracy compared to the baselines.\n\n**Difference from Kaushik et al. and Veitch et al.**: Thank you for raising this point, which deserved more attention in the paper. Since reviewer J1UG also asked about comparison with similar papers to the one you mentioned, we included this in our general response and elaborate here.  \nThe focus of our work is on studying estimation methods for Counterfactual Data Augmentation (CDA) and explaining their improved performance from the view of sample complexity. Compared to previous work on CDA, Kaushik et al. 20 (and subsequent work in Kaushik et al. 21) do not propose scalable estimation methods and rely on manual rewriting of text by humans. Veitch et al. 21 propose a regularizer that promotes conditional independence, instead of CDA. We claim and show empirically, that in large-scale problems, such methods fail due to challenging optimization and statistical complexity.  We are not aware of other works that adapt causal estimation methods for the purpose of data augmentation.  \nAs for formal results, Kaushik et al. 21 treat a rather specific data-generating process, and the most general treatment we are aware of is by Wang & Veitch 23 (follow up to Veitch et al. 21). These works discuss properties of CDA under infinite data, as opposed to our finite sample point of view. Specifically, our formal setting is an instance of the “purely spurious” problems defined by Wang & Veitch, where they show CDA achieves a min-max optimal model provided infinite data (as a side note, the assumptions for our estimation methods are not specific to cases where the graph in Fig. 1 holds, and hence they can be utilized in other such “purely spurious” settings). In contrast, we make an argument for the *sample complexity* of *approximate* CDA (i.e. augmentation might not be perfect).   To this end, we identify a baseline that in our problem setting is also min-max optimal provided infinite data, and draw a comparison between the sample complexities of the two methods to explain the performance we observe in practice. Finally, our main empirical contribution is favorable out-of-distribution (OOD) generalization in a complex, large-scale, real world problem of medical note classification, which further distinguishes our work.\n\n**How often models “hallucinate”**:\nEstimating this reliably warrants manual inspection of many pairs of original and generated text to determine in each instance whether the LLM filled in data that did not exist in the original text. Since our final goal is not a perfectly accurate counterfactual estimation but rather favorable downstream classification accuracy, we did not perform an extensive evaluation to determine how frequently this occurs. Instead we found examples of such incidents during our manual experimentation with the models, and verified that it is indeed a possible failure (this also makes sense considering we provide example texts with certain properties), which can be amended by matching on auxiliary data. \nFurthermore, note that because of the focus on downstream prediction accuracy, our bound implies that inaccurate estimations of the counterfactual text can damage performance, regardless of whether these accuracies are “hallucinations” or other types of mistakes (e.g. omitting data from the note). Therefore our example in section 4 is merely one instance of possible mistakes that can be prevented by additional context.\n\n**Details on experiments**:\nFollowing your comments, we added an experiment with equal-sized observational samples and augmented samples. Preliminary results on this new experiment, which sheds light on the effect of augmentation while controlling for sample sizes, are in the attached PDF. We should obtain the full results (across all runs and tasks) during the discussion period and will update them. \nRegarding the source of randomness across 5 runs, these correspond to different random seeds.\n\nFinally, thank you for a careful reading of the synthetic experiments. For Figure 4, lower values of $\\lambda$ result in larger corruptions. Hence the 3 top curves, which correspond to CDA, show the best performance for $\\lambda=0.4$ (lowest corruption in this simulation) and worst performance for $\\lambda=0.2$. So as intuition suggests, and opposite to your interpretation in the review, larger corruption results in inferior performance. We think the confusing part is that a *lower* value of $\\lambda$ corresponds to a *higher corruption*. We will better emphasize this in the description of the figure, thank you for turning our attention to this.\n\n**Additional references**: Wang & Veitch 23 The Causal Structure of Domain Invariant Supervised Representation Learning"}, {"review_id": "nUUE2KDdJ5", "reviewer": "Reviewer_8nR3", "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.", "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?\n2. Following the first question, how critical is the selection of LLM in this work?\n3. How critical is the quality of the auxiliary data? Also, would it be difficult to identify informative auxiliary data in practice?\n4. Regarding the writing,\n- line 134: Should the term under argmax be c \\in [K] instead of y \\in [K]?\n- Would it be more clear if the authors specify that the employed LLM is GPT4 in the main text instead of Appendix?", "limitations": " The authors adequately address the limitation.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "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.\n2. 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.\n3. 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.\n2. 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.\n3. 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\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper and recognizing positive points on the importance of the topic we study, the achieved empirical performance and writing quality. In what follows we aim to address each question and concern in your review.\n\n**Limited novelty**:\nThe novelty of our work, besides the favorable performance (w.r.t to all baselines) on a large scale real-world task, is:\n1. Introducing causal methods for producing counterfactually augmented data that substantially improves OOD performance. We note that merely “introducing auxiliary data” is not enough to achieve the performance we observe, and the auxiliary data must be used in a disciplined manner (in our method, we are guided by causal inference methodology, i.e. matching and diff-in-diff).\n2. Our theoretical analysis departs significantly from previous work on counterfactual data augmentation by considering sample complexity aspects instead of analyzing performance given infinite data.\n3. The proposed method and analysis open up the possibility for using any applied causal method for generating training data. We highlight two different approaches in this paper: matching and diff-in-diff. However, in principle, other methods from the econometrics literature, such as instrumental variables and synthetic controls, can be leveraged for better OOD performance using our approach and we expect future work to significantly extend ours.\n\n**Baselines**: \nFollowing our model for the problem in Fig. 1, we choose baselines that are appropriate for this task since they (a) are proposed in a paper that studies this anti-causal prediction problem, and (b) achieve state-of-the-art performance on it.\n\nWe have experimented with other baselines on the medical notes tasks and following your comment we will include results on GroupDRO and IRMv1 in the paper. The results can be found in the PDF attached to the rebuttal. IRMv1 underperforms the baseline from our original submission and GroupDRO is comparable, while CATO (A) has the best performance.\n\nWe refrained from including these baselines in the original submission, since previous work already documented their limited success on most large scale problems (e.g. [Rosenfeld et al. 20, Kamath et al. 21, Gulrajani and Lopez-Paz 20]) and we wanted to avoid making this the focus of our work. The medical notes task is far larger in scale than datasets where these methods have shown significant gains, and devising methods that show gains in large scale problems was among our main motivations for this work. Following your comment, we see why it is important to include these results, and it led us to revise our decision. Thank you for helping us improve this aspect of the paper.\n\n**LLM-induced biases**:\nThis is a great point! We have also included a discussion on this in our general response.  \nYour comment about LLMs varying in their generative abilities is on point. We’ve experimented with several models in the medical notes task and will include our results on Bio-BERT, Sentence-BERT and GPT4 in our revised paper. We find that models trained on synthetically augmented data using each of these LLMs results in improved OOD performance w.r.t the baselines.  \nConnecting this to your second question about biases of off-the-shelf LLMs, it is possible that LLMs introduce biases into our problem, inherited from their own training data and these biases may be specific to each LLM. This requires further study, however from our manual examination we found their quality satisfactory (see Appendix C for generation examples) and that all of them helped in OOD generalization which is our main evaluation metric (note that inaccurate counterfactual estimations in our case are only harmful if they adversely affect downstream predictive performance).\nNonetheless, we believe that further studies are required to characterize the possible biases of LLMs in the context of our work, for instance by comparisons of the counterfactual estimation made by LLMs with a rewrite of the text made by the human whose style we are trying to emulate (over a small curated dataset).\n\n**Quality of auxiliary data**:\nThis is also an important point. As in any causal matching problem, having reliable auxiliary data is crucial for selecting control examples. If the auxiliary data does not capture any information or contains wrong information we won’t be able to correctly match examples.\nLooking into our experimental analysis, we show that in two problems with radically different auxiliary data (complete patient history in the clinical narratives, and three review indicators in CeBAB) we get better performance than very strong baselines.\n\nWe will add additional analysis where we corrupt the auxiliary data in both real-world experiments (we have already done that for the synthetic data).\n\n**Comments on writing**:\nThanks for turning our attention to the typo in line 134, the term refers to the Bayes-optimal classifier, so it should be $y\\in{[L]}$ (L being the number of classes) instead of $y\\in{[K]}$.\nWe will also move details on the LLMs we use into the main paper.\n\n**References**:  \n1. Rosenfeld et al. 20, The Risks of Invariant Risk Minimization\n2. Kamath et al. 21, Does Invariant Risk Minimization Capture Invariance?\n3. Gulrajani & Lopez-Paz 20, In Search of Lost Domain Generalization"}]

Author rebuttal:
We thank the reviewers for their efforts and valuable comments, which led us to make important revisions to the paper, particularly by adding two additional baselines to our experiments (IRMv1 and GroupDRO), clarifying the connections of our work to previous papers on invariant learning and counterfactual augmentations, and adding a discussion on limitations of LLMs and possible biases they can introduce.

We are grateful for the appreciation of the novelty of our formal treatment of finite sample effects (reviewers J1UG, Erys), the superior strength of our method on real-world experiments (reviewers Erys, 8nR3), the novelty of using LLMs for counterfactual augmentation (reviewers MXTx, Erys), and for finding our paper very interesting (reviewers MXTx, YqAy, 8nR3). Moreover, we are happy to hear that all reviewers found our work to be clear and well-written.

Below we summarize our responses to the reviewers’ major concerns and review the major changes we introduced following the reviewers' comments.

1. **Additional baselines**: Reviewers YqAy, 8nR3 were curious about the performance of popular methods in OOD generalization that were excluded from our baselines. Following our model for the problem, we chose baselines that are appropriate for this task since they (a) are proposed in a paper that studies this anti-causal prediction problem, and (b) achieves state-of-the-art performance on it. Following reviewers comments, we have added experiments with two other baselines, GroupDRO and IRMv1, on the medical notes tasks. These are arguably the two most popular approaches in the literature, and the results can be found in the PDF attached to the rebuttal. As can be seen in our newly-added experiments, our method consistently outperforms these methods.

2. **Previous work on Counterfactual Data Augmentation (CDA)**: Reviewers J1UG, Erys wanted clarification about the connection of our work to previous prominent papers (e.g. Kaushik et al. 2021, Joshi & He 2022, Veitch et al. 2021). We emphasize that the focus of our work is on studying estimation methods for CDA and explaining their improved performance from the point of view of sample complexity. In comparison to previous work on CDA, these works do not propose scalable estimation methods, and we are not aware of other works that adapt causal estimation methods for the purpose of data augmentation. They also discuss properties of CDA algorithms under infinite data, as opposed to our finite sample point of view. Finally, our main empirical contribution is favorable out-of-distribution (OOD) generalization in a complex, large-scale, real world problem of medical note classification, which further distinguishes our work.

3. **Limitations of LLMs**: Reviewers MXTx, 8nR3 thought that the paper could benefit from additional discussion on the potential limitations and biases of LLMs in the context of our problem. The following points summarize a discussion that will be added to our revised paper.
* *LLM generation quality*: We acknowledge that LLMs vary in their ability to generate realistic text. It is possible that LLMs introduce biases into our problem, inherited from their own training data. This requires further study, however from our manual examination we found their quality satisfactory (see Appendix C for generation examples) and that OOD generalization also improved for models trained on the augmented data they generate. We'll include this analysis and results with different LLMs, namely Bio-BERT, Sentence-BERT and GPT4, in the revised paper.

* *Counterfactual approximation*: Other than generation quality, the additional challenge in using LLMs for CDA is our ability to elicit a good approximation to the counterfactual text. Our methods rely on principles from causal inference to advance disciplined approaches for this task. While further studies are required (e.g. systematically comparing small sets of manual re-writes of texts to the elicited LLM output), we view our work as a promising first step in this direction, which we expect to be significantly extended and improved in future work.

* *Effect of biases on OOD generalization*: Since we focus on OOD generalization, the limitations and possible biases mentioned above must be weighed within this context. Namely, we should bear in mind that even though generation may be biased, this bias is only harmful when it affects the generalization of a downstream classifier, and this is what we evaluate. Further, in OOD generalization we consider cases where the training data is biased in the first place, and training a standard predictive model also results in a biased solution. Hence we must weigh risks and limitations of alternative solutions vs. those of LLMs.

**Overview of major changes**:
1. We added two widely-used baselines (IRMv1 and GroupDRO), and an experiment with modified sample sizes (suggested by reviewer Erys). Results are in the attached PDF.
2. We made precise our connection to relevant previous works, highlighting our novelty in discussing finite sample properties of our approach and the usefulness of our method for approximating counterfactuals.
3. We added a section to discuss the possible biases that LLMs may introduce into our problem, emphasizing that further studies are required. We also note that when our goal is OOD generalization, and training data may already be biased, we need to weigh the biases of LLMs vs. alternatives.

**References**:
1. Kaushik et al. 2021 Explaining the efficacy of counterfactually augmented data.
2. Veitch et al. 2021 Counterfactual invariance to spurious correlations: Why and how to pass stress tests.
3. Joshi & He 2022 An Investigation of the (In)effectiveness of Counterfactually Augmented Data

Meta-review: The paper presents an augmentation technique to improve generalization of ML models. Reviewers highlighted a concern that the method critically depends on LLMs for generating the augmentations, but otherwise acknowledged the paper's contribution, empirical experiments, and its clarity in stating necessary assumptions. I would recommend the authors to revise their paper based on the feedback, in particular, 1) discuss the limitations of relying on LLMs; 2) how this paper tackles the limitations of CF data augmentation as noted in prior work (Joshi et al.).
Decision: Accept (poster)