CAETC: Causal Autoencoding and Treatment Conditioning for Counterfactual Estimation over Time

Nghia D. Nguyen; Pablo Robles-Granda; Lav R. Varshney

cs.LGMar 12, 2026

CAETC: Causal Autoencoding and Treatment Conditioning for Counterfactual Estimation over Time

Authors: Nghia D. Nguyen, Pablo Robles-Granda, Lav R. Varshney

Organizations: Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign · VinUni-Illinois Smart Health Center, VinUniversity · AI Innovation Institute, Stony Brook University · Department of Electrical and Computer Engineering, University of Illinois Urbana-Champaign

Abstract

Counterfactual estimation over time is important in various applications, such as personalized medicine. However, time-dependent confounding bias in observational data still poses a significant challenge in achieving accurate and efficient estimation. We introduce causal autoencoding and treatment conditioning (CAETC), a novel method for this problem. Built on adversarial representation learning, our method leverages an autoencoding architecture to learn a partially invertible and treatment-invariant representation, where the outcome prediction task is cast as applying a treatment-specific conditioning on the representation. Our design is independent of the underlying sequence model and can be applied to existing architectures such as long short-term memories (LSTMs) or temporal convolution networks (TCNs). We conduct extensive experiments on synthetic, semi-synthetic, and real-world data to demonstrate that CAETC yields significant improvement in counterfactual estimation over existing methods.

Figures & tables

Figure 1: Causal graph for time-dependent confounding over

\mathcal{H}_{t}

Figure 2: A. Prior works concatenate the representation

\Phi(\mathcal{H}_{t})

with treatment

A_{t+1}

to predict next outcomes at

F^{Y}

. We make a distinction to model treatment

A_{t+1}

as a transformation on the representation before being forwarded to respective heads. B. The history

\mathcal{H}_{t}

is encoded into representation

\Phi(\mathcal{H}_{t})

before being forwarded to respective heads for decoding. The treatment, outcome, and time-varying covariates estimators

F^{A}

F^{Y}

and

F^{X}

reconstruct

A_{t}

\bm{Y}_{t}

and

\bm{X}_{t}

. C. Simultaneously, balancing is applied to

\Phi(\mathcal{H}_{t})

by maximizing the entropy of the treatment balancer

F^{B}

. D.

A_{t+1}

-specific transformation is applied to the

\Phi(\mathcal{H}_{t})

before being decoded to the next outcomes by

F^{Y}

. E. The conditioning layer

F^{C}

is encouraged to learn treatment-specific transformations of

\Phi(\mathcal{H}_{t})

Figure 3: To handle input mismatch of the history

\mathcal{H}_{T_{0}}=\{\bm{V},A_{t},\bm{Y}_{t},\bm{X}_{t}\}_{t=1}^{T_{0}}

and autoregressively decoded sequences

\{\bm{V},A_{t},\bm{Y}_{t}\}_{t=T_{0}+1}^{T_{0}+\tau}

, we replace the future time-varying covariates

\{\bm{X}\}_{T_{0}+1}^{T_{0}+\tau}

with a learnable vector

\bm{M}

Figure 4: RMSEs for NSCLC fully synthetic data on random trajectories with increasing levels of time-dependent confounding on the training dataset. Each bar represents the RMSE for 5-step-ahead predictions. Lower is better.

Figure 5: RMSEs for NSCLC fully synthetic data on no-confounding test set (

\gamma=0

) with increasing levels of time-dependent confounding on the training dataset. Training on

\gamma=0

and testing on

\gamma=0

represents the performance upper bound for each method. Due to time-dependent confounding, training on

\gamma\neq 0

and testing on

\gamma=0

is expected to reduce performance. Lower is better.

	$\tau=1$	$\tau=2$	$\tau=3$	$\tau=4$	$\tau=5$	$\tau=6$	$\tau=7$	$\tau=8$	$\tau=9$	$\tau=10$	Avg
LSTM	.363 $\pm$ .028	.417 $\pm$ .025	.463 $\pm$ .021	.506 $\pm$ .020	.546 $\pm$ .017	.581 $\pm$ .016	.615 $\pm$ .016	.649 $\pm$ .015	.682 $\pm$ .016	.715 $\pm$ .017	.554
RMSN	.408 $\pm$ .029	.471 $\pm$ .032	.530 $\pm$ .030	.587 $\pm$ .030	.641 $\pm$ .030	.690 $\pm$ .030	.737 $\pm$ .028	.782 $\pm$ .025	.825 $\pm$ .022	.867 $\pm$ .023	.654
CRN	.365 $\pm$ .028	.424 $\pm$ .027	.476 $\pm$ .023	.524 $\pm$ .021	.569 $\pm$ .019	.608 $\pm$ .018	.647 $\pm$ .018	.685 $\pm$ .018	.723 $\pm$ .020	.760 $\pm$ .022	.578
CT	.324 $\pm$ .017	.411 $\pm$ .024	.499 $\pm$ .035	.582 $\pm$ .040	.660 $\pm$ .041	.733 $\pm$ .043	.808 $\pm$ .046	.882 $\pm$ .046	.941 $\pm$ .048	.999 $\pm$ .053	.684
CAETC LSTM	.330 $\pm$ .021	.393 $\pm$ .019	.442 $\pm$ .014	.487 $\pm$ .009	.528 $\pm$ .006	.564 $\pm$ .007	.600 $\pm$ .011	.634 $\pm$ .014	.668 $\pm$ .017	.702 $\pm$ .020	.535
CAETC TCN	.322 $\pm$ .016	.391 $\pm$ .019	.447 $\pm$ .023	.495 $\pm$ .027	.540 $\pm$ .029	.582 $\pm$ .030	.621 $\pm$ .032	.657 $\pm$ .032	.692 $\pm$ .032	.726 $\pm$ .032	.547

Table 1: RMSEs on MIMIC-III semi-synthetic data with counterfactual random trajectories. Lower is better.

	$\tau=1$	$\tau=2$	$\tau=3$	$\tau=4$	$\tau=5$	Avg
LSTM	6.80 $\pm$ 0.02	7.32 $\pm$ 0.05	7.61 $\pm$ 0.05	7.82 $\pm$ 0.05	7.99 $\pm$ 0.07	7.51
RMSN	7.29 $\pm$ 0.07	7.80 $\pm$ 0.07	8.10 $\pm$ 0.10	8.33 $\pm$ 0.16	8.54 $\pm$ 0.21	8.01
CRN	6.81 $\pm$ 0.04	7.35 $\pm$ 0.05	7.64 $\pm$ 0.06	7.86 $\pm$ 0.07	8.04 $\pm$ 0.10	7.54
CT	6.70 $\pm$ 0.04	7.37 $\pm$ 0.06	7.71 $\pm$ 0.07	7.93 $\pm$ 0.07	8.10 $\pm$ 0.04	7.56
CAETC LSTM	6.72 $\pm$ 0.05	7.27 $\pm$ 0.05	7.56 $\pm$ 0.06	7.76 $\pm$ 0.06	7.91 $\pm$ 0.08	7.45
CAETC TCN	6.69 $\pm$ 0.06	7.24 $\pm$ 0.05	7.53 $\pm$ 0.04	7.72 $\pm$ 0.03	7.88 $\pm$ 0.06	7.41

Table 2: RMSEs on MIMIC-III real-world data. Lower is better.

CAETC LSTM	$\tau=1$	$\tau=2$	$\tau=3$	$\tau=4$	$\tau=5$	Avg
Full	7.74 $\pm$ 0.57	9.01 $\pm$ 0.84	9.28 $\pm$ 0.78	9.09 $\pm$ 0.8	8.42 $\pm$ 0.72	8.71
$\delta^{C}=0$	7.75 $\pm$ 0.57	9.09 $\pm$ 0.86	9.45 $\pm$ 0.87	9.37 $\pm$ 0.98	8.79 $\pm$ 1.04	8.89
$\delta^{C}=0$ $\delta^{E}=0$	8.01 $\pm$ 0.91	9.58 $\pm$ 1.31	10.06 $\pm$ 1.44	10.05 $\pm$ 1.61	9.45 $\pm$ 1.49	9.43

Table 3: Ablation with the same setting of Fig. 5 with

\gamma=8

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Time-varying covariates

heart rate

red blood cell count

sodium

mean blood pressure

systemic vascular resistance

glucose

Appendix

Table 4: Selected time-varying covariates and static covariates from the MIMIC-Extract processed data.

Table 5:

Table 6:

Table 7:

=\big[\mathcal{L}^{RY}(\theta_{F^{Y}},\theta_{\Phi})+\mathcal{L}^{Y}(\theta_{F^{Y}},\theta_{F^{C}},\theta_{\Phi})\big]+\delta^{A}\big[\mathcal{L}^{RA}(\theta_{F^{A}},\theta_{\Phi})+\mathcal{L}^{C}(\theta_{F^{A}},\theta_{F^{C}},\theta_{\Phi})\big]+\delta^{X}\mathcal{L}^{RX}(\theta_{F^{X}},\theta_{\Phi})+\delta^{E}\mathcal{L}^{E}(\theta_{\Phi})

Appendix

Algorithm 1 Training algorithm for CAETC

Input :

\{A_{t^{\prime}}\}_{t^{\prime}=1}^{t}

\mathcal{H}_{t}

LSTM (# layers)

Linear Layer

ELU Activation

Output :

\Phi(\{A_{t^{\prime}}\}_{t^{\prime}=1}^{t})

\Phi(\mathcal{H}_{t})

Appendix

Table 8: Architecture for

\Phi

(RMSN)

Input :

\mathcal{H}_{t}

LSTM (# layers)

Linear Layer

ELU Activation

Output :

\Phi(\mathcal{H})_{t}

Appendix

Table 9: Architecture for

\Phi

(CRN)

Input :

\mathcal{H}_{t}

3x Transformers (# layers)

Output :

\Phi(\mathcal{H})_{t}

Appendix

Table 10: Architecture for

\Phi

(CT)

Method	Components	Hyperparameters	NSCLC range	MIMIC III range (semi-synthetic & real)
LSTM & CRN	Encoder & Decoder	# layers	1	1,2
		# hidden units	16, 32, 48	32, 48, 64
		Dropout rate	0.0, 0.1
	Encoder	# epochs	100	150
		Learning rate	0.001, 0.0001
		Batch size	64, 128

Appendix

Table 11: Hyperparameter search range for every method & dataset. If multiple components (e.g., encoder & decoder) are in the same cell, the hyperparameter is shared for those components (e.g., encoder and decoder have the same number of layers) for the same run. Otherwise, the hyperparameter is component-specific for the same run. We run a random grid search with 10 combinations for every method for fair comparison. The final results reported are averaged over 3 random seeds.

Causal Effect Estimation

Causal effect estimation measures how outcomes would change under interventions, using experiments or observational data and explicit assumptions.

See all 264 papers

CEPAE: Conditional Entropy-Penalized Autoencoders for Time Series Counterfactuals
Feb 17, 2026Tomàs Garriga, Gerard Sanz, Eduard Serrahima de Cambra +1Causal Counterfactual Generation Causal Inference
Novartis · Barcelona Supercomputing Center
Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction
Jun 4, 2026Amirhossein Zare, Amirhessam Zare, Herlock Rahimi +2Structural Causal Models Counterfactual Prediction
School of Medicine, Tehran University of Medical Sciences, Tehran, Iran. · Department of Electrical and Computer Engineering, Yale University, New Haven, CT, USA. · School of Medicine, Shiraz University of Medical Sciences, Shiraz, Iran. +2
Causal State-Space Model for Causal Inference: Estimating Longitudinal Individual Treatment Effects
Aug 8, 2026Abisoye Abidakun, Mingjun Zhong, Georgios LeontidisCausal Effect Estimation Selective SSMs
Department of Computing Science, University of Aberdeen, Aberdeen, UK · UiT The Arctic University of Norway

Feb 17, 2026J/K move · Enter open · S save

CEPAE: Conditional Entropy-Penalized Autoencoders for Time Series Counterfactuals

Tomàs Garriga, Gerard Sanz, Eduard Serrahima de Cambra, Axel Brando

Novartis · Barcelona Supercomputing Center

Causal Counterfactual Generation Causal Inference

Open paper AskPDF arXiv

The ability to accurately perform counterfactual inference on time series is crucial for decision-making in fields like finance, healthcare, and marketing, as it allows us to understand the impact of events or treatments on outcomes over time. In this paper, we introduce a new counterfactual inference approach tailored to time series data impacted by market events, which is motivated by an industrial application. Utilizing the abduction-action-prediction procedure and the Structural Causal Model framework, we first adapt methods based on variational autoencoders and adversarial autoencoders, both previously used in counterfactual literature although not in time series settings. Then, we present the Conditional Entropy-Penalized Autoencoder (CEPAE), a novel autoencoder-based approach for counterfactual inference, which employs an entropy penalization loss over the latent space to encourage disentangled data representations. We validate our approach both theoretically and experimentally on synthetic, semi-synthetic, and real-world datasets, showing that CEPAE generally outperforms the other approaches in the evaluated metrics.

Abstract

Causal Effect Estimation

Explore similar work

CEPAE: Conditional Entropy-Penalized Autoencoders for Time Series Counterfactuals

Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

Causal State-Space Model for Causal Inference: Estimating Longitudinal Individual Treatment Effects