SCOT: Improved Temporal Counterfactual Estimation with Self-Supervised Learning

We appreciate the reviewer for providing positive feedback on our work. Here we want to restate our main contributions over existing works such as RMSN and CRN as follows: (1) As far as we know, 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. (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. We have also added the above discussions in Related Work (Appendix A) in our revised version.

Thank you for providing feedback and suggestions to improve the submission. Here we list our responses to your major concerns.

1. “Contribution is incremental. Applying existing random augmentations with the component-wise contrastive loss is quite straightforward while the main part of method seems to come from Causal Transformer”.

2. “What are the exact improvements over Causal Transformer and why?”

(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.

3. “Explain the settings of “cold-start”, “zero-shot transfer”. How this differs from basic potential outcomes evaluation (e.g., in-sample vs out-of-sample evaluation).”

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.

4. “Especially wrt to the analogy to unsupervised domain adaption it is important to describe the exact setting more specifically”

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).

5. “For reproducibility, the code for the experiments should be provided.”

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

I thank the authors for their extensive answers, which made many aspects of the paper more clear to me. In principle, I think this is an interesting paper with nice experimental evaluation and reasonable motivation (even though it could be still improved at some points). Hence, I am raising my score to 5. However, I would like to emphasize that, while nicely validated, the contribution of the paper itself is only marginal and not sufficient for an ICLR paper for the following reasons:

1. Contribution: Architecture. It is important to emphasize that the main parts of the proposed architecture have been originally proposed by [1]. The architectural contribution of this work here is the featurewise attention only. Here, it is important to mention that (i) featurewise attention in general has also been used for treatment effect estimation in previous works (e.g., [2]), and (ii) the gains of the feature-wise attention blocks are only marginal compared to the original temporal attention blocks from the causal transformer of [1].

2. Contribution: Self-supervised learning for treatment effect estimation and link to UDA. I think that self-supervised learning for more expressive and robust expression representation learning for time-varying treatment effect estimation definitely is an useful approach. However, also here the approach lacks novelty in the proposed points: (i) the authors state to provide a novel view of the connection between UDA and estimation of potential outcomes (Sec. 3.4). But there exist already works discussing this connection (e.g., [3, 4]) in more detail. (ii) Then, the authors simply leverage already existing results for contrastive learning and unsupervised domain adaptation. Also, there exists already closely related work for self-supervised contrastive learning for unsupervised domain adaptation of time series data (e.g., [5]) benefitting from this connection. (iii) The applied data augmentations are not novel [6]. Here would be an opportunity to use customized augmentations grounded in the causal graph of the treatment effect setting. (iv) The applied component-wise contrastive loss is quite straightforward per intuition but not quite clear why here an improvement over the simple contrastive loss is expected. Also, the improvements of this component seem to be not that significant compared to the other ablations in table 3. (Additionally, the code should be provided for such a submission for reproducibility and comparison to the existing works.)

For now, my questions are answered. We will carefully consider the above points, as well as the rebuttal / revised paper, in the discussion period!

[1] Melnychuk, Valentyn, Dennis Frauen, and Stefan Feuerriegel. "Causal transformer for estimating counterfactual outcomes." International Conference on Machine Learning. PMLR, 2022.

[2] Zhang, Yi-Fan, et al. "Exploring transformer backbones for heterogeneous treatment effect estimation." arXiv preprint arXiv:2202.01336 (2022).

[3] Johansson, Fredrik D., et al. "Generalization bounds and representation learning for estimation of potential outcomes and causal effects." The Journal of Machine Learning Research 23.1 (2022): 7489-7538.

[4] Johansson, Fredrik, Uri Shalit, and David Sontag. "Learning representations for counterfactual inference." International conference on machine learning. PMLR, 2016.

[5] Ozyurt, Yilmazcan, Stefan Feuerriegel, and Ce Zhang. "Contrastive Learning for Unsupervised Domain Adaptation of Time Series." International Conference on Learning Representations. 2022.

[6] Woo, Gerald, et al. "CoST: Contrastive Learning of Disentangled Seasonal-Trend Representations for Time Series Forecasting." International Conference on Learning Representations. 2022.

We greatly appreciate the reviewer for carefully evaluating our responses and for providing comprehensive feedback. Still, here we would like to reillustrate the novelty of our work. In short, our SSL-based method is the first to address the counterfactual outcome estimation problem for time series when both treatment bias and feature distribution shifts exist.

We are ready to share our code upon request once the double-blind review stage ends.

We would also like to kindly request the reviewer to update the score in the original response to reflect the change in the review.

Contributions in the model architecture.

We would like to emphasize that our proposed encoder architecture considers both temporal and feature-wise interactions for the temporal counterfactual outcome estimation task.

• (a) The temporal attention as well as the relative positional encoding are first proposed in [7] and [8] respectively, and they are the only two parts our model shares with [1].

• (b) Our model utilizes the temporal attention and the feature-wise attention alternatingly to capture both temporal and feature-wise interactions. Instead, [2] only considers feature-wise interactions in a static setting.

• (c) We added the ablation study results showing the effect of temporal/feature-wise attention blocks:

• Semi-Synthetic MIMIC-III:

• M5:

Contributions in SSL for counterfactual outcome estimation.

• (a) We would like to clarify that the novelty of our view connecting UDA and potential outcome estimation is that we provide a bound of the counterfactual outcome estimation error when the feature extractor is optimized using SSL. Theorem 3.1 shows that when the feature extractor minimizes the contrastive loss, the error of counterfactual outcome estimation is bounded by the error of outcome estimation on factual data. It justifies our steps that (1) first train the feature extractor with SSL loss and (2) then train both the feature extractor and the predictor with factual outcome estimation loss. Meanwhile, both the contrastive loss and the factual outcome estimation loss can be easily generalized to multi-valued or continuous treatments by selecting S-Learner-like model architectures.

Instead, while [3,4] also discusses the connection between UDA and potential outcome estimation, they derive a bound composed of both factual outcome estimation and explicit distributional distance regularization, which leads to regularization-based methods like BNN[4] and CFR[3]. Moreover, both [3,4] use the assumption that the treatment is binary, and it requires extra steps to extend the distributional distance regularization term beyond binary treatments in practice [3].

• (b) While there exist works addressing UDA of time series forecasting/classification with SSL, as we have discussed in related work, our work is still novel in that we are not solving the prediction problem, but the counterfactual outcome estimation problem. Here both treatment bias and feature distribution shift exist, and we are the first to address both issues with SSL and provide competitive empirical results. While the effect of SSL on feature distribution shift has been widely discussed in the standard UDA setting, we further provide theoretical justification explaining why SSL also works for treatment bias by taking the UDA view.

Thank you for providing feedback and suggestions to improve the submission. Here we list our responses to your major concerns.

1. “What do symbols (cross, needle, and dashed line) in Figure 1 mean?” 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.
2. “The necessity of component-wise contrastive loss.” 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.
3. “For the outcome predictor, $\bar{\mathbf{H}}_t$ already includes treatment sequences, what is the motivation and justification for remodeling it using 2 convolution blocks?” 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})$.
4. “Assumption B.1 is too strong.” 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.
5. “Figure 3: Pretrained is worse than the rest in separation, but there is not much difference among the remaining, right?” 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.
6. “Reporting in table 3 sounds manipulated.” 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.

Thank you for your thoughtful and helpful review and suggestions! We have addressed your concerns and questions in our replies based on the review. We are looking forward to having an in-depth and fruitful discussion so that we can clarify any further confusion or concerns.