Enhancing Counterfactual Estimation: A Focus on Temporal Treatments

1 Introduction

Accurate estimation of counterfactuals over time is crucial for evaluating the temporal effects of different treatment strategies, which can optimize medical decision-making and significantly impact healthcare [Yazdani and Boerwinkle, 2015]. While randomized controlled trials are the gold standard for causal inference [Hariton and Locascio, 2018], estimating counterfactuals from observational data is gaining attention due to the high costs and ethical constraints associated with conducting these trials in real-world settings.

Recent neural network techniques have advanced this field, particularly by focusing on integrating treatments within network architectures, supported by evidence in static causal inference settings. For example, [Shalit et al., 2017] develop a two-head network architecture for binary treatment observational data, influencing many subsequent studies [Shi et al., 2019; Hassanpour and Greiner, 2020; Johansson et al., 2022]. For continuous treatments, [Nie et al., 2020] introduce the Varying Coefficient Network (VCNet), enabling the prediction network to function continuously with treatment, thus strengthening the impact of treatment on predictions.

In longitudinal settings, current and historical treatments, viewed as time series, often jointly influence future patient outcomes through complex temporal interactions. For example, understanding the intricate temporal interactions between antibiotics is crucial for optimizing antibiotic usage and minimizing the development of antibiotic resistance [Roemhild et al., 2022]. However, existing methods often overlook the role of treatments in their model design. For example, Counterfactual Recurrent Network (CRN) [Bica et al., 2020] and Causal Transformer (CT) [Melnychuk et al., 2022] learn representations from historical information and then concatenate them with current treatments as input to a feedforward neural network for counterfactual estimation. This simple concatenation of current treatments with historical information embedded in representations may overly simplify interactions between temporal treatments, thus undermining the model's ability to capture the complex dynamics that evolve over time.

Another challenge in counterfactual estimation when observational data as time series is the complex confounding bias introduced by time-varying confounders. CRN and CT attempt to mitigate this issue by learning a balanced representation that excludes current treatment assignment information. However, in real-world applications, current treatments are often closely related to historical treatments. Our research demonstrates that adopting such a balanced representation can lead to the loss of historical treatment information, which may adversely affect counterfactual estimation.

The challenges identified make it difficult for existing methods to effectively and comprehensively utilize temporal treatment information. Our study models the evolution of patient states as a dynamical system. Subsequently, using the Koopman operator [Koopman, 1931], which linearizes nonlinear systems to effectively model complex dynamics, we propose a novel model named the Counterfactual Temporal Dynamics Network via Neural Koopman Operators (CTDNKO). Figure 1 shows CTD-NKO's integration of two Recurrent Neural Network (RNN) modules. One RNN is tasked with learning the system state, while the other combines current and historical treatment information to learn a causal Koopman operator. This architecture capitalizes on the ability of Koopman operators to model complex nonlinear systems, thus enhancing the capacity of capturing causal dynamics induced by temporal treatments. Additionally, CTDNKO learns an intrinsic Koopman operator to capture inherent system dynamics, such as the circadian rhythm of heart rate [Massin et al., 2000].

Furthermore, CTD-NKO introduces a novel weighting method to mitigate confounding bias by aligning the joint and marginal distributions of the system state and the current treatment. This approach avoids the loss of historical treatment information caused by learning a balanced representation that excludes current treatment assignment information, further improving counterfactual estimation performance. CTD-NKO offers several key advantages over existing methods, including unified treatment representation, separation of causal and non-causal dynamics, and enhanced expressiveness in modeling complex temporal dependencies (see Appendix A for detailed discussion). Overall, our main contributions are threefold:

• CTD-NKO improves the capability of counterfactual estimation over time by effectively learning complex interactions among temporal treatment using Koopman operators.
• We theoretically point out that using treatment-invariant balanced representations to mitigate confounding bias may lead to the loss of historical treatment information. To avoid this issue, we alternatively propose a novel weighting method.
• Experimental results on both synthetic and real-world datasets demonstrate that CTD-NKO achieves state-ofthe-art performance and efficiency.

2 RelatedWork

Counterfactual estimation over time. Current leadingedge methodologies for analyzing time-varying outcomes utilize advancements in deep neural networks. Notable examples include RMSN [Lim et al., 2018], CRN [Bica et al.,

2020], G-Net [Li et al., 2021], and CT [Melnychuk et al., 2022]. RMSN incorporates two propensity networks and uses a training approach based on Inverse Probability of Treatment Weighting (IPTW) for its prediction models. G-Net enhances the conventional G-computation technique via a deep learning framework. Both CT and CRN, on the other hand, focus on creating balanced representations that effectively predict outcomes while not being predictive of current treatment allocations. These methods often learn historical and current treatment information separately, which is not conducive to exploring complex dynamics, especially the causal dynamics driven by temporal interactions between treatments. A recent study [Kacprzyk et al., 2024] introduces a method based on ordinary differential equations (ODE) rather than neural networks for estimating counterfactuals over time. This method offers improved interpretability and the ability to handle irregularly sampled data, opening new avenues for research. However, its reliance on ODE discovery may impose limitations due to strong assumptions, e.g., the functional form of the ODE involved. A more comprehensive survey of relevant literature on causal inference can be found in Appendix B.

Koopman operator. Koopman operator [Koopman, 1931] explores how nonlinear dynamics can be linearized through an infinite-dimensional operator on Koopman space, often approximated by Dynamic Mode Decomposition (DMD) [Brunton et al., 2016]. Recent advancements in integrating Koopman operator with machine learning, particularly through the use of Deep Neural Networks (DNNs) [Takeishi et al., 2017; Morton et al., 2019; Yeung et al., 2019; Han et al., 2020; Fan et al., 2022], have greatly enhanced the ability to directly derive measurement functions from data. These innovations encompass the utilization of DNNs to construct Koopman invariant subspaces and the dynamic adaptation of Koopman operators to accommodate evolving system dynamics, marking a significant departure from conventional static approaches. For example, [Takeishi et al., 2017] propose a data-driven approach to learn Koopman invariant subspaces (LKIS), enhancing the ability of DNNs for Koopman spectral analysis. [Brunton et al., 2022] explore various methods, including data-driven techniques such as DNNs, that utilize predefined functions to improve the Koopman operator's ability to model and analyze nonlinear dynamical systems by approximating their underlying linear structure.

However, our work differs from these studies in its focus. Rather than developing novel theories related to Koopman operator, we are more inclined towards its application. This aligns with recent works that employ Koopman operators to address temporal domain adaptation. For instance, KNF [Wang et al., 2023] employs predefined measurement functions for learning a Koopman operator and an attention map, targeting temporal distribution changes in forecasting. Concurrently, Koopa [Liu et al., 2023] introduces an inventive architecture, merging deep residual structures with Koopman operators to enhance efficiency and performance in managing non-stationary time series. We observe that current temporal counterfactual estimation methods often overlook the interaction between treatments over time in their model design. Our paper applies Koopman operators to encapsulate sequential treatment data over time, distinguishing it from existing methods by effectively addressing this oversight.

3 Background

3.1 ProblemFormulation

Consider an i.i.d. observational dataset D, which contains

detailed information of N patients. Mathematically, it can be represented as D = {x(i) t , a(i) t , y(i) t T (i) t=1 ∪ v(i)}N i=1 . For each patient indexed by i, the observations include timevarying covariates X(i) t ∈ X, treatments received A(i) t ∈ A, and outcomes Y(i) t ∈ Y at discrete time steps T (i) . Additionally, patients' static covariates (e.g., gender and age) are denoted as V(i) ∈ V. To simplify the notation, the patientspecific superscript (i) will be omitted when it does not affect the context understanding.

Based on the potential outcome framework proposed by Robin [Rubin, 1978] and its extension to accommodate time-varying treatments by Robins and Hernan [Robins and Hernan, 2008], we aim to estimate time-varying counterfactual outcomes, following previous studies [Bica et al., 2020; Lim et al., 2018; Li et al., 2021]. Let the patient's historical information be ¯Ht = ( ¯Xt, ¯At-1, ¯Yt,V), where ¯Xt = (X1, · · · ,Xt), ¯Yt = (Y1, · · · ,Yt), and ¯At-1 = (A1, · · · ,At-1). Our goal is to estimate the potential outcome Yt+τ [¯at:t+τ-1] after a sequence of treatments ¯at:t+τ-1 = (at, · · · , at+τ-1), given the patient's historical information ¯Ht, i.e.,

The identifiability of treatment effects from observational data relies on assumptions outlined in prior research [Melnychuk et al., 2022; Lim et al., 2018]: consistency, sequential ignorability, and sequential overlap. These are detailed in Appendix C.

3.2 Koopman Operator

A discrete-time dynamical system is represented by ut+1 = F(ut), where ut denotes the system state and F describes the dynamics. Identifying direct transitions between states can be challenging due to nonlinearity or noise [Eivazi et al., 2021; Morton et al., 2018]. Koopman theory [Koopman, 1931], however, allows for the projection of the state ut into a measurement function space (a.k.a. Koopman space) g, managed by an infinite-dimensional linear operator K, such that:

Koopman operators can effectively model complex nonlinear dynamics [Koopman, 1931; Yeung et al., 2019; Li and Jiang, 2021], leveraging linear operators to advance understanding of dynamical systems.

4 Methods

The CTD-NKO, illustrated in Figure 1, comprises two primary RNN modules. One RNN is tasked with learning the system state, while the other is dedicated to learning the Koopman operator that governs state evolution. Subsequently, the model employs the evolved state for counterfactual prediction. We will next detail the problem modeling and the specifics of our method. Code available at https://github.com/wangxin0126/CTD-NKO IJCAI.

4.1 Modeling Patient State Dynamics via Koopman Theory

In this study, we model patient state evolution as a dynamical system, with transitions intricately linked to historical data, especially prior treatments. These transitions usually display complex dynamics due to temporal treatment interactions. To capture this complexity, we utilize Koopman theory in our modeling approach.

Inspired by the synergy between Koopman theory and machine learning [Takeishi et al., 2017; Morton et al., 2019; Fan et al., 2022; Wang et al., 2023; Liu et al., 2023], we employ DNNs alongside a set of measurement functions G := [g1, · · · , gn], each mapping R → R, to learn the dynamics of complex systems. These functions include nonlinear mappings such as exponential functions, enhancing the DNNs' ability to capture the inherent nonlinearity of dynamical systems [Wang et al., 2023; Kutz et al., 2016; Brunton et al., 2022]. Specifically, an encoder Φr is used to extract the latent state vector rt = Φr( ¯Ht) ∈ Rn from historical data. We then apply these measurement functions to project the learned state into the Koopman space:

where upper indices indicate vector components. For clarity, we denote the system state in the Koopman space by st = Ψ( ¯Ht) = G(Φr( ¯Ht)) ∈ S, where Ψ is the composition of G and Φr .

In counterfactual estimation for longitudinal data, the evolution of a system can be naturally divided into two types: causal and non-causal. Causal evolution primarily arises from the interaction between the treatment sequence and other historical information, while non-causal evolution reflects the inherent patterns of the system itself. For instance, in the absence of medication, a patient's heart rate or blood pressure typically exhibits periodic variations. As such, we define causal Ktc and intrinsic Ki Koopman operators to capture these evolutionary processes, allowing us to describe the system's evolution in the Koopman space as follows:

4.2 Counterfactual Estimation

Consistent with prior studies [Wang et al., 2023; Liu et al., 2023], we utilize DNNs to derive a matrix representation of the Koopman operator. For causal evolution, we employ an encoder Φc to generate the low-dimensional representation Φc( ¯Ht, At), which is subsequently transformed as an n × n matrix representing the causal Koopman operator ˆKtc . In our implementation, both Φr and Φc are implemented using Long Short-Term Memory units (LSTM), with the parameters denoted as θΦ. For non-causal dynamics, a separate matrix forms the intrinsic Koopman operator ˆKi. Given Equation 4, we compute the estimated system state ˆst+1 in the Koopman space for the next time step.

Subsequently, this allows for counterfactual predictions based on ˆst+1. CTD-NKO adopts an autoregressive recursive strategy [Chevillon, 2007; Taieb and Atiya, 2015] for multistep-ahead prediction, which is also employed in G-Net [Li et al., 2021]. Therefore, during the training process, we need to predict the output and time-varying covariates for the next step. To perform output prediction, we define a feedforward neural network GY , parameterized by θY , to decode the expected output from ˆst+1. We use the Mean Squared Error (MSE) to define the following loss function:

To predict the covariates, we design a module FX that consists of two feedforward neural networks, GX and JX, with parameters represented by θX. Similar to the output prediction, we use GX to decode the expected covariates from ˆst+1. Some covariates may change at a slower pace, such as cholesterol levels. Inspired by the gating mechanism in GRUs, we utilize JX to design a smoothing mechanism that adapts to this change trend:

where ρ = Sigmoid(JX(ˆst+1)) is used to regulate the smoothing degree, achieving a balance between historical observations and predicted values. Similarly, we define the following loss function:

Furthermore, to encourage the model to learn Koopman operators correctly, we define the following loss function based on the 'true' st+1 learned by the encoder:

4.3 Balancing viaWeighted Factual Loss

The aforementioned modeling approach facilitates accurate next-step outcome predictions. Let Kt = Ktc +Ki ∈ ΩK, and let f : S ×ΩK → Y and L represent the prediction function and the loss function, respectively. The marginal error that we aim to minimize is then defined as:

which is consistent with the evaluation method proposed by [Melnychuk et al., 2022] and reflects the model's performance in estimating counterfactual distributions. Given the unobservability of counterfactual outcomes, we are constrained to estimate only the factual error ϵF from the observed data:

However, due to confounding bias, the distribution during counterfactual evaluation often diverges from the factual distribution, suggesting that using factual loss as a surrogate for marginal loss could introduce biases. To address this issue, [Bica et al., 2020] and [Melnychuk et al., 2022] propose to learn a treatment-invariant balanced representation Φ( ¯Ht) through an encoder Φ. This aims to equalize P(Φ( ¯Ht)) and P(Φ( ¯Ht)|At), decoupling the representations of historical data from current treatments, thereby reducing confounding. However, this approach may compromise the precision of counterfactual estimates, leading us to propose Lemma 4.1 for further clarification.

Lemma 4.1. Let Zt = Φ( ¯Ht), and let I(·; ·) denote the mutual information. Suppose I( ¯At-1;At) > 0. When completely eliminating the correlation between the representation of the historical information and the current treatment, i.e., I(Zt;At) = 0, the information about ¯At-1 contained in Zt must be lossy relative to the historical information ¯Ht:

Lemma 4.1 suggests that using treatment-invariant balanced representations leads to the loss of historical treatment information, which may have a negative impact on accurately predicting future treatment outcomes. To address this, we propose an alternative approach to mitigate confounding bias, allowing the weighted-adjusted ϵF to be used as an estimate of ϵM. The key idea of our method is to minimize the distance between the weighted joint distribution and the marginal distribution of the system state St and the current treatment At. To achieve this, we introduce the Integral Probability Metric (IPM), a measure of the distance between two probability distributions Q and P :

where ζ is the concatenation of St and At, andM is a family of functions m : S × A → R.

Theorem 4.2. Assume Ψ be a one-to-one representation function. When the weight w satisfies IPM(Pw(St,At), P(St)P(At))=0, then w can equate factual and marginal errors, i.e.,

Theorem 4.2 indicates that by minimizing the distance between the weighted distribution Pw(St,At) and P(St)P(At), we can obtain a set of weights (over time) that make the adjusted factual error close to the marginal error. In this study, we choose the family of 1-Lipschitz functions, where the employed IPM is known as the wasserstein distance [Villani and others, 2009]. For detailed proofs, please see Appendix D.

Algorithm 1 Pseudocode of Training CTD-NKO

Ensure: Optimized parameters θ, EMA parameters θEMA

4.4 Training and Inference

Algorithm 1 presents the pseudocode for the training process of CTD-NKO, with inputs including the observed data D and other necessary parameters, and outputs being the optimized model parameters. For detailed parameter settings, please refer to Appendix H. During training, we adopt the Exponential Moving Average (EMA) strategy [Tarvainen and Valpola, 2017] to obtain more reliable results, following the CT study [Melnychuk et al., 2022]. The first line of Algorithm 1 initializes the EMA parameters θ0EMA , and in Line 15, the parameters are iteratively updated as follows:

where j denotes the iteration number and β denotes the exponential smoothing factor.

In Lines 3-5 of Algorithm 1, we update the weights every c epochs. The weight calculation is based on Theorem 4.2, which involves minimizing the IPM distance between the weighted joint distribution and the marginal distribution of the system state representation and the current treatment. We obtain samples from the joint distribution and the marginal distribution by applying the actual treatment and the shuffled treatment to the observed data, respectively. Then, we learn weights by minimizing the wasserstein distance between the weighted joint distribution and the marginal distribution. For specific details, please refer to Appendix G. The model parameters are updated in Lines 8-13 of Algorithm 1, with the aim of minimizing the following objective function:

Here, λX denotes a predefined weight coefficient, and λKp denotes a weight coefficient that gradually increases with the training epochs. This design takes into account that the learning of st may not be accurate enough in the early stages of training; therefore, the importance of LK(θΦ, t) is reduced during the early training phase. We implement CTD-NKO using the Pytorch Lightning framework and employ the Adam algorithm [Kingma and Ba, 2014] for gradient optimization. After training, CTD-NKO performs one-step-ahead predictions and uses an autoregressive approach for multi-stepahead predictions.

5 Experiments

In this section, we conduct a series of experiments to validate the effectiveness of the proposed CTD-NKO. Following the standard workflow of the Counterfactual Estimation Benchmark [Melnychuk et al., 2022], we compare CTD-NKO with existing models on both synthetic and real-world datasets. We then analyze the efficiency of the baseline methods and CTDNKO on various data settings. Finally, we investigate the role of different components of CTD-NKO through ablation studies.

In this study, we select several state-of-the-art models from the recent literature on estimating time-varying counterfactuals to serve as comparative baselines. These include neural network-based models like RMSN [Lim et al., 2018], CRN [Bica et al., 2020], G-Net [Li et al., 2021], and CT [Melnychuk et al., 2022], as well as the non-neural networkbased model INSITE [Kacprzyk et al., 2024]. To guarantee a fair comparison, we perform hyperparameter tuning for these baseline methods (refer to Appendix H for details).

5.1 Counterfactual Estimation Performance Comparison

Experiments with FS-Tumor Dataset

Data. The FS-Tumor dataset has been widely adopted in previous studies evaluating counterfactual estimation over time, such as [Bica et al., 2020; Melnychuk et al., 2022; Lim et al., 2018; Kacprzyk et al., 2024]. In the dataset, a parameter γ in the biomathematical model of tumor controls the time-varying confounding factors. An increase in the parameter γ signifies a greater influence of historical data on treatment allocation, implying a more pronounced presence of confounding biases. For more details, please refer to Appendix F

Results. In this study, we select the range of γ from 0 to 4. It is worth noting that γ = 0 represents a scenario where the treatment is assigned completely at random, indicating the absence of confounding bias. Figure 2 illustrates the comparison of τ-step-ahead prediction results, where τ ∈ {2, 4, 6}. For more comprehensive experimental results, please refer to Appendix F. The results in Figure 2 demonstrate that CTDNKO exhibits the best prediction accuracy in the majority of cases, with the lowest Root Mean Square Error (RMSE). This advantage becomes more pronounced as the confounding bias increases, i.e., when γ is larger.

Experiments with MIMIC-III Datasets

Data. MIMIC-III is a comprehensive database that encompasses electronic health records of patients in the intensive care unit and has been widely utilized to evaluate the performance of various models in complex real-world medical settings. In accordance with the related studies [Hatt and Feuerriegel, 2021; ?; ?], we construct a Real-World MIMIC (RW-MIMIC) dataset by selecting 25 dynamic patient covariates and 3 static features from the MIMIC-III database. Our primary focus was to investigate the impact of two common treatments, i.e., vasopressors and mechanical ventilation, on patients' blood pressure.

As a real-world data source, MIMIC-III does not contain information on counterfactual outcomes. To control for confounding bias and obtain counterfactuals for evaluation while analyzing high-dimensional patient trajectories, we construct a semi-synthetic dataset based on MIMIC-III, termed SSMIMIC, following the methodology described in [Melnychuk et al., 2022]. This dataset is generated considering treatment effects, endogenous dependencies, and exogenous dependencies, building upon the research methods of [Schulam and Saria, 2017]. For more details on these two datasets, please refer to Appendix F

Results. Table 1 and Table 2 present the τ-step-ahead prediction results on the RW-MIMIC and SS-MIMIC datasets, respectively. The results demonstrate that CTD-NKO consistently exhibits superior performance on both datasets, achieving the lowest RMSE and a relatively small standard deviation. These experimental findings suggest that CTD-NKO can effectively manage complex long-term dependencies and may be well-suited for applications that closely mirror realworld complexities. Notably, we do not compare INSITE on the MIMIC-III dataset due to the challenges posed by its high-dimensional, time-varying covariates. For more detailed reasons, please refer to Appendix G.

5.2 Model Efficiency Analysis

In practical applications, both operational efficiency and predictive performance are crucial. Hence, we evaluate the CTD-NKO against neural network-based models across several datasets in three key aspects: the τmax-step-ahead prediction RMSE, training speed, and peak GPU memory usage. The term 'peak GPU memory usage' refers to the maximum amount of memory utilized by a GPU during training. This metric is instrumental in assessing a model's efficiency and scalability [Ikuzawa et al., 2016; Nie et al., 2022; Bergner et al., 2023], thereby enabling optimal hardware utilization and cost-effectiveness in applications.

Figure 3 presents the comparative results: On the FSTumor, RW-MIMIC, and SS-MIMIC datasets, compared to the state-of-the-art model CT, CTD-NKO reduces training time per epoch by 69.3%, 96.6%, and 89.1% respectively, while the peak GPU memory usage is only 36.4%, 16.6%, and 12.6% of that consumed by the CT model. The performance advantage of the CT model stems from its adoption of the more powerful Transformer architecture; however, this also leads to larger memory consumption, particularly when handling complex data. In contrast, models based on LSTM architecture, while less demanding in terms of memory, tend to exhibit slightly inferior predictive performance. The CTDNKO model, applying the Koopman theory and utilizing a straightforward LSTM architecture, achieves a good balance between performance and computational cost. Furthermore, CTD-NKO, similar to G-Net, utilizes an autoregressive recursive strategy for multi-step-ahead predictions. This strategy, focusing solely on encoder training, boosts efficiency compared to the CRN's encoder-decoder method, which requires training multiple components.

5.3 Ablation Study

To assess the significance of various components within the CTD-NKO model, we conduct ablation studies to compare the predictive performance of the full model against variants lacking specific components across different γ settings on the FS-Tumor dataset, as presented in Table 4. Specifically, 'w/o weighting' denotes the model's use of uniform weights by omitting the weighting balancing module; 'w/o global operator' indicates that only the causal Koopman operator is utilized for modeling the temporal state transitions.

When the time-varying confounding bias is larger, the disparity between the counterfactual and factual distributions increases. The experimental results demonstrate that under these conditions, the learned balancing weights of CTD-NKO

can mitigate this issue, thus leading to a more pronounced performance improvement. Additionally, the global Koopman operator enhances the model's performance, with its impact growing more substantial at farther prediction steps τ and larger confounding bias levels. Finally, we compare the proposed weighting method with the domain confusion loss strategy introduced by CT on CTD-NKO and test it on FSTumor (γ = 4). The results in Table 3 demonstrate that our weighting strategy outperforms learning balanced representations, which is consistent with our theoretical analysis.

6 Conclusion

Our work emphasizes the importance of effectively modeling temporal treatment interactions and preserving historical treatment information in counterfactual estimation over time. We propose the Counterfactual Temporal Dynamics Network via Neural Koopman Operators (CTD-NKO), which models temporal counterfactual reasoning as an evolutionary process of a nonlinear dynamical system, capturing complex temporal treatment interactions and inherent rhythmic dynamics. Additionally, we introduce a novel weighting method to mitigate confounding bias while reducing the loss of historical treatment information. These designs enable CTD-NKO to outperform state-of-the-art methods in terms of performance and efficiency, providing a powerful tool for improving medical decision-making and optimizing treatment strategies.

Acknowledgements

This research is supported by National Key RD Program of China (No. 2021ZD0111700), National Nature Science Foundation of China (No. 62137002, 62176245), Special Foundation for Science and Technology Innovation and Entrepreneurship of CCTEG NO. 2020-2-TD-CXY006, Key Science and Technology Special Project of Anhui Province (No. 202103a07020002). We thank the anonymous reviewers for their constructive comments that help improve the manuscript.