Survival analysis or time-to-event studies focus on modeling the time of a future event, such as death or failure, and investigate its relationship with covariates or predictors of interest. Specifically, we may be interested in the causal effect of a given intervention or treatment on survival time. A typical question may be: will a given therapy increase the chances of survival of an individual or population? Such causal inquiries on survival outcomes are common in the fields of epidemiology and medicine [15, 41] . As an important current example, the COVID-19 pandemic is creating a demand for methodological development to address such questions, specifically, when evaluating the effectiveness of a potential vaccine or medical treatments outside randomized controlled trial settings.
Traditional causal survival analysis is typically carried out in the context of a randomized controlled trial (RCT), where the treatment assignment is controlled by researchers. Though they are the gold standard for causal inference, RCTs are usually long-term engagements, expensive and limited in sample size. Alternatively, the availability of observational data with comprehensive information about patients, such as electronic health records (EHRs), constitutes a more accessible but also more challenging source for estimating causal effects [17, 21] . Such observational data may be used to augment and verify an RCT, after a particular treatment is approved and in use. Moreover, the wealth of information from observational data also allows for the estimation of the individualized treatment effect (ITE), namely, the causal effect of an intervention at the individual level. In this work, we develop a novel framework to estimate the ITE for survival or time-to-event outcomes from observational data.
instance, patients who are severely ill are likely to receive more aggressive therapy, however, their health status may also inevitably influence survival. Traditional survival analysis neglects such bias, leading to incorrect causal estimation. Second, the exact time-to-event is not always observed, i.e., sometimes we only know that an event has not occurred up to a certain point in time. This is known as the censoring problem. Further, censoring might be informative depending on the individuals' characteristics, thus proper adjustment is required for accurate causal estimation.
Traditional causal survival-analysis approaches typically model the effect of the treatment or covariates (not time or survival) in a parametric manner. Two commonly used models are the Cox proportional hazards (CoxPH) model [12] and the accelerated failure time (AFT) model [39] , which presume a parametric relationship between the covariates and survival probability. Further, proper weighting for each individual has been employed to remove confounding bias from parametric models [1, 2, 18] . Such parametric specification makes these models interpretable but compromises their flexibility, and makes it difficult for them to be adapted for high-dimensional data or to capture complex interactions among covariates. Fortunately, recent advances in machine learning, such as representation learning or generative modeling, have enabled causal inference methods to handle highdimensional data and to characterize complex interactions effectively. For instance, there has been recent interest in tree-based [11, 38] and neural-network-based [33, 42] approaches. When estimating the ITE, some methods propose to regularize the transformed patient covariates or representations to have balanced distributions across treatment groups, thus removing the confounding bias and improving ITE prediction. However, most approaches employing machine learning techniques for counterfactual inference deal with continuous or binary outcomes, instead of time-to-event outcomes. Hence, a generalization to the context of survival analysis is needed.
To the best of our knowledge, this work constitutes the first attempt to leverage representation learning to estimate ITE for survival outcomes in observational studies. We develop a framework to predict event times from a low-dimensional (latent) transformation of the original covariate space. To address the specific challenges associated with causal survival analysis, we make the following contributions:
• We develop an optimization objective incorporating proper adjustments for informative censoring, as well as a balanced regularization term bounding the generalization error for ITE prediction. For the latter, we repurpose a recently proposed bound [33] for our time-to-event scenario. • We propose a generative model for event times, to relax restrictive survival parametric assumptions and allow for more flexible modeling. Our approach can also provide nonparametric uncertainty quantification for ITE predictions. • We provide survival-specific evaluation metrics, including a new nonparametric hazard ratio estimator, and discuss how to perform model selection for survival outcomes. The proposed model demonstrates superior performance relative to the commonly used baselines in real-world and semi-synthetic datasets. • We introduce a survival-specific semi-synthetic dataset and demonstrate an approach for leveraging prior randomized experiments in longitudinal studies for model validation.
We first introduce the basic setup for performing causal survival analysis in observational studies. Suppose we have N units, with N 1 units being treated and N 0 in the control group (N = N 1 + N 0 ). For each individual, we have covariates X, which can be heterogeneous, e.g., a mixture of categorical and continuous covariates which, in the context of medicine, may include labs, vitals, procedure codes, etc. We also have a treatment indicator A, where A = 0 for the controls and A = 1 for the treated, as well as the outcome (event) of interest T . Under the potential outcomes framework [31] , let T 0 and T 1 be the potential event times for a given subject under control and treatment, respectively. In practice we only observe one realization of the potential outcomes, i.e., the factual outcome T = T A , while the counterfactual outcome T 1−A is unobserved.
In the context of survival analysis, the problem becomes more difficult because we do not always observe the exact event time for each individual, but rather the time up to which we are certain that the event has not occurred; specifically, we have a (right) censoring problem, most likely due to the loss of follow-up. We denote the censoring time as C and censoring indicator as δ ∈ {0, 1}. The actual observed time is Y = min(T A , C), i.e., the outcome is observed (non-censored) if T A < C and δ = 1. . Covariates X = x are mapped into latent representation r via deterministic mapping r = Φ(x). The potential outcomes are sampled from t a ∼ p(T A |X = x) for A = a via stochastic mapping h A (r,˜ ), where randomness is induced with a flow based transformation,˜ , of a simple distribution p( ), i.e., uniform or Gaussian. (b) and (c) show the proposed causal graphs for non-informative and informative censoring, respectively.
In this work, we are interested in the expected difference between the T 1 and T 0 given X for a given unit (e.g., subject), which is commonly known as the individualized treatment effect (ITE). Specifically, we wish to perform inference on the conditional distributions of T 1 and T 0 , i.e., p(T 1 |X) and p(T 0 |X), respectively, as shown in Figure 1 (a). In practice, we observe N realizations of (Y, δ, X, A) for observed time, censoring indicator, covariates and treatment indicator, respectively; hence, from an observational study the dataset takes the form
. Below, we discuss several common choices of estimands in survival analysis.
We begin by considering survival analysis in the absence of an intervening treatment choice, A. Let F (t|x) P (T ≤ t|X = x) be the cumulative distribution function of the event (failure) time, t, given a realization of the covariates, x. Survival analysis is primarily concerned with characterization of the survival function conditioned on covariates S(t|x) 1 − F (t|x), and the hazard function or risk score, λ(t|x), defined below. S(t|x) is a monotonically decreasing function indicating the probability of survival up to time t. The hazard function measures the instantaneous probability of the event occurring between {t, t + ∆t} given T > t and ∆t → 0. From standard survival function definitions [24] , the relationship between cumulative and hazard function is formulated as
From (1) we see that f (t|x) P (T = t|X = x) = λ(t|x)S(t|x), is the the conditional event time density function [24] .
We now introduce the binary treatment A, and are interested in its impact on the survival time. For ITE estimation, we are interested in the difference between the two potential outcomes T 1 , T 0 . Let S A (t|x) and λ A (t|x) denote the survival functions and hazard functions for the potential outcomes T A , i.e., T 1 and T 0 . Several common estimands of interest include [43, 36] :
• Difference in expected lifetime:
The inference difficulties associated with the above estimands from observational data are two-fold. First, there are confounders affecting both the treatment assignment and outcomes, which stem from selection bias, i.e., the treatment and control covariate distributions are not necessarily the same. Further, we do not have direct knowledge of the conditional treatment assignment mechanism, i.e., P (A = a|X = x), also known as the propensity score. Let ⊥ ⊥ denote statistical independence. For estimands to be identifiable from observational data, we make two assumptions:
i.e., no unobserved confounders or ignorability, and (ii) overlap in the covariate support 0 < P (A = 1|X = x) < 1 almost surely if p(X = x) > 0. Second, the censoring mechanism is also unknown and may lead to bias without proper adjustment. We consider two censoring mechanisms in our work, (i) conditionally independent or informative censoring: T ⊥ ⊥ C|X, A, and (ii) random or non-informative censoring: T ⊥ ⊥ C. Note that for informative censoring, we also have to consider potential censoring times C 1 and C 0 and their conditionals p(C 1 |X) and p(C 0 |X), respectively. Figure 1 shows causal graphs illustrating these modeling assumptions.
To overcome the above challenges and adjust for observational biases, we propose a unified framework for counterfactual inference that accounts for survival outcomes. Specifically, we repurpose the counterfactual bound in [33] for our time-to-event scenario and introduce a non-parametric approach for survival outcome predictions based on normalizing flows [29] . Below we formulate a theoretically grounded and unified approach for estimating (i) the encoder function r = Φ(x), which deterministically maps covariates x to their corresponding latent representation r ∈ R d , and (ii) two stochastic time-to-event generative functions, h A (·), to implicitly draw samples from both potential outcome conditionals t a ∼ p h,Φ (T A |X = x), for A = {1, 0}, and where t a indicates the sample from p h,Φ (T A |X = x) is for A = a. Further, we formulate a general extension that accounts for informative censoring by introducing two stochastic censoring generative functions, ν A (·), to draw samples for potential censoring times c a ∼ p ν,Φ (C A |X = x). The model-specifying functions, {h A (·), ν A (·), Φ(·)}, are parameterized via neural networks. See the Supplementary Material (SM) for details. Figure 1a summarizes our modeling approach.
We wish to estimate the treatment effect from intervention do(A = a) [28] , where the potential outcomes, event times, are sampled from distributions parameterized by functions
where (3) follows from (2) because identifiability holds, provided that X is a sufficient set from A into T [28] , consistent with the causal graphs in Figure 1b and 1c. We obtain (4) from (3) via the strong ignorability assumption, i.e., {T 0 , T 1 } ⊥ ⊥ A|X and 0 < P (A = a|X = x) < 1, and the consistency assumption, i.e., T = T A |A = a. A similar argument can be made for informative censoring based on Figure 1c , so we can also write c a ∼ p ν,Φ (C A |X = x). Given (4), model functions {h A (·), Φ(·)} and ν A (·) for informative censoring can be learned by leveraging standard statistical optimization approaches, that minimize a loss hypothesis L given samples from the empirical distribution (y, δ, x, a) ∼ p(Y, δ, X, A), i.e., from dataset D. Specifically, we write L as
where h,Φ (t a , y, δ) is a loss function that measures the agreement of t a ∼ p h,Φ (T A |X = x) (and c a ∼ p ν,Φ (C A |X = x) for informative censoring) with ground truth {y, δ}, the observed time and censoring indicator, respectively.
For some parametric formulations of event time distribution p h,Φ (T A |X = x), e.g., exponential, Weibull, log-Normal, etc., and provided the censoring mechanism is non-
, which implies that the conditional event time density and survival functions can be calculated in closed form from transformations {h A (·), Φ(·)} of x. See the SM for parametric examples of (5) accounting for informative censoring.
We further define the expected loss for a given realization of covariates x and treatment assignment a over observed times y (censored and non-censored), and the censoring indicator δ as [33] . For a given subject with covariates x and treatment assignment a, we wish to minimize both the factual and counterfactual losses, L F and L CF , respectively, by decomposing (5) into L = L F + L CF as follows
Let u P (A = 1) denote the marginal probability of treatment assignment. We can readily decompose the losses in (6) according to treatment assignments. The decomposed factual
, and similarly, the decomposed counterfactual
CF . In practice, only factual outcomes are observed, hence, for a non-randomized non-controlled experiment, we cannot obtain an unbiased estimate of L CF from data due to selection bias (or confounding). Therefore, we bound L CF and L below following [33] .
is the latent distribution for group A = a, and α > 0 is a constant. Then, we have:
The integral probability metric (IPM) [27, 35] measures the distance between two probability distributions p and q defined over M , i.e., the latent space of R. Formally,
represents a class of real-valued bounded measurable functions on M [33] . Therefore, model functions {h a (·), Φ(·)} can be learned by minimizing the upper bound in (7) consisting of (i) only factual losses under both treatment assignments and (ii) an IPM regularizer enforcing latent distributional equivalence between the treatment groups. Note that if the data originates from a RCT it follows (by construction) that
Below we formulate an approach for estimating functions h A (·) and ν A (·) for synthesizing (sampling)
that can be used as likelihood for censored observations, they are restrictive, and have been shown to generate unrealistic high variance samples [9] . So motivated, we seek a nonparametric likelihoodbased approach that can model a flexible family of distributions, with an easy to sample approach for event times t a ∼ p h,Φ (T a |X = x). Normalizing flows [29] have been applied successfully to generative modeling due to their tractability and high sample quality [23] . Typically, a simple distribution p( ), i.e, Gaussian or uniform, is transformed via a sequence of invertible transformations z = z L • . . . z 2 • z 1 , into an arbitrarily complex distribution. We conveniently model the event time generation process where the source of randomness is obtained from an invertible planar flow parameterized by
. While we can apply L invertible planar transformations for more complex structured distributions, empirically, we did not see quantifiable performance benefits, therefore we model a single layer transformation as
is the dimensionality of the normalizing flow; each component of is drawn independently from Uniform(0, 1), and˜ h may be viewed as a skip connection with stochastic in . Further, h A (r,˜ h ) and Φ(x) are time-to-event generative and encoding functions, respectively, parameterized as neural networks. For simplicity the dimensions of r and are set to d, however, they can be set independently if desired. In practice, we are interested in generating realistic event-time samples; therefore, we account for both censored and non-censored observations by adopting the objective from [9] , formulated as
where the first term encourages sampled event times t a to be close to y, the ground truth for observed events, i.e., δ = 1, while penalizing t a for being smaller than the censoring time when δ = 0. Further, the expectation is taken over samples (a minibatch) from empirical distribution p(Y, δ, X, A). Informative censoring We model informative censoring similar to (9) but mirroring the censoring indicators to encourage accurate censoring time samples c a for δ = 0, while penalizing c a for being smaller than y for δ = 1 (observed events). Specifically, we set an independent planar flow like in (8) but parameterized by {U ν , W ν , b ν } and censoring generative functions
(10) Further, we introduce an additional time-order-consistency loss that enforces the correct order of the observed time relative to the censoring indicator, i.e., c a < t a if δ = 0 and t a < c a if δ = 1, thus TC (h, ν, Φ) = E (δ,x,a)∼p(δ,X,A), ∼p( ) [δ · (max(0, t a − c a )) + (1 − δ) · (max(0, c a − t a ))] (11) Note that TC (h, ν, Φ) does not depend on the observed event times but only on the censoring indicators. Finally, we write the consolidated NFR-Informative censoring (NFR-INFO) loss by aggregating (9), (10) and (11) as
Learning Model functions {h A (·), Φ(·), ν A (·)} are learned by minimizing the bound (7), via stochastic gradient descent on minibatches from D, with the factual loss (9) for non-informative censoring (NFR) and (12) for informative censoring (NFR-INFO). Further, for the IPM regularization loss in (7), we optimize the dual formulation of the Wasserstein distance, via the regularized optimal transport [37, 14] . Consequently, we only require α −1 ζ h,Φ (x, a) to be 1-Lipschitz [33] and α is selected by grid search on the validation set using only factual data (details below).
We propose a comprehensive evaluation approach that accounts for both factual and causal metrics. Factual survival outcome predictions are evaluated according to standard survival metrics that measure diverse performance characteristics, such as concordance index (C-Index) [16] , mean coefficient of variation (COV) and calibration slope (C-slope) [8] . See the SM for more details on these metrics. For causal metrics, defined below, we introduce a nonparametric hazard ratio (HR) between treatment outcomes, and adopt the conventional precision in estimation of heterogeneous effect (PEHE) and average treatment effect (ATE) performance metrics [19] . Note that PEHE and ATE require ground truth counterfactual event times, which is only possible in (semi-)synthetic data. For HR, we compare our findings with those independently reported in the literature from RCT data.
Nonparametric Hazard Ratio In a medical setting, the population hazard ratio HR(t) between treatment groups is considered informative thus has been widely used in drug development and RCT [41, 26] . For example, HR(t) < 1, > 1, or ≈ 1 indicate population positive, negative and neutral treatment effects at time t, respectively. Moreover, HR(t) naturally accounts for both censored and non-censored outcomes. Standard approaches for computing HR(t) rely on the restrictive proportional hazard assumption from CoxPH [12] , which is constituted as a semi-parametric linear model λ(t|a) = λ b (t) exp(aβ). However, the constant covariate (time independent) effect is often violated in practice (see Figure 2b ). For CoxPH, the marginal hazard ratio between treatment and control can be obtained from regression coefficient β learned via maximum likelihood without the need for specifying the baseline hazard λ b (t):
So motivated, we propose a nonparametric, model-free approach for computing HR(t).
Definition 1 We define the nonparametric marginal Hazard Ratio and its approximation,ĤR(t), as
where for HR(t) we leveraged (1) and (13) to obtain (14) and S (t) dS(t)/dt. The nonparametric assumption for S(t) makes the computation of S (t) challenging. Provided that S(t) is a monotonically decreasing function, for simplicity, we fit a linear function S(t) = m · t + c, and set S (t) ≈ m. Bias from S (t) can be reduced by considering more complex function approximations for S(t), e.g., polynomial or spline. Further, we leverage the model-free population point-estimate-based nonparametric Kaplan-Meier [22] estimator of the survival functionŜ PKM (t) in [8] to marginalize both factual and counterfactual predictions given covariates x. The approximated hazard ratio,ĤR(t), is thus obtained by combining the approximationsŜ PKM a (t) and m a . See the SM for additional details. A similar formulation for the conditional, HR(t|x), can also be derived. Note that for some parametric formulations, HR(t|x), can be readily evaluated because f (t a |x) and S(t a |x) are available in closed form.
and γ(·) is a deterministic transformation. In our experiments, γ(·) is the average over samples from t a ∼ p h,Φ (T A |X = x). Alternative estimands, e.g., thresholding survival times γ(T A ) = I{T A > τ }, can also be considered as described in Section 2.1.
The population treatment effect estimation error is defined as 
We describe the baselines and datasets that will be used to evaluate the proposed methods (NFR and NFR-INFO). Pytorch code including the new semi-synthetic dataset (see below) will be made publicly available. Throughout the experiments, we use HR(t) in (13) for CoxPH based methods and (14) for all others. The bound in (7) is sensitive to α, thus we propose approximating proxy counterfactual outcomes {Y CF , δ CF } for the validation set, according to the covariate Euclidean nearest-neighbour (NN) from the training set. We select the α that minimizes the validation loss L = L F + L CF from the set (0, 0.1, 1, 10, 100).
Baselines We consider the following competitive baseline approaches: (i) propensity weighted CoxPH [32, 6, 30] ; (ii) IPM (7) regularized AFT (log-Normal and Weibull) models; (iii) an IPM (7) regularized deterministic semi-supervised regression (SR) model with accuracy objective from [9] , as a contrast for the proposed stochastic predictors (NFR and NFR-INFO); and (iv) survival Bayesian additive regression trees (Surv-BART) [34] . For CoxPH, we consider three normalized weighting schemes: (i) inverse probability weighting (IPW) [20, 7] , where IPW i = aî ei + 1−ai 1−êi ; ii) overlapping weights (OW) [13, 25] , where OW i = a i · (1 −ê i ) + (1 − a i ) ·ê i ; and iii) the standard RCT uniform assumption. A simple linear logistic modelê i = σ(x i ; w), is used as an approximation, e i , to the unknown propensity score P (A = 1|X = x). See the SM for a details of the baselines. Datasets We consider the following datasets: (i) FRAMINGHAM, is an EHR-based longitudinal cardiovascular cohort study that we use here to evaluate the effect of statins on future coronary heart disease outcomes [5] ; (ii) ACTG, is a longitudinal RCT study comparing monotherapy with Zidovudine or Didanosine with combination therapy in HIV patients [15] ; and (iii) ACTG-SYNTHETIC, is a semi-synthetic dataset based on ACTG covariates. We simulate potential outcomes according to a Gompertz-Cox distribution [4] with selection bias from a simple logistic model for P (A = 1|X = x) and AFT-based censoring mechanism. The generative process is detailed in the SM. selection and censoring bias, naive approaches would result in a counter-intuitive treatment effect from statins. However, there is severe confounding from covariates such as age, body mass index, diabetes, history of peripheral artery disease (PAD), history of myocardial infarction (MI), history of stroke, etc., that influence both treatment likelihood and survival time. Table 3 , demonstrates that NFR-INFO is clearly the best performing approach. Specifically, its HR(t), reverses the biased observational treatment effect, to demonstrate positive treatment from statins, which is consistent with prior large RCT longitudinal findings [41] . Qualitative Results Figure 2a demonstrates that NFR-INFO matches the ground truth population hazard, HR(t), better than alternative methods on ACTG-SYNTHETIC data. See the SM for ACTG and FRAMINGHAM. Figure 2b shows sub-population log hazard ratios for four patient clusters obtained via hierarchical clustering on the individual log hazard ratios, log HR(t|x), of the test set of FRAMINGHAM data. Interestingly, these clusters stratify treatment effects into: positive (2), negative (1 and 3), and neutral (4) sub-populations. Moreover, the estimated density of median log HR(t|x) values in Figure 2c illustrates that nearly 70% of the testing set individuals have log HR(t|x) < 0, thus may benefit from taking statins. Further, we isolated the extreme top and bottom quantiles, HR(t|x) < 0.024 and HR(t|x) > 1.916, respectively, of the median log HR(t|x) values for the test set of FRAMINGHAM, as shown in Figure 2c . After comparing their covariates, we found that individuals with the following characteristics may benefit from taking statins: young, male, diabetic, without prior history (cardiovascular disease, PAD, stroke or MI), high body max index, cholesterol, triglycerides, fasting glucose, and low high-density lipoprotein. Note that individuals with contrasting covariates experience adverse effects from taking statins. There seem to be consensus that diabetics and high-cholesterol patients benefit from statins [10, 40] . See the SM for more qualitative results.
We have proposed the first unified counterfactual inference framework for survival analysis. Our approach adjusts for bias from two unknown sources, namely, confounding due to covariate dependent selection bias and censoring mechanism (informative or non-informative). We demonstrate superior performance in both survival outcome predictions and treatment effect estimations than competitive alternatives across three diverse datasets, including a semi-synthetic dataset which we introduce. Moreover, we formulate a model-free nonparametric hazard ratio metric for comparing treatments effects or leveraging prior randomized real-world experiments in longitudinal studies.
We proposed a unified framework for individualized treatment effect estimation for survival outcomes from observation data. Our approach is an inexpensive alternative to time consuming RCTs, albeit, requires proper adjustment of biases due to confounding and censoring. Though we tailored our study to medical applications, i.e., drug development, risk profiling, clinical trials, etc., we believe the framework is general and can be extended to other applications that characterize the future event time given covariates. It is our intent to encourage researchers in the survival analysis and causal inference communities to consider the problem of counterfactual inference for survival outcomes.
Additionally, our proposed method can be used for social good by adjusting for selection bias and censoring in high stake decision making, i.e, financial, healthcare, recidivism risk profiling. We hope our end users will use this framework responsibly, thus we propose a transparent hybrid approach in decision making with input from both this framework and experienced domain professionals. 
The general likelihood-based loss hypothesis that accounts for informative censoring is formulated as:
where (2) follows from the conditional independence (informative censoring) assumption T ⊥ ⊥ C|X, A. For some parametric formulations of event p h,Φ (T A |X = x) and censoring p ν,Φ (C A |X = x) time distributions, e.g., exponential, Weibull, log-Normal, etc., then − h,Φ (t a , c a , y, δ) is the closed-form log-likelihood, where:
where {S h,Φ (·), G ν,Φ (·)} and {f h,Φ (·), e ν,Φ (·)} are survival and density functions respectively.
Several common estimands of interest include [17, 16] :
• Difference in expected lifetime:
In our experiments, we consider both the hazard ratio and difference in expected lifetime. The difference of expected lifetime is expressed in terms of both survival functions and expectations:
where (5) follows from standard properties of expectations and (6) from 1 − F (t|x) = S(t|x) and 0 −∞ F (t|x)dt = 0. Below we formulate an approach for estimating the individualized and population hazard ratio.
From standard survival function definitions [10] , the relationship between survival and hazard function is formulated as
We propose a new nonparametric hazard ratio, model-free estimator for computing HR(t).
Definition 1 We define the nonparametric marginal Hazard Ratio and its approximation,ĤR(t), as
where for HR(t) we leveraged (7) and S (t) dS(t)/dt. For the estimatorĤR(t), provided that S(t) is a monotonically decreasing function, for simplicity, we fit a linear function S(t) = m · t + c and set S (t) ≈ m. Further, we leverageŜ PKM (t) in [4] , defined as the model-free population point-estimate-based nonparametric Kaplan-Meier [9] estimator. We denote J distinct and ordered observed event times (censored and non-censored) by the set
whereŜ PKM A (t 0 ) = 1, and I(b) represent an indicator function such that I(b) = 1 if b holds or I(b) = 0 otherwise. Further, γ(·) is a deterministic transformation for summarizing T A , in our experiments, γ(·) = median(·), computed over samples from t a ∼ p h,Φ (T A |X = x). Note from (9), we marginalize both factual and counterfactual predictions given covariates x.
A similar formulation for the conditional, individualized HR(t|x), can also be derived, where the cumulative density F A (t|x) = 1 − S A (t|x), is estimated with a Gaussian Kernel Density Estimator (KDE) [14] on samples from the model, t a ∼ p h,Φ (T A |X = x). Then we have:
where, S (t|x) dS(|t|x)/dt is also approximated with fitting a linear function S(t|x) = m · t + c, and setting S (t|x) ≈ m. Note that for some parametric formulations, HR(t|x), can be readily evaluated because f (t a |x) and S(t a |x) are available in closed form.
Concordance Index C-Index (also related to receiver operating characteristic) is a widely used survival ranking metric which naturally handles censoring. It quantifies the consistency between the order of the predicted times or risk scores relative to ground truth. C-Index is evaluated on point estimates, we summarize individualized predicted samples from NFR and NFR-INFO, i.e., t a = median {t s } 200 s=1 , where t s is a sample from the trained model.
Calibration Slope Calibration quantifies distributional statistical consistency between model predictions relative to ground truth. We measure population calibration by comparing population survival curves from model predictions against ground truth according to [4] . We desire a high calibrated model, with calibration slope of 1, while a slope < 1 and slope > 1 indicates underestimation or overestimation risk, respectively.
The coefficient of variation (COV) σµ −1 , the ratio between standard deviation and mean, quantifies distribution dispersion. A COV > 1 and < 1 indicates a high or low variance distribution, in practice, we desire low variance distribution. We use Mean COV 
Cox proportional hazard (CoxPH) CoxPH assumes a semi-parametric linear model λ(t|a) = λ b (t) exp(aβ) , thus the hazard ratio between treatment and control can be obtained without specifying the baseline hazard λ b (t):
A simple logistic modelê i = σ(x i ; η), is used to approximate the unknown propensity score P (A = 1|X = x). Methods that adjust for selection bias (or confounding) learn β by maximizing a propensity weighted partial likelihood [13, 2, 12 ]
We consider three normalized weighting schemes for w , namely, (i) inverse probability weighting (IPW) [8, 3] , where IPW i = aî ei + 1−ai 1−êi , (ii) overlapping weights (OW) [6, 11] , where OW i = a i · (1 −ê i ) + (1 − a i ) ·ê i , and (iii) the standard RCT Uniform assumption. Note that this modeling approach requires fitting over the entire dataset, thus has no inference capability.
Accelerated Failure Time (AFT) We implement IPM regularized neural-based log-Normal and Weibull AFT baselines. Both approaches have a desirable closed form S h,Φ (t a |x), thus enabling maximum likelihood based estimation, where
The log-Normal mean and variance parameters are learned such that, log t a = µ h,Φ (h(r, a)) + , where ∼ N (0, σ 2 h,Φ (h(r, a))) and r = Φ(x). Further, we learn the Weibull scale and shape parameters, where t a = λ h,Φ (h(r, a)) · (− log U ) (k h,Φ (h(r,a))) −1 and U ∼ Uniform(0, 1). We regularize (13) with the IPM loss, for maximum likelihood optimization.
Semi-supervised regression (SR) To demonstrate the effectiveness of our flow-based uncertainty estimation approach we contrast NFR with a deterministic accuracy objective from [5] , where t a = h(r, a) and:
where (14) is regularized according to the IPM loss.
Survival Bayesian additive regression trees (Surv-BART) Surv-BART [15] is a nonparametric tree-based approach for estimating individualized survivalsŜ(t (j) a |X = x) (defined at pre-specified J time-horizons) from an ensemble of regression trees. Note, Surv-BART does not adjust for both selection bias and informative censoring. While, we fit two separate models based on factual treatment and control data, causal metrics are estimated with both factual and counterfactual predictions.
The ACTG-SYNTHETIC, is a semi-synthetic dataset based on ACTG covariates [7] . We simulate potential outcomes according to a Gompertz-Cox distribution [1] with selection bias from a simple logistic model for P (A = 1|X = x) and AFT-based censoring mechanism. Below is our generative scheme:
, U ∼ Uniform(0, 1) log C ∼ Normal(µ c , σ 2 c ) Y = min(T A , C) , δ = 1 if T A < C, else δ = 0 , where {β A , α A , λ A , b, a, η, µ c , σ c } are hyper-parameters and {µ AGE , µ CD40 } are the means for age and CD40 respectively. This semi-synthetic dataset will made publicly available.
See Table 1 for additional quantitative comparisons on ACTG dataset. 
We detail the architecture of neural-based methods, namely, baselines (AFT-log-Normal, AFT-Weibull, SR) and our proposed methods (NFR and NFR-INFO). All methods are trained using one 
A , ν
A }, take as input the concatenated [h (1) A ,˜ h ] and [ν (1) A ,˜ ν ] respectively. Finally, we set the planar flow dimensions for both {˜ ν ,˜ h } to 100.