OrthoGen: A Generative Orthogonal Learner for Time-Varying Treatments

Tomàs Garriga; Valentyn Melnychuk; Konstantin Hess; Eduard Serrahima de Cambra; Axel Brando; Gerard Sanz; Stefan Feuerriegel

cs.LGOct 7, 2026

OrthoGen: A Generative Orthogonal Learner for Time-Varying Treatments

Authors: Tomàs Garriga, Valentyn Melnychuk, Konstantin Hess, Eduard Serrahima de Cambra, Axel Brando, Gerard Sanz, Stefan Feuerriegel

Organizations: Novartis · Barcelona Supercomputing Center · Universitat Polit`ecnica de Catalunya · LMU Munich · Munich Center for Machine Learning

Abstract

Estimating conditional distributional potential outcomes (CDPOs) over time is important in medicine (e.g., to estimate patient-specific risks under different treatment sequences). However, this task is challenging because of time-varying confounding, yet existing adjustment strategies for this task are limited. In this paper, we aim to learn CDPOs under time-varying treatments using flexible generative models. Our contributions are two-fold. (1) We introduce a tailored adjustment strategy for our setting, namely, generative recursive g-computation. Our adjustment strategy recursively propagates full conditional outcome distributions rather than conditional means, modeling the variables of interest directly rather than full trajectories. Building on our adjustment strategy, we formulate simple generative learners for CDPO estimation. However, these learners can be sensitive to nuisance estimation errors, which motivates an orthogonal learner. (2) We thus introduce OrthoGen, a Neyman-orthogonal and doubly robust generative learner. Importantly, we show that OrthoGen further achieves rate double robustness and quasi-oracle efficiency under suitable conditions. Our learners are flexible and can be instantiated with different generative backbones (e.g., normalizing flows and diffusion models). Across experiments with synthetic, semi-synthetic and real-world datasets, we find that OrthoGen is highly effective. To the best of our knowledge, we are the first to propose a generative orthogonal learner for estimating CDPOs under time-varying treatments.

Figures & tables

Figure 1: Task: Estimating the CDPOs under different treatment sequences.

Method	Adjustment	Backbone	Modeling target	Neyman-orthogonality
Wu et al. (2024)	IPTW	CVAE / DM	✗ Marginal DPO	✗
Mu et al. (2025)	IPTW	DM	✓ CDPO	✗
PI ( ours )	G.R. g-comp.	Model agnostic	✓ CDPO	✗
RA ( ours )	G.R. g-comp.	Model agnostic	✓ CDPO	✗
OrthoGen ( ours )	DR	Model agnostic	✓ CDPO	✓
Notes: G.R. g-comp. = generative recursive g-computation; CVAE = conditional variational autoencoder; DM = diffusion model; DPO = distributional potential outcome; PI = plug-in; RA = regression-adjusted; DR = doubly robust.

Table 1: Comparison of related approaches for direct learning of distributional potential outcomes over time.

Figure 2: Illustration of generative recursive g-computation for

\tau=3

. Left: Continuation densities are fitted through backward recursion. The terminal density is fitted from observed outcomes among patients receiving the target treatment. At each earlier treatment step

\delta

, pseudo-outcomes are sampled from the fitted successor density and used to fit the preceding density. Right: Generating the pseudo-samples that will be used to fit the continuation nuisances at earlier steps.

	Method	$\bm{\tau=1}$	$\bm{\tau=2}$	$\bm{\tau=3}$	$\bm{\tau=4}$	$\bm{\tau=5}$
NFs	PI	$\mathbf{0.308\pm 0.004}$	$0.235\pm 0.007$	$0.350\pm 0.015$	$0.249\pm 0.010$	$\mathbf{0.344\pm 0.015}$
	RA	$0.320\pm 0.008$	$0.232\pm 0.008$	$0.347\pm 0.015$	$\mathbf{0.247\pm 0.008}$	$0.348\pm 0.018$
	IPTW	$0.325\pm 0.009$	$0.240\pm 0.002$	$0.373\pm 0.015$	$0.305\pm 0.019$	$0.518\pm 0.062$
	OrthoGen	$0.312\pm 0.005$	$\mathbf{0.229\pm 0.009}$	$\mathbf{0.342\pm 0.011}$	$0.267\pm 0.015$	$0.401\pm 0.051$
DMs	PI	$\mathbf{0.306\pm 0.013}$	$0.312\pm 0.025$	$0.506\pm 0.038$	$0.482\pm 0.015$	$0.612\pm 0.014$
	RA	$0.411\pm 0.030$	$0.404\pm 0.018$	$0.610\pm 0.032$	$0.537\pm 0.025$	$0.651\pm 0.019$

Table 2: Synthetic data.

	Method	$\bm{\tau=1}$	$\bm{\tau=2}$	$\bm{\tau=3}$	$\bm{\tau=4}$	$\bm{\tau=5}$
NFs	PI	$1.13\pm 0.39$	$1.46\pm 0.70$	$1.59\pm 0.21$	$1.78\pm 0.67$	$2.55\pm 0.58$
	RA	$1.22\pm 0.33$	$1.22\pm 0.55$	$1.47\pm 0.23$	$1.59\pm 0.44$	$2.22\pm 0.52$
	IPTW	$0.78\pm 0.11$	$0.67\pm 0.20$	$1.45\pm 0.57$	$0.92\pm 0.22$	$2.06\pm 0.87$
	OrthoGen	$\mathbf{0.49\pm 0.09}$	$\mathbf{0.47\pm 0.13}$	$\mathbf{0.91\pm 0.13}$	$\mathbf{0.73\pm 0.22}$	$\mathbf{1.23\pm 0.45}$
DMs	PI	$0.56\pm 0.06$	$0.59\pm 0.13$	$0.81\pm 0.32$	$0.93\pm 0.36$	$0.93\pm 0.36$
	RA	$0.63\pm 0.08$	$0.74\pm 0.20$	$0.99\pm 0.38$	$1.09\pm 0.39$	$1.09\pm 0.41$

Table 3: Semi-synthetic data.

	Method	$\bm{\tau=1}$	$\bm{\tau=3}$	$\bm{\tau=5}$
NFs	PI	$0.338\pm 0.005$	$0.417\pm 0.006$	$0.459\pm 0.014$
	RA	$0.329\pm 0.004$	$0.405\pm 0.007$	$0.453\pm 0.015$
	IPTW	$0.319\pm 0.006$	$0.410\pm 0.004$	$0.457\pm 0.012$
	OrthoGen	$\mathbf{0.307\pm 0.007}$	$\mathbf{0.387\pm 0.007}$	$\mathbf{0.442\pm 0.013}$
DMs	PI	$\mathbf{0.314\pm 0.008}$	$0.428\pm 0.005$	$0.500\pm 0.007$
	RA	$0.330\pm 0.007$	$0.446\pm 0.005$	$0.506\pm 0.005$

Table 4: Real-world data.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value	Role
Transformer blocks	$1$	Causal encoder block
Model dimension $d_{\mathrm{model}}$	$30$	Token and attention representation width
Attention heads	$3$	Head width $d_{\mathrm{model}}/3=10$
Feed-forward dimension	$20$	Hidden width inside the encoder block
Dropout probability	$0.1$	Applied to the input embedding and encoder block
Readout dimension	$20$	Dimension of the context passed to the downstream head

Appendix

Table 5: Fixed hyperparameters of the Transformer history encoder. None of these values was tuned specifically for our experiments.

Family	Hyperparameter	Values considered
NF	Conditioner width	$\{20,32,64\}$
	Spline bins	$\{5,10,15,20\}$
	Learning rate	$\{3\!\times\!10^{-4},10^{-3}\}$
	Batch size	$\{32,64\}$
	Weight decay	$\{0,10^{-4}\}$
	Context noise $\sigma_{C}$	$\{0,0.01,0.05\}$

Appendix

Table 6: Consolidated hyperparameter search space for the terminal nuisance. Only parameters that were varied are shown.

\sigma_{C}

and

\sigma_{Y}

denote training-only Gaussian noise applied to the history context and standardized outcome, respectively.

Dataset	Family	Selected architecture	Selected optimization and regularization
Semi-synthetic MIMIC-III	NF	Width $20$ ; $10$ bins	Learning rate $10^{-3}$ ; batch $64$ ; weight decay $10^{-4}$ ; $100$ epochs; $(\sigma_{C},\sigma_{Y})=(0.01,0)$
	DM	$S=100$ ; denoiser width $20$	Learning rate $3\!\times\!10^{-4}$ ; weight decay $10^{-4}$ ; $\sigma_{C}=0.01$
Real-world MIMIC-III	NF	Width $20$ ; $10$ bins	Learning rate $10^{-3}$ ; batch $64$ ; weight decay $10^{-4}$ ; $100$ epochs; $(\sigma_{C},\sigma_{Y})=(0.01,0)$
	DM	$S=100$ ; denoiser width $20$	Learning rate $3\!\times\!10^{-4}$ ; weight decay $10^{-4}$ ; $\sigma_{C}=0.01$
Fully synthetic	NF	Width $20$ ; $10$ bins	Learning rate $10^{-3}$ ; batch $64$ ; weight decay $10^{-4}$ ; $30$ epochs; $(\sigma_{C},\sigma_{Y})=(0,0)$
	DM	$S=150$ ; denoiser width $20$	Learning rate $3\!\times\!10^{-4}$ ; weight decay $10^{-4}$ ; $\sigma_{C}=0$

Appendix

Table 7: Selected generative-model hyperparameters for each dataset.

Dataset	NFs	DMs
Fully synthetic	$56$ min	$122$ min
Semi-synthetic MIMIC-III	$110$ min	$176$ min
Real-world MIMIC-III	$27$ min	$63$ min

Appendix

Table 8: Average end-to-end wall-clock runtime per replicate for a representative treatment horizon

\tau=3

. Each replicate includes nuisance estimation, fitting of PI, IPTW, RA, and OrthoGen , and evaluation.

Explore similar work

Are Good Generators Good Decision-Makers? Policy Learning for General Interventions via Retargeted Counterfactual Generation
Jun 5, 2026Raphael C Kim, Jingsen Zhu, Ramin Zabih +1Counterfactual Generation Causal Inferences
Cornell Tech, Cornell University, New York, NY · Department of Biostatistics, Department of Population Health, New York University Grossman School of Medicine, New York, NY
Doubly Robust Adaptive Conformal Inference for Causal Effects Under Temporal Dependence
Jun 29, 2026Andreas Koukorinis, Ricardo SilvaConformal Inference Causal Inferences
Department of Computer Science University College London · Department of Statistical Science University College London
DoTime: A Synthetic Benchmark Generator for Interventional and Counterfactual Time Series
Jul 29, 2026Dennis Thumm, Billy Tim Anthony, Ying ChenCausal Inferences Synthetic Benchmark
National University of Singapore, Singapore, Singapore

Jun 5, 2026, stat.MLJ/K move · Enter open · S save

Are Good Generators Good Decision-Makers? Policy Learning for General Interventions via Retargeted Counterfactual Generation

Raphael C Kim, Jingsen Zhu, Ramin Zabih, Michele Santacatterina

Cornell Tech, Cornell University, New York, NY · Department of Biostatistics, Department of Population Health, New York University Grossman School of Medicine, New York, NY

Counterfactual Generation Causal Inferences Latent Feature Interventions Auto Research

Open paper AskPDF arXiv

Generative models are increasingly used to support decision-making in complex systems, where interventions may be joint and high-dimensional, and outcomes are high-dimensional. However, using generators for these decision-making settings are challenged by three problems. First, they are often trained on noisy logs with limited intervention data. Second, generator learning can be noisy across different environments. Third, generators are not optimized for the decisions they support. We introduce policy learning via retargeted counterfactual generation, which trains a generator for the decisions it supports in three steps. We (1) learn a doubly-robust, invariant counterfactual generator for high-dimensional interventions and outcomes, (2) conduct policy-learning based on its rollouts, and (3) retarget the generator toward the learned policy's interventions and relearn the policy, so the generator is accurate where decisions are made. Theoretically, our generator's excess counterfactual risk has a doubly robust product remainder, and retargeting removes the worst-case density ratio between the logged and learned policies from the regret. We demonstrate the efficacy of our approach across synthetic data, Cell Painting images of SARS-CoV-2-infected cells, and physical video simulations.

Abstract

Explore similar work

Are Good Generators Good Decision-Makers? Policy Learning for General Interventions via Retargeted Counterfactual Generation

Doubly Robust Adaptive Conformal Inference for Causal Effects Under Temporal Dependence

DoTime: A Synthetic Benchmark Generator for Interventional and Counterfactual Time Series