G-Transformer for Conditional Average Potential Outcome Estimation over Time

Authors: Konstantin Hess, Dennis Frauen, Valentyn Melnychuk, Stefan Feuerriegel
Primary area: causal reasoning
Keywords: causal inference, potential outcomes, treatment effects, G-computation, time-varying confounding, medicine

Abstract:
Estimating potential outcomes for treatments over time based on observational data is important for personalized decision-making in medicine. Yet, existing neural methods for this task either (1) do not perform proper adjustments for time-varying confounders, or (2) suffer from large estimation variance. In order to address both limitations, we introduce the G-transformer (GT). Our GT is a novel, neural end-to-end model which adjusts for time-varying confounders, and provides low-variance estimation of conditional average potential outcomes (CAPOs) over time. Specifically, our GT is the first neural model to perform regression-based iterative G-computation for CAPOs in the time-varying setting. We evaluate the effectiveness of our GT across various experiments. In sum, this work represents a significant step towards personalized decision-making from electronic health records.

Decision: reject
Reviewer scores: 5, 8, 3, 5

Reviews and discussion:
<thread>
creation date:
1730316288525

thread_type:
official_review

<thread_object>
creator: reviewer_5auu

<content>
summary:
The paper proposes the G-transformer (GT) approach, which combines the G-computation formula and a multi-input transformer to estimate conditional average potential outcomes (CAPO) over time, given time-varying confounding and treatment observational data. GT iteratively generates pseudo-outcomes by reformulating the G-computation formula as recursive conditional expectations. Experimental results on synthetic and semi-synthetic datasets demonstrate that the proposed approach outperforms the baselines in terms of average root mean squared error (RMSE).

strengths:
- The paper addresses the important and impactful real-world problem of estimating CAPO over time under time-varying confounding and treatment from observational data.
- The proposed GT, iteratively generates pseudo-outcomes by reformulating the G-computation formula as recursive conditional expectations, which seems interesting.
- Experimental results on synthetic and semi-synthetic datasets demonstrate that the proposed approach outperforms the baselines in terms of RMSE, especially under conditions of high time-varying confounding and longer CAPO estimation horizons.
- The code and data are publicly available, making it easy to reproduce the results.

weaknesses:
- The novelty is quite limited; the proposed approach could be considered a simple extension of the causal multi-input transformer (Melnychuk et al., 2022) with pseudo-outcome predictions.

- While the paper claims to be the first to use pseudo-outcomes for CAPO estimation over time, several works have already explored the use of pseudo-observations in causal inference and survival analysis, including [1,2,3]. I encourage the author(s) to discuss these works and highlight key differences with their proposed approach.

- The paper asserts that related works, such as G-Net, which properly adjust for time-varying confounding, result in high estimation variance without providing substantial experimental or theoretical results to support this claim. Only RMSE values with error bars are provided (Tables 2 and 3), which seems insufficient given that this is a central claim of the GT approach.

- The experimental results seem underwhelming. Only one semi-synthetic dataset is considered, and RMSE is provided to quantify performance. I encourage the author(s) to explore evaluations tied to meaningful clinical measures, including precision in the estimation of heterogeneous effects, average treatment effects, and qualitative analysis measures of over- or underestimation of outcomes.

- Clarity: The notation and writing need improvement for clarity and readability. The use of both $U$ and $H$ to denote the random variables $\{Y_t, X_t, A_t\}$ does not seem necessary. The paper seems to refer to $H$ throughout, and it's unclear why $U$ was introduced. The definitions for $\bar{Y}$, $\bar{X}$, and $\bar{A}$ are not provided. The recursion could be made more explicit by combining Equations (8) and (9) and providing a simple example, e.g., $\tau = 3$. The paper's use of $\delta$, $\tau$, and $t$ for indexing time needs further simplification. Additionally, the definition for $\tau$ should be provided in the *Setup* section.


**References**
- [1] Andersen et al. (2017), "Causal inference in survival analysis using pseudo-observations", Statistics in medicine.
- [2] Andersen and Perme (2010), "Pseudo-observations in survival analysis", Statistical methods in medical research.
- [3] Chien-Lin et al. (2022), "Causal inference for recurrent event data using pseudo-observations", Biostatistics.

**Minor**
- Table 1: I encourage the author(s) to include citations to related works for improved readability.
- Notation overload:  $\phi$  is used to parameterize both functions $g(\cdot  )$ and $z( \cdot)$.
- Typos: Line 140 should be  $U_{t+\tau}$ remove the 'minus'.
- I encourage the author(s) to increase the font size for tables and figures to be consistent with the main text.
- The definition for $\delta$ in Line 244 is not consistent with that in Line 267.

questions:
- It seems that the GT approach has a clear advantage over existing approaches when $\tau \ge 2$. Could you provide scenarios where estimating CAPO over time ($\tau \ge 2$) has clear clinical benefits?
- Could you quantify the variance of the baseline approaches, for example, using the coefficient of variation?
- Could you provide additional results on the precision in the estimation of heterogeneous effects, average treatment effects, and qualitative analysis measures of over- or underestimation of outcomes?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730437898412

thread_type:
official_review

<thread_object>
creator: reviewer_i9bm

<content>
summary:
This work proposed a framework to estimate conditional average potential outcome over time using a combination of transformer and iterative pseudo-outcome regression. The framework is tested in experiments on a synthetic dataset based on tumor growth model and a semi-synthetic dataset based on MIMIC-III, and is shown to outperform several benchmark models on these two datasets.

strengths:
1. Paper is well written, problem formulation, architecture descriptions and the presentation of learning procedures are all clearly stated and very easy to follow.

2. The results look compeling. Table 2 and 3 showed noticable improvement of the proposed model over the compared benchmarks.

weaknesses:
1. The novelty of this work is the most important concern to me. The combination of transformer and temporal-difference learning (i.e. pseudo-outcome regression as called by the authors) are already proposed in [1] for the estimation of CAPO (on top of that, [1] also went one step further and proposed the adjustment for marginal causal parameters). Besides, the heterogeneous token transformer architecture in [1] appears to be a more clean and intuitive solution as opposed to the proposed architecture in my judgement. The proposed transformer architecture is rather trivial and does not significantly differ from [2]. At the very least, the author should compare their architeture to [1], show experimental results and properly state their contributions compared to prior work.

2. The paper has a very weak description of prior work related to pseudo-outcome regression in my opinion. Not only is the exact same procedure temporal-difference learning proposed before as stated above, it is widely used in causal inference and reinforcement learning known as sequential regression [3] and Q-learning [4]. The omission of these works by the authors made it look like pseudo-outcome regression is a novel contribution of their own.

[1] Shirakawa, Toru, et al. "Longitudinal Targeted Minimum Loss-based Estimation with Temporal-Difference Heterogeneous Transformer." International Conference on Machine Learning. PMLR, 2024.

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

[3] van der Laan, M. J. and Gruber, S. Targeted Minimum Loss Based Estimation of Causal Effects of Multiple Time Point Interventions. The International Journal of Biostatistics, 8(1), 2012.

[4] Chebotar, Yevgen, et al. "Q-transformer: Scalable offline reinforcement learning via autoregressive q-functions." Conference on Robot Learning. PMLR, 2023.

questions:
My suggestions correspond to the two weaknesses above. I think comparison to prior work that already proposed the combination of transformer and temporal-difference learning/pseudo-outcome regression is absolutely crucial; literature review related to pseudo-outcome regression ([1][3][4] and other related work) should be added, and contribution of this work should be properly and correctly stated.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730619140990

thread_type:
official_review

<thread_object>
creator: reviewer_5wft

<content>
summary:
This study proposes G-Transformer to estimate conditional average potential outcomes (CAPOs) over time while addressing two major limitations of the previous studies: (1) adjustments for time-varying confounders; (2) and large estimation variance. Particularly, the authors propose a viable solution to G-computation, iterative pseudo-outcome regression, using Transformer.

strengths:
The study is motivated and the problem is well-defined by the research gap of the previous studies. Specifically, the authors address the limitations of the previous studies regarding G-computation to estimate CAPOs with neural networks by providing a viable solution of iterative pseudo-outcome regression, although the neural network architecture itself is similar to the existing study (Melnychuk et al., 2022).

The author provides sufficient and proper supplementary materials to explain the details of the study.

weaknesses:
The reviewer did not find any significant weaknesses.

The reviewer thinks a more principled and diverse way to evaluate the quality of the generated pseudo-outcomes will be desirable. While the experiment conducted in E.2 could empirically show the learning step is sensitive to the noise of the generated pseudo-outcomes, it might not be sufficient to show the quality of the generated pseudo-outcomes as it is titled. An ablation study that separates each step (i.e., another step is frozen when one step is trained) might be useful for a controlled experiment to measure the sensitivity of the performance in terms of how much the model is trained to generate the pseudo-outcomes.

questions:
Table 2 and 3 have low readability due to small font sizes.

The reviewer thinks Supplementary E.2 regarding the quality of pseudo-outcome is worth being mentioned in the main text
</content>
</thread_object>

</thread>
<thread>
creation date:
1730654305324

thread_type:
official_review

<thread_object>
creator: reviewer_aabk

<content>
summary:
The paper studies conditional average potential outcome estimation over time. The authors propose a model called G-Transformer for addressing two primary challenges in the field: adjustment for time-varying confounders and low variance in estimates. The authors validate their approach through experiments on synthetic and semi-synthetic data, demonstrating improvements over current baselines.

strengths:
- The paper provides a clear, detailed description of model architecture.

- The authors conduct extensive experiments with synthetic and semi-synthetic datasets, demonstrating GT's effectiveness under different levels of confounding settings.

- The code is provided.

weaknesses:
- The contribution of this paper is somewhat limited. 
  * From the methodology perspective, the proposed method's improvement over existing approaches appears marginal. The primary contribution is exploring the transformer backbone in treatment effect estimation, which has been widely investigated in existing work (static setting [1] and dynamic setting [2]).
  * From the application perspective, the paper claims that their model can help personalize decision-making from patient trajectories in medicine. However, it lacks discussion on how the model's estimates would concretely translate to real-world decision support. Furthermore, the assumptions on which the model relies (e.g., positivity, ignorability) are restrictive and may not hold robustly in practical, real-world scenarios.

- While the paper claims to reduce estimation variance, it lacks a theoretical analysis of variance bounds. Given GT’s objective to minimize estimation variance, a more rigorous theoretical justification or a derivation of variance bounds would significantly strengthen this claim.

-  The paper does not consider the existence of unobserved (and potentially time-varying) confounders, which can introduce biases and affect CAPO estimation in real-world EHR settings.

- The study uses a semi-synthetic dataset derived from MIMIC-III. It would be valuable to see the model's performance on fully real-world datasets to assess its utility in practical clinical decision-making problems.

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

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

questions:
See weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732149437439

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers

comment:
Thank you very much for the constructive evaluation of our paper and your helpful comments! We addressed all of them in the comments below. Furthermore, we uploaded an updated version of our paper and marked all significant changes to the original version in **blue color**. 

Our **main improvements and clarifications** are the following:



* **Clarification of our contribution over existing methods:** We spell out the novelty of our work more clearly: our work proposes a **novel end-to-end training procedure for estimating CAPOs over time**. In other words, we do **not** propose a new architecture but **an entirely new training procedure in an end-to-end-architecture**. To this end, we revised our Sections 2, 3 and 4, where we reiterate that existing neural methods either **(i)** completely lack proper adjustments or **(ii)** have problematic adjustment strategies. Further, we clarify that our transformer-based instantiation is only **one** possible choice, and point to our **G-LSTM ablation study in our Supplement G.1**. Importantly, our G-LSTM outperforms all baselines (including more advanced architectures like the causal transformer from Melynchuk et al 2022), which shows that the source of gain is our new end-to-end training algorithm.
* **Related work:** \
We revised our related work **Section 2**:  \
**(i)** We added more explanations on why methods for estimating average potential outcomes have a completely different purpose and are **not** applicable in our setting. \
We further added an **extended related** work to our paper in a **new Supplement A**: **(ii)** Therein,  we discuss methods from survival analysis based on pseudo-outcomes and clarify why they are also **not** applicable to our setting.  \
**(iii)** Finally, we provide a detailed discussion on how Q-learning in the reinforcement literature and G-computation relate and highlight important differences.
* **Difference to existing methods:** We added a new **Section 4.4** where we discuss the differences to our method in greater detail. We explicitly point to the limitations of existing neural approaches and how our method improves upon them.
* **Clarification on low-variance claims:** We clarified our low-variance claims. In particular, we did the following:
**(i)** We **removed** our previous “low-variance” terminology throughout the paper to avoid ambiguities. \
**(ii)** We showed mathematically that IPW adjustments as employed by RMSNs construct pseudo-outcomes with **larger variance** than pseudo-outcomes constructed by G-computation in **Proposition 3**. \
**(iii)** We further explained that G-Net requires estimating **high-dimensional probability distributions** which may lead to large estimation variance (see our **revised Section 4.4 and Supplement F)**.

* **Additional experiments based on real-world data:** We added new experiments based on real-world data in a **new Supplement H**. We use the MIMIC-III dataset, which gives measurements from intensive care units aggregated at hourly levels and we predict the effect of vasopressors and mechanical ventilation on diastolic blood pressure. Our GT has state-of-the-art performance. Thereby, we show that our GT can easily handle difficulties in predicting patient outcomes in real-world datasets.  \
This supports our claim that our end-to-end training algorithm is highly effective.
* **Discussion about medical applicability:** We provide a discussion of how our GT is applicable in medical practice in a **new Supplement B**. Therein, we also discuss why our assumptions are standard in the literature, why they are weaker than for methods in the static setting, and why they are plausible  (see our also **revised Section 3**).
* **Coefficient of variation:** We report an additional metric for our main experiments, namely, the coefficient of variation (see our **new Supplement I). **The results show that** our GT provides the most stable estimates** of the CAPOs.

We highlight all changes in the **updated PDF** in **blue color**. We specifically incorporated all changes (marked with **Action**) into the updated version of our paper. Given these improvements, we are confident that our paper provides valuable contributions to the causal machine learning literature and is a good fit for ICLR 2025.
</content>
</thread_object>

</thread>


Metareview:
metareview:
In this paper, the authors propose a framework to estimate the CAPO in temporal settings based on a G-transformer and pseudo-outcome regression. The paper is well written and sound, providing synthetic and semi-synthetic experiments. The authors provided a revised version of their paper with extensive changes.

However, Reviewer i9bM mentioned a lack of novelty compared to [1]. The authors have argued that the goal of that paper is limited to APO, but the reviewer suggested a misunderstanding of [1], which also performs CAPO. The authors did not engage further. Based on my cursory overview of [1], reviewer i9bM is correct, and I believe that comparing with [1] is a key step in demonstrating the efficiency of the proposed method. In addition, I would suggest extending the experiments to include real-world datasets.

[1] Shirakawa, Toru, et al. "Longitudinal Targeted Minimum Loss-based Estimation with Temporal-Difference Heterogeneous Transformer." International Conference on Machine Learning. PMLR, 2024.