SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 a patient's visit and each may be written by a different caregiver. Prediction is made using the release note from the hospital whose embedding consists our features  $\mathbf{x}$ . For simplicity let us consider a single note  $\mathbf{x}_{\mathrm{pre}}$  taken prior to  $\mathbf{x}$ . Difference-in-difference [19, 78, 79] estimation of causal effect is based on the parallel-trends, or constant effect assumption that two units  $i,j$  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  $\mathbf{m}$  includes  $c_{\mathrm{pre}}$ , the caregiver assigned pre-treatment.
Assumption 2 (constant effect). Let  $\mathbf{x}_{i,pre}$  be the pre-treatment features for unit  $i$ , and assume  $\mathbf{m}_i$  includes the pre-treatment attribute  $c_{i,pre}$ . There exists a function  $\rho : [K] \times \mathcal{M} \to \mathcal{X}$  such that  $\mathbf{x}_i(c) = \mathbf{x}_{i,pre} + \rho(c, \mathbf{m}_i)$ .
Under this assumption, to calculate  $\mathbf{x}_i(c)$  we can use any unit  $j$  for which  $\mathbf{m}_i = \mathbf{m}_j$  and has  $c_{j} = c$  to estimate  $\rho (c,\mathbf{m}_i) = \mathbf{x}_j - \mathbf{x}_{\mathrm{pre},j}$ . The resulting estimation procedure is given in algorithm 1(B) and illustrated in section 4.2.
Before empirically evaluating our methods, we discuss alternatives for learning robust classifiers in our setting, and how their properties fair compared to counterfactual augmentation.
4.3 Why
Bother with Counterfactual Data Augmentation?
Reasoning about counterfactuals with problem-specific domain knowledge is a considerable challenge, and it is interesting to see 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 (see e.g. Makar et al. [25], Shimodaira [80]).
Reweighting baseline. Intuitively, re-weighting samples from the uncorrelated distribution  $P(Y,C) = P(Y)P(C)$  by setting for each example  $i$  a weight  $w_{i} = P_{\mathrm{train}}(Y = y_{i})P_{\mathrm{train}}(C = c_{i}) / P_{\mathrm{train}}(Y = y_{i},C = c_{i})$  and
Panel A: Matching patients using auxiliary data
Panel B: Generating counterfactual discharge summaries
Figure 2: Generating counterfactual clinical notes for patients using auxiliary data with Algorithm 1(A).
minimizing the weighted empirical risk:
$$
\hat {\mathcal {R}} _ {\mathbf {w}} ^ {\ell} (h) = \frac {1}{m} \sum_ {i \in [ m ]} w _ {i} \ell \left(h (\mathbf {x} _ {i}), y _ {i}\right).
$$
It can be proved that at the limit of infinite data the method learns a min-max optimal hypothesis, as it also effectively minimizes  $\mathcal{R}_{P_{\epsilon}}^{l}$  (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.3
Comparing sample complexities. To make this statement precise, we can apply the bounds from Cortes et al. [81] and compare them with an upper bound that we 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  $d_{\alpha,\mathrm{train}}(Y,C) = \left[\sum_{y \in [L], c \in [K]} P_{\mathrm{train}}^{\alpha}(Y = y, C = c) / P_{\mathrm{train}}^{\alpha - 1}(Y = y)P_{\mathrm{train}}^{\alpha - 1}(C = c)\right]^{\frac{1}{\alpha - 1}}$ , and we may derive the following bound for a hypothesis  $h \in \mathcal{H}$  and any  $\delta \in [0,1]$ :
$$
\mathcal {R} _ {P _ {\perp}} ^ {\ell} (h) \leq \widehat {\mathcal {R}} _ {\mathbf {w}} ^ {\ell} (h) + \sqrt {\frac {2 d _ {2 , \text {t r a i n}} (Y , C) \cdot \log (1 / \delta)}{N}} + \frac {d _ {\infty , \text {t r a i n}} (Y , C) \cdot \log (1 / \delta)}{N}. \tag {2}
$$
A complementary lower bound on  $\widehat{\mathcal{R}}_{\mathbf{w}}^{l}(h)$  can also be derived based on results in Cortes et al. [81]. To compare this with counterfactual augmentations, denote our augmentation model by  $\tau : \mathcal{X} \times \mathcal{M} \to \mathcal{X}^K$ , which is some measurable function whose output's  $c$ -th coordinate is the counterfactual estimate w.r.t. caregiver  $c$ , i.e.  $\hat{\mathbf{x}}(c) = \tau_c(\mathbf{x}, \mathbf{m})$ . The following statement quantifies the relation between the accuracy of  $\tau(\cdot)$  in approximating counterfactuals and the classification accuracy of a model learned from the augmented data, via minimization of  $\widehat{\mathcal{R}}_{\mathrm{aug}}^\ell(h)$  in eq. (1).
Lemma 2. Consider a prediction problem with a spuriously-correlated attribute (definition 1), a measurable function  $\tau: \mathcal{X} \times \mathcal{M} \to \mathcal{X}^K$ , and let  $d_1(P, Q)$  denote the total variation distance between two distributions  $P, Q$ . Further let  $h^*$ ,  $h_{aug}^*$  denote the optimal hypotheses w.r.t  $\mathcal{R}_{P_\perp}^{\ell_{01}}$ ,  $\mathcal{R}_{aug}^{\ell_{01}}$  respectively and let  $\lambda_{aug} = \left[R_{P_\perp}^{\ell_{01}}(h_{aug}^*) - R_{P_\perp}^{\ell_{01}}(h^*)\right]$ . For any hypothesis  $h \in \mathcal{H}$ , and any  $\delta \in (0, 1)$  it holds that with probability at least  $1 - \delta$  over the draw of the training set,
$$
\mathcal {R} _ {P _ {\perp}} ^ {\ell_ {0 1}} (h) \leq \widehat {\mathcal {R}} _ {a u g} ^ {\ell_ {0 1}} (h) + \sqrt {\frac {\log (1 / \delta)}{N}} + K ^ {- 1} \cdot \sum_ {c \in [ K ]} d _ {1} (\tau_ {c, *} (P _ {t r a i n} (X, M)), P (X (c))) + \lambda_ {a u g}.
$$
The divergence  $d_{1}(\tau_{c,*}(P_{\mathrm{train}}(X,M)),P(X(c)))$  is a distance between the true distribution over counterfactual instances  $P(X(c))$  and our augmented data  $\tau_{c,*}(P_{\mathrm{train}}(X,M))$ . Divergences other than total-variation can be used, resulting in tighter bounds, e.g. see Ben-David et al. [82]. As we generate better counterfactuals this divergence decreases, and it can also be shown that  $h^*$  and  $h_{\mathrm{aug}}^*$  coincide. Hence  $\lambda_{\mathrm{aug}}$  vanishes and the bound scales with  $N^{-\frac{1}{2}}$ , resulting in a gain of factor  $d_{2,\mathrm{train}}(Y,C)$  over the upper bound on  $\widehat{\mathcal{R}}_{\mathbf{w}}^{\ell_{01}}(h)$  in Equation (2). We discuss the details in the appendix, and in Section 5 we show this empirically through simulations.
Takeaways and additional baselines. We emphasize that the counterfactual datapoints should not be interpreted as "more data" in the sense of i.i.d training examples, they rather embody knowledge 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  $P_{\perp}$ . 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-Mean Discrepancy (MMD) from Makar et al. [25] who show improved sample complexity under a linear hypothesis class.
<sup>3</sup>We remark that other works discuss the potential benefits of data augmentation for identification in other problem settings, e.g. [72, Thm. 9] and [17].
4The notation  $\tau_{c,*}(\cdot)$  denotes the pushforward measure. We note that in our implementation  $\tau_{c}$  is data dependent and we ignore this dependence to enable a simple analysis.
We empirically study the following questions: (1) Can CATO enhance OOD performance of downstream 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 of counterfactuals?
These questions seek to establish causally-motivated augmentations as a practical approach for improving 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 an overview of the tasks we experiment with.
Input (x)	Label (y)	ID Data	OOD Data	Spurious Feature (c)	auxiliary data (m)
Clinical Narratives	Condition Prediction
Note Segmentation
Demographic Traits	MIMIC-III	i2b2-2010
partner data
i2b2-2006	Caregiver ID	Medications,
Lab Results,
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 et 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 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 generalization that are used in the literature. The augmentation approaches are compared here to 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 segmentation and demographic traits identification<sup>6</sup>, for which we have both ID and OOD data. We utilize several electronic health records (EHR) datasets. We train on MIMIC-III [86], a widely-used medical dataset containing over 2 million notes from 38,597 adult patients, 49,785 hospital admissions, and 3,500 healthcare professionals between 2001 and 2012. MIMIC-III is commonly used in NLP research for clinically-related tasks and for pre-training language models for the medical domain [87]. When available, we use i2b2 2006 and 2010 competitions as our held-out hospital dataset. In the note segmentation task, we use private held-out data.
Generating notes from counterfactual caregivers. To generate augmentations, we select caregivers with multiple patients and notes for more than one patient. For each caregiver-patient pair where both their last progress note and discharge summary were written by that caregiver<sup>7</sup>, we match them to similar patients having the same initial caregiver but a different one for their discharge summary. In matching, we select patients with similar medications and lab results (denoted as patient's
5 Appendix B includes results where the Baseline model is also BioBERT, SentenceBERT or GPT3.
See Appendix B for results on the demographic traits identification task.
7During a patient's stay, progress notes capture its current state. When leaving the hospital, a discharge summary is written.
auxiliary data  $m$  in Table 1). We then generate counterfactual discharge summaries for matched patients using Algorithm 1(A) and train the model using original data and generated counterfactuals.
Figure 3 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.
Figure 3: Results (F1 averaged across 5 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 a 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 augmentations improve performance slightly. In the OOD setting, all OOD methods help (reweighting, MMD, IRM, GroupDRO, CATO (A)), but our causally-motivated augmentation approach is substantially 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 approach improves over the vanilla PubMed BERT model, but is 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 condition 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 as CeBAB-Spurious, where there's a spurious correlation between training and deployment.
To construct CeBAB-Spurious, we leverage the availability of both the original and perceived rat
Method	CeBAB	CeBAB-Spur.
Observational	0.85	0.64
Reweighting	0.84	0.68
Naive Aug.	0.80	0.62
Conditional Aug.	0.84	0.70
CATO (B)	0.84	0.75
Table 2: Accuracy on CeBAB and CeBAB-Spurious. CATO (B) outperforms all baselines when we introduce a spurious correlation.
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 sentiment. 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 information. To incorporate this variability, we introduce a new attribute called food-mention to signify the presence of food-related information in a review. If the perceived food rating is either negative or positive, we assign a value of 1 to the food-mention attribute; otherwise, it is set to 0. We subsample the data such that there is a 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 to 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 to better OOD generalization, while naive data augmentation hurts model performance. In line with the sample complexity argument in Section 4, conditional augmentation effectively doesn'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), we generate synthetic data for a binary classification problem where  $K = 8$  (cardinality of  $C$ ). We sample  $\tilde{P}(C \mid Y)$  to simulate varying degrees of spurious correlations. Then we draw  $\mathbf{x} = [\mathbf{x}^*, \mathbf{x}_{\mathrm{spu}}]$  from a Gaussian distribution,
$$
\mathbf {x} _ {i} = \left[ \begin{array}{c} \mathbf {x} _ {i} ^ {*} \\ \mathbf {x} _ {\mathrm {s p u}, i} \end{array} \right] \sim \mathcal {N} \left(\left[ \begin{array}{c} \boldsymbol {\mu} _ {y _ {i}} \\ \boldsymbol {\mu} _ {c _ {i}} \end{array} \right], \left[ \begin{array}{c c} \sigma^ {2} \mathbf {I} _ {d ^ {*}} & 0 \\ 0 & \sigma_ {\mathrm {s p u}} ^ {2} \mathbf {I} _ {\mathbf {d} _ {\mathbf {c}}} \end{array} \right]\right).
$$
In this case  $\hat{\mathbf{x}}_i(c)$  is obtained by adding  $\mu_c - \mu_{c_i}$  to  $\mathbf{x}_{\mathrm{spu},i}$ . To corrupt our augmentation, we instead add  $\xi_{i}\left(\mu_{c} - \mu_{c_{i}}\right)$  where  $\xi_{i}$  is drawn from a truncated Gaussian centered at  $\lambda \in (0,1)$ . We train models with a fixed sample size (in the appendix we also examine varying sample sizes and additional types of corruption) and evaluate the trained models' accuracy on  $P_{\perp}$  to examine the interplay between spurious correlation strength (measured by mutual information  $I(Y;C))$ , and counterfactual augmentation quality. As can be seen in Figure 4, corruptions degrade performance under stronger spurious correlations, though a strong corruption is required for reweighting to become preferable.
In this work, we have presented a data augmentation approach based on the causal structure of auxiliary data for improving OOD generalization, specifically focusing on text classification tasks. However, our approach is not without
Figure 4: OOD accuracy  $(1 - \mathcal{R}_{P_{\perp}}^{l_{01}}(h))$  and  $Y, C$  correlation strength  $(I(Y;C))$ . Lower values of  $\lambda$  correspond to stronger corruptions of the augmentations. Even with substantial corruption ( $\lambda = 0.2$ ) and strong correlation, augmentations outperform baselines.
limitations. The validity of our assumptions, the specification of the causal graph and the quality 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 causally-motivated data augmentation methods like ours can help address challenges in developing robust and reliable machine learning systems, particularly in safety-critical applications.