Generalization bounds and algorithms for estimating the effect of multiple treatments and dosage
Authors: Alexis Bellot
Anish Dhir
Giulia Prando
Venue: Submitted to ICLR 2023 (2023)
Keywords: Treatment effect estimation

Abstract: Estimating conditional treatment effects has been a longstanding challenge for fields of study such as epidemiology or economics that require a treatment-dosage pair to make decisions, but may not be able to run randomized trials to precisely quantify their effect. In the context of representation learning, there is an extensive literature relating model architectures with regularization techniques to solve this problem using observational data. However, theoretically motivated loss functions and bounds on generalization errors only exist in select circumstances, such as in the presence of binary treatments. In this paper, we introduce new bounds on the counterfactual generalization error in the context of multiple treatments and continuous dosage parameters, which subsume existing results. This result, in a principled manner, guides the definition of new learning objectives that can be used to train representation learning algorithms. We show empirically new state-of-the-art performance results across several benchmark datasets for this problem, including in comparison to doubly-robust estimation methods.

Under review as a conference paper at ICLR 2023
GENERALIZATION BOUNDS AND ALGORITHMS FOR ES- TIMATING THE EFFECT OF MULTIPLE TREATMENTS AND DOSAGE
Anonymous authors Paper under double-blind review
ABSTRACT
Estimating conditional treatment effects has been a longstanding challenge for fields of study such epidemiology economies that require treatment-dosage pair to make decisions, but may not be able to run randomized trials to precisely quantify their effect. This may be due financial restrictions or ethical consid- erations. In the context of representation learning, there is an extensive literature relating model architectures with regularization techniques to solve this problem using observational data. However, theoretically motivated loss functions and bounds on generalization errors only exist in selected cireumstances, such as the presence of binary treatments. In this paper, we introduce new bounds on the counterfactual generalization error the context multiple treatments and continuous dosage parameters, which subsume existing results. This result, in a principled manner, guides the definition new learning objectives that can be used to train representation learning algorithms. We show empirically of-the-art performance results across several benchmark datasets for this problem, including in comparison doubly-robust estimation methods.
state-
INTRODUCTION
Treatment effect estimation is the problem of predicting the effect of an intervention (e.g. treatment- dosage pair) on an outcome of interest guide meaiin.. The challenge for prediction models is to learn this map from observational data, which is formally generated from a different structural causal model which treatment assignment varies according to an individual's s covariates, instead of being fixed by the decision-male Counterfactuals define the outcome that would have been observed had the assigned treatment been different. For concreteness, consider designing policy for the administration of chemotherapy regiments; not all cancer patients in the available data equally likely to be offered the same type and dosage, with varied factors, e.g. age, wealth, etc., involved in the decision-makee making process. Evaluating new treatment combination for given patient is data point that invariably under-tepeesenned in the empirical distribution of the data.
Treatment effect estimation is studied under wide range of assumptions, including experimental de- signs that feature ignorability (Imbens, 2000; Imai & Van Dyk, 2004), multiple treatments, sequential decision-making making problems, and different generative models encoded in general causal graphs (Pearl, 2009). There is a growing literature on several parts of this problem in the field of machine learning that attempts define loss functions that are conducive learning representations of covariates predictive of both observed and counterfactual outcomes. Existing methods could be generally categorized by the theoretical guarantees that inspire training objectives, driven either by bounds for the generalization eтror or by Noubly-robustness. guarantees. In the first line of research, Shalit et al. (2017); Johansson et (2020) showed the binary treatment setting that the counterfactual error, that is not computable from data by design, could be instead bounded by the in-sample error plus term that quantifies the difference in distributions between treated and untreated populations, leading to differentiable loss function that can be used to train expressive neural networks. Several papers used this insight to investigate different neural network architectures for this problem. For example, Johansson al. (2016) proposed to use separate feed-forward prediction heads on top of common representation, Zhang et al. (2022) use transformers, De Brouwer et al. (2022); Seedat et al. (2022) neural differential equations. In turn, doubly-robust estimators combine expressive function

---

Under review as a conference paper at ICLR 2023
approximators and inverse probability weighting leveraging statistical non-parametric asymptotic guarantees of both estimators (Funk et al., 2011; Kennedy, 2016; 2020). In particular, when the direct estimate of the outcome is biased, such when using nonparametric or high-dimensional regression, the doubly robust estimator weights the model residuals by inverse propensity weights in order remove the bias. Its convergence and consistency for treatment effect estimation requires only that one of the estimators is consistent. In principle, any consistent function approximator could be used, which in the context of neural networks has led several adaptations of loss functions and architectures. For example, Shi et al. (2019) adapted the architecture of Johansson et al. (2016) for this purpose introducing targeted regularization, and Nie et al. (2020) proposed varying coefficient networks in the context of continuously-valee vallued dosage parameters. In both cases, however, the authors provide guarantees for population average treatment effect estimation, in contrast with conditional average treatment effect estimation.
Despite the generality of these results, no guarantees and no theoretically motivated loss functions exist for learning representations for counterfactual estimation in the general setting of multiple treatment types and/or continuous treatment values or dosages. The challenge in the context of representation learning is that there is notion of treatment group as each individual gets assigned potentially different and unique treatment value. Lack of overlap in finite samples and subsequently large estimation variance for counterfactual predictions that adjustments for distributional differences are, in principle, not applicable. In particular, the intuition for reducing variance by regularization deviates from previous proposals (that regularize representations of covariates to match distributions among groups with different treatment types (Shalit et al.. 2017)) as potentially infinite set of counterfactuals for each individual must be considered. Even the analysis of multiple categorical treatments is currently an open question as, while pairwise comparisons between treatment specific distributions could be implemented in principle, it is not computationally tractable to do so in practice. At this moment, only heuristic neural network architectures for this problem have been proposed, including Dose Response networks that consist of multi-task layers for dosage sub-intervals defined on top of common representation (Schwab et al., 2020), variants of generative adversarial networks (Bica et al., 2020), and varying coefficient networks (Nie et 2020).
exacerbated in this setting the extreme
In this paper, we investigate the design of representation learning- algorithms for predicting (conditional average) treatment effects in the context of multiple treatments and continuous dosage parameters. Our analysis starts by extending definitions of loss and generalization error to this broader setting, over all possible treatment-dosage pairs. We then show by using the definition of integral probability metrics that the generalization error can be bounded by a term that is computable from data and that involves the factual error and a term that quantifies the statistical dependence between the pair of coatment-ddsagee random variables and observed confounders. In principle, any treatment space on which we can deñnn a probability measure is consistently accounted for, which gives well- defined bounds on the generalization error for treatments with multiple types and continuous values, and in particular, our bound includes as a special case existing guarantees in the binary treatment case (Shalit et al., 2017). This bound suggests new training objectives for learning representations conducive to counterfactual estimation. Moreover, such objectives are tractable: both avoiding combinatorial numbers of pairwise comparisons and avoiding binning dosage values into different sub-intervals. A further contribution we make is to design extensive numerical comparisons that compare both methods driven by bounds the generalization error (that typically target conditional average treatment effects) and methods driven by doubly-robust guarantees (that typically target average treatment effects). Moreover, we do so independently of the adopted neural architecture which provides the first analysis of different objectives for the problem of treatment effect estimation with multiple, continuously-vat valued treatments. We hope these results can give some insight into the trade-offs of different approaches to this problem and demonstrate the ability of representation learning techniques to tackle wider ranging scenarios within treatment effect estimation.
2 BACKGROUND
We start by introducing the notation and definitions used throughout the paper. In particular, we use capital letters for random variables (X), small letters for their values (z), bold letters for sets of variables (X) and their values (x), and 5 for the spaces where they are defined (Ωx) if not explicitly stated. To simplify notation, we consistently use the shorthand P(x) to represent probabilities or
2

---

Under review as a conference paper at ICLR 2023
densities P(X x) and similarly P(y x) to represent P(Y l x). For three sets of variables X, Y, Z the conditional independence statement "X is conditionally independent of Y given Z z" is written as X I.
We use the semantics of the Rubin- Jubin-Neyman potential outcomes framework, see c.g. Section in Rubin (2005). We assume that for individual with observed covariates fxx, and tuple T = -W. S) defining the treatment type out of & distinct treatments W E Ωw = [wI..... wxx and dosage parameter S S2y I. there is a corresponding potential outcome Y, that would have been observed had the assigned treatment been 7 = t. With observational data only one of these potential outcomes is observed for each unit depending on the treatment assignment. We will refer to the unobserved potential outcomes as counterfactuals. Let Y denote the observed outcome and Or {(w, 8) w Sw: 8 E TR] denote the set of all treatment options. The goal to derive estimates of the expected potential outcomes for given set of input covariates: [[[, x for any value of and x. Under the following standard assumptions (Rubin, 2005), it is well understood that the treatment effect between two selected treatment options i and f2 reduces to a contrast of conditional distributions, presented in Prop. below.
Assumption 1 (Unconfoundedness). The treatment assignment, t (w, 8) Dry, and potential outcomes, Y, are conditionally independent given the covariates x, i.e. LT) X.
Assumption 2 (Overlap). For any E Ωx such that P(x) > O, we have > P(I x) > 0. each t e Ω2.
Assumption 3 (Consistency).. The observed outcome is the potential outcome, as a function of any t
treatment, when the treatment is set to the observed exposure, i.e. Y = Yt ifT = any
Dr.
Proposition! 1 (Identifiability). Under assumptions I and 3, and any tı, b2 E Q2g
E[ - Yt, x]= E x,4 E | x, 12],
which is composed entirely of observable quantities and can be estimated from data given Assump. 2.
We refer to the quantity E[Y,, Yr, x] as the conditional (or individual the conditioning set identifies unit) treatment effect (CATE), and the E[Y,,-+,, -Y,,] as the average, population, treatment effect (ATE). Our results rely on defining representation functions Φ : f2x - DR. where DR is the representation space, that preserve unconfoundedness and overlap, and the identifiability of the treatment effect. For this purpose, it is sufficient to assume ó to be injective'.
-
Corollary (Identifiability given representation). Under the assumption that the representation is injective, 1 > P(t ((x)) > 0 and Y, ((X), that is unconfoundedness and overlap hold conditional on features (xx)
Without loss of generality we will assume that QR is the image of f2x under c. We will write P also denote the distribution induced by 6 over QR and let h QR x Qr fyy be a prediction function defined over PR. Next, we define two complimentary loss functions: one is the standard machine learning loss, which we will call the factual error on the estimation at the observed treatment type and dosage tuple, and the other is the counterfactual error, as an average error over all other treatment assignment options, made for an individual with a particular treatment type and dosage tuple.
Definition 1. For given loss function L: By f2y Rt, the expected factual and counterfactual losses of and o at treatment t E 2r are defined as,
Cx(1) = S JoY
S C(y, (w)
(2)
Ccr(1) Rppr [ ( C(y, M(())
Jax Jny
P(xl())xdyi.
The counterfactual error defines the average error made for the counterfactual prediction at treatment tuple (w, s) on all individuals that are observed to be assigned a different treatment t. This
'The remark been made that injectivity of representation difficult to enforce (Zhang et al., 2020: Johansson et al., 2019). An algorithmic solution, discussed by Zhang et al. (2020), is to include a decoder from the representation to the input domain and reconstruction loss in the training objective to encourage solutions with invertible latent representations. A reconstruction loss and encoder-decnder architecture can be included on top of the regularization terms proposed in this paper.
3

---

Under review as a conference paper at ICLR 2023
definition extends the binary treatment to assess the quality of counterfactual predictions at t Stp Similarly, we define an average measure of factual and counterfactual performance over all possible treatment options E S2y.
Definition 2. The average factual and counterfactual error over all treatment options are defined.
Cr ve-S. Jan (r(())(()dt, Lor J. (c)(())(t)dt.
Next we define the enror made on the estimation of a counterfactual contrast for given pair of treatments, instead of an average over all counterfactual treatment options.
Definition 3. Let the treatment effect between two different treatments tuples t2 ?r be given by ((t))ta(()) = E[, x E[ x]. The error in treatment effect estimation is then defined as,
(())))) = [ ((T(L)))))) (x)
Jox
where P())))) Qx Sy denotes its estimate.
3 REPRESENTATION LEARNING FOR COUNTERFACTUAL ESTIMATION
As is apparent in the presence of multiple treatments and continuously-valued dosages, there is no notion treatment group as each individual gets assigned a potentially different and unique treatment value. The intuition for reducing variance by regularization deviates from previous proposals potentially infinite set of counterfactuals for each individual must be considered Shalit et al. (2017). The following theorem shows that the average counterfactual error defined in Def. 2 can be bounded by terms that are explicitly computable from observational data.
Theorem (Bound on average counterfactual generalization error). Under the assumption that injective, it holds that,
LCF & LF X sup [{ g(r,t) (P(r))()( - P(r, ((),rt))r
gest, Jor JoR
Og defines a space of functions 9 On D2 expressive enough include Say f(yi, (()x X as a function of Φ(x) and where À > O depends on the choice of representation function ø.
This theorem states that the average counterfactual error is upper-bounded by the factual error plus a term that quantifies the dependence between treatment tuple T and covariates X. As the treatment tuple contains multiple treatment types w, as well as continuous dosages s, this single bound is valid for multiple treatment values well continuous dosages.
Bias Variance tradeoff Counterfactual estimation would be unbiased by minimizing factual losses Cp by Prop. but the variance in the estimation of counterfactuals for treatment-dosage pairs that are not heavily represented in observational data will be high. This will contribute larger generalization error and captured in the supremum in the second term of Eq. (6). In particular, the supremum quantifies an imbalance in the association of and R. by using distributional distances between joint distributions and the product of marginals. PP(r) (()) ()) P(r, 1) is large if not all treatment and feature combinations are evenly represented in the data. This observation recovers an interesting intuition Q, is chosen to be expressive enough. The observation being
that suppen, Say San g(r,t) (P()) - P(r, dr)dt = 0 if and only if the representation is independent of treatment assignment, i.e. (X)LLT. This extreme case leads lower variance as counterfactuals for treatment-dosage tuple have the same effective sample size that of the observational data. The hyperparamter À controls the tradeoff between the bias and the variance of the counterfactuals. Two choices for P. consider are the space of functions in universal Reproducing Kernel Hilbert Space (RKHS) with characteristic kernels (Sriperumbudur et al., 201 1), which recovers the the space of Lipschitz functions with Lipschitz constant bounded by which recovers the Wasserstein distance (Villani, 2009).
Hilbert Schmidt Independence Criterion (Gretton et al., 2007), and
4

---

Under review as a conference paper at ICLR 2023
Binary treatment One insight from Thm. 1 is that bias in the treatment assignment in the context a general treatment choices, such as multiple treatment types or continuously-vallee treatments, takes the form of high statistical dependence between random variables, that is more general than differences between distributions. In particular, differences in distributions between treatment groups defined by Shalit et al. (2017) in the binary treatment case can be formulated as statistical dependence between random variables. The following corollary recovers the generalization bound of Shalit al. (2017) as special case.
Corollary 2. Let 22T {0, 13. Then, by Thm. 1,
CCF & LF + X sup geStg .RR
g(r) (P(r T T =) P(r T 0))dr
(7)
and is equivalent to (Shalit et al., 2017, Lemma 1).
We show next similar result that gives generalization bounds for the treatment effect comparing two specific treatment options, instead of average over all possible counterfactual options, that may be of interest in applications specifically comparing two treatment options.
Theorem 2. Let t1.t2 t2 P2r be two treatment tuples to compared. Then,
22 < LF(ti) + sup gest
g(r) (P(() - PPrr T t)))dr + LF(22)
[
sup gest,
( - P(r T =2))2r) ov, -Ye,, (8)
where oy, and oy,, stand for the variance of the random variables Ye, and Yi,, respectively, under the distribution P(x).
3.1 ARCHITECTURES AND ALGORITHMS FOR COUNTERFACTUAN ESTIMATION
This section discusses the architectures of the representation and prediction functions used, well as training objectives to leverage the generalization bound in Thm. 1. The training objective that we define can be used with any neural network architecture that parameterizes representation function on f2x - Qn and a separate prediction function hg On x 22 - By with sets parameters and 0 respectively.
Following the discussion in Sec. 3, we learn representation & and prediction function to minimize trade-off between predictive accuracy and imbalance in the representation space using the following objective:
N Σ
y ha hx(0)(x(1)) o(1))) + -IMM,, (o(X),
min O.
where 7 is hyperparameter, the number of samples, and IPMn, (6(X,,,
supgest, Seep San g(r,t) (Pt) P(r, t)) is the integral probability metric for a chosen space of functions S2g
Concretely, we wish to increase the predictive accuracy while making the representation as indepen- dent of the treatment possible. We consider the Hilbert Schmidt Independence Criterion (HSIC) and the Wassertein distance as choices for the integral probability metric. In practice, the HSIC can be approximated with finite data sample using (Gretton et al., 2007, Eq. (3)). For the Wasserstein distance, we simulate a sample with joint distribution P(r) P(t) by randomly permuting the observed treatment-dosage pair across individuals to generate a sample (r) ( : 1,...,,,, NN). where a: : (1,...,,, N) - {1....,N N} is bijective function. The original data ((r(m) ((0)) : 1,..., N) drawn from the distribution P(r, l). The empirical distributions are compared using the arguments in (Cuturi & Doucet, 2014). Both these regularization terms are differentiable and all parameters can be updated using stochastic gradient descent.
Each treatment type w corresponds to a separate prediction network head, i.c. he := (h (w)) weD while the representation layer is common across all treatment types. In particular, this implies that each
5

---

Under review as a conference paper at ICLR 2023
(K
:
(v...)
IPMa.
Figure 1: Sketch of the architecture.
sample (x(x), mo(() ,(()),(() (0)) used update only the prediction network h1 (w(")) corresponding to the observed treatment (((), while all samples are used to update the representation layer Φg- A sketch of this training routine is given in Fig.
The following network architectures for the prediction functions wesly have been proposed in the literature.
(
Dose Response Networks (DRNet) Schwab et al. (2020) propose Dose Response Networks for predicting the effect of dosage task network with a shared set of layers and multiple task-specific heads. In this context, the range of dosage values is split into separate bins and each of them is associated with separate head. Each task-specific network in addition takes the dosage value as input, but crucially the parameterization of the prediction function common all dosages belonging the same sub-interval. For example, the range of dosage values for treatment type w could be divided into sub-inter inteittervals, thus using 5
outcome of interest. The architecture takes the form of multi-
task-specific heads (i(w) (hg (h(w..)) ... : hg (((,,))) hgg (w.,d) SRR 22 fl í 1, ,...., 5. To some extent this approach accounts for the heterogeneity in the dose-response function but remains limited by the binning choice and may be vulnerable to abrupt changes in the prediction on the dosage values
(w.1)
that separate two bins, as demonstrated by Nie et al. (2020).
Varying Coefficient networks (VCNet) Varying Coefficient networks (Nie et al., 2020) are proposed for dose-response estimation, but a multi-task architecture can be designed case. In particular, the authors define the parameters for each prediction network h, := B(aa)
special
PR Ωy to be functions (s)) (s),s). . Od, (s)) of dosage themselves, where de is the total number of parameters. Each scalar parameter 0: Ug - IR is given by linear combination (s(s) -ل= (),) of polynomial basis functions ( f defined the space of dosage values Ug. The coefficients {asst : I,... dø, / = 1,... L} define the trainable parameters and the map hg (r, (11) (r, s) is differentiable with respect to {as.1
(ar)
I,..., do...,do,1 ... L). For example, DRNets are recovered by choosing { [/ to be a picce-wise constant functions spline basis of the form 1 (s; & & 8,) with different 8;, 85. More general choices can be made, such as B-splines, that lead to continuous dose response curves. The influence of the dosage parameter is different to that of a covariate and thus ensures dosage information is not lost in high-dimensional representations, which in practice has been shown to lead to better counterfactual prediction performance.
EXPERIMENTS
This section conducts controlled experiments on synthetic and semi-synthetic datasets previously used in the literature. Overall, we found that simulation results support generalization guarantees with different architectures benefiting from the proposed regularization strategy using both the HSIC and Wasserstein distances.
4.1 BASELINES AND METRICS
We consider several baselines for comparison, including different neural network architectures without regularization and with doubly-robust regularization techniques. In particular, consider standard multilayer perceptron (MLP) that optimises the (factual) squared error loss objective to learn the weights of the network, standard VCNet (Nie al., 2020), and DRNet (Schwab et al., 2020). In

---

Under review as a conference paper at ICLR 2023
the context of doubly-robust optimization, Shi et al. (2019); Nie et al. (2020) propose to learn a joint representation Φ(x) that is conducive to both counterfactual h1 SR x Ωr - f2y and propensity score estimation l22 IR. 22 by using loss function that trades-off the two objectives, e.g.,
1 Σ
y 3 -1 (()))) 2
a CrossEntropy (12(0(x(1)) (())))
(10)
h and hy are consistent estimators of the outcome and propensity scores respectively, as well as satisfy the non-parametric estimating equation,
N Σ
p(y(") (m)) xm)); Pr. h22 O,
(11)
where e denotes a perturbation term that optimized and where (in the binary treatment case for simplicity).
ply, t, x; h1, b2, c) (x, 1) (x, 0) + (11 1- h=(x))
1 1-11 (y (x, (12)
hx((x)
then the resulting estimator will have desirable asymptotic properties for average treatment effect (Shi et al., 2019; Kennedy, 2016). We consider iı parameterized by both VCNets and DRNets. Algorithms trained to minimize Eq. (10) are denoted VCNet-PS, DRNet-PS, and algorithms trained to minimize both Eqs. (10) and (11) (also known Targeted Regularization),
referred as VCNet-TR, DRNet-TR. TR. Finally, we consider Generalized Propensity Scores (GPS) (Imbens, 2000; Imai & Van Dyk, 2004) that fit linear model using inverse propensity scores. Our proposed methods are labeled DRNet-HSIC. DRNet-Waas -CNet-HSIC, and VCNet- Wass, which combine existing architectures with the proposed regularization methods. We include details on network architectures, hyperparameters optimisation and computational time in Appendix C.
For performance comparisons, we consider the Mean Integrated Squared Error (MISE),
MISE
11 N n= wesl
Σ Σ
( s(mo))))
(13)
where we use the notation (() and g(())
(w.e) (w.e) for the true and predicted outcome for individual given treatment-dosage pairs (w, s) E Qrr, and the expectation is taken with respect to the dosage
parameter, i.e. E (w.ss) Sis ((w.s)) P(s)ds. Intuitively, MISE calculates how well an algorithm estimating individual level dose response and thus accounts for the heterogeneity in treatment response. In contrast, the Average Mean Squared Error (AMSE) evaluates population average counterfactual prediction by taking sums and integrals before comparisons between predicted and true outcomes. We define and evaluate AMSE in Appendix D.
(n)
4.2 DATASETS
The nature the treatment-effects estimation problem does not allow for meaningful evaluation on real-world datasets. This is simply because never observe counterfactual for a given unit. There are, however, established synthetic and semi-synthetic datasets that have been used by Schwab et al. (2020); Bica al. (2020); Nie al. (2020). Following these proposals we use,
Fully synthetic. A data generating mechanism with a total of 6 randomly generated covariates and a single treatment with dosage ranging from 0 to I that involve complex functions for both treatment assignment and outcome function, as defined by Nie et al. (2020).
HHDP-continuous The original semi-synthetic IHDP dataset from Hill (2011) contains binary treatments with 747 observations on 25 covariates. We adapt this dataset to the continuous dosage context by changing the treatment assignment and outcome function. We generate these in similar way Nie et al. (2020).
News. The News dataset consists of 3000 randomly sampled news items from the NY Times corpus (Newman, 2008), which originally introduced as a benchmark in the binary treatment setting. We generate a continuously-valued treatment and corresponding outcome in a similar way as Bica al. (2020).
7

---

Under review as a conference paper at ICLR 2023
1.6
Syrihhet
IHDP continaus
VCNee: Wa55
DRM Wass
144
VChes-MtiC
14-
1.2
1.0
1.0
0.0
0.
04
Regutarser
Regu
(a) Synthetic.
(b) CHDD-ccontinuos..
Figure Oft- sample MISE error versus IPM regularization, relative to error 7 (no regularisation), on 50 realizations Synthetic (a) and IHDP-continnous (b) datasets. Average values (dot markers) and one standard deviation (shaded areas) are shown.
In each of our experiments we generate 50 independent realizations from each of the above datasets (20 for News), with samples split into train/validation/ee set with ratios 0.6/0.2/0.2. Further details on the data generating mechanisms, as well as about networks architecture, hyper-parameters tuning and training times provided in Appendix C.
4.3 EPFECTIVENESS OF REGULARISATION
Our first experiment tests the effectiveness of the proposed regulariser by evaluating counterfactual prediction performance as function of that determines the influence of the independence constraint feature space in Eq. (9).
We consider both DRNets and VCNets architectures, with both HSIC and Wasserstein regularizers on the Synthetic and IHDP-continuout datasets. Fig. 2 compares MISE performance results for these models with varying values of ? relative to 7 0 (without regularisation). Both datasets include confounding factors which induce bias or imbalance in the treatment assignment T for different covariate subgroups X. On both plots we observe that the proposed regularization term (with increasing y > 0 relative to y = 0) confers an advantage to training with regularization term that explicitly corrects for this imbalance for the purpose of predicting counterfactuals. The gain of some 7 > 0 is consistent across different neural network architectures and across different datasets, which illustrates generalization guarantees but also shows that some form of regularisation may broadly be applicable in practice.
4.4 PERFORMANCE COMPARISONS
In this section we conduct wide-range comparison against the benchmark prediction algorithms using the three data generating mechanisms described in Section 4.2. Table reports average values and standard deviations of V'ISSE over 50 (20) realizations of Synthetic and IHDP-continuous (News) datasets. On average, the proposed regularization technique, using either the HSIC or Wasserstein distances between distributions, outperforms all other regularization techniques on both choices of neural network architecture. Several trends are interesting to discuss in more detail.
Existing representation learning algorithms that optimise doubly-robust objectives optimal. The results show that, in terms of the MISE, our regularisation based on counterfactual generalisation outperforms doubly robust methods. This can be explained by the fact that doubly robust methods have guarantees when estimating average treatment effects, and not individual or conditional treatment effects. The proposed regularization techniques, with guarantees for coun- terfactual generalization error, instead, are designed for good performance in conditional average treatment effect estimation and often substantially outperform in terms of MISE. We believe that this discrepancy is due to the doubly robust methods discarding information that helps predict the individual outcome, resulting in a worse MISE performance. This also emphasizes the fact that estimating average counterfactuals and individual counterfactuals Indeed, the cross-entropy term in Eq. (10) encourages the representation to retain information that is predictive of the treatment; hence, it encourages the discarding of information that is predictive of
always
require different objectives.
8

---

Under review as a conference paper at ICLR 2023
Synthetic
IHDP-continuous
News
GPS
2.80 (0.51)
4.91 (0.87)
MLP
(0.09)
0.74 (0.07)
1.05 (0.07)
DRNet
0.34 (0.06)
0.60 (0.06)
0.84 (0.04)
DRNet-PS
0.43 (0.07)
0.66 (0.15)
0.84 (0.04)
DRNei-TR
0.41 (0.03)
2.07 (3.54)
0.82 (0.05)
DRNet-HSIC
0.39 (0.06)
0.52 (0.05)
0.87 (0.04)
DRNet-Wand Wass
0.38 (0.07)
0.51 (0.08)
0.86 (0.04)
VCNet
0.33 (0.03)
0.56 (0.09)
0.85 (0.05)
VCNet-PS
0.62 (0.42)
0.59 (0.11)
0.98 (0.09)
VCNet-TR 0.42 (0.12)
0.64 (0.39)
0.99 (0.06)
VCNet-HSIC 0.28 (0.04)
0.56 (0.09)
0.87 (0.05)
VCNet-Wass
0.38 (0.03)
0.55 (0.10)
0.81 (0.05)
Table I: Average values and standard deviations (within brackets) of /MISE across 50 (20) realizations of Synthetic, CoDPicontisuous (News) datasets, Bold notation highlights the best- best-performing algorithm on cach dataset.
the outcome but not the treatment, which is simply noise when predicting the treatment. On average there is also significant gain by considering more expressive neural network architectures, for instance DRNet outperforms MLP and VCNet outperforms DRNet all metries and data generating mechanisms. Finally, we note that GPS requires matrix inversion which was not feasible to compute on the nigh-dimensional News dataset.
CONCLUSION
In this paper, we investigate the task of estimating the conditional average causal effect of dosage from combination of observational data and assumptions the causal relationships in the underlying system. When these assumptions hold, we give new bounds on the counterfactual generalization error in the context of multiple treatment types and continuously-valuue dosage parameters that subsume generalization guarantees from the binary treatment case. Using this result, we provide new learning objectives that can be used guide the training of representation learning algorithms. We show empirically new state-of-the-art performance results across several benchmark datasets for this problem. To our knowledge, this is the first paper exploring representation learning and regularization for conditional average counterfactual estimation in the context of multiple, continuous valued treatments a principled manner. We hope these results can demonstrate the ability representation learning techniques to tackle wider ranging scenarios within treatment effect estimation.
REFERENCES
Ioana Bica, James Jordon, and Mihaela van der Schaar. Estimating the effects of continuous- interventions using generative adversarial networks. Advances in Neural Information Processing Systems, 33:16434-16445, 2020.
Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, and Whitney Newey. Double/deeiaaseddnemma machine learning of treatment effects. American Economic Review, 107(5):261 2017.
Marco Cuturi and Arnaud Doucet. Fast computation of wasserstein barycenters. In International conference on machine learning, pp. 685-693. PMLR, 2014.
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. 4722- PMLR, 2022.

---

Under review as a conference paper at ICLR 2023
Michele Jonsson Funk, Daniel Westreich, Chris Wiesen, Til Stürmer, M Alan Brookhart, and Marie Davidian. Doubly robust estimation of causal effects. American jourrnaljjouraal of epidemiology. 173(7): 761-767, 2011.
Daniel Greenfeld and Uri Shalit. Robust learning with the hilbert-schmidr independence criterion. In International Conference Machine Learning, pp- 3759-3768. PMLR, 2020.
Arthur Gretton, Kenji Fukumizu, Choon Hui Teo, Le Song, Bernhard Schölkopf, and Alexander J Smola. A kernel statistical test independence. NIPS. 2007.
Miguel A Hernán and James M Robins. Causal inference, 2010.
Jennifer Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1):2172240, 2011.
Keisuke Hirano and Guido W Imbens. The propensity score with continuous treatments. Applied Bayesian modeling and causal inference from incomplete-dus perspectives, 226164:73-84, 2004.
Kosuke Imai and David A Van Dyk. Causal inference with general treatment regimes: Generalizing the propensity score. Journal of the American Statistical Association, 99(467):854-866 2004.
Guido W Imbens. The role of the propensity score in estimating dose-response functions. Biometrika, 8763),06--70,, 2000.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference on machine learning, pp. 3020-3029. PMLR, 2016.
Fredrik D) Johansson, David Sontag, and Rajesh Ranganath. Support and invertibility in domain- invariant representations. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 527-536. PMLR, 2019.
Fredrik Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. arXiv preprint arXiv:2001.07426, 2020.
Edward H Kennedy. Semiparametric theory and empirical processes in causal inference. In Statistical causal inferences and their applications in public health research, pp. 141-167. Springer, 2016.
Edward H Kennedy. Optimal doubly robust estimation of heterogeneous causal effects, arXiv preprint arXiv:2004.14497, 2020.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
David Newman. Bag of words data set. UCI Machine Learning Respository, 2008
Lizhen Nie, Mao Ye, Dan Nicolae, et al. Venet and functional targeted regularization for learning causal effects of continuous treatments. In International Conference on Learning Representations, 2020.
Judea Pearl. Causality. Cambridge university press, 2009.
Donald Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association. 100(469):322-335 331-331, 2005.
Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 5612-5619, 2020.
Nabecl Seedat, Fergus Imrie, Alexis Bellot, Zhaozhi Qian, and Mihaela der Schaar. Continuous- time modeling of counterfactual outcomes using neural controlled differential equations. arXiv preprint ar arXiv:2206.083 2022.
10

---

Under review as a conference paper at ICLR 2023
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: general- ization bounds and algorithms. In International Conference on Machine Learning, pp. 0085-0055. PMLR, 2017.
Claudia Shi, David Blei, and Victor Veitch. Adapting neural networks for the estimation of treatment effects. Advances in neural information processing systems, 32, 2019.
Bharath K Sriperumbudur, Kenji Fukumizu, and Gert RG Lanckriet. Universality, characteristic kernels and rkhs embedding of measures. Journal of Machine Learning Research, 12(7), 2011.
Mark Van Der Laan and Daniel Rubin. Targeted maximum likelihood learning. The international journal of biostatistics, 2(1), 2006.
Cédric Villani. Optimal Iransport: old and new, volume 338. Springer, 2009.
Yao Zhang, Alexis Bellot, and Mihaela Schaar. Learning overlapping representations for the estima-
tion of individualized treatment effects. In International Conference on Artificial Intelligence and Statistics, pp. 1005-1014. PMLR, 2020.
Yi-Fan Zhang, Hanlin Zhang, Zachary Lipton, Li Erran Li, and Eric P Xing. Can transformers be strong treatment effect estimators? arXiv preprint 220x:22002..1336, 2022.

---

Under review as a conference paper at ICLR 2023
A RELATED WORK ON DOUBLY ROBUST ESTIMATION OF THE AVERAGE TREATMENT EFFECT
Thm. I suggests that the imbalance in the distribution of X across treatment dosage pairs is relevant for the expected generalization error of fitted models. Estimators inspired from the semi-parametric literature, known as doubly robust estimators (Van Der Laan & Rubin, 2006; Chemoozukov et al., 2017), instead try to optimize average treatment effects (ATE), e.g. Ex x] - Ex x]. by constructing prediction function h1 fix x Stg Ωy, a propensity score function h2 Ox f2r, and perturbation term 6, satisfying the non-parametric estimating equation,
N
1
Σ ( x(m), mm) hr, h2, E) = O,
(14)
N
where (in the binary treatment case for simplicity).
ply, t, x; h1, h2, = h1(x,))
0) +
(v - (i(x,))) -. (15)
(x)(x) 1- h2(x))
(x(x,0 () is an estimator of [E[ x]. while (x(x) ) is an estimator of the probability of treatment P(t| x) and e Qr is a perturbation term that is optimized. In the literature, a common estimation approach is to rely on (task-agnostic) fitted models h1 and h2. and then choose so that this equation is satisfied. If h and hy are consistent estimators of the outcome and propensity scores respectively, as well as satisfy Eq. (14), the resulting estimator of the ATE will have desirable asymptotic properties (Shi et al., 2019; Kennedy, 2016). However, as these guarantees are on the average treatment effects, they do not necessarily guarantee accurate estimates of conditional treatment effects.
In the context of neural networks, Shi et al. (2019); Nie et al. (2020) propose to learn joint representation ((x) that is conducive to both counterfactual h S2 x Ωr - fy and propensity score estimation h2 SR. + Utr by a using a loss function that trades-off the two objectives, e.g..
N
2 + a CrossEntropy (ha(s(x(1)) tlm))) (n)) - (())))
(16)
N
as in (Nie et al., 2020, Eq. (1)) or (Shi et al., 2019, Eq. (2.2)). The motivation is that: "If the average treatment effect is identifiable conditioning on the propensity score [...] it suffices to adjust for only the information in x that is relevant for predicting the treatment", see (Shi et al., 2019, Theorem 2.1). Intuitively, the cross entropy term in Eq. (16) encourages the representation to retain information that is predictive of the treatment. Hence, it encourages the discarding of information that is predictive of the outcome but not the treatment, which is simply noise when predicting the treatment.
Variables that affect the outcome and not treatment are referred to as effect modifiers in the literature, e.g. (Hernán & Robins, 2010). By definition, the treatment effect varies across different conditioning sets of these effect modifiers. As effect modifiers are responsible for the heterogeneity of treatment effects, it is necessary to condition on them to obtain accurate conditional treatment effects. Thus, to compute conditional average or "individualized" treatment effects such representations may be too restrictive because they tend ignore effect modifiers.
In contrast, our regularizer penalizes the dependence between the representation and the treatment distributions explicitly. Loosely speaking we discard covariate information predictive of treatment but outcome information is retained. Hence, our regularizer should preserve these effect modifiers leading to more accurate estimates of conditional treatment effects. We conclude that, in general, optimal average treatment effects does not necessarily imply optimal conditional average treatment effects as measured by expected losses Definitions and 2 We verify this intuition in our experiments.
Definitions 1 and also involve averages but makes head to head comparisons between observed outcomes and predicted outcomes for each individual (x, t) in the term Jay L(ye, t)) (x)xy (which are then averaged across individuals) instead of averaging predicted counterfactuals across the whole population before comparison with average true outcomes across different dosage levels.
12

---

Under review as a conference paper at ICLR 2023
B PROOFS
Theorem 1 (Generalization bound for the average counterfactual error). Under the assumption that one to one, holds that,
LCF & LF + A- sup
g(r,t) (P(r))((() - P(r, t))
(17)
gelg
3rr URn
2g defines a space of functions 9 UR x Z2r R expressive enough to include
Say Cly,, (() (/ng))dd/////> X as a function of v(x) and where À depends on the choice of representation function Φ.
Proof. Let & SR fix be the inverse of á and let (x,, Say t)) dy.e The following derivations show the claim.
Jay Jox ()xx [ ( ((t)))
LCF LF
Jor S5x
[ (n, (eo(x,)))(((x)PPP))))) (P(x))(t) P(x, txd
Jo Jix
(v)(((()) ((()()))P() P()(())
Jor Jon
(v)(((()).)) (P(r) P(r, Je 'rdt drdt
Jrr Sn
[ geft, JDr JR
< A sup
g(r,t) (P(r) P(r, ((rxt))drdti
For the third equality, the distribution over QR Q2 can be obtained by the standard change of variables formula, using the determinant of the Jacobian of (br),, denoted Js giving P(e(r), f) = P(r, Je (which cancels with the inverse Jacobian that appears after the change of variables in the differential term). The last inequality comes from the assumption that 4,, f e 2g., which is justified and extensively discussed in Shalit et al. (2017).
To prove Thm. 2, we will use the following lemma.
Lemma For convenience, we write m(t, x) E[Y; x] and we define its estimate given prediction function f OR x Qx - Ω2y by ftt, x). If L is the square loss, it then holds that,
[ [ C(n,
CCF(1)
(18)
J0x Jrr
[ (ve- f(x, ((x P(y) Pxx)
Jox Jay
(19)
Jo [ (f(x,t) m(x, l P(x)dxxx
(20)
Jay
I (m(x,t) t) yd) 2 P(v Px)
Jax Sty
(21)
(f(x,t) (m) -).) Px))
(22)
Jxx Jsy
((xxt)-m(xxt))) P(x)) σy,
(23)
Jax
The third term the third equality evaluates to zero because m(x, t) Sey Y/P(y1 xx)dy and we have defined the variance of Y, with respect to the distribution P(x) as = Say Soy (m(x, 2 x) dx
13

---

Under review as a conference paper at ICLR 2023
Theorem 2 (Generalization bound for selected treatment tuples t and tz). Let t1,t Sr be two treatment tuples to be compared. Then,
[ (P(r) P(r T = tr) + Cr(t2)
221 & Cr(t1) + sup
gelta
RR
+ sup
g(r) - (P(r) P(r T = t2))drr -oy,, -org (24)
gestg
Sun
where oY,, and oy,, stand for the variance of the random variables Yr, and Yız, respectively, under the distribution P(x).
Proof.
((11.22)
(f(11,x) ((1,,) m(t1, xl(2
(25)
Sxx
2 (f(t1, m((ttt x))² x)))))((x)(sss) +
(f(t2, m(tx,x))
(26)
Jix
.0x
2(Ccr((1) -oyi, 2(((r((12) -ori))
(27)
2 CF(t1) + sup geft,
g(r) (0(((((()))))))))))))))) P t1))dr
-Y+,
(28)
SuR
2 (p((2) sup
g(r) (P(r) P(r[T 62))
-Y12
(29)
gell,
RRR
The first inequality holds by the fact that (a + b)2 < 2a + 2b for any a, b RR The second equality holds by Lemma and the last inequality holds by the same arguments used Theorem
C EXPERIMENTAL DETAILS
C.1 DATA GENERATING MECHANISMS
This section describes the data generating mechanisms used our experiments.
Synthetic. We generate synthetic data similar to Nie et al. (2020). With covariates x E R6 all drawn from uniform distribution between 0 and generate the continuous dosages and outcomes as follows.
10 10 (msxn(mma(aa) Ø1, 22, x3)) + max(I3, 25)3 + sin(0.5za)(1 [1 exp(x4 0.5zz))+
x
1 (Ii 2252
(30)
2 (sn) + 2xg
6.5 + (0,
8/x, (2 0.5)) (22
4 max( sin(xa) V(0,0.25))
(31)
1 1 2x3
where 8 = =1 exp(-x))
IHDP Continuous. The IHDP dataset contains 25 covariates with binary treatments and continuous outcomes (Hill, 2011). Disregarding original treatments and outcomes, we use the covariates to generate new continuous dosages and outcomes to test method. We follow the data generating
14

---

Under review as a conference paper at ICLR 2023
procedure of Nie et al. (2020), namely:
8|x
201
+
2 max Ts. n6)
+ tanh
5 5 /e/
J; - 02
4 (N0 0.25),
(32)
+ J2
0.2 + min R6, 266
[
x/x,
sin (3T3)
tanh 5 [ Iz
C1
+ exp
0.5 + min (12, 11,, 13, ))
0.2(x1 - I6)
+ +NN(0, 0.25),
1.2
[J]
(33)
Lies Zi
41 = Fp(x)
(34)
[
Eier Tij [1]
22 Ep(x)
(35)
where
= (1 + exp( (08 -3)) (4, 7, 9, 10, 11, 12, 13, 14, 15).
= (16, 17, 18, 19, 20, 21, 22, 23, 24, and J
News. This dataset contains words sampled from 5000 news articles (Newman, 2008). The covariates are word counts. We generated continuous dosage and outcomes by following the data generation method listed in Bica et al. (2020). We first sample three vectors v ~ A(0, 1), with V; ///l for i 1,2,3. Then, dosages are drawn from a Beta distribution:
max (1,
2x V2
(2, 9);
(36)
v
Finally, outcomes are sampled according to:
v² - 0.3
8' = exp
(37)
y = 2 (maxx min((, + 20x V3 (4(s 0.5)3) ) sin
T3 ( 2
+ +NO(0, 0.25).
(38)
C.2 ARCHITECTURES AND TRAINING DETAILS
In both VCNet and DRNet, the representation part of the network ó,, and the prediction heads he have two layers each, with 50 hidden units and ReLU activations. Following Nie et al. (2020), we use B-spline with degree 2 and knots placed at (1/3, 2/3) for VCNet and 5 regression heads for DRNet.
For the MLP model, we use a 4-layers network to represent similar power of approximations to ensure fair comparison. We optimise the networks using Adam (Kingma & Ba, 2014) with weight decay of 0.005 for regularisation and a batch size of 1000. Learning rate chosen within the set (0.01, 0.005, 0.001, 0.0005, 0.0001, 0.00005) using the procedure outlined in C.3. Each data set is split into a train/ vatin//aiiatio//test set with ratios 0.6/0.2/0.2. To avoid overfitting, we stop the training if the validation loss did improve after 50 epochs.
Propensity score regularization (-PS methods) In addition to the representation net ó f2x - OR and to the prediction net h1 QR Ω2g Gl, propensity score regularized methods also include separate head hy Pg + Zr. Parameters are tuned by minimizing the loss
N
Cps(ó, hy, h2) = N
1
Σ
((m) (((6(x(m)) + a CrossEntropy (hz(0(x( E(m))
(39)
In our experiments, he is modelled through a softmax layer over grid of 10 bins. Average treatment effects are then estimated by considering an additional perturbation term following Shi et al. (2019) and Nie et al. (2020). is treated as a hyperparameter and chosen within the set (0.5, 1) using the procedure detailed in C.3. The implementation in practice follows the publicly available code of Nie al. (2020).
15

---

Under review as a conference paper at ICLR 2023
Synthetic
IHDP-continuous
News
DRNet
55.82 (0.36)
49.82 (0.21)
82.91 (1.76)
DRNet-PS
55.83 (0.30)
50.09 (0.17)
83.47 (2.00)
DRNei-TR
87.71 (0.22)
79.65 (0.58)
125.17 (2.77)
DRNet- HSIC
83.28 (1,47)
75.86 (1.44)
134.47 (1.74)
DRNet WASS
63.57 (0.42)
59.67 (0.30)
99.20 (1.64)
VCNet
32.73 (0.15)
29.71 (0.34)
77.51 (0.23)
VCNet-PS
32.91 (0.21)
29.79 (0.21)
76.13 (0.26)
VCNeI-TR 51.82 (0.24)
43.82 (0.17)
114.03 (0.30)
VCNet-HSIC 59.43 (0.61)
53.88 (0.94)
128.23 (0.76)
VCNeI-WASS 41.69 (0.29)
38.64 (0.29)
91.83 (0.52)
Table 2: Computational times (in seconds) required for 2000 epochs of training. Averages and standard deviations (within brackets) 10 runs (5 for News dataset) are reported.
Targeted regularization (-TR methods) Methods labeled -TR use the functional targeted regular- ization approach presented in Nie et al. (2020) which optimizes the loss function
((m))) (-))- h1(0(x(())) N
B
N Σ
eN (N(((())
CrRlØ, h2) pos(o.hn.hh)
12( h2(((x(()))
(40)
where (.)+) Σj-, (-) is modelled through JN spline basis functions 1by of degree The number of basis might change with the sample size N. Following Nie al. (2020), we select the learn- ing rate for ex (-).). 3 and the number of spline knots within the sets (0.001, 0.0001), {20, 10, 5) / and (5, 10, 20). respectively. Again, the implementation in practice follows the publicly available code of Nie al. (2020).
IPM regularization (-HSIC and Wass methods) IPM regularized methods minimize the pro- posed loss Equation (13) in the main body of this paper, where the 7 is selected within the 10) using the procedure in C.3. The implementation of Wasserstein dis-
(10*/6 i 18, 17,, ...,..., 9, tance regulariser follows the one available at https: : / iiggtttt thub. com/ cl ini ca 1ml / / rnet/ / blob/ //aattateaaae cfr///r//////////e /ut//uuta il, py #L 6. HSIC regulariser is computed according (Greenfeld Shalit, 2020, Eq. 3), using two RBF kernels with length-scales (0.05, 0.1, 0.5, 1, 5, 10, 50, 100, 500}.
Generalised Propensity Score (GPS) (2004).
We use our implementation following Hirano & Imbens
C.3 HYPER-PARAMETERS TUNING
We use grid-search to tune the hyper-parameter. Namely, we generate dataset for each hyperpa- rameters setting, randomly splitting into a train/test set with ratio of 0.8/0.2 and we choose the hyperparameters values giving the best MISE test score.
C.4 RUN TIME COMPARISONS
Table 2 reports the computational time (in seconds) required by the algorithms compared in the experimental section for 2000 training epochs. These results are machine and algorithm- specific but do serve as relative comparison of run times for different neural network architectures and regularization techniques. In general, Wassertstein IPM regularisation is more computationally efficient than TR and IPM regularisation through HSIC metric.
16

---

Under review as a conference paper at ICLR 2023
Synthetic
IHDP-continuous
News
MIIS
VAMSE
VMIISE
VAMSE
VMISSE
v VAMSE
GPS
2.80 (0.51)
2.75 (0.52)
4.91 (0.87)
4.88 (0.88)
MLP
0.72 (0.09)
0.62 (0.11)
0.74 (0.07)
0.59 (0.05)
1.05 (0.07)
0.66 (0.15)
DRNet
0.34 (0.06)
0.20 (0.04)
0.60 (0.06)
0.42 (0.05)
0.84 (0.04)
0.37 (0.09)
DRNet-PS
0.43 (0.07)
0.25 (0.06)
0.66 (0.15)
0.43 (0.12)
0.84 (0.04)
0.37 (0.09)
DRNei-TR
0.41 (0.03)
(0.02)
2.07 (3.54)
0.68 (0.70)
0.82 (0.05)
0.30 (0.10)
DRNet-HSIC
0.39 (0.06)
0.22 (0.02)
0.52 (0.05)
0.35 (0.06)
0.87 (0.04)
0.44 (0.07)
DRNet-Want
0.38 (0.07)
0.21 (0.02)
0.51 (0.08)
0.34 (0.08)
0.86 (0.04)
0.43 (0.07)
VCNet
0.33 (0.03)
0.13 (0.06)
0.56 (0.09)
0.33 (0.09)
0.85 (0.05)
0.31 (0.11)
VCNet-PS
0.62 (0.42)
0.32 (0.30)
0.59 (0.11)
0.31 (0.12)
0.98 (0.09)
0.25 (0.13)
VCNet-TR
0.42 (0.12)
0.19 (0.11)
0.64 (0.39)
0.31 (0.13)
0.99 (0.06)
0.4 (0.08)
VCNet-HSIC
0.28 (0.04)
0.10 (0.03)
0.56 (0.09)
0.33 (0.08)
0.87 (0.05)
0.35 (0.11)
VCNet-Wass
0.38 (0.03) 0.11(0.03)
0.55 (0.10)
0.33 (0.08)
0.81 (0.05)
0.29 (0.1)
Table 3: Average values and standard deviations (within brackets) of /MIS and v VAMSE across 50 (20) realizations of Synthetic, IHDP coDP-connisuuus (News) datasets. Bold notation highlights the best-performing algorithm each dataset.
D FURTHER EXPERIMENTS
This section includes experiments comparing performance with respect Average Mean Squared Error (AMSE), in addition to Mean Integrated Squared Error (MISE),
MISE
1
E
(g) -g(n (w)
N
k n-1
wefl
N
2
AMSE Σ wefl
(g (n) (w, (w 8))
where (w, s) and ((") (w, s) stand for the true and predicted outcome for individual n given treatment-dosage pairs (w,,) 8) Dr, and Bg(s) Snn g(s) P(s)ds. The AMSE calculates the accuracy of the population level dose response.
As the doubly robust methods get rid of effect modifiers, that are useful for accurate predictions, but have theoretical guarantees for the average treatment effects, expect these methods to get better AMSE. On the other hand, as the regularizers proposed in this work provide guarantees on the counterfactual error, expect models trained with these achieve better MISE score.
In Table we show mean performance for both MISE (as in the main body of this paper) and AMSE with the objective contrast the proposed methods, designed for optimal conditional average counterfactual prediction, i.e. MISE, and doubly-robust methods designed for population average counterfactual prediction, i.e. AMSE. Overall, we note that the proposed regularization technique, using either the HSIC or Wasserstein distances between distributions, is competitive across all datasets and metrics, including AMSE.
Across all datasets, there is a clear trend showing that regularizing for optimal generalization performance terms of the MISE with HSIC leads to good population average performance as well, as measured by AMSE. Doubly-robust methods (-PS,-TRE are designed for optimality in estimation of the AMSE, and across all datasets they are either optimal or competitive compared to all other algorithms. It is interesting to note that the performance achieved by the proposed regularisation techniques (HSIC, Wasserstein) are very close to the optimum AMSE while, in contrast, doubly robust methods often perform significantly worse in terms of MISE than the optima achieved by HSIC and Wasserstein regularisation. This further confirms the intuition presented in Appendix A: estimating average counterfactuals and individual counterfactuals require different objectives.
17

---

Under review as a conference paper at ICLR 2023
Moreover, in terms of AMSE, it still holds that neural network architectures with better expressiveness to model heterogeneous dose-response curves perform better on all datasets.
18


Reviews:
[{"review_id": "y4cPE72qgxB", "reviewer": "Reviewer_RujR", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper seeks to define appropriate representations of the coviariates for treatment effect estimation in the setting of multiple treatments with continuous dosage parameters. Previous work has looked at this problem in the binary treatment setting. As in previous work, the paper defines a standard ML loss and a counterfactual loss. Bounds for the counterfactual loss are derived in terms of the standard ML loss and the IPM between the joint distribution of the treatment and representation and the product of marginals for a sufficiently expressive function class. The key intuition is to increase predictive accuracy while making the representation as independent as possible of the treatment. Experimentally, different function classes for the IPM and different network architectures are implemented and compared. ", "summary_of_the_review": "While the experimental results are impressive, I think I need to be convinced that this is not just a trivial extension of the binary treatment setting. The paper as it stands doesn't make this case very strongly. \n", "strength_and_weaknesses": "The paper extends the generalization bounds for treatment effect estimation for the multiple treatments with continuous dosage setting. While novel, it doesn't seem like this extension is nontrivial. The bounds are of the same form. The paper would be much stronger if the differences and contributions of the paper are highlighted over just mentioning that they extend the binary treatment results. The experimental validation however, is thorough and the proposed methods and regularizers are better than existing approaches.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "tcOUbDB_sf", "reviewer": "Reviewer_BdyF", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors develop an approach for conditional average treatment effect estimation for interventions that have an arbitrary dimension and type (i.e., could be continuous valued or binary). They do so by extending previous work by Johansson/Shalit on binary treatments. \n", "summary_of_the_review": "The paper tackles an important issue but it falls a bit short of suggesting a truly novel approach that deals with challenges that naive extensions of binary treatment effect estimation are prone to. ", "strength_and_weaknesses": "Strengths: \n\nBecause the majority of work in causal literature focuses on binary interventions, this paper represents an important and understudied problem in the field of causality. It is likely to be of interest to a wide audience. \n\nWeaknesses: \n\nThe main limitation that I see with this work is that it represents the most straightforward extension of Johansson/Shalit’s work on binary treatments, without addressing the inherent challenges that come with the high dimensionality of the intervention variable. Specifically:\n\n1. One of the core limitations that prevents extensions of methods designed for binary treatment effect estimation the high dimensional treatment setting is that in this setting the overlap assumption is more likely to be violated. This is the kind of challenge that I would expect any paper on high dim treatment to tackle head on, so it was rather disappointing that the authors did not.  \n\n2. Another core limitation preventing extensions from binary to categorical extensions (in the framework presented by Johansson/Shalit et al) is the fact that naive extensions would lead to multiple heads and extreme sample splitting (meaning each head is updated by a very small number of samples). However, this is the direction that the authors choose to take. Can the authors comment on how they expect their model to perform as k increases? This is particularly hard to assess given that it seems like all the experiments seem to have been done on datasets with k=1? At least that was my understanding from reading the appendix in Nie et al. Can the authors confirm that none of the experiments were done on k>1? \n\n3. Of note is that those two limitations have been addressed elsewhere (see [Kaddour et al](https://arxiv.org/abs/2106.01939)). Can the authors comment on how their approach compares (both theoretically and empirically) to work by Kaddour et al? It seems to me that Kaddour et al address the two main limitations of high dim/arbitrary type of interventions that this work does not address. \n\n\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Mxxa43txw-f", "reviewer": "Reviewer_Z4Tv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper provides a generalization bound for a treatment effect on the treated (ETT) with a representation learning function and develops a learning principle based on the proposed bounds.", "summary_of_the_review": "Even if the paper is well-written, I think the paper needs intuitive and insightful explantions about Theorem 1. ", "strength_and_weaknesses": "## Strength\n1. This paper is well-written. \n2. This paper is self-contained.\n3. This paper contains extensive simulation results.\n\n## Weakness\n1. Some theories can be more formalized. Details are provided in the next section.\n2. Intuitive/insightful explanations about the proposed results are missing.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "9URHYz2sD6", "reviewer": "Reviewer_zcLw", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper addresses the problem of representation learning for counterfactual estimation under non-binary treatments. The authors observe that counterfactual error can be reduced by ensuring independence between treatments and covariates. Using this observation the authors propose a loss and show improved empirical results with respect to other proposed learners based on neural networks. ", "summary_of_the_review": "As I stated above, I believe that the novelty and significance of this paper is extremely limited by the lack of contextualization within the larger literature on causal inference for general valued treatments. I would encourage the authors to better contextualize the results of their work and provide comparison against relevant methods from both inside and outside the representation learning community. ", "strength_and_weaknesses": "Strengths: \n* Counterfactual estimation for general treatment regimes is an important problem\n* The authors do a nice job of motivating the loss functino\n\nWeaknesses:\n* The paper does not seem to engage with the literature on learning for continuous and multivalued treatments. For example, \"Regret Minimization for Causal Inference on Large Treatment Space\" from ICML 2021 makes nearly the same observation as this paper and employs HSIC for dependence minimization. Additional examples include \"Double Debiased Machine Learning Nonparametric Inference with Continuous Treatments\", which addresses the problem of learning continuous treatment effects using machine learning, and the balancing weight literature which explicitly expresses the risk of the causal estimate in terms of dependence between treatment and covariates. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject