NOFLITE: Learning to Predict Individual Treatment Effect Distributions
Authors: Toon Vanderschueren
Jeroen Berrevoets
Wouter Verbeke
Venue: Accepted by TMLR ()

Abstract: Estimating the effect of a treatment on an individual's outcome of interest is an important challenge in various fields, such as healthcare, economics, marketing, and education. Previous work in machine learning has focused on estimating the expected value of the treatment effect. However, effective personalized decision-making requires more than just the treatment expected effect; it requires knowing the entire treatment effect distribution. Knowing this distribution allows analyzing the treatment's expected utility or quantifying the uncertainty regarding a treatment's effect. This information is essential for prescribing optimal treatments. The ability of a model to predict accurate individual treatment effect distributions is captured by its likelihood. In light of this, we propose a novel neural architecture, NOFLITE, that uses normalizing flows to directly optimize this likelihood, while simultaneously learning flexible estimates of the individual treatment effect distribution. Experiments on various semi-synthetic data sets show that NOFLITE outperforms existing methods in terms of loglikelihood. Moreover, we illustrate how the predicted distributions can enable an in-depth analysis of the treatment effect and more accurate decision-making.

Published in Transactions on Machine Learning Research (11/2023)
NOFLITE: Learning to Predict Individual Treatment Effect Distributions
Toon Vanderschueren KU Lewven University of Antwerp
ton..uonerschherrnkuuuuuuo..
Jeroen Berrevocts
jb2384 @cam.ac......mm..
University of Cambridge
Wouter Verbeke
wouter. nouter.uerbekee@ulleee
KU Lewven
Reviewed OpenReview: https: opwprevieeeeveeee...mmmoe net/ forum? id= EjgopDaiG
Abstract
Estimating the effect of treatment tant challenge in various fields, such as healthcare, economics, marketing, and education. Previous work in machine learning has focused on estimating the expected value of the treatment effect. However, effective personalized decision-making requires more than just the treatment expected effect; requires knowing the entire treatment effect distribution. Knowing this distribution allows analyzing the treatment's expected utility or quantifying the uncertainty regarding treatment's effect. This information is essential for prescribing optimal treatments. The ability of model to predict accurate individual treatment effect distributions is captured by its likelihood. In light of this, propose novel neural ar- chitecture, NOFLITE, that uses normalizing flows to directly optimize this likelihood, while simultaneously learning flexible estimates of the individual treatment effect distribution. Experiments on various semi-synthotic data sets show that NOFLITE outperforms existing hoth in terms of loglikelihood. Moreover, we illustrate how the predicted distributions can enable an in-depth analysis of the treatment effect and more accurate decision decision-making.
individual's outcome of interest is an impor-
1 Introduction
Knowing how certain treatment action will affect instance's outcome of interest is of great importance in various domains, such as healthcare (Berrevoets et al., 2020), marketing (Devriendt et al., 2021), education (Olaya et al., 2020), and economics (Vanderschueren et al., 2023a). A wide variety of existing work has looked at using machine learning (ML) for estimating the individual treatment effect, to help decision-makers optimize treatment assigument at an individual level. Existing work on treatment effect estimation ML has proposed novel approaches based on a variety of different MT. algorithms, including neural networks (Johansson et al., 2016; Shalit al., 2017; Zhang al., 2020: Curth & van der Schaar, 2021a), Gaussian processes (Alaa & van der Schaar, 2017), and decision trees (Hill, 2011; Rzepakowski & Jaroszewiez, 2012; Wager & Athey, 2018), as well as general meta-learners (Künzel et al., 2019; Curth & van der Schaar, 2021b).
In spite of this growing body of literature, existing work has almost exclusively focused on accurately esti- mating the expected value of the treatment effect. We argue, however, that more comprehensive approach is needed: to effectively support dacision-making we require accurately modeling the effect 's entire distribution. Such an approach is essential for adopting ML for treatment docision-making practico. First, unlocks a wide range of descriptive statistics, enabling a more detailed analysis of the treatment effect. For exam- ple, it allows for reasoning about uncertainty of events resulting from treatment: e.g., to get uncerainty

---

Published in Transactions on Machine Learning Research (11/2023)
p(r)
TU(T)
p(T)
(((T)
E (u((T)
B(u(T))
Different utilities... and resulting expectations
Different treatments...
...and expected utilities
Optimizing treatment decisions requires knowing the treatment effect distribution. Predicting the treatment effect distribution p(r) allows for assessing treatment's utility by pairing it with a utility function w(T) and obtaining its expected utility [(u(T)). Consequently, the individual meett effect distribution instrumental to analyzing treatment utilities (on the left) and comparing treatment preferences (on the right). Left: For a given treatment and the corresponding treatment effect distribution p(r), we can compare different utility functions u(r) and their expected utilities. Right: Similarly, for a given utility function u(T); different treatment effect distributions p(r) incur different. expected utilities. These types of analyses are not possible using only the individual's expected treatment effect IE(T), illustrating the importance of estimating the entire treatment effect distribution p(T) for personalized decision- decision-making.
Figure 1:
intervals or estimate the probability of the treatment's effect being positive. This way, the treatment effect distribution subsumes other common estimands which focus on a single attribute of this distribution, such as the mean, median, or other quantiles. Second, the trealment effect distribution is essential for deciding upon the optimal treatment. By pairing the treatment effect distribution with a utility function, we can obtain the expected utility (sce Figure for a graphical illustration). Utility functions have been influential in variety of fields, including economics (Hall, 1920; Kahneman & Tversky, 1979), game theory (Nash Jr, 1950), insurance (Spence & Zeckhauser, 1978), design (Thurston, 1991), and healthcare (Pliskin et al., 1980).
Contributions. The goal in this work is to estimate individual treatment effect distributions. To obtain good estimates of this distribution, we require learning a model with a high likelihood from observational data. This work addresses treatment effect estimation through this lens and, in doing so, makes three contributions. 1. To learn models with high likelihood, we propose novel neural architecture, NOFLITE, which employs normalizing flows to learn flexible estimates of the treatment effect distributions. 2. We propose an end-to-end training strategy to directly maximize the metric of interest-the model's likelihood- while regularizing to account for treatment assignment bias. This way, we aim to learn an unbiased model with high likelihood given the (unknown) test distribution from observational training data. 3. We evaluate our method empirically and compare performance to other state-of-the-art approaches using three semi- synthetic data sets. Contrary existing work, this evaluation is centered around the model's likelihood.
2 Related work
Our work builds upon existing work in (individual) treatment effect estimation and normalizing flows. In this section, discuss the most closely related literature for each category in turn.
2.1 Treatment effect estimation
A wide variety of methodologies have been proposed for estimating the expected individualized treatment effect (ITE). For complete overview, refer to Yao et al. (2021). Even though estimating the distribution of the ITE has received far less attention, several methods that were originally proposed for estimating the expected value are nevertheloss capable of lcarning distributional estimands. These methods include Bayesian Additive Regression Trees (BART) (Hill, 2011) and Causal Multi-task Gaussian Processes (CMGP) (Alaa & van der Schaar, 2017), as well as a variety of metalearners - general combinations of supervised ML methods (Künzel et al., 2019). Particularly relevant our work are recent approaches based on generative neural networks. On the one hand, generative adversarial networks have been used learn the counterfac- tual distribution and deal with confounding bias, for binary treatments (GANITE; Yoon et al., 2018) and
2

---

Published in Transactions on Machine Learning Research (11/2023)
Method
Neural?
Likelihood?
Confounding?
Reference
BART
Hill (2011)
CMGP
X
Alaa & van der Schaar (2017)
CEVAR
V
x
Louizos et al. (2017)
GANITE
x
Yoon et al. (2018)
CCN
x
Zhou et al. (2022)
NOFLITE
/
(Ours)
Table I: Literature table. We compare NOFLITE with existing methodologies for estimating individual treatment effect distributions. Three dimensions are considered: (1) whether a neural network is used, (2)
whether the model's likelihood Le is optimized directly, and (3) whether it adjust for confounding.
multiple treatments with continuous dosages (SCIGAN: Bica et al., 2020). On the other hand, variational autoencoders have been used to adjust for hidden confounder using (noisy) proxy (Louizos et al., 2017).
To the best of our knowledge, the only work explicitly looking at learning individual treatment effect distri- butions is the recently proposed Collaborating Causal Networks (CCN) (Zhou et al., 2022). As opposed to NOFLITE, CCN does not learn by maximizing the model's likelihood, but rather uses collaborating networks (Zhou al., 2021), where one network is trained to estimate the cumulative distribution function and an- other model estimates its inverse. An additional, more subtle distinction that they focus on estimating and evaluating the distribution of the potential outcomes, instead of the treatment effect distribution. A structured comparison of existing methodologies and the proposed NOFLITE is provided in Table 1.
In this work, we rely on generative neural network. More specifically, we estimate the treatment effect distribution using normalizing flows, a flcxible type of generative model that transforms a distribution to a simple prior through a series of learned transformations. Similar in spirit to our work is the recent Interventional Normalizing Flows (Melnychuk al., 2022). They propose to use normalizing flows to estimate the density of the average treatment effect resulting from an intervention, building upon the theoretical results of Kennedy et al. (2023). In contrast, aim to predict distributions at an individual level.
Although our work focuses on the static setting, other work has looked at forecasting treatment outcomes over time (Bica et al., 2019; Vanderschueren et al., 2023b). In this setting, CF-ODE has recently been proposed for learning uncertainty estimates for treatment outcomes over time (De Brouwer et al., 2022).
The fact that the expectation of the treatment effect insufficient for decision-making has motivated previous work in related areas. Most closely related ours are methods for predicting the treatment effect distribution at the population level, instead of the individual level (DiNardo et al., 1995; Chernozhukov et al., 2013; Firpo & Pinto, 2016; Kennedy et al., 2023; Melnychuk et al., 2022). Other work has been done on different, but related problems, such as estimating quantile treatment effects (Chernozhukov & Hansen, 2005; Firpo, 2007), finding the treatment regime that optimizes the quantile treatment effect (Wang et al., 2018), conformal inference treatment effects (Lei & Candès, 2021; Alaa et al., 2023), or bounding the potential outcomes to support docision-making makinn (Makar et al., 2020). More generally, other related work in machine lcarning aims to estimate confidence intervals or distribution properties for more comprehensive off-policy evaluation (e.g. Jiang & Huang, 2020; Chandak et al., 2021 b:a; Xu et al., 2022; Zhang et al., 2023; Wu et al., 2023) or learns to optimize the distribution of returns of reinforcement learning agent (Bellemare et al,, 2017; 2023).
2.2 Normalizing flows
Our approach builds upon a type of deep generative model called normalizing flows. Normalizing flows offer distinct advantages over other types of generative models like generative adversarial networks (Goodfellow et al., 2014) or variational autoencoders (Kingma & Welling, 2014). The key benefit of using normalizing flows is that they allow for exact evaluation of the density. Consequently, the model can be directly
3

---

Published in Transactions on Machine Learning Research (11/2023)
optimized for the metric we are interested in: the model's likelihood. We provide brief introduction to normalizing flows below. For a more detailed overview, we refer to Papamakarios et al. (2021).
A normalizing flow is an invertible mapping y + 2 from the empirical/origiin latent space Z (Rezende & Mohamed, 2015; Dinh et al., 2017). During training, the flow learns to map the empirical distribution p(y) to a known (simple) prior distribution p(z), typically a Gaussian distribution. The mapping g consists of series of invertible transformations g(y) = 91 o gv(y) with parameters 0 learned by neural network. This way, the density can be obtained using the change of variables formula:
data space ) to a
ply) = pz (B(v)) det
(1)
by
Using this formulation, we can evaluate the model's density exactly. Consequently, we can optimize the mapping& g to directly maximize the model's likelihood. After training, inference can be done by sampling from the simple prior p(z) and transforming the sample based on the inverse flow g (z).
Normalizing flows have been successfully applied in a wide range of tasks, such as gencrating images (Dinh et al., 2014; 2017; Kingma & Dhariwal, 2018), audio (Oord el al., 2018; Kim et al., 2018; Tran el al., 2019; Prenger et al., 2019), and graphs (Madhawa et al., 2019), well as reinforcement learning (Haarnoja et al., 2018; Ward et al., 2019; Schroccker et al., 2019) and energy forecasting (Dumas et al., 2022; Ge ct al., 2020).
Different, families of transformations have been proposed in the literature on normalizing flows. In this work, we use d sigmoidal flows (Huang et al., 2018). Our choice is motivated by sigmoidal flows offering a flexible transformation the resulting flow is a universal approximator - and their excellent empirical performance. Deep sigmoidal flows use a transformation g that is a strictly monotonic neural network. The parameters of this network are given by conditioner network. To comply with the monotinicity requirement, the learned transformer parameters are restricted to strictly positive weights and strictly monotonic activation functions- more specifically, a sigmoid activation. Although the inverse transformation g-1 is not known analytically, it can be approximated numerically. There is a wide variety of other flow transformations that could potentially also be used with NOFLITE, including gaussianizations flows (Chen & Gopinath, 2000; Meng et al., 2020), residual flows (Chen el al., 2019), or neural spline flows (Durkan et al., 2019).
3 Problem Formulation
Notation. We describe our problem setting using the Soyman-Rubile potential outcomes framework (Ru- bin, 2005). Let each instance be defined by covariates X E X C Rd binary treatment indicator T T f0, 1), and an outcome Y y C R. Let the potential outcomes y("), y()) E IR be de- fined as the outcomes that would be observed given treatment T = 0 and 1
Goal. Given an individual's covariates we are interested in predicting its individual treatment e p(T), ie., the difference between both potential outcomes T = y()) - (0)),, Most existing work aims to learn an individual's expected treatment effect:
B(r) = E(Y() - (())) x))
(2)
We argue, however, that this expectation itself is insufficient for pocision-making in many applications. Instead, we need to learn the distribution of the individual treatment effect:
p(r) = P(y0) - (() x = x) P(Y x
x= =x) P X x)).
(3)
This distribution can be used to optimize an individual's 's treatment. For example, a mecision-maker can use it to assess probabilistic statements, such as the probability of the treatment effect being strictly positive p(T 0). Alternatively, the ITE distribution can be used to evaluate treatment decisions by pairing it with (personalized) utility function to obtain the expected utility: B(u(r) = u(T)p(F) dr, see Figure 1.
When using the term Individual Treatment Effect (ITE), refer to the instauce's measured covariates included in X. Note, however, that these covariates need not completely describe this individual and, bocause of this, may refer to multiple individuals. The only requirement is that X satisfies the ignorability assumptions (1-3). Because of this distinction, earlier work has argued that is more precise denote T as the Conditional Average Treatment Effect. (CATE), see Vegetabile (2021).
4

---

Published in Transactions on Machine Learning Research (11/2023)
Our goal is to obtain good estimates of the treatment effect distribution, i.c., to obtain model A E Q with a high likelihood p(r(0) = - p(T,|0) = =T= p(y, (1) y (0)) given a (hypothetical) test set Duest {1 " y/ ) containing both counterfactuals. Compared to point estimates such the mean squared error, the likelihood incorporate uncertainty and capture the entire data distribution. This way, it provides more comprehensive measure of model performance and facilitates more robust decision- decision-making. maniing
Data and assumptions. We assume access to observational dataset Dtrain ((1 yi )))) sampled from the joint distribution TYY, Learning model for estimating the individual treatment effect distribution from this data is challenging for several reasons. We only observe one factual outcome (() in practice, while the other, counterfactual outcome y(l-t) never observed. Consequently, the treatment effect itself is nover observed, which is known as the fundamental problem of causal inference (Holland, 1986). Additionally, because the data are observational, treatments were assigned by a (potentially unknown) policy based on instance covariates. Therefore, it is necessary adjust for confounding order obtain unbiased estimates of (() and, consequently, p(T). A final challenge is that need to learn each individual's entire distribution based on only one sample for each individual by leveraging data from similar individuals.
To identify the individualized treatment effect from observational data, we require the following standard assumptions (Roscnbaum & Rubin, 1983; Rubin, 2005):
Assumption (Consistency) An instance's 's observed outcome given treatment is equal to its potential outcome:
Assumption 2 (Overlap) Each instance has strictly positive probability of receiving each treatment: 0 << P(T =t X z) < 1,, V,ET, E.
Assumption 3 (Unconfoundedness) Potential outcomes are independent of the treatment given the co- variates: (Y y() 11 T | X.
4 NOFLITE: Estimating ITE distributions using normalizing flows
To tackle the problem formulated above, we propose NOFLITE': neural architecture using normalizing flows for estimating individual treatment effects distributions. A high-level overview of the architecture is shown in Figure 2. NOFLITE consists of two parts. The first part learns to predict simple, conditional prior p(zla, (), in this case a Gaussian distribution parametrized by (H, The second part learns to transform this prior to a more complex posterior distribution of the potential outcome p(y|æ, t). The entire model is trained end-to-end by directly maximizing its likelihood, while regularizing to deal with confounding. Both the model architecture and training procedure are explained in more detail below.
4.1 Architecture
NOFLITE'S architecture consists of two parts (see Figure 2): (1) an encoder f, i.e., neural network that estimates simple, parametrized prior distribution, and (2) normalizing flow g- to transform this prior to a more complex posterior distribution. In the following, we provide detailed description of both modules.
4.1.1 Encoder: learning conditional priors
The first part of the model a neural network that encodes the input (a,t) as simple prior distribution p(2 1). This prior can be seen as first approximation of the empirical distribution. We use a normal distribution defined by parameters μ and σ: (X,T) + (M, E). Depending on the application, other distributions could be used, such as uniform or log-normal distribution. The only requirements are that the distribution is defined by a finite set of parameters and that its likelihood can be computed analytically.
2 All code available http:://gitiie http://gthhaa.ccomttonod/////T
5

---

Published in Transactions on Machine Learning Research (11/2023)
Encoder
Normalizing flow
p(2)
91 ook (y(()
p(((())
y(0
9p o 91 (z)
Conditional prior
Conditional priors p()
Normalizing flow (conditional on z and t)
Potential outcome distributions p(Y(()=))
parameters 14,
Figure 2: NOFLITE architecture overview. We visualize the S-learner configuration of our method. Opti- mization: NOFLITE learns (1) a conditional prior p(z|æ, t) for each instance, and (2) an invertible mapping g(y) from that conditional prior to the empirical distribution pzy))>z), possibly conditioned on x and/or t. Both the encoder and normalizing flow use neural networks and are trained with gradient descent. Inference: Based on input (x,t), t), the conditional prior p(z|ær is estimated. Samples are drawn from this prior and transformed using the inverse flow g- to estimate the distribution of each potential outcome pz))
To deal with confounding, x is first transformed to balanced representation d. This is achieved by learning representation in which the distributional distance between the treated and non-treated representations is minimized (Johansson et al., 2016; Shalit et al., 2017). The training procedure is discussed in detail below.
4.1.2 Normalizing flow: learning complex posteriors
Although the encoder could be trained on its own to maximize the likelihood of the conditional prior, the model's hypothesis space would in that case be limited to simple parametric distributions, such as the normal distribution. Therefore, to augment the representational capacity of our model, the second part of the network uses normalizing flow that gradually transforms the simple prior distribution p(2(t,z) x) into more complex posterior distribution (*) [z, t) that can match more complex potential outcome distributions. The number of flow Iransformations & in g is a hyperparameter that can be tuned. This number can be used to tune the complexity of the transformation and, therefore, the complexity of the posterior distribution.
Normalizing flows are an active field of research and, as such, new flow transforms continue to be proposed. In principle, any type of normalizing flow is compatible with NOFLITE, as long some requirements fulfilled. First, they have to be compatible with a univariate distribution (e.g. coupling layers (Dinh et al., 2017) are not suitable). Second, we have to be able to compute both g and g- Nevertheless, these requirements leave variety of flow transformations (e.g., Durkan et al., 2019; Chen et al., 2019), with several potential candidates that have been shown to be universal approximators (e.g., Huang et al., 2018; Meng et al., 2020).
In this work, we instantiate g with a sigmoidal flow (Huang et al., 2018). This flow type offers great flexibility and is a universal approximator. The transformation is defined by monotonic transformer network, whose parameters estimated by second conditioner network. By constraining the outputs of the conditioner network, the monotonicity of the transformer network is ensured. Specifically, the transformer network has strictly positive weights and strictly monotonic activation functions. The conditioner network can take different factors into account, such as covariates x and/or treatment f. Although the inverse mapping is not known analytically, we can use an interpolation scarch to find it mumercalll.
4.1.3 Metalearner configuration
We implement our model in two metalearner configurations. For the S-learner configuration, the encoder concatenates the balanced representation ó with the treatment t and uses this combination as inputs to predict the parameters of the conditional prior p(z(x,t). For the T-learner configuration, both treatments have individual output heads after the shared balanced representation á. Similarly, we define two corre- sponding metalearners for the normalizing flow. The S-learner shares the flow across treatments, potentially conditioning the flow transformation on the treatment-depending on the flow type. The T-learner learns a separate flow per treatment. Either one might be more suitable, depending on the data generating process; we this is a hyperparameter that be tuned.

---

Published in Transactions on Machine Learning Research (11/2023)
4.2 Optimization
The model is trained end-to-end simultaneously estimate accurate treatment effect distributions by max- imizing the likelihood, while regularizing deal with confounding:
CROFLITE CLL + ACMMD
(4)
with hyperparameter I- trading-off likelihood maximization and balancing.
To calculate the model likelihood, the encoder and normalizing flow cooperate and meet in the middle, i.c., in the latent space Z. On the one hand, the encoder maps the inputs (z, () to prior distribution ((u,o). On the other hand, the normalizing flow maps the outcome y to a latent 2. The goal of both components is to cooperate and maximize the likelihood of 2 given N(,, o). This results in the following training objective:
LT. = log pz (B(v)) + log det
(g(a/) 8ie
The first term, the likelihood of gy). can be computed analytically given the Gaussian prior from the encoder. The second term, the log determinant of the Jacobian, can be calculated from the flow transformation. This way, using a normalizing flow allows computing the likelihood exactly and, therefore, optimizing directly. In practice, we minimize the negative oog-likelhhood LLL, which is equivalent to maximizing this log- ligliikl
We learn balanced representation to deal with confounding (Johansson al., 2016; Shalit al., 2017). This is achieved by minimizing the distributional distance between the representations of different treatments. Specifically, we use the linear Maximum Mean Discrepancy (MMD) (Gretton et al., 2012):
LMMD
=
1
Σ ((xs)
Σ M(n))
(6)
n2 SEDo
11
JED.
where Do and Dj denote the control and treatment group with no and elements (per batch), respectively.
4.3 Inference
Inference is done in two steps. For each instance, samples are drawn from its prior distribution 2 N (u(2,, t), (I, These samples are then transformed to samples from the posterior using the inverse mapping: 1 = g (z). After this process is completed for both treatments, a sample for the ITE is obtained by taking the difference between each sampled potential outcome: 4 = g(2) - g(0)
5 Results
In this section, we evaluate our proposed approach and compare it with several benchmarks. Our main goal is to assess whether NOFLITE learns to predict accurate individual treatment effect distributions from different types of observational data sets. More specifically, our experiments aim to answer three questions. (1) Does NOFLITE predict accurate individual trealment. effect distrinutions? This is our main question of interest. We evaluate this with the loglikelihood. (2) Does the predicted distribution allow for a more detailed analysis of individual the trealmend effect, based on statistics derived from this distribution? This evaluated by looking at the predicted distribution's expected value (using the precision in estimation of heterogeneous treatment effects), as well as the derived confidence intervals (using the intorsection-ove-uniii and empirical coverage). (3) Does the predicted distribution enable accurate dacision-making when paired with utility function? We evaluate this by looking at the accuracy of the recommended treatments.
The remainder this soction starts by describing our experimental sctup, including data sets and benchmark methodologies in Section 5.1 and evaluation metrics in Section 5.2. The empirical results are presented in Section 5.3. More information on hyperparameter tuning for NOFLITE is provided in Appendix A.2.
7

---

Published in Transactions on Machine Learning Research (11/2023)
5.1 Data and benchmarks
Evaluating individualized estimates of the treatment effect is challenging, because only outcome is observed and, because of that, we do not have access to the ground truth. To overcome this challenge, we follow existing work and evaluate based on semi-synthetic data. More specifically, in this work, we evaluate NOFLITE using commonly used benchmark data sets: IHDP, EDU, and News. These are introduced briefly in the following; Appendix A.I provides more detailed information on the data generating processes.
IHDP. The Infant Health and Development Program (THDP; Hill, 2011) is a semi-syuthetic data set that is commonly used to evaluate machine learning models for causal inference. Covariates are based on a real-world randomized experiment in which some infants (n= = 747) were targeted with child care and home visits. The resulting outcome simulated based on these covariates (ær E [R22) and the treatment, with the resulting treatment effect following normal distribution. Although this dataset has recently been criticized (Curth et al., 2021), we include it due to it being a widely-used benchmark.
The Education data set (EDU; Zhou et al., 2022) measures the effect of providing mother with adult education benefits on their children's learning. The data is simulated based on covariates E R using two (non-linear) neural networks, with added Gaussian (t 0) or exponential noise (t= 1). For both potential outcomes, the noise level depends on one of the covariates; more specifically, a single binary variable. Confounding is introduced through covariate- baared propensities and removing well-balanced subjects.
EDU.
News. The News data set (Johansson et al., 2016) shows the effect of reading an article on either mobile or desktop (t) on the reader's experience (y), based on the article's content in word counts (z). The data is simulated using a topic model =(x), which used define two centroids in the topic space: 2 (desktop) e
and of (mobile). Device assignment and reader experience are both simulated based on the similarity of the article's topic ((x) to the rontroids zf and 2f. The resulting data is very high-dimensional (d odd=3,m 3,477..
Benchmark methodologies. We compare NOFLITE with several other machine learning models capable of learning distributions. First, compare against Causal Multi-task Gaussian Processes" (CMGP; Alaa & van der Schaar, 2017). Second, compare against other methods relying on generative neural networks. The Causal Effect Variational Autoencoder (CEVAE; Louizos et al., 2017) uses variational autocncoder to model latent confounders given noisy proxies. GANITE (Yoon et al., 2018) uses generative adversarial networks to deal with selection bias and to obtain a probabilistic estimate of the treatment effect. Finally, we compare against Causal Collaborating Networks (CCN; Zhou et al., 2022).
5.2 Performance metrics
The main goal in this work is to predict accurate individual treatment effect distributions. Therefore, the main metric of interest is the model's loglikelihood, which allows for a comprehensive assessment of each instance's predicted distribution, by quantifying how likely the test data is given the predicted distributions. We estimate the loglikelihood as follows. For each instance i, we sample 500 samples from the model based on the covariates zy. Then, for each instance, we fit a Gaussian kernel density estimator (d,, A) using these samples and estimate the loglikelihood of the true treatment effect according to this kernel density estimator log p(T, kd;; 0)). The loglikelihood is averaged over all instances:
LL
Σ 72
logp(T: Ta kdde
(7)
The loglikelihood evaluates the predicted distributions globally. As argued in the motivation, statistics derived from the distribution can be used to facilitate decision-making. Therefore, to facilitate more holistic assessment of performance, we present additional metrics that analyze specific properties of the
For the high dimensional News data set, we take the first 100 principal components avoid excessive training times. This number was tuned based on validation set..
8

---

Published in Transactions on Machine Learning Research (11/2023)
predicted distribution. First, we use the square root of the precision in estimation of heterogencous effects (VPEHE, see Hill, 2011) to evaluate the accuracy of the expected value of the individual treatment effect:
PEHE = 12 Σ (1 (u('vv)")
1000 02 0.00) -0
1 n
7 +2)².
(8)
Moreover, we evaluate the empirical coverage (Cov) of the estimated 90% confidence interval ĆI, i.e., the probability that an observed sample of the ITE lies within the estimated interval:
1 - 1( ((TECC)
Cov =
(9)
1=1
For the 90% CI, the empirical coverage should ideally be 0.90 with small variance between iterations. For some data sets, we know the distribution that was used to generate the data. In these cases, we can additionally compare the prodicted and true confidence intervals. We propose to evaluate their overlap using the inlerseation-over-uiii (IoU) comparing the predicted and ground truth confidence intervals:
IOU CINCI
(10)
CIUCT
This metric is bounded between O and 1. The worst value, O, indicates an emply intersection. The best value, 1, indicates an intersection equal to the union and, consequently, an estimated confidence interval that is equal to the ground truth.
Finally, we evaluate the quality of decisions made based on the predicted distributions, for a given utility distribution. The decision to treat or not based on whether the expected utility positive. For IHDP, we use u(T) (T- given that the average treatment effect will be 4 on average (Hill, 2011). For EDU, we use u(T) (1),,,, as the treatment effect was empirically observed to be approxinately on average. The predicted optimal treatment decision is compared with the theoretical optimal based on the ground truth distribution using the accuracy (Acc). For the News data, we cannot do this analysis, as the true distribution is not known. For all metrics and models, we use 500 samples (per instance) for evaluation.
5.3 Empirical results
We compare the different models for the IHDP, EDU, and News data sets in Table 2. In terms of loglikelihood, NOFLITE obtains the best performance out of all methods under consideration for each data set. These findings demonstrate NOFLITE's ability to learn accurate individual treatment effect distributions from a variety of observational data sets and associated data generating processes. For the IHDP data set, NOFLITE slightly outperforms the next-best model, CMGP, Both methods use a Gaussian prior, which matches IHDP's data gencrating process. On the EDU data set, the data generating process is more complex and requires more flexiblity. Indeed, the Gaussian prior of CMGP results in relatively worse performance and the more complex models, such as NOFLITE and CCN, perform better. This illustrates how the normalizing flows can allow for a more flexible model of posterior when required. Finally, the good performance on the News data set illustrates NOFLITE's ability of learning from high dimensional data. As opposed to CMGP, which needs PCA preprocessing, it can also handle this high dimensional data out-of-the-box.
When we look at the metrics evaluating statistics derived from the treatment effect distribution, i.e. the PEHE and IoU, there is no clear winner overall: either CMGP, CCN or NOFLITE result in the best per- formance. Nevertheless, for all metrics, NOFLITE's performance is competitive with the best performing methodology for cach data set. This illustrates how optimizing the likelihood and learning distributions re- sults in good performance for the metrics evaluating properties of this distribution. Nevertheless, if there is one particular metric interest (e.g., PEHE), other objectives than NOFLITE's ' loglikelihood might be prefer- able (e.g., the mean squared crror). Morcover, we observe that NOFLITE obtains relatively high accuracy for both the IHDP and EDU data sets. This illustrates that pairing the NOFLITE'S predicted distributions with utility function enables qualitative mecision

---

Published in Transactions on Machine Learning Research (11/2023)
LL (t)
VPEHE (L)
IoU (t)
Coe (90%)
Ace (t)
CMGP
0.73
0.77
0.88
0.90
BART
2.16
2.22
0.63
0.83
0.84
GANITE
2.32
2.01
0.01 10.00
0.55
CEVAE
-2.97
5.71
0.25
0.98 10.00
0.57
CCN
-2.1620.08
1.46.0.17
0.66...02
0.84
0.86.
NOFLITE
1.90
=.00
0.75
0 1888
0.90=o.0x
(a) IHDP (n= 747; d 25)
LL (t)
VPEHE (1)
IoU (1)
Cov (90%)
Acc (1)
LL (1)
PPHEE (4)
Cov (90%)
CMGP
0.22 10.74.0.0
0.56go.co
0.91 10.00
0.70gn.00
CMGP - -2.29
2.21
0.95
BART
10.00
2.53
0.562o.co
0.89
2.69
BART -2 -2.4
2.71
2.97
GANITE
1.26
0.37...co
0.466go.0
0. 0.722..01
GANITE
18.91
0.01
CEVAE
=..7s
2.20
0, 2.25
1. 1.00
0. 49
CEVAE
2.83
3.74
0.97
CCN 1. 1.65...0
0.31
0. 0..44u..z
0.87
0. 0.76
CCN
2 2.25
2.23+o.04
0,84
NOFLITE
0.26
0.64
0.89 1892..mm on
0.76go.01
NOFLITE 2. 2.1520.02
2.18 on
0.93
(b) EDU
32)
(c) News (n= 5-000;
Table 2: Empirical results. We compare NOFLITE against a variety of existing methods for three data sets: (a) IHDP, (b) EDU, and (c) News. For each metric, arrows indicate whether lower (4) or higher (t) value is better; the ideal coverage is 90%. We show the best result in bold. We highlight the loglikelihood ( (LL) ) in gray to emphasize that this is our main metric of interest, as it evaluates the quality of the predicted distributions. For all data scts, GANITE achieves loglikelihood of less than -00, indicatod by a dash (-).
Finally, we illustrate how NOFLITE can be used for practical applications. We do this by training our model on an ileration of the News data set and showing the model's output for few selected test instances, see Figure 3. Figures 3a to 3e show the ITE distribution and related statistics based on samples from the learned model. This not only allows for visualizing the estimated distribution p(T:) and its associated expected value E(T), but also for assessing the uncertainty using the 90% confidence interval or a boxplot. Moreover, due to the News data set being semi-synthetic, we can compare the ground truth Tobserved with the estimated distribution. Finally, the model can be used to consider the treatment decision for an instance by, e.g., looking at the probability of its treat effect being strictly positive. Figure 3d shows the heterogeneily in distributions of different instances, both in terms of expected value and shape.
6 Conclusion
Estimating an instance's individual treatment effect distribution is essential requirement for personalized decision-making. To this aim, we presented NOFLITE, a. flexible neural method for estimating treatment effect distributions that directly optimizes the metric of interest for this task: the model's likelihood. By leveraging normalizing flows, the model is not constrained to any particular parametric distribution, but can instead trade off a simple normal distribution with a more complex posterior, depending on the data. Experiments on a varicty of data sets demonstrated NOFLITE's good performance practice and underlined its excellent representational capacily, as illustrated by its ability obtain high likelihoods for a variely of data set sizes, dimensionalities, and data generating processes.
Future work in normalizing flows could benefit our method, as novel flow transformations will be compatible with our method and help address polential limitations of our work. First, depending on the flow transfor- mation that is used, training or inference can be slow in normalizing flows. Second, NOFLITE introduces a variety of hyperparametors. Therefore, novel flow transformation with less hyperparameters might benefit adoption of our method in practice. Related to these points, it would be interesting to incorporate advances in other types of generative models, such as diffusion models (see e.g. Yang et al., 2022), in future work.
10

---

Published in Transactions on Machine Learning Research (11/2023)
0.30
1977 Pirt
F(F)
Tateree 5055 C7
0.10
H+
908
HH 908 CT
003
0.00
10
0.0
20
23
0.00 10
(a) Instance 1
(b) Instance 2
(c) Instance 3
(d) Heterogeneity
Figure 3: NOFLITE illustration. We visualize NOFLITE's output for selected test instances for one iteration of the News data set. Individual treatment effect distributions: Figures (a-c) visualize one particular instance's ' treatment effect distribution and related statistics based on samples from the model. Distribution hetero- geneity: Figure (d) compares the predicted individual treatment effect distributions of several instances.
Although estimating individual treatment effect distributions can be very valuable, it constitutes a challening problem -sspccillly when working with high-dimensional, observational data. For many applications, the available data may be limited. Although NOFLITE can be used without flow transformations, its flexibility comes from the normalizing flows, which require data to be trained succesfully. Although NOFLITE's hyper- parameters can be tuned by looking the fit on a validation set, validating causal inference models itself is challenging problem (Curth & van der Schaar, 2023). Therefore, we consider it promising area for future work to analyze NOFLITE from a theorctical perspective and come up with performance guarantecs based statistical learning theory.
Additionally, our method relies on the standard ignorability assumptions in causal inference. Judging the fcasibility of these assumptions is impossible based on data alone and requires the judgment of domain experts. There is a growing body of work looking at learning treatment effects under violations of these assumptions (e.g., under hidden confounding, see Kallus al. (2019); Oprescu et al. (2023)). Another interesting direction for future work is to analyze performance of our method in settings where the ignorability assumptions are violated, and to extend our methodology to account for these violations.
Acknowledgments
We would like to thank Thomas Van Hout and the anonymous reviewers for insightful comments and dis- cussions on earlier drafts of this paper. Toon Vanderschueren is supported by the Research Foundation Flanders (FWO PhD Fellowship 1117322N).
References
Ahmed Alaa, Zaid Ahmad, and Mark van der Laan. Conformal meta-learners for prodictive inference of individual treatment effects. aixiv preprint or Xix:2230 14895, 2023.
Ahmed M Alaa and Mihaela van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. Advances in Neural Information Processing Systems, 30, 2017.
Marc G Bellemare, Will Dabney, and Rémi Munos, A distributional perspective on reinforcement learning. In International conference on machine learning, pp. 449-458. PMLR, 2017.
Mare G Bellemare, Will Dabney, and Mark Rowland. Distributional reinforcement learning. MIT Press, 2023.
Jeroen Berrevocts, James Jordon, Ioana Bica, Mihacla van der Schaar, et al. Organite: Optimal transplant donor organ offering using an individual treatment effect. Advances in newral information processing systems, 33:20037 33-20037-20050, 2020.
11

---

Published in Transactions on Machine Learning Research (11/2023)
Ioana Bica, Ahmed M Alaa, James Jordon, and Mihacla van der Schaar. Estimating counterfactual treat- ment outcomes over time through adversarially balanced representations. In International Conference Learning Representations, 2019.
Toana Bica, James Jordon, and Mihaela van der Schaar. Estimating the effects of contimuous--aluued inter- ventions using generative adversarial networks. Advances in Neural Information Processing Systems, 33: 2020.
Lukas Biewald. Experiment tracking with weights and biases, 2020. URL https:///ww www.w/wwww.....m..m.... com/ Soft- available from wandb.com.
Yash Chandak, Scott Niekum, Bruno da Silva, Erik Learned-Miller, Emma Brunskill, and Philip S Thomas. Universal off-policy evaluation. Advances in Neural Information Processing Systems, 34:27475-274 2021a.
Yash Chandak, Shiv Shankar, and Philip S Thomas. High-confidence off-policy (or counterfactual) variance estimation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 6939 6947, 2021b.
Ricky TQ Chen, Jens Behrmann, David K Duvenaud, and Jörn-Henrik Jacobsen. Residual flows for invertible generative modeling. Advances in Newral Information Processing Systems, 32, 2019.
Scott Chen and Ramesh Gopinath. Gaussianization. Advances in neural information processing systems, 13, 2000.
Victor Chernozhukov and Christian Hansen. An iv model of quantile treatment effects. Econometrica, 73 (1):245 261, 2005.
Victor Chernozhukov, Tván Fernández-Val. and Blaise Melly. Inference on counterfactual distributions. Econometrico, 81(6):2205-2260 2013.
Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint arXi::151..07289, 2015.
Alicia Curth and Mihacla van der Schaar. On inductive biases for heterogencous treatment effect cstimation. Advances in Neural Information Processing Systems, 34:15883- -18844, 2021a.
Alicia Curth and Mihacla van der Schaar. Nonparametrie estimation of heterogeneous treatment effects: From theory to learning algorithms. In International Conference on Artificial Intelligence and Statistics, pp. 1810-1818. PMLR, 2021b.
Alicia Curth and Mihaela der Schaar. In search of insights, not magic bullets: Towards demystification of the model selection dilemma in hetcrogencons treatment effect estimation. 2023.
Alicia Curth, David Svensson, Jim Weatherall, and Mihacla van der Schaar. Really doing great at estimating cate? critical look al ml benchmarking practices treatment effect estimation. In Thirly-ffthh conference neural information processing systems datasets and benckmarks track (round 2), 2021.
Edward De Brouwer, Javier Gonzalez, and Stephanie Hyland. Predicting the impact of treatments over time with uncertainty aware neural differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 4705 4722. PMLR, 2022.
Floris Devriendt, Jeroen Berrevocts, and Wouter Verbeke. Why you should stop predicting customer churn and start using uplift models. Information Sciences, 548:497-515, 2021.
John DiNardo, Nicole Fortin, and Thomas Lemieux. Labor market institutions and the distribution of wages, 1973-1992: A semiparametric approach, 1995.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv preprint or Xie: 10.8516, 2014.
12

---

Published in Transactions on Machine Learning Research (11/2023)
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. 2017.
Jonathan Dumas, Antoine Wehenkel, Damien Lanaspeze, Bertrand Cornélusse, and Antonio Sutera. doep generative model for probabilistic energy forecasting in power systems: normalizing flows. Applied Eneryy. 305:117871, 2022.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. Advances in neural information processing systems, 32, 2019.
Sergio Firpo. 2007.
Efficient semiparametric estimation of quantile treatment effects. Econometrica, 75(1):259-24
Sergio Firpo and Cristine Pinto. Identification and estimation of distributional impacts of interventions using changes in incquality mcasures. Journal of Applied Econometrics, 31(3):457-486, 2016.
Leijiao Ge, Wenlong Liao, Shouxiang Wang, Birgitte Bak-Jensen, and Jayakrishnan Radhakrishna Pillai. Modeling daily load profiles of distribution network for scenario generation using flow-based generative network. JEEE Access, (177587-77597 2020.
Ian Goodfellow, Jean Pouget-Abadic, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. kernel two-sample test. The Journal of Machine Learning Research, 13(1):723 773, 2012.
Tuomas Haarnoja, Kristian Hartikainen, Pieter Abbeel, and Sergey Levine. Latent space policies for hicrar- chical reinforcement learning. International Conference on Machine Learning, pp. 1851 1860. PMLR, 2018.
Duke Hall. Principles of economics. 1920.
Jennifer L Hill Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statisties, 20(1):217-24) 2011.
Paul W Holland. Statistics and causal inference. Journal of the American Statistical Association, 81(396):
945 960, 1986.
Chin- Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. International Conference on Machine Learning, pp. 2078 2087. PMLR, 2018
Nan Jiang and Jiawei Huang. Minimax value interval for off-policy evaluation and policy optimization. Advances in Newral Information Processing Systems, 33:2747-2785 2020.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference machine learning, pp. 3020 3029. PMLR, 2016.
Daniel Kahneman and Amos Tversky. Prospect theory: An analysis of decision under risk. Econometrion, 47(2):263 292, 1979.
Nathan Kallus, Xiaojic Mao, and Angela Zhou. Interval ostimation of individual-leeel causal effects under unobserved confounding. The 22nd international conference on artificial intelligence and statistics, pp. 2281-2290. PMLR, 2019.
EH Kemnedy, S Balakrishnan, and LA Wasserman. Semiparametric counterfactual densily estimation. Biometrika, pp. asad017, 2023.
Sungwon Kim, Sang-gil Lee, Jongyoon Song, Jachyeon Kim, and Sungroh Yoon. Flowavenet: A generative flow for raw audio. arXiv preprint (to::811.0215:. 2018.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. 2015.
13

---

Published in Transactions on Machine Learning Research (11/2023)
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In International Conference Learning Represendations, 2014.
Durk P Kingmn and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions. Advances in neural information processing systems, 31, 2018.
Sören R Künzel, Jasjeet S Sekhon, Peter Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the national academy of sciences, 116(10):4156- 4165, 2019.
Lihua Lei and Emmanuel J Candès. Conformal inference of counterfactuals and individual treatment effects. Journal of the Royal Statistical Society Series B: Slatistical Methodology, 83(5)9111 938, 2021.
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep lalent-variable models. Advances in neural information processing systems, 30, 2017.
Kaushalya Madhawa, Katushiko Ishiguro, Kosuke Nakago, and Motoki Abe. Graphnvp: An invertible flow model for generating molecular graphs. arXiv preprint arxiv:1905. 11600, 2019.
Maggic Makar, Fredrik Johansson, John Guttag, and David Sontag. Estimation of bounds on potential outcomes for decision making. In International Conference on Machine Learning, pp. 6661-6661661111 6671. PMLR, 2020.
Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Normalizing flows for interventional densily estimation. arXiv preprint arXio:2209.0633 2022.
Chenlin Meng, Yang Song, Jiaming Song, and Stefano Ermon. Gaussianization flows. In International Conference on Artificial Intelligence and Statistics, pp. 4336-4345. PMLR, 2020.
John F Nash Jr. The bargaining problem. Econometrica: Journal of the econometric society, pp. 155 162, 1950.
Diego Olaya, Jonathan Vásquez, Sebastián Maldonado, Jaime Miranda, and Wouter Verbeke. Uplift mod- eling for preventing student dropout in higher education. Decision support systems, 134:113320, 2020.
Aaron Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukcuoglu, George Driess- che, Edward Lockhart, Luis Cobo, Florian Stimberg, et al. Parallel wavenet: Fast high-fidelity speech synthesis. In International conference on machine learning, pp. 3918-388 1926. PMLR, 2018.
Miruna Oprescu, Jacob Dorn, Marah Ghoummaid, Andrew Jesson, Nathan Kallus, and Uri Shalit. B-learner: Quasi-oracle bounds on heterogeneous causal effects under hidden confounding. 2023.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshmi- narayanan. Normalizing flows for probabilistic modeling and inference. Journal of Machine Learning Research, 22(57):1-66 64, 2021.
Joseph S Pliskin, Donald S Shepard, and Milton C Weinstein. Utility functions for life years and health status. Operations research,
224, 1980.
Ryan Prenger, Rafael Valle, and Bryan Catanzaro. Waveglow: A flow-based generative network for speech synthesis. In ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Process- ing (ICASSP), pp. 3617 3621. IEEE, 2019.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International con- ference on machine learning, pp. 1530-1533 PMLR, 2015.
Paul R. Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studics for causal effects. Biometrika, 70(1):41-5) 1983.
Jonas Rothfuss, Fabio Ferreira, Simon Walther, and Maxim Ulrich. Conditional density estimation with neural networks: Best practices and benchmarks. arXiv preprint 1Xx:::90330005, 2019.
14

---

Published in Transactions on Machine Learning Research (11/2023)
Donald Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Stalistical Associalion, 100:322-33 331, 02 2005. doi: 10.2307/27590541
Piotr Rzepakowski and Szymon Jaroszewiez. Decision trees for uplift modeling with single and multiple treatments. Knowledge and Information Systems, 32:303 327, 2012.
Yannick Schroecker, Mel Vecerik, and Jon Scholz. Generative predecessor models for sample-efficient imita- tion learning. In International Conference Learning Representations, 2019.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pp- 3076-300000 3085. PMI.R, 2017.
Michael Spence and Richard Zeckhauser. Insurance, information, and individual action. In Uncertainty in Economics, pp. 333 343. Elsevier, 1978.
Deborah L Thurston. A formal method for subjective design evaluation with multiple attributes. Research in engineering Design, 3(2):105-122, 1991.
Dustin Tran, Keyon Vafa, Kumar Agrawal, Laurent Dinh, and Ben Poole. Discrete flows: Invertible gener-
alive models of discrete data. Advances Newral Information Processing Systems, 32, 2019.
Toon Vanderschueren, Robert Boute, Tim Verdonek, Bart Baesens, and Wouter Verbeke. Optimizing the preventive maintenance frequency with causal machine learning. International Journal of Production Economics, 258: 108798, 2023a.
Toon Vanderschueren, Alicia Curth, Wonter Verbeke, and Mihaela van der Schaar. Accounting for informa- tive sampling when learning to forecast treatment outcomes over time. 2023b.
Brian G Vegetabile. On the distinction between' conditional average treatment effects' (cate) and' individual treatment effects' (ffects'iite) under ignorability assumptions. arXiv preprint arxiv:2108 04939, 2021.
Stefan Wager and Susan Athey. Estimation and inference of hetcrogencous treatment effects using random forests. Journal of the American Statistical Associalion, 113(523):1222 1242, 2018.
Lan Wang, Yu Zhou, Rui Song, and Ben Sherwood. Quantile-optimal treatment regimes. Journal of the American Statistical Association, 113(523):1243 -25144,41111,1, 2018
Patrick Nadeem Ward, Ariella Smofsky, and Avishek Joey Bose. Improving exploration in soft-actor-critic with normalizing flows policies. arXiv preprint arXiv:1906 190.0170,,,7711, 2019.
Runzhe Wu, Masatoshi Uchara, and Wen Sun. Distributional offline policy evaluation with predictive error guarantees. Xrxiv preprint ar arXiv:2302.094 09456, 2023.
Tengyu Xu, Zhuoran Yang, Zhaoran Wang, and Yingbin Liang. unifying framework of off-poliey general value function evaluation. Advances in Neural Information Processing Systems, 35:13570-133 -15583, 2022.
Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Yingxia Shao, Wentao Zhang, Bin Cui, and Ming-Hsua Hsuan Yang. Diffusion models: A comprehensive survey of methods and applications. arXiv preprint ar Xrxii:2200.000 00796, 2022.
Liuyi Yao, Zhixuan Chu, Sheng Li, Yaliang Li, Jing Gao, and Aidong Zhang. A survey on causal inference.
ACM Transactions on Knowledge Discovery from Data (TKDD), 15(5):1-46, 2021.
Jinsung Yoon, James Jordon, and Mihaela Van Der Schaar. Ganite: Estimation of individualized treatment effects using generative adversarial nets. In International Conference on Learning Representations, 2018.
Yao Zhang, Alexis Bellot, and Mihacla Schaar. Learning overlapping representations for the estimation of individualized treatment effects. In International Conference on Artificial Intelligence and Statistics, pp. 1014-1014. PMLR, 2020.
15

---

Published in Transactions on Machine Learning Research (11/2023)
Yingying Zhang, Chengehum Shi, and Shikai Luo. Conformal off-policy prediction. In International Confer- ence on Arlificial Intelligence and Statistics, pp. 2751-2768 PMT.R, 2023.
Tianhui Zhou, Yitong Li, Yuan Wu, and David Carlson. Estimating uncertainty intervals from collaborating networks. The Journal of Machine Learning Research, (2(1):11445-116)) 2021.
Tianhui Zhou, William E Carson IV, and David Carlson. Estimating potential outcome distributions with collaborating causal networks. Transactions on Machine Learning Research, 2022.
A Appendix
A.1 Data sets and associated data generating processes
This section gives more extensive details on the different semi-synthelic data sets used in this work. We describe each data set and the underlying data generating process that were used to simulate it.
IHDP (n = 747, = 25; Hill, 2011). We use data based on the standard response surface B. In this setting, potential outcomes are generated as
y: exp((x; + A)8) +
and y = x3 - w +
(11)
where E ~~ (0, 1), W is a (fixed) offset matrix with each value equal to 0.5, is a sparse coefficient vector with each clement sampled from (0,0.1,0.2 0.2, 0.3, 0.4) with corresponding probabilities (0.6, 0.1, 0.1, 0.1, 0.1). For each iteration of the data set, is set to ensure that the conditional average trealment effect for the treated (CATT) and conditional average treatment effect on the controls (CATC) both equal 4 on average. (Shalit et al., 2017).
We use the 100 replications from
EDU (n 8,627, d = 32; Zhou al., 2021). For each treatment group, a neural network fyt is trained based on the observed outcomes. The potential outcomes are then simulated
(0) fyro (I;) + (2 x 223) and r = (1; + (2 xx3 )1,
(12)
with en V(0,0.55* and ε exp(2). 223 refers to the 23rd covariate, which is binary covariate indicating whether the instance's mother has received previous education.
News (n 5,000, d = 3,477; Johansson et al., 2016). First, based on topic model two trcatment- specific cent roids =6 and zf are estimated in the topic space z(x). Each potential outcome is then generated as the similarity between 2(x;)) and of as
r!" C(I(x)) + 4, Tx) +
with scaling factor C and 8 1)0 Treatment assignment is modeled as p(te 1/x;) ex)(x==(=x)T=f))
oope=ztaa)aas))ooooeeeee with k = 10.
A.2 Hyperparameter optimization
We add more information on the chosen hyperparameters Table 3. Given the absence of the ground truth in observational data, hyperparameter tuning was done using wandb (Biewald, 2020) using the potential outcome's 's loglikelihood, PEHE, and loU of validation set. Additionally, to tune the hyperparameter X for the de-biasing term, we use an inverse propensity weighted loglikelihood, allowing us to assess performance under covariate shift. Training is done using gradient descent with the Adam optimizer (Kingma & Ba, 2015). We use several types of regularization: L1, f2, and noise regularization (Rothfuss et al., 2019). The encoder networks uses exponential linear units (ELU) as activation functions (Clevert et al., 2015).
The optimal hyperparameters additionally give some insights into how NOFLITE can deal with a diversity of data generating processes. For instance, the best metalearner was the T-learner for all data sets. The
16

---

Published in Transactions on Machine Learning Research (11/2023)
Hyperparameter
IHDP
EDU
News
General
Metalearner
T
T
T
Encoder
Hidden layers balancer
2
1
3
Hidden layers encoder shared Hidden layers encoder separate Hidden neurons encoder
3
0
0
2
2
3
8
8
32
Flow
Number of flow transformations k
0
4
1
Flow type
SigmoidX
SigmoidX
Hidden neurons transformer
4
2
Hidden neurons conditioner
16
32
Hidden layers conditioner
2
1
Training settings
Learning rate
5e-4
5e-4
5e-4
Batch size
128
512
128
Training steps
5,000
5,000
10,000
Regularization λě
le-3
0
5c-4
Regularization Arą
5e-4
1e-3
5e-3
Ammdd
le-2
le-2
Noise regularization
0
0
1
Noise regularization y
5e-1
1e-1
5e-1
Truncation probability
0
0
1e-2
Table 3: Hyperparameter tuning. We show the optimal hyperparameters for the different data sets. The flow type 'SigmoidX' refers to the deep sigmoidal flow (DSF) of Huang et al. (2018) conditioned on the balanced representation Φ of x. During inference, drawing samples from truncated normal distribution was observed to be more stable for the News data.
potential outcomes were generated scparately for all data sets under consideration, which corresponds more closely to the T-learner's hypothesis space. Moreover, the complexity of the distribution can be tweaked depending on the data generating process: NOFLITE does not any flows for IHDP, uses only flow for News, and uses four flow transformations for EDU, These settings match the increasingly more complex distributions being used in the corresponding the data generating processes.
17


Reviews:
[{"review_id": "idnRkJIgXW", "reviewer": "Reviewer_gFHo", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* This is a simple (in a good way) idea to combine ideas from generative modeling and the literature on causal neural networks to obtain distribution estimates. \n* The performance of the proposed estimator is quite good compared to other neural network based approaches. \n* Well written (very nice clarity in both task description and solution) and organized. \n* I can see how the proposed method would be easily extensible and could provide a nice basis for future work.\n\nWeaknesses:\n* The inclusion of the MMD constraint in this setting is motivated as it is in the prior literature (to enforce balance). Given the probabilistic framing for the rest of the paper it could be useful to have a larger discussion about how the MMD constraint should be interpreted in the probabilistic context (e.g., the discussion in Kanagawa, et al. proposition 6.1 (https://arxiv.org/pdf/1807.02582.pdf). \n* It would be useful to have a comparison against non-neural network based estimators as well, especially BART which is mentioned in related work but not compared to. \n* The paper has no theoretical results. While this isn't necessarily a reason to prevent acceptance it would be nice if there were some discussion on the properties of the proposed estimator.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a neural network based model for estimating individual treatment effect distributions. The authors approach this problem by using two headed network, using normalizing flows in order to allow for more complex distributions, and imposing a balance constraint using maximum mean discrepancy in order to have validity of the estimated counterfactuals. The core of the idea is to combine probabilistic modeling with normalizing flows and ideas from estimation of the mean potential outcomes using neural networks. Experiments are provided which demonstrate the efficacy of the proposed method.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "As I mention above, I would like to see a stronger motivation for the balance constraint from a probabilistic perspective, an extended set of simulation results that compare against more \"classical\" benchmarks (which remain very strong and competitive baselines). It would also be nice if the authors discussed theoretical properties along the lines of Johannson, et al. (Learning Weighted Representations for Generalization Across Designs), though I understand if the authors feel this is beyond the scope of the paper.", "broader_impact_concerns": "Individual effect estimation has immediate impact to optimization and recommendation algorithms that various actors may use to alter of manipulate the behavior of people and other systems. While this work does not provide any immediate ethical quandaries past what currently exist in the prior art it is important to be mindful of.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "HHVw763PBs", "reviewer": "Reviewer_KX1e", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\nThe use of Normalizing Flows to construct the approximate treatment effect distribution is an intriguing and reasonable idea. I feel that this contribution is a specific strength of this paper.\n\nI agree with the position provided by the authors that estimating the full treatment effect distribution enables more informed decision making. I am encouraged to see more work that focuses on these type of efforts. I do however have some concerns (listed below) about the way this claim is presented and justified however.\n\nThe proposed NOFLITE is adequately grounded among the relevant literature. It is made clear by the authors where prior methods may not fully satisfy accurate effect estimation. I appreciated the detail provided in Section 2 about Normalizing Flows, I do feel that more could be written to explain how the inverse flow $g^{-1}(z)$ is executed. That would be a really helpful point of information to support the full conceptual grounding of the proposed modelling approach.\n\nThe construction of the NOFLITE approach is easy to understand and clearly articulated at the outset of Section 4. Figure 2 is really informative! The one aspect of this that could be improved is discussing how the encoder handles confounding. More detail is warranted there to explain the balanced representations used. (Also, the terms in Section 4.1.2 \"The flow transformation uses a monotonic transformer network that is parametrized by a conditioner network...\" are very jargony without any context about what they may mean. Perhaps these could be expanded upon slightly? And Section 4.1.3 seems unnecessary in the main body of the paper without supporting context in Section 2.)\n\n## Weaknesses\n\nThere are not very clear descriptions of what the actual claims and intended contributions of this paper are, outside of the the generic framing of estimating the full treatment effect distribution with a neural architecture that harnesses normalizing flows to construct the distributional estimate. Many of the statements in the introduction are vague and imprecise, appealing to a generic intuition that the estimation of a treatment effect distribution and its presented consequences are already accepted. For example, the sentence: \"By pairing the treatment effect distribution with a utility function, we can obtain the expected utility.\" is fine but doesn't completely convey why utility functions are important or useful nor expresses why expected utility is a desirable metric. The following sentence provides really weak justification. Appealing to the use of the metric in other fields isn't a sufficient explanation or description of why it's a good idea here. There are similar problems with how the \"first\" stated advantage of estimating the effect distribution is presented. What do the expanded set of descriptive statistics actually provide? Why is this useful? I think more precise statements and an expanded introduction could strengthen the paper and clearly set out the major claims being put forward by the authors.\n\nFigure 1 would be improved by visualizing how $\\mathbb{E}[p{\\tau}]$ fails to adequately capture the true treatment effect. Additionally, it's unclear what the different type of utility functions may measure (Fig 1 (left)) and why they would be under consideration. This isn't sufficiently discussed in the text of the paper. For example, what does it mean to have a \"personalized utility function\"? Finally, Figure 1 should be situated closer to Section 3. It isn't very helpful to be placed so early (as currently presented, this may change if there were illustrations of how the expected treatment effect falls short in tandem with a more detail Introductory section).\n\nUnfortunately, I fear that the authors missed an important comparative baseline approach in the CF-ODE of de Brouwer, et al (2022, AISTATS) which satisfies the very same considerations proposed as important dimensions of consideration in developing models for estimating treatment effect distributions (Neural, Likelihood, Confounding). The CF-ODE model satisfies all three considerations, similar to NOFLITE, and also enables inference over continuous time settings when estimating the treatment effect distributions. This is an unfortunate oversight and I would highly recommend that the authors include this method in their experimental benchmarking so the claims of NOFLITE's performance could be fully centered among the full literature. It would be really interesting to see how the distributions provided by NOFLITE's use of Normalizing Flows would compare to those generated by the Neural ODE based CF-ODE model. A comparative analysis that could be done is in measuring each distribution's calibration for each possible treatment (see Sun and Yu, 2022 for various calibration metrics).\n\nIt's unclear why it's necessary to have a monotonicity requirement in the Normalizing Flow. It would be helpful to more clearly lay out the assumptions and constraints used in the formulation of NOFLITE.\n\nIt's unclear whether comparing the empirical baselines to NOFLITE via the model's likelihood is totally fair given that the baselines are not optimized to maximize a likelihood? Some clarification from the authors in the text on this point, justifying the use of this comparative metric would be really helpful.\n\nMore discussion should be included at the beginning of Section 3 to illustrate/explain why the expectation is \"insufficient for decision-making in may domains\". It's not appropriate to vaguely present these types of arguments without more information. \n\nThe major and possible disqualifying weaknesses, such that the major claims of the paper are unsupported are:\n**There is no convincing demonstration that supports the major claim of the paper, that a full treatment effect distribution is necessary for adequate decision making. The empirical analysis rests on three metrics and does not discuss the implications of downstream use. On top of this, the improvements seen by NOFLITE are marginal at best and are not consistent across all metrics.**\n\n**There is zero demonstration or analysis to support the suggested strengths of the full distribution, more coherent statistical analysis _and_ the use of \"personalized utility functions\".** <-- This is a major disappointment because the use of utility functions is a neat idea and I was disappointed that there was no supporting evidence for the concepts and ideas put forward at the beginning of the paper. \n\n##### Suggested Literature\n- `Edward De Brouwer, Javier Gonzalez, and Stephanie Hyland. Predicting the impact of treatments over time with uncertainty aware neural differential equations. AISTATS 2022`\n- `Sophia Sun and Rose Yu. Copula conformal prediction for multi-step time series forecasting. arXiv preprint arXiv:2212.03281, 2022`", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a method, based on Normalizing Flows, to estimate the full distribution of treatment effects. Claims are made that this full distribution enables a more accurate analysis of the treatment effect than simply estimating the expectation of that distribution. The paper optimizes the model that estimates this full distribution based on the log likelihood and is compared to recent generative approaches to estimating the treatment effect on three semi-synthetic datasets.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "## Major\n\nBetter clarifying details about the specific motivations of the approach.\n\nClarify how the log likelihood is an adequate comparative metric across the various baselines.\n\nInclusion of CF-ODE as an empirical baseline.\n\nAdditional analysis of the statistics enabled by the full distribution.\n\nAnalysis and discussion of how utility functions can be used to modify and/or improve the analysis of the estimated treatment effect distributions.\n\n## Minor\n\nMore detail about how the Normalizing Flows are used in practice. \n\nFigure 2 typo in inversion of the flow ($g_1$ is not inverted).\n\nReframe the metrics PEHE and IoU to discuss their use in evaluating the learned/estimated distribution.", "broader_impact_concerns": "None that I can ascertain.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vgMA76knSf", "reviewer": "Reviewer_fyxB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Contribution:\n\nStrengths:\n+ Well-written. The main ideas are presented clearly and the paper includes nice figures/illustrations to convey key ideas.\n+ Estimating the distributions of conditional treatment effects is an important but very hard problem. \n\nproblems: \n- The term “Individual treatment effect” has the potential to be misleading. I think it’s better to be more precise and call these heterogeneous/conditional treatment effects or at the very least have a discussion of how the word “individual” refers to a measured covariate vector. \n- Empirical results: IoU metric for evaluating the confidence intervals seems to suggest there is poor coverage although the metric is hard to evaluate. It would be more helpful to present empirical coverage results— e.g., does the 90% CI actually contain the true ITE 90% of the time. \n- No theoretical guarantees. While normalizing flows have shown success in domains where evaluation is straightforward (e.g., image generation), causal inference is different since in real-world settings we never observe the ground-truth outcomes needed for evaluation. It’s not clear that the empirical results are sufficient to establish that this method will perform well in real-world contexts. Such claims would be bolstered by theoretical guarantees. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes a normalizing flow based method to estimation the distributions of conditional treatment effects under standard (but strong) assumptions of no measured confounding, consistency, and overlap. They propose two meta-learner configurations for the architecture, one (S-learner) which learns one flow for both treatments and the other (T-learner) that learns a separate flow for each treatment. They present empirical results using semi-synthetic design on three datasets, Infant Health Development Program and an Education dataset, and a News dataset. Their empirical analysis investigates how their method compares to existing methods in terms of log-likelihood, PEHE, and a intersection-over-union measure of the estimated confidence intervals. \n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "recommendations:\n- Add discussion about whether it’s reasonable to make the required assumptions and provide references to work in this space that relaxes those assumptions (e.g., Interval Estimation of Individual-Level Causal Effects Under Unobserved Confounding by Nathan Kallus, Xiaojie Mao, and Angela Zhou)\n- Provide standard coverage empirical results (as described above)\n- Add discussion on the difficulty of high-dimensional conditional estimation and discuss implications for proper use of the method (e.g., sample sizes required etc)\n- Use more precise language like \"conditional or heterogeneous treatment effect estimation\" instead of \"individual treatment effect estimation\".\n", "broader_impact_concerns": "For work that has the potential to be used in high-stakes domains and settings where evaluation is challenging, it's really important to be careful about the language used and to be transparent about the assumptions and limitations of the proposed method. As noted above, it would be important to update the paper to clarify what is meant by \"individual treatment effect estimation\" and to devote more discussion to the important assumptions required by the method. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies a problem relevant to the TMLR community and approaches it from a new perspective. After the substantial revision and discussion between authors and reviewers, all reviewers leaned toward accepting the paper.  

However, the main motivating idea—that the average effect is not enough for good decision making—has many more links to previous work than acknowledged in the related work section. For the final version of the paper, I recommend adding a paragraph, or substantially updating the second paragraph in related works, contextualizing the manuscript in the closely related literature on reinforcement learning, off-policy evaluation, and conformal prediction of treatment effects. To name a few *recent* examples, 
* Wang, L., Zhou, Y., Song, R., & Sherwood, B. (2018). Quantile-optimal treatment regimes. Journal of the American Statistical Association, 113(523), 1243-1254.
* Chandak, Y., Niekum, S., da Silva, B., Learned-Miller, E., Brunskill, E., & Thomas, P. S. (2021). Universal off-policy evaluation. Advances in Neural Information Processing Systems, 34, 27475-27490.
* Bellemare, M. G., Dabney, W., & Rowland, M. (2023). Distributional reinforcement learning. MIT Press.
* Xu, T., Yang, Z., Wang, Z., & Liang, Y. (2022). A Unifying Framework of Off-Policy General Value Function Evaluation. Advances in Neural Information Processing Systems, 35, 13570-13583.
* Wu, R., Uehara, M., & Sun, W. (2023). Distributional Offline Policy Evaluation with Predictive Error Guarantees. arXiv preprint arXiv:2302.09456.
* Zhang, Y., Shi, C., & Luo, S. (2023, April). Conformal Off-Policy Prediction. In International Conference on Artificial Intelligence and Statistics (pp. 2751-2768). PMLR.

I am confident that the authors can find even more examples of aiming to estimate the full distribution of treatment effects / policy rewards / utility functions.
Decision: Accept (poster)