Treatment Rule Optimization Under Counterfactual Temporal Point Processes with Latent States

Authors: Xinye Chen, Gegenduolan Zhang, Yinghao Fu, Chao Yang, Shuang Li
Primary area: interpretability and explainable AI
Keywords: counterfactual reasoning, temporal point processes, latent confounder, rule learning

Abstract:
In high-stakes areas like healthcare, retrospective counterfactual analysis—such as evaluating what might have happened if treatments were administered earlier, later, or differently—is vital for refining treatment strategies. This paper proposes a counterfactual treatment optimization framework using temporal point processes to model outcome event sequences. By sampling potential outcome events under new treatment decision rules, our approach seeks to optimize treatment strategies in a counterfactual setting. To achieve accurate counterfactual evaluation of new decision rules, we explicitly introduce latent states into the modeling of temporal point processes. Our method first infers the latent states and associated noise, followed by counterfactual sampling of outcome events. This approach rigorously addresses the complexities introduced by latent states, effectively removing biases in the evaluation of treatment strategies. By proving the identifiability of model parameters in the presence of these states, we provide theoretical guarantees that enhance the reliability and robustness of the counterfactual analysis. By incorporating latent states and proving identifiability, our framework not only improves the accuracy and robustness of treatment decision rules but also offers actionable insights for optimizing healthcare interventions. This method holds significant potential for improving treatment strategies, particularly in healthcare scenarios where patient symptoms are complex and high-dimensional.

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

Reviews and discussion:
<thread>
creation date:
1729170809019

thread_type:
official_review

<thread_object>
creator: reviewer_vw1g

<content>
summary:
The paper proposes a new method for counterfactual prediction and policy learning in a continuous time setting. The authors leverage a probabilistic approach and train their model using a variant of the EM algorithm. The model is evaluated on both simulated and real-world data.

strengths:
- The paper targets an important problem, relevant in many disciplines such as healthcare or economics
- The proposed method is flexible and the probabilistic formulation potentially allows for uncertainty quantification
- The method allows for both counterfactual prediction and policy learning

weaknesses:
- The problem formulation is a little sloppy. For example, in Eq. (7) you use the do operator but never define an appropriate SCM or potential outcomes to work with. Also, notation sometimes seems confusing and redundant, e.g., treatment times and events are first denoted with $t, m_a$ and later with $x, \tau$.
- The related work on counterfactual prediction should be expanded. There are various methods for both discrete and continuous time that have not been cited. Just to list a few: https://papers.nips.cc/paper_files/paper/2018/hash/56e6a93212e4482d99c84a639d254b67-Abstract.html, https://arxiv.org/abs/2204.07258, https://arxiv.org/abs/2002.04083, https://arxiv.org/abs/2310.17463, https://arxiv.org/abs/2407.05287
- It would be helpful if the authors could spill out their contribution as compared to previous work more clearly. In their related work, the authors claim that their method can handle unobserved confounders which is not true for previous work. However, the latent variables $z$ considered in their papers are not confounders (see Fig. 1).
- Related to the above, the novel ideas in this paper seem rather limited. The main contribution seems to be the specification of the latent variable model. The learning algorithm (inner vs outer loop) seems to correspond to a classical policy learning algorithm (where first the policy value is estimated and then the policy is updated) combined with a standard EM algorithm
- In Eq. (7), the authors only condition on the latent states z but not on observed variables (e.g., outcome or treatment history). This does not match with the motivation of the paper (e.g., in the introduction: "Given the observational treatment and outcome trajectories, can we modify specific treatment actions...?")
- The model relies on a variety of parametric assumptions which can restrict flexibility. 
- If I understand correctly, the model does not account for (observed) time-varying confounders. These are common in medicine and often available in electronic health records. Taking into account such time-varying confounding necessitates the usage of specialized adjustment mechanisms to correctly estimate counterfactual outcomes/ treatment effects (see here: https://arxiv.org/abs/2407.05287). Ignoring these confounders can severely limit the method's practical applicability.
- The method relies on various assumptions that are not spelled out explicitly, e.g., consistency (no spillover effects), positivity, and ignorability (no unobserved confounders).
- Identifiability: Theorem 1 assumes that the latent variables $z$ are identified in order to identify the model parameters. How can this be ensured in practice? Furthermore, I somewhat doubt the validity of the theoretical arguments as the authors never make use of assumptions necessary to identify counterfactuals or treatment effects (see the point above).

questions:
- Why does the proposed method work with "meta-rules" instead of learning an optimal policy in a fully data-driven manner? It seems like the aim is to "hard-code" certain aspects of the policy. I think this should be justified e.g., with references to relevant literature.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729570474030

thread_type:
official_review

<thread_object>
creator: reviewer_um1n

<content>
summary:
The paper offers a solution for finding optimal treatment strategies in the presense of discrete-time latent states, continuous-time treatment and outcome processes. The paper represents treatment and outcome processes as marked point processes and model them through Hawkes process given past history and latent states. It then treat the latent states and unknowns from models as parameter of interest and apply an EM algorithm to identify.

strengths:
The Hawkes process modeling and the assoicated EM algorithm are intuitive and natural. The paper is contributing into the field of little consideration and can be novel.

weaknesses:
1. The parametric and highly linear modeling of intensity process can be highly questionable. As in a data-rich environment in ICU with great complication and unknown influences across latent states, treatment, and outcome, I found such a modeling untrustworthy. The concern is further aggrevated when one needs to choose the number of latent states by themsselves (and maybe also timing?), in your real-data application. This concern can be somewhat relieved in two ways:
 a. Justify the usage of such modeling in the literature to show this can be a common and acceptable practice, especially in the presence of latent states and continuous-time treatment and outcomes.
 b. In simulation, investigate scenarios under which the modeling assumption (3) is lightly, moderately, and highly violated and investigated  the performance of your framework.
 c. Can you explore the possibility of relaxing modling assumptions, for example, maybe only the latent states part need to be linear while the part involving history of treatment and outcomes can be non-parametric? I believe this is not too much to ask, given that in a similar continuous-time setting, a complete nonparatric identification can be achievable without any parametric assumptions, for example, Continuous-time targeted minimum loss-based estimation of intervention-specific mean outcomes, HC Rytgaard, TA Gerds, MJ van der Laan The Annals of Statistics 50 (5), 2469-2491.

2. We did not observe the latent states. This not only means its value but also the number of latent states and when they happened. The authors seem to be providing this information in the simulation to their estimator. But in practice, how do choose and decide the number of latent states and when they happened? This can be very difficult and unrealistic to assume.
 a. Can you specify how to select the number of latent states adaptively?
 b. Explain how you determined the number and timing of latent states in their real data analysis
 c. Conduct a sensitivity analysis varying the number of latent states
 
3. It is commonly known that EM algorithm that this paper heavily relies on, suffers from convergence issue and identification issue. Therefore, Theorem 1 is one key ingredient for the paper to be useful. I think this concern about EM algorithm can be relieved by a clear presentation of assumptions in Theorem 1 (instead of citing assumptions from Bonnet et al. (2023). You need an extensive discussion on these assumptions:
 a. When they hold or fails;
 b. How to interpret them both theoretically and in practice (e.g., in your real data application);
 c. Are you able to verify them in your simulations?
 d. Are we able to test them or conduct sensitivity analysis on them?

questions:
Clarify a confusion:
1. One point I am quite doubtful and thus lowering my scoring is. Why did you only model the intensity process for the timing of treatment and outcome but not the mark? In Equation (8) to (9), the likelihood of m_j suddenly disappeared. I found that highly counterintuitive because the dosage of treatment (m_t, j) and disease progression (m_o, j) shall influence the intensity process greatly. Also, their own distribution shall contribute to the likelihood as well.
 a. If marks are not modeled, explain the rationale and discuss potential limitations
 b. Consider extending the model to explicitly incorporate marks in the intensity and likelihood
2. The last point brings up this question, if you did not model the mark of the outcome process, how did you compute the gradient update on Page 6 line 311?
3. Can you clarify what the baseline means in Figure 3? I searched "baseline" in the paper and still is confused.
4. In Step 1: Initialization in Section 6, how to initialize treatment decision rule parameters? Randomly?
5. In Section 3.2, you used x to represent dosage, did you switch to m_{d, k} in Section 6?
6. Page 3 line 151, do you mean z(t) = [z_k]_{k <= t}?

Address a limitation:
1. In the simulation, apart from with latent states but using EM, without latent states by ignoring latent states and apply MLE, one can add a third one by pretending we observe latent states and directly apply your decision rule optimization algorithm in Section 6 (is this "baseline"?)
</content>
</thread_object>

</thread>
<thread>
creation date:
1729994623766

thread_type:
official_review

<thread_object>
creator: reviewer_y7ml

<content>
summary:
The paper proposes a counterfactual treatment optimization framework based on temporal point processes for evaluating and optimizing treatment decision rules in high-stakes areas like healthcare. The framework consists of an outer loop for exploring and optimizing treatment decision rules, and an inner loop for evaluating these rules via counterfactual sampling of symptom events. To address challenges posed by latent states, a two-stage procedure is introduced - first inferring latent states and associated noise to mitigate biases, and then conducting counterfactual sampling. Identifiability of model parameters in the presence of latent states is theoretically proven, enhancing robustness. The framework aims to refine existing treatments and generate insights for optimizing patient outcomes.

strengths:
1. Tackling a crucial problem of treatment optimization under latent confounding in a counterfactual manner.

2. Integrating temporal point processes to model event sequences rigorously.

3. The paper establishes strong theoretical guarantees through comprehensive identifiability proofs. This mathematical rigor ensures the reliability of counterfactual evaluations and strengthens the practical applicability of their approach.

weaknesses:
1. The notation of the do-operator seems a bit confusing. For example, on line 212, it would be helpful if the meaning of $do(H_a(T) = H'_a(T')|{f_1, \ldots, f_D}, {z_1, \ldots, z_K})$ could be mathematically illustrated.

2. The authors did not compare their approach with other off-policy optimization methods in the experiments, which might raise some doubts about the effectiveness of their algorithm. 

3. While the paper provides detailed reward optimization results for synthetic experiments (Section 7.1), presenting similar quantitative analysis for the MIMIC-III database (e.g., convergence plots and reward comparisons similar to Figure 3) would help readers better appreciate how the proposed framework translates to real-world healthcare scenarios.

questions:
1. What if the number of latent factors K is unknown, or the latent states are not categorical - will the identifiability result still establish under such relaxed assumptions?
2. How can the framework handle scenarios where meta-rules are not well-defined or unavailable?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730738555495

thread_type:
official_review

<thread_object>
creator: reviewer_srx6

<content>
summary:
The paper develops a new method for time-varying off-policy learning based on confounded irregularly-sampled data. The authors assume observational data is generated by a Hawkes process with discrete latent variables, which capture hidden confounding. Then, given the observational data, they (1) employ an expectation-maximisation (EM) algorithm to infer the latent variables, (2) generate counterfactual outcomes, and (3) perform gradient-based policy learning. The method was evaluated based on one synthetic and one real-world experiment.

strengths:
The paper aims to tackle a challenging problem of causal inference and reinforcement learning: time-varying off-policy learning based on the irregularly sampled, confounded observational data and there are not that many related methods tailored for this setting.

weaknesses:
I find it very hard to acknowledge the contribution of the paper due to several major flaws:
1. It was unclear, what causal assumptions the paper relies on. Specifically, does one need to assume a non-Markovian time and time-varying potential outcomes framework [1],  or some Markovian decision process (MDP)? Also, additional identifiability assumptions are required due to the presence of hidden confounding [2-3]. I am not even sure, whether the paper works with the interventional quantities (e.g., potential outcomes after the interventions in the future), or counterfactual quantities (e.g., the outcomes after intervention on the past treatments). The latter one, for example, requires even stronger assumptions [4] (e.g., invertibility of the latent noise of the outcome).  Therefore, it is hard to judge whether the proposed method is sound or even correct.
2. The authors did not provide any derivations for the identification of the target causal quantity, i.e., policy value after the intervention (Theorem 1 relates to the identifiability of the latent variables, which is not the same as policy value identification).   
3. Lack of baselines. Although a lot of relevant methods are mentioned in the related works section, none are provided in the synthetic benchmark. 
4. Poor quality. Multiple notions and notations are not properly introduced or defined throughout the paper. For example, $\phi(t-s)$, (Eq. 3); triggering functions; $x$, (line 208); etc. Thus, it is hard to understand the theoretical claims of the paper.

As all the above-mentioned problems cannot be easily fixed during the rebuttal, I recommend rejecting this work.

  References:  
- [1] Robins, J. M. and Hern  ́an, M. A. Estimation of the causal effects of time-varying exposures. CRC Press, Boca Raton, FL, 2009.
- [2] Milan Kuzmanovic, Tobias Hatt, and Stefan Feuerriegel. Deconfounding temporal autoencoder: estimating treatment effects over time using noisy proxies. In Machine Learning for Health, pp. 143–155. PMLR, 2021.
- [3] Ioana Bica, Ahmed Alaa, and Mihaela Van Der Schaar. Time series deconfounder: Estimating treatment effects over time in the presence of hidden confounders. In International conference on machine learning, pp. 884–895. PMLR, 2020.
- [4] Hızlı, Çağlar, et al. "Causal Modeling of Policy Interventions From Sequences of Treatments and Outcomes." arXiv preprint arXiv:2209.04142 (2022).

questions:
- Why are $H_0$ and $H_1$ two distinct nodes on the causal diagram in Fig. 1? Shouldn’t $H_0$ be a part of $H_1$?
-  What mixture distribution is meant in line 152? Isn’t a mixture of categorical variables also a categorical variable?
- What is the intuition of the counterfactual interventions in the past (namely both conditioning and intervening on the same treatments) from the practical point of view? Should regular future interventions be sufficient for developing optimal treatment regimes?
</content>
</thread_object>

</thread>


Metareview:
metareview:
Optimizing sequential policies is of much importance, and definitely a major AISTATS topic. The continuous-time modeling and counterfactual sampling show potential, but reviewers and myself continued to struggle in appreciating the novel aspects and the fundamental needs of counterfactual sampling. Despite the references given, following the finer details wasn't all that clear - for instance, counterfactual sampling requires some sort of cross-world inference, while more standard approaches for causal RL (including Q-learning with causal assumptions) wouldn't require that level of untestable modeling. Maybe there are clearer ways of presenting it, but currently even equations like (9) use non-standard notation like $p(z | do(x | y))$ (if I understood it correctly, the intention would be $p(z |y, do(X = f(y)))$? I wasn't sure) 

A promising set of ideas, but we concluded that the manuscript will have much to gain by another round of preparation.