SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Data Augmentations for Improved (Large) Language Model Generalization
Amir Feder  $^{*1,2}$ , Yoav Wald  $^{*3}$ , Claudia Shi  $^{1}$ , Suchi Saria  $^{3}$  and David Blei  $^{1}$
<sup>1</sup> Columbia University, <sup>2</sup> Google Research, <sup>3</sup> Johns Hopkins University
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.
The reliance on spurious correlations is a significant challenge for Machine Learning (ML) safety as it can lead to performance degradation of deployed models. Spurious correlations are prevalent in various applications such as medical imaging [1, 2], text classification [3], and risk prediction systems [4]. Failures due to spurious correlations occur under distribution shift [5-7], which may result from differences in data recording protocols, shifts in the underlying population being monitored, or the way the ML tool is being used. In this paper, we focus on text classification and explore how using language models in a domain-informed way can help us avoid reliance on spurious correlations.
Consider a scenario where we want to make robust predictions about patients' conditions, probability of readmission, etc., using clinical narratives written in hospitals [8-10]. In this setting, a common issue arises due to clinical practice, where patients with certain conditions are 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 unintentionally rely on the style to make predictions. This leads to poor generalization on unseen hospitals, i.e. failure to generalize out of distribution(OOD), due to changes in clinical practice [7]. However, collecting 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, we propose leveraging available auxiliary data (e.g., time, document type, demographics) and incorporating knowledge about the causal structure of the problem to build a more robust classifier. For example, in the note classification task, we can use our knowledge that some auxiliary data, such as the patient's current state, can affect doctor assignment, to improve the classifier's robustness.
*Equal Contribution. Correspondence to amir.feder@columbia.edu
Causal inference often makes use of such auxiliary data and has now been used in a variety of ways to improve OOD generalization [6, 11-14]. Data augmentation methods have demonstrated impressive performance in these tasks as well [15-17], and with recent improvements in generative models, forming additional principles to incorporate domain knowledge into data augmentations seems like a promising path forward.
In this work we pursue this and develop causally-driven 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 LLMs towards counterfactual estimation. Our contributions are:
1. Through extensive experiments, we show how the use of language models in a manner that is informed by causal knowledge improves model robustness in challenging safety-critical tasks in healthcare. Furthermore, our findings are reinforced by experiments that incorporate semi-synthetic scenarios, and simulations where there are ground-truth counterfactuals.
2. We formalize counterfactual data augmentation in a prediction setting as a method to deconfound the target and a spuriously correlated attribute. We show how deconfounding improves OOD generalization. In a setting where sample complexities for alternative methods (reweighting and invariance penalties) can be derived, we show favorable generalization bounds for accurately performed data-augmentation.
3. Our data-augmentation 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 LLMs and other generative models, can be a 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).
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 different 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 learning [13, 21-26]. These methods have been applied in a 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 guarantees. 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, a popular and straightforward approach is 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 augmentation [37]. Constructing counterfactual instances that involve perturbations to confounding 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 counterfactual data augmentation in text involve manual editing by humans, heuristic keyword replacement, or automated text rewriting [37, 39, 41-50]. Manual editing is accurate and effective [38, 51] but expensive, hence our goal is to make counterfactual data augmentation scalable, demanding smaller
human effort. Keyword-based methods can be limited in coverage and difficult to generalize across languages [52]. Generative approaches offer a balance of fluency and coverage [53], but generating meaningful counterfactuals is challenging [54]. Our work departs from previous techniques by using causal auxiliary data structure and LLMs 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 identifying this information, and standardized datasets and competitions exist for this purpose [56-60]. Best performing approaches have leveraged transformer architectures both for token-level classification 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 NLP systems [67, 68]. In our experiments, we follow recent papers in clinical NLP addressing challenges of degraded performance across different hospitals [69-71].
To formally analyze how counterfactual data augmentation helps OOD generalization, we consider a setting where the label is spuriously correlated with a known attribute. This setting has been used previously to study learning with "shortcuts" [25] and spurious correlations [29]. We note that our approach is applicable and valid under additional settings and causal graphs (e.g. "purely spurious" problems defined in Wang and Veitch [72]) and we elaborate on this at ??. The data generating process used here motivates counterfactual data augmentation 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 (see section 4.3).
Consider a classification problem with  $L$  classes, where the label  $Y$  is correlated with a certain attribute  $C$  in the training data and this correlation may change arbitrarily at test time (denoted by a red edge  $C \leftrightarrow Y$  in fig. 1). In our medical notes example,  $C$  is the caregiver writing the note and  $Y$  is the underlying condition we wish to diagnose. We denote the number of caregivers in our training data by  $[K]$ . For a given loss function  $\ell : \mathbb{R}^L \times [L] \to \mathbb{R}$  and distribution  $P$ , we denote the expected loss of a hypothesis  $h : \mathcal{X} \to \mathbb{R}^L$  by  $\mathcal{R}_P^\ell(h)$  and its expected accuracy by  $\mathcal{R}_P^{\ell_{01}}(h)$ . The data-generating process is depicted by the causal model in fig. 1, for our motivating example of clinical notes classification  $X$  is a vector representation of the clinical note and  $X^*$  is 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: Prediction problem with a spuriously correlated attribute.
Definition 1. We denote the set of distributions induced by interventions on a causal model with the structure in fig. 1 by
$$
\mathcal {P} = \left\{P (X \mid X ^ {*}, C) P (X ^ {*} \mid Y) P (Y) \tilde {P} (C \mid Y): \tilde {P} (C \mid Y = y) \in \Delta^ {K - 1} \forall y \in [ L ] \right\},
$$
where all distributions other than  $\tilde{P}(C \mid Y)$  are fixed. In a prediction problem with a spuriously correlated attribute, the learner is provided with a set  $\{(\mathbf{x}_i, y_i, c_i)\}_{i=1}^N$  sampled i.i.d from  $P_{train} \in \mathcal{P}$ . We assume that  $X^* = e(X)$  almost surely for some  $e: \mathbb{R}^d \to \mathbb{R}^{d^*}$ .
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  $\tilde{P}(C \mid Y)$  do not change the conditional distribution  $P(Y \mid X^{*})$ . Therefore an optimal solution that does not rely on  $C$  is  $h^{*}(\mathbf{x}) = \arg \max_{y \in L} P(Y = y \mid e(\mathbf{x}))$ . In clinical note classification,  $X^{*}$  represents all the information in the note about the patient conditions, unsullied by the writing style of caretaker  $C$ . To obtain  $h^{*}(\mathbf{x})$  we will rely on risk minimization w.r.t a distribution where  $Y$  and  $C$  are uncorrelated.
3.1 Learning Robust Classifiers when Counterfactuals are Available
Consider the unconfounded distribution  $P_{\perp} \in \mathcal{P}$  that is given by intervening on  $C$ , setting it independent of  $Y$  and uniformly distributed,  $\tilde{P}(C \mid Y) = P_{\mathrm{unif}}(C)$ . An optimal classifier under  $P_{\perp}$  has the following min-max optimality guarantee.
Lemma 1. For the prediction problem in definition 1, the Bayes optimal classifier under the unconfounded distribution  $P_{\perp} \in \mathcal{P}$  where  $C$  is uniformly distributed and independent of  $Y$  is  $h^{*}(\mathbf{x}) = \arg \max_{y \in [K]} P_{\perp}(Y = y \mid X^{*} = e(\mathbf{x}))$ . It is a minimizer of  $\min_{h: \mathcal{X} \to [L]} \max_{P \in \mathcal{P}} \mathcal{R}_P^{\ell_{01}}(h)$  and  $\mathcal{R}_P^{\ell_{01}}(h^{*}) = \mathcal{R}_{P_{\perp}}^{\ell_{01}}(h^{*})$  for all  $P \in \mathcal{P}$ .
Hence we would like to minimize risk w.r.t  $P_{\perp}$  and we cannot do that directly by via ERM since our training data is sampled from  $P_{\mathrm{train}} \neq P_{\perp}$ . Instead we consider risk minimization over an augmented dataset that contains counterfactual instantiations of our training data under different values of  $C$ .
Minimizing  $\mathcal{R}_{P_1}$  via Counterfactual Data Augmentation. Returning to 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 if each patient had been seen by all possible caregivers  $c \in [K]$  and each caregiver had written their own version of the note  $\mathbf{x}_i(c)$ . Given these counterfactual clinical notes, we seek a hypothesis that minimizes the average loss over all such possible scenarios, denoted by  $\widehat{\mathcal{R}}_{\mathrm{aug}}^{\ell}(h)$ .
Definition 2. Consider a prediction problem with a spuriously-correlated attribute (see Definition 1). For a given example  $\mathbf{x}_i$ , we denote its counterfactual with attribute value  $c \in [K]$  as derived from the corresponding causal model, by  $\mathbf{x}_i(c)$ . For estimates of the counterfactuals  $\{\hat{\mathbf{x}}_i(c)\}_{i \in [N], c \in [K]}$  and a hypothesis  $h \in \mathcal{H}$ , the counterfactually augmented empirical risk is
$$
\widehat {\mathcal {R}} _ {a u g} ^ {\ell} (h) = \frac {1}{N K} \sum_ {i \in [ N ], c \in [ K ]} \ell \left(h \left(\hat {\mathbf {x}} _ {i} (c)\right), y _ {i}\right). \tag {1}
$$
We use approximate counterfactuals  $\hat{\mathbf{x}}_i(c)$  in our definition to highlight that in practice we cannot obtain a precise estimate of  $\mathbf{x}_i(c)$ . In the ideal case where  $\hat{\mathbf{x}}_i(x) = \mathbf{x}_i(c)$ , the expected loss  $\mathcal{R}_{\mathrm{aug}}^{\ell}(h)$  where  $N \to \infty$ , satisfies  $\mathcal{R}_{\mathrm{aug}}^{\ell}(h) = \mathcal{R}_{P_{\perp}}^{\ell}(h)$ . This follows by a simple derivation and it is part of a claim we give later in Lemma 2. Hence obtaining this dataset is useful for our goal of minimizing risk under  $P_{\perp}$ . Our main challenge is then to derive effective approximations for counterfactuals such as clinical notes under alternative writing styles.
4 Assumptions and Algorithms for Estimating Counterfactuals
Perfectly capturing writing style is a strong assumption. Even if we could perfectly model writing styles, we only observe a limited set of variables - the actual notes  $x$ , 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  $c$  and  $\tilde{c}$ , where a note  $\mathbf{x}_i$  was written by  $c_i = \tilde{c}$ . We want to estimate what  $\mathbf{x}_i(c)$ , the note caregiver  $c$  would have written, might look like. To this end we will build a model  $\tau_c(\cdot)$  that takes data and generates a note in caregiver  $c$ 's style. Now suppose caregiver  $c$  usually sees patients with high blood pressure and always includes blood pressure values in notes, while  $\tilde{c}$  rarely does. A naive model estimating  $\hat{\mathbf{x}}_i(c) = \tau_c(\mathbf{x}_i)$  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 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 a proof for completeness. We set the distribution over  $C$  in  $P_{\perp}$  as uniform for simplicity, the derivation for non-uniform distributions is analogous.
4.1 Identification of the Counterfactual Distributions
To make effective use of this data, we suggest that the input to the model  $\tau_c: \mathcal{X} \times \mathcal{M} \to \mathcal{X}$  will include a baseline text to be edited and auxiliary data  $\mathbf{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 1 (Strong ignorability). For all  $P \in \mathcal{P}$  it holds that  $X(c) \perp C \mid M$ , and for all values of  $\mathbf{m} \in \mathcal{M}$ ,  $P(\mathbf{m}) > 0$ .
Under this assumption, we can rewrite the counterfactual distribution with the observed distribution,
$$
P (X (c)) = \int P (X (c) \mid M = \mathbf {m}) P (M = \mathbf {m}) d \mathbf {m} = \int P (X \mid C = c, M = \mathbf {m}) P (M = \mathbf {m}) d \mathbf {m}.
$$
However, in practice, we do not observe many samples from  $P(X \mid C = c, M = \mathbf{m})$ , making it a poor approximation for the counterfactual distribution. We address this by using counterfactual data augmentation [37]. Formally, we assume that for all possible counterfactual distributions  $c \in [K]$  there exist a function  $\tau_{c}$  that maps from the observed distribution  $P(X \mid M = \mathbf{m})$  to the target counterfactual distribution  $P(X(c) \mid M = m)$ .
We approximate the loss under the counterfactual distributions through the empirical loss produced by data augmentation. That is, for a hypothesis  $h \in \mathcal{H}$
$$
\mathbb {E} _ {P (X (c))} [ \ell (h (\mathbf {x}), y) ] \approx \frac {1}{N} \sum_ {i \in [ N ]} \tau_ {c} (\mathbf {x} _ {i}, \mathbf {m} _ {i}).
$$
Note that whenever the text in the training set is already written by caregiver  $c$ , i.e.  $c_{i} = c$ , we will simply keep the original text  $\mathbf{x}_i$
Evaluation of Augmented Distribution. The right hand-side of the above equation is a Monte-Carlo estimator of the distribution of augmented notes, which averages the distributions  $\tau_{*,c}(P_{\mathrm{train}}(X,M))$  over all caregivers  $c \in [K]$ . The distribution  $\tau_{*,c}(P_{\mathrm{train}}(X,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  $\tau_{c}$ .
High-quality counterfactual estimation, as measured by small distributional divergence between our estimator and the target distribution, will help in lowering the upper bound on the risk  $\mathcal{R}_{P_{\perp}}^{\ell}(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  $\tau_{*,c}(P_{\mathrm{train}}(X,M))$  is obtained simply by running training data through  $\tau_{c}$ , while a sample from  $P(X(c))$  can be obtained either by adjusting for  $M$ , or we can obtain a sample from  $P(X|C = c,M = \mathbf{m})$  for each value of  $\mathbf{m}$  and compare that to a sample obtained by augmenting validation data where  $M = \mathbf{m}$ . In both cases two-sample tests can be applied and obtain estimates of divergences between the two distributions. That is 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 = \mathbf{m})$  for certain values of  $\mathbf{m}$  and  $c$ .
We now describe the estimation methods that obtain  $\tau_{c}$ . 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  $\mathbf{x}_i(c)$ . Our framework for estimating counterfactuals  $CATO$  (Causal-structure Driven Augmentations for Text OOD Generalization) involves the use of an LLM to model the conditional probability distribution of text. Counterfactuals are formed by matching similar auxiliary data examples or manipulating texts' vector representations, as described below.
Prompting with matched examples. Our first estimation method in Algorithm 1(B) draws insights from matching [76]. We construct a prompt for an LLM, that given an original text  $\mathbf{x}$  and a set of
Algorithm 1 CATO
Input: Training set  $\{(\mathbf{x}_i,y_i,c_i,\mathbf{m}_i)\}_{i = 1}^N$
Hypothesis class  $\mathcal{H}$
Version  $\in \{(A),(B)\}$
Optional pre-treatment data  $\{(\mathbf{x}_{\mathrm{pre},i})\}_{i = 1}^{N}$
Output: A hypothesis  $h_{\mathrm{aug}}(\mathbf{x})$
1: if Version  $= (A)$  then
2: Get  $\tau_{c}(\mathbf{m},\mathbf{x})$  with preprocess (A)
3: Get  $\hat{\mathbf{x}}_i(c) = \tau_c(\mathbf{x}_{i,\mathrm{pre}},\mathbf{m}_i)\forall i\in [N]$
4: else
5: Get  $\tau_{c}(\mathbf{m},\mathbf{x})$  with preprocess (B)
6: Get  $\hat{\mathbf{x}}_i(c) = \tau_c(\mathbf{x}_i,\mathbf{m}_i)\forall i\in [N]$
7: end if
8: return  $h_{\mathrm{aug}} \in \mathcal{H}$  that minimizes  $\widehat{\mathcal{R}}_{\mathrm{aug}}^{\ell}$ .
Pre-process CATO (A)
Assume:  $\mathbf{m}$  includes the label  $y$  and pretreatment attribute  $c_{\mathrm{pre}}$ , among other auxiliary data. We are given  $\{\mathbf{x}_{j,\mathrm{pre}}\}_{j = 1}^{N}$ .
1: Set  $\pmb {\rho}(c_j,\mathbf{m}_j) = \mathbf{x}_j - \mathbf{x}_{j,\mathrm{pre}}$  for  $j\in [N]$
2: return  $\tau_{c}(\mathbf{x},\mathbf{m})\coloneqq \mathbf{x}_{\mathrm{pre}} + \rho (c,\mathbf{m})$
Pre-process CATO (B)
Assume:  $\mathbf{m}$  includes the label  $y$  among other auxiliary data.
1: return prompt  $\tau_{c}(\mathbf{x},\mathbf{m})$  that rewrites  $\mathbf{x}$  in the style of matching examples with attribute  $c$ , i.e.  $\{\mathbf{x}_j:(\mathbf{m}_j,c_j) = (\mathbf{m},c)\}$ .
context notes, asks the LLM to rewrite  $\mathbf{x}$  in their style. Now given text  $\mathbf{x}$  with auxiliary data  $\mathbf{m}$  that we wish to estimate with counterfactual value  $c$  (i.e. writing style),  $\tau_{c}(\mathbf{x},\mathbf{m})$  runs this prompt with context notes whose auxiliary data is similar to  $\mathbf{m}$  and their attribute value equals the desired  $c$ .