Paper: SCOT: Improved Temporal Counterfactual Estimation with Self-Supervised Learning (ICLR-2024)

Reviewer comment (Novelty):
The main weakness of the paper seems to be that contribution is only incremental. Applying existing random augmentations with the component-wise contrastive loss is quite straightforward while the main part of method seems to come from the work of [1].

Author response:
**(Responses to both 1 and 2)** Here we want to restate our main contributions as follows:

(1) To the best of our knowledge, our work is the first to adapt self-supervised learning (SSL) in temporal counterfactual outcome estimation. Instead of simply putting SSL in our task, we also provide theoretical insights for the contribution of SSL in Section 3.4. We demonstrate that the counterfactual outcome estimation problem can be viewed as an unsupervised domain adaptation (UDA) problem if we regard the outcome estimation of each treatment value as a separate task. Therefore, we can adapt the generalization error of SSL in UDA settings for counterfactual estimation settings.

(2) In addition to SSL, we also propose a new encoder architecture combining both temporal and feature-wise attention to capture both temporal dependencies and feature interactions in the observed history. 

(3) Our work is significantly different from Causal Transformer (CT) [1] in both methodology and empirical results.
- (a) The encoder of CT only considers attention across groups of covariate/treatment/outcome variables, while our proposed method considers a finer-grained level of feature interactions via attention across each input variable.
- (b) CT utilizes the adversarial domain confusion loss to learn a balanced representation of history to alleviate treatment bias and generalize to counterfactual cases. Adversarial training involves 2 conflicting objectives: the learned representations should be able to predict outcomes but not treatments, which improves the difficulty of training. We empirically find that such a loss cannot lead to robustly improved performance. In Table 1, the complete CT model performs slightly better than its variant trained without the adversarial loss (named CT(ERM)) only in the zero-shot transfer setup on M5 data and worse in the other settings. Different from CT, we provide a novel view of the connection between UDA and counterfactual estimation (Section 3.4) and choose SSL, which is already proven to be effective in UDA, to improve generalization to counterfactual outcome estimation. We also verify our conclusion in the ablation study. In Table 3, we verify that SSL brings performance improvement of 4.4% and 0.2% on Semi-synthetic MIMIC-III and M5 respectively.

Reviewer comment (Writing):
The motivation of substantial parts of the work is often not clear and could be improved. E.g., what are the exact improvements over [1] and why are these expected to work better. Also, the used terminology for describing the setting (“cold-start”, “zero-shot transfer”, “source/target domain”) is a bit confusing and its not quite clear how this differs from basic potential outcomes evaluation (e.g., in-sample vs out-of-sample evaluation). Especially with respect to the analogy to unsupervised domain adaptation it is important to describe the exact setting more specifically (in Sec. 2 as well as in Sec 3.4., e.g., distinguish between the domain shift between the distributions conditioned on the treatments induced by the confounding bias vs. the domain shift stemming from the unconditional feature distribution shifts in source and target domains). Hence, source and target domains should be clearly defined in the Problem Formulation Sec. 2.

Author response:
**(Responses to both 1 and 2)** Here we want to restate our main contributions as follows:

(1) To the best of our knowledge, our work is the first to adapt self-supervised learning (SSL) in temporal counterfactual outcome estimation. Instead of simply putting SSL in our task, we also provide theoretical insights for the contribution of SSL in Section 3.4. We demonstrate that the counterfactual outcome estimation problem can be viewed as an unsupervised domain adaptation (UDA) problem if we regard the outcome estimation of each treatment value as a separate task. Therefore, we can adapt the generalization error of SSL in UDA settings for counterfactual estimation settings.

(2) In addition to SSL, we also propose a new encoder architecture combining both temporal and feature-wise attention to capture both temporal dependencies and feature interactions in the observed history. 

(3) Our work is significantly different from Causal Transformer (CT) [1] in both methodology and empirical results.
- (a) The encoder of CT only considers attention across groups of covariate/treatment/outcome variables, while our proposed method considers a finer-grained level of feature interactions via attention across each input variable.
- (b) CT utilizes the adversarial domain confusion loss to learn a balanced representation of history to alleviate treatment bias and generalize to counterfactual cases. Adversarial training involves 2 conflicting objectives: the learned representations should be able to predict outcomes but not treatments, which improves the difficulty of training. We empirically find that such a loss cannot lead to robustly improved performance. In Table 1, the complete CT model performs slightly better than its variant trained without the adversarial loss (named CT(ERM)) only in the zero-shot transfer setup on M5 data and worse in the other settings. Different from CT, we provide a novel view of the connection between UDA and counterfactual estimation (Section 3.4) and choose SSL, which is already proven to be effective in UDA, to improve generalization to counterfactual outcome estimation. We also verify our conclusion in the ablation study. In Table 3, we verify that SSL brings performance improvement of 4.4% and 0.2% on Semi-synthetic MIMIC-III and M5 respectively.

Our evaluations focus on the out-of-sample data, following the common practice of previous works [1-3]. In addition to the out-of-sample setting, we further consider scenarios where distributions of history of training/test samples are different. Under this setting, source/target domains refer to the distribution of history in training/test data respectively.

We have revised the draft to distinguish between the domain shift from the confounding bias and that from the unconditional feature distribution shifts. 
The source/target domains in Section 3.4 are used as analogical names to describe the labeled and unlabeled subsets regarding a treatment value in the unsupervised domain adaptation framework to help us analyze the error of counterfactual outcome estimation, and are different from the definitions we use in Section 2. In Section 2, source/target domains describe the different distributions of the observed history between train/test data (i.e. the feature distribution shifts).
For synthetic and semi-synthetic data, we always evaluate with data after removing treatment bias, while for real-world data we can only evaluate with factual examples. In both cases, settings “zero-shot transfer” and “data-efficient transfer” add feature distribution shifts between train and test, while the “standard supervised learning” setting has no unconditional feature distribution changes (results are reported in Table 6 in Appendix G).

Reviewer comment (Reproducibility):
For reproducibility, the code for the experiments should be provided.

Author response:
**“For reproducibility, the code for the experiments should be provided.”**

We will release the code and data as soon as our work is published.

Reviewer comment (Novelty):
The paper contribution should be spelled out better. If the authors claim their main contribution is around UDA, why not use a better experimental setting (and more common datasets that align with UDA-ts)? Why not review the UDA for time-series stream in the related work? Why build on SOTA methods for UDA-ts? If the authors say that their main contribution is around self-training, why not use a proper method that is more grounded in the causal graph (given that the methods for data augmentation are largely 1:1 copies from Woo et al)?

Author response:
**“If the authors say that their main contribution is around self training ,why not use a proper method that is more grounded in the causal graph?”**

One of our main contributions is to adapt self-supervised learning (SSL) for counterfactual outcome estimation over time. We provide theoretical insights for the contribution of SSL in Section 3.4 with the help of a UDA view of the counterfactual outcome estimation task, and we verify its contribution in ablation studies. Our results reveal that SSL-learned representations of history help alleviate the negative impact of the correlation between **history** and **future treatments**. We agree that digging into the internal causal graph of the observed history and developing data augmentation methods grounded in the internal graph would be a promising direction.

Reviewer comment (Theory):
In Figure 3, the representations trained only on the self-supervised objective seem already unpredictive of the treatment. This indicates that either the dataset does not contain a high-level of confounding by default or the self-supervised learning can already balance the representations. If the latter is possible, a (theoretical) intuition why this works would be interesting, especially to support the use of self-supervised learning in this setting.

Author response:
**“This indicates that either the dataset does not contain a high-level of confounding by default or the self-supervised learning can already balance the representations. If the latter, possible (theoretical) intuition why this works would be interesting, especially to support the use of self-supervised learning in this setting.”**

The reason for using self-supervised learning (SSL) is to leverage its ability to learn generalizable representations, which is important for both (i) counterfactual outcome estimation and (ii) cold-start scenarios. For (i), we provide a novel view to connect the counterfactual outcome estimation problem to the unsupervised domain adaptation (UDA) problem in Section 3.4 and show that the generalization error bound of counterfactual outcome estimation can be derived from the error bound of target domain data in the UDA problem. For (ii), the cold-start scenario is a more direct case of UDA, and using SSL again benefits the performance with its generalizable representations. 

As we have discussed, existing works such as CRN and CT use adversarial losses to learn the balanced representations: representations should maintain the ability to predict outcomes but not treatments, which loses part of the history information and limits the outcome estimation ability, as well as increasing the training difficulties. In contrast, the SSL stage in our method enables generalization to counterfactual data without the cost of adversarial training. To evaluate the ability of treatment prediction of learned representations, we train a linear probe on representations and use the average accuracy of predicting each type of treatment as the metric. We use the same architecture of SCOT, extract the output of the encoder on the target domain data as the representations, and compare their accuracies in the following settings (i) trained with outcome prediction losses on source domain data without self-supervised learning; (ii) trained with self-supervised learning losses only; (iii) trained with self-supervised learning losses first and supervised outcome prediction losses later on source domain data (the zero-shot setting). The accuracies of treatment prediction are (i) 70.5% (ii) 71.4% and (iii) 71.1%. The results show that the self-supervised learning stage even slightly improves the treatment prediction ability of representations. This also verifies the conclusion in existing work [5] that the improved generalization of contrastive pre-training is not the result of invariant representations as conventional UDA methods.

Reviewer comment (Theory):
In general, the paper could adapt more to the terminology of counterfactual vs interventional distribution of the causal ladder of Pearl. E.g., in Figure 1, the “cold-start / UDA case” is an interventional question (what will happen if treatment A will be prescribed) and not a counterfactual question (what would have happened if instead of treatment B treatment A had been prescribed).

Author response:
**“In general, the paper could adapt more to the terminology of counterfactual vs interventional distribution of the causal ladder of Pearl.”**

First, we would like to emphasize that for cold-start cases (zero-shot setting), both (1) treatment bias in training and (2) feature distribution shifts exist. With the existence of (1), estimating the outcome of all possible treatments from observed data with treatments prescribed with biases toward history is a counterfactual problem.

In addition, as we have stated in Section 2, we follow the commonly used potential outcome framework extended to time-varying data, where we estimate the **average/expectation of treatment outcomes over a certain population conditioning on the observed history** instead of the outcomes of individuals described in the causal ladder of Pearl.

We have also seen the terminology “counterfactual” widely used in existing works [1-3] addressing the same treatment outcome estimation over time tasks and we keep it for consistency.

Reviewer comment (Experiments):
It is assumed that in general, the rMSE is evaluated on unconfounded test data for the semi-synthetic datasets. In Sec 4.2., does the data in the target domain leveraged for fine-tuning still contain confounding bias such that the final test dataset still has a distribution shift with respect to confounding to this data, or is the confounding bias removed here too?

Author response:
**“Does the data in the target domain leveraged for fine-tuning still contain confounding bias?”**

For test data, we always remove the confounding bias in both source and target domains if the data is synthetic/semi-synthetic. For the real-world dataset, since we have no access to the ground-truth outcomes of arbitrary treatments, we have to use factual data in evaluation and it might contain confounding bias.

Reviewer comment (Presentation):
Figure 1 is not intuitive. What do these symbols (cross, needle and dashed line) mean?

Author response:
In Figure 1, we use an example from the healthcare domain to illustrate the treatment outcome estimation task: given the observed history of vital signs of a patient, we want to estimate the future outcomes after giving the patient one or multiple steps of treatment. Here the symbol cross/needle refers to the value of a binary treatment (giving the medicine/doing nothing). Dashed lines refer to the results of counterfactual outcome estimation.

Reviewer comment (Novelty):
As mentioned above, I would challenge the necessity of component-wise constrastive loss because it brings little improvement. If this is one of the main contributions.

Author response:
We introduce the component-wise contrastive loss as a part of the self-supervised contrastive learning method. In Section 3.4 we also provide theoretical insights that the contrastive learning benefits the model in generalizing to counterfactual outcome estimation. In Table 3 we observe that the component-wise contrastive loss contributes to a 2.6% performance gain on Semi-Synthetic MIMIC-III data and 0.7% with M5 data. Notice that Semi-Synthetic MIMIC-III allows us to evaluate the exact counterfactual estimation performance, while the M5 is from real-world and thus we can only evaluate with factual data, the former scenario weighs more on the need for generalization to counterfactual estimation and demonstrates a larger improvement. Moreover, we observe that the overall variance of results on both datasets decreases after including the component-wise contrastive loss (0.0926->0.0869, 0.0715->0.0609).
While the improvement from component-wise contrastive loss is smaller compared to the contributions from the encoder and the loss, previous works like Causal Transformer [1] also demonstrate that the improvement from the design that helps generalize to counterfactual estimation (for Causal Transformer it is the counterfactual domain confusion loss) is also relatively small. In Table 1 of [1], the performance of the complete model (CT(ours)) and the variant without counterfactual domain confusion loss (CT(alpha=0)) is comparable in all horizons.

Reviewer comment (Experiments):
Some parts are unclear and can be confusing. For the outcome predictor, Ht already includes treatment sequences, what is the movitation and justification for remodeling it using 2 convolution blocks? Does it imply that the latent representation of Ht could not capture the treatment sequences well? If not, can you include this in the ablation study?

Author response:
We would like to clarify that $\bar{\mathbf{H}}\_t$ only includes the treatment sequences in the observed history, while the decoder takes future steps of treatment as input. In other words, Ht encodes historical treatments, while the decoder predicts based on both history and future treatments. As shown in Figure 2 (a) and (c), if we estimate outcomes of future \tau steps after time t, Ht encodes the historical sequence of treatment from time $1 \dots t-1$, while the decoder uses $(a_t, …, a_{t+\tau-1})$.

Reviewer comment (Theory):
I think the Assumption B.1 is too strong. Correct me if I am wrong, but I believe it removes uncertainty. However, when given a treatment, there is still probabilities, being either high or low, of different outcomes. In that sense, B.1 is much stronger than the no unmeasured counfounding that people usually use.

Author response:
We would like to clarify that Assumption B.1 (consistency) connects the potential outcome and the observed outcome, which is the fundamental assumption in causal inference to ensure identifiability. Instead of removing uncertainty, it implies that the variation within treatments would not result in a different outcome [2]. In addition, all three assumptions (B.1, B.2, B.3) are commonly used in existing works to ensure identifiability [1,3,4] in time-varying data.

Reviewer comment (Evaluation):
Figure 3: Pretrained is worse than the rest in separation, but there is not much difference among the remaining, right?

Author response:
In terms of the separation of outcome values, we observe in Figure 3 that pretrained embeddings are less separated since the pretraining stage does not involve the supervised training with outcome labels. After tuning the pretrained embeddings with supervised outcomes, the embeddings are more predictive of outcomes and thus are better separated w.r.t outcome values.

Reviewer comment (Experiments):
I think the reporting in table 3 sounds manipulated - in each dataset, you cherrypick the best-performing supervision strategy as "SCOT", and the worse-performing variants as "ablation study". This really looks unprofessional to me. I would rather fix one as SCOT across all datasets.

Author response:
As we describe in Section 3.3, we treat the weights of stepwise losses as hyperparameters and can be selected based on the validation errors. Since we cannot enumerate all possible choices of weights, we consider three simple schemes introduced in Section 4.3 (uni/inv/sq.inv) in our experiments and select the one with the lowest validation error in each experiment.

Reviewer comment (Experiments):
As mentioned in the strengths, I think the main contributors (from the table 3) is the encoder design, as you stack two blocks of encoders (of course you have multiple layers as well) to capture feature-wise and temporal information. This can bring more room for discussion: i) is the performance gain from model complexity (if you use much more parameters), or one of the encoders? ii) in ablation study, why don't you remove one of the encoder blocks (just stacking feature-wise block or temporal block) and see how it performs? That will be more interesting.

Author response:
(i) In Table 5 in the appendix, we report the number of parameters of baselines and our model. We control the model complexity so that the number of parameters in our model is comparable to or even fewer than the baselines. (ii) We provide the ablation study results of only keeping the temporal/feature-wise attention blocks as follows:

Semi-Synthetic MIMIC-III:

||$\tau=1$|2|3|4|5|6|Avg|Gain(%)|
|-|-|-|-|-|-|-|-|-|
|SCOT |0.2266+-0.0249|0.4501+-0.0893|0.5406+-0.0987|0.5964+-0.1020|0.6344+-0.1040|0.6637+-0.1052|0.5186+-0.0869|(-)|
|Temporal blocks only|0.2729+-0.0409|0.4711+-0.0836|0.5607+-0.0937|0.6160+-0.0995|0.6553+-0.1038|0.6831+-0.1043|0.5432+-0.0856|4.53%|
|Feature-wise blocks only|1.1210+-0.0827|1.1287+-0.0855|1.1755+-0.1076|1.2055+-0.1251|1.2296+-0.1418|1.2547+-0.1566|1.1858+-0.1155|56.3%|

M5:

||$\tau=1$|2|3|4|5|6|Avg|Gain(%)|
|-|-|-|-|-|-|-|-|-|
|SCOT |6.4054+-0.0547|6.9328+-0.0634|7.2428+-0.0700|7.4585+-0.0580|7.7012+-0.0627|7.8278+-0.0651|7.2614+-0.0609|(-)|
|Temporal blocks only|6.4085+-0.0538|6.9547+-0.0535|7.2673+-0.0453|7.4825+-0.0388|7.7167+-0.0380|7.8328+-0.0430|7.2771+-0.0409|0.2%|
|Feature-wise blocks only|6.8805+-0.0333|7.6298+-0.0212|7.9706+-0.0254|8.1215+-0.0298|8.3989+-0.0411|8.5303+-0.0435|7.9219+-0.0311|8.3%|

Reviewer comment (Experiments):
How do you deal with missingness in the time series? This is common in healthcare dataset.

Author response:
In the Semi-Synthetic MIMIC-III dataset, we use the extracted features from [5]. [5] uses the so-called “Simple Imputation” described in its ”Data Pre-processing” section: missing values are first forward filled and then set to individual-specific mean if there are no previous values. If the variable is always missing for a patient, we set it to the global mean. In general, datasets can always be imputed to eliminate missing values. We also added the description of processing missing values in Appendix E.

Reviewer comment (Experiments):
Since we talked about pretraining and transfer learning, instead of showing t-sne plots which hardly differentiates NCS, CS and Transfer, I would rather want to know the model performance of pretraining without finetuning.

Author response:
We first want to clarify that the supervised training stage on source domain data after pretraining is necessary for the outcome estimation task - the decoder outputting outcome estimation requires the supervised training stage. Our analysis of the generalization of our model to counterfactual estimation is also based on such a self-supervised pretraining -> supervised training scheme.

After this step, we have three different settings depending on the domain of test data as well as whether we further finetune the model. (1) Non-Cold-Start (NCS): directly evaluate the model on source domain data; (2) Cold-Start (CS): directly evaluate the model on target domain data; (3) Transfer: finetune the model with a few samples from the target domain and evaluate on target domain data.

We already report the performance in NCS, CS, and Transfer settings in Table 6 (in Appendix G), Table 1, and Table 2 respectively. While we are not fully certain what results the “pretraining without finetuning” is referring to, we assume Table 1 fits the description to some extent: the cold-start setting has no fine-tuning stage on target domain data. We hope the reviewer can clarify the question further.

Reviewer comment (Evaluation):
This is minor but in Table 3 fold 2, the best one for MIMIC is actually sq. inv.

Author response:
We thank the reviewer for pointing out the mark issue and we have fixed it in the revised version.