DTRNet: Precisely Correcting Selection Bias in Individual-Level Continuous Treatment Effect Estimation by Reweighted Disentangled Representation
Authors: Mengxuan Hu
Zhixuan Chu
Sheng Li
Venue: Accepted by TMLR ()

Abstract: Estimating the individual-level continuous treatment effect holds significant practical importance in various decision-making domains, such as personalized healthcare and customized marketing. However, most current methods for individual treatment effect estimation are limited to discrete treatments and struggle to precisely adjust for selection bias under continuous settings, leading to inaccurate estimation. To address these challenges, we propose a novel Disentangled Representation Network (DTRNet) to estimate the individualized dose-response function (IDRF), which learns disentangled representations and precisely adjusts for selection bias. To the best of our knowledge, our work is the first attempt to precisely adjust for selection bias in continuous settings. Extensive results on synthetic and semi-synthetic datasets demonstrate that our DTRNet outperforms most state-of-the-art methods. Our code is available at \href{https://github.com/xuanxuan03021/DTRNet_final_2}{DTRNet}.

Published in Transactions on Machine Learning Research (11/2024)
DTRNet: Precisely Correcting Selection Bias in Individual- Level Continuous Treatment Effect Estimation by Reweighted Disentangled Representation
Mengxuan Hu
qly7su@virginia. edu
School of Data Science University of Virginia
Zhixuan Chu Zhefiang University
shiruanchu@cju.ede wdw..uuu..uuuu...m..moooo...
Sheng Li School of Data Science University of Virgimia
shengli@virginia.co edu
Reviewed on OpenReview: https: // opwrrevvee.e....e...mm net/ forum? id= 12TfzA9bXu
Abstract
Estimating the individual-level continuous treatment effect holds significant practical impor- tance in various decision-making domains, such as personalized healthcare and customized marketing. However, most current methods for individual treatment effect estimation limited to discrete treatments and struggle to precisely adjust for selection bias under con- timuous settings, leading to inaccurate estimation. To address these challenges, we propose a novel Disentangled Reprosentation Network (DTRNet) to estimate the individualized dose-response function (IDRF), which learns disentangled representations and precisely adjusts for selection bias. To the best of our knowledge, our work is the first attempt to precisely adjust for selection bias in continuous settings. Extensive results on synthetic and semi-synthetic datasets demonstrate that our DTRNet outperforms most state-of-the-are methods. Our code is available at [DTRNet
1 Introduction
In various fields, from medicine to marketing, estimating the causal effects of continuous treatments at individual level is not just academic exercise; it is crucial for decision making. Take precision medicine example: The central question often focuses on determining the "optimal dosage of medicine to achieve the optimal outcome for a given patient". Therefore, understanding the causal relationship between continuous treatments and outcomes can help develop customized medication regimens tailored to individual patients.
When estimating the effect of individual treatment (TTE), two predominant challenges arise: the inability to observe counterfactual outcomes and the presence of selection bias. For example, when a specific dose of treatment is assigned to patient, only the factual outcome corresponding to that dose is observed, leaving counterfactual outcomes unobserved for other doses, Furthermore, unlike randomized controlled trials where treatments are assigned at random, the dosage patient reccives in practice may depend on certain patient-specific features (e.g., older individuals may receive higher dosages more frequently). This dependency can introduce selection bias, thereby compromising the accuracy of counterfactual outcome estimation. For example, bocomes challenging to accurately estimate the treatment. effoct of higher dosages in younger populations. Beyond conventional causal inference techniques such as stratification methods (Pearl 2009) and matching methods (Abadie al. 2004), recent research has harnessed representation learning

---

Published in Transactions on Machine Learning Research (11/2024)
for counterfactual prediction and selection bias mitigation (Johansson et al. 2016 Shalit et al. 2017 Chu
et al. 2021 2022a by Schwab et al.; 2020 Curth & van der Schaar 2021 Bellot et al. 2022 Wang et al.
2022
Acharki et al. 2023 Huang et al. 2024 These approaches involve learning latent representations
from covariates, balancing these representations across treatment groups to climinate selection bias, and
subsequently estimating counterfactual outcome on these balanced representations (Shalit et al., 2017
Despite prior endeavors, several crucial challenges still remain unresolved. Most existing studies Yao
et al. 2021 Acharki et al.) 2023 Huang et al., 2024 are limited to discrete treatment settings. These methods cannot be easily extended to continuous treatment settings due to the infinite number of unobserved counterfactual outcomes and the difficulty of adjusting for selection bias in an infinite set of treatments. @ Existing methods struggle to precisely adjust for selection bias. Many previous approaches resort to a simplistic and brute-force method, which balances the entire representations, leading to inaccurate estimation. However, we argue that not all the information in the latent representations should be balanced to adjust for selection bias [Groenland] 2008 Chu et al. 2020 2023a). For example, although confounder factors in representations bring selection bias, they also contribute to outcome predictions (Hassanpour & Greiner 2019ah. Balancing instrumental factors in the representations is also theoretically implausible, as they are related to the treatment assignment and independent of all confounders (Wu ct al. 2022): hence, they should not be identical across treatments. Attempting to balance them can lead to increased bias and variance in causal effect estimation (Myers et al. 2011). Therefore, the entire representations of input covariates should not be indiscriminately balanced.
Although some follow-up methods have attempted to address the first challenge by dividing continuous treatments into bins and assigning one network head to each bin (Schwab et al. 2020 Acharki et al.) 2023), using varying coefficient structure (Nie et al., 2021), or designing specific loss for continuous treatments (Bellot et al. 2022 Zhu et al. 2024a), these methods do not guarantee any precise adjustment for selection bias (i.e., the second challenge). On the other hand, although Hassanpour & Greiner (2019b) demonstrated that disentangled representations can be leveraged to preciscly correct for selection bias, it is based on the assumption of binary treatment and cannot be readily extended to continuous settings (i.e., the first challenge). To the best of our knowledge, there no research that simultaneously solves these two problems, namely generating appropriate disentangled representations that preciscly adjust for selection bios to estimate the continuous treatment effect at the individual level, which is defined as the Individualized Function. (IDRF).
To address these challenges, we propose novel method named Disentangled Representation Network (DTRNet).. Specifically, we follow the convention assumption in causal inference WW et al.) 2020 Kuang et al,
2017 Hassanpour & Greiner 2019a) that covariates are determined by three latent factors: (1) Instrumental factors. (2) Confounder factors. (3) Adjustment factors. (as shown in Fig. [ a). DTRNet first learns
disentangled representations for each factor, providing the opportunity to precisely adjust for bias by using only the relevant representations instead of the entire representations. Then, we precisely adjust for selection bias by adopting re-weighting function and predict outcomes based on the represcntations of confounder and adjustment factors through a varying coefficient network, which enables continuous treatment effect. estimation. A rigorous theoretical proof supporting the debiasing ability of our re-weighting function is also provided. Our contributions are summarized as follows.
We propose a novel method, DTRNet,, which learns disentangled representations for unbiased cont treatment effect estimation at the individual level. These disentangled representations provide robust foundation for precisely adjusting for selection bias by using only the relevant representations rather than the entire set.
The effectiveness of DTRNet''s precise bias adjustment is supported by theoretical proofs and does not require prior knowledge of the treatment distribution, enhancing its practical applicability.
We have conducted extensive experiments to validate the effectiveness and disentangling capability of our model. The results demonstrate that our method performs well on both synthetic and semi-synthetic datasets, with each component contributing to its advanced performance. Additionally, the results indicate that the learned disentangled components accurately capture the corresponding factors.
2

---

Published in Transactions on Machine Learning Research (11/2024)
Treatment Prediction
X
lep
((x)
Laiise $
Re-weighhing & Outcome Prediction
A
Y
w(4, (())))))
(xx)
dise
T
Independient Regularization
Y
Lind
(xx)
Figure 1: Causal graph and framework of our DTRNet. Figure (a) shows the causal graph involving covariates (X); treatment (7), outcome (Y), instrumental factors (T), confounder factors (A), adjustment factors (T). The solid line represents causal relations, and dot lines denote affiliations. Figure (b) shows the framework of DTRNet. Three contracted neural networks are utilized to obtain the deep representations of the three factors. Then [(x,) and A(x;) are concatenated to predict the distribution of 4- Aax,), and M(x)) are used to predict outcomes through varying coefficient network 90(r,), while T(;; attempts to encode little information about treatment.
2 Related Work
2.1 Potential Outcomes Framework
The potential outcomes framework, as proposed by [Rubin] 1974), widely used to define the individual treatment effect (TTE). Specifically, to make the definition clear, we use binary treatments (T' and T 0) to illustrate, which can be further extended to multiple treatments by comparing their potential outcomes. For each individual 2₁, there two potential outcomes: Y!T 1) and Y,T 0) corresponding to two possible treatments (T and T = 0). As only one treatment can be taken for each individual, only one of the potential outcomes can be observed (observed outcome), and the remaining unobserved outcome is the counterfactual outcome. Therefore, one major challenge in estimating ITE is to infer the counterfactual outcome. After obtaining the counterfactual outcome, the ITE calculated as follows.
ITE, = Y,T = Y(T=
(1)
Morcover, estimating the ITE is also challenged by the presence of selection bias. Some covariates may act as common causes for both the treatment and the outcome, leading to spurious correlations between them, which affect the accuracy of the treatment estimation.
2.2 Machine Learning for ITE Estimation
As mentioned above, selection bias and counterfactual prediction are two major challenges for ITE. To address these issues, various methods have been proposed (Alaa van der Schaar
2017
Hansen 2008 Chipman
ct al. 2010 Hansen 2008 Wager Athey 2018 Chu Li 2023) Yao al. 2021
Acharki et al. 2023
Moreover, several state-of-the-art methods use deep representation learning models estimate ITE based on
treatment-invarian representations (Johansson et al. 2016 Shalit et al. 2017
Louizos et al. 201 Yoon
et al. 2018 Yao et al.
2018 Bellot et al.) 2022 Wang al. 2022 Acharki et al.
2023
Guo et al.
2023
Nagalapatti ct al. 2024
Sui ot al.] Huang et al. 2024 Zhu et al. 2024b Chu et al.
2023b
2024
Huang &
Zhang 2023 Zhu et al.
2022 Jiang el al. 2023
Gao el al. 2024 Pogodin et al.
2022
Specifically, the
discrepancy loss between the deep representations of the treatment group and control group is used to balance
the distribution of the two groups to adjust for selection bias. Subscquently, one network head for cach treatment built on the deep-balanced representations estimate ITE. However, due to the commonly used
two-head design in existing work, these models cannot be easily generalized to continuous treatment. settings.
3

---

Published in Transactions on Machine Learning Research (11/2024)
To address this limitation and estimate the average dose-response function (ADRF) under continuous
setting, several methods have been proposed (Schwab et al. 2020 Nie et. al. 2021 Shi et al. 2019 Bellot
et al. 2022 Jesson et al. 2022 Nagalapatti et al., 2024 Schwab et al. (Schwab et al. 2020 introduced a modification of TARNET (Shalit et al. 2017). called Dose Response Networks (DRNet), which divided continuous dosage inlo equally sized strata and assigned a head to each of them. To further achieve continuity of the ADRF, Nie et al. NNe et al., 2021 proposed a varying coefficient neural network (VCNet). Instead of using a multi-head model, they used varying coefficient prediction head whose weight depends on the treatment t, which improves the expressiveness of the treatment effect. In addition, generative adversarial models (Zhang et al.), contrastive learning based models (Zhu et al. 2024a) have also been proposed to estimate the ADRF. However, the methods fail to accurately adjust for selection bias, as they do not make any adjustments to eliminate the bias, or resort. to a simple and unsophisticated approach to balance the entire representations.
net (GAN)-based models (Bica et al. 2020), transformer-
Hassanpour et al. Hassanpour & Greiner (2019a) proposed that discntangled representations of covariates enhance the capture of their underlying factors, thereby improving ITE estimation performance. Instead of balancing the entire deep representations, they only applied discrepancy loss on the representations of adjustment factors and used ro-weighting methods to adjust for selection bias brought by confounders. As a result, no confounding variables are discarded. However, the method can only work with binary treatments. While subsequent research efforts (Curth & van der Schaar 2021; Chauhan et al. 2023) have introduced more refined disentanglement representations, it is worth noting that these advancements are predominantly tailored to binary settings.
Unlike existing work, our DTRNet attempts to address the two limitations mentioned above by generating appropriate disentangled representations for three underlying factors and precisely adjusting for selection bias to estimate IDRF.
3 Methodology
In this section, we begin by defining the problem and providing an overview of DTRNet. Next, we delve into the functions and explanations of each component of our DTRNet. Following this, we provide theoretical substantiation that our devised re-weighting function is able to precisely adjust for selection bias. Lastly, we discuss the advantages of our approach in comparison to other sate-of-the-ate models.
3.1 Problem Setting
St, 1 denote a dataset of size N, where 3i, yr, 4 are independent realisations of random variables X,Y with support VX T = 1]), respectively. We refer to X E RTT as covariates, which contain information about features of an unit. Y represents the outcome, and T represents the continuous treatment in the range from 0 to 1 Nie et al. 2021 Our goal is lo estimate the Individualized Dose-Response Function (IDRF) with continuous treatments by eliminating selection bias.
Let D= = (1,.
Definition 3.1. (Individualized Dose-Responee Function (IDRF)). IDRF measures the treatment effect for an individual under continnous treatments, which can be defined as:
2, p(t, x) E|Y =
(2)
We assume that covariates are generated from three underlying factors: (1) instrumental factors (T(x)) that are associated with the treatment but not with the outcome except through the treatment, (2) confounder factors (Axz)) that are associated with both the treatment and outcome, and (3) adjustment factors (T(z)) that are predictive of the outcome but not associated with the treatment. Therefore, the treatment assignment is affected by instrumental factors and confounder factors, while the outcome is affected by confounder factors and adjustment factors. This assumption is commonly used in previous research to disentangle the covariales for precise information extraction (Yoo et al. 2021 Wu et al. 2020), an illustrative example can be found in Appendix π The underlying relationship can be illustrated using a causal graph shown Fig. [a). It is important to note that IDRF identification is made under the following convention assumptions (Wu et al.
2020 Kuang et al. 2017
Hassanpour & Greiner 2019a).
4

---

Published in Transactions on Machine Learning Research (11/2024)
Assumption 3.2 (Stable Unit Treatment Value Assumption (SUTVA)). There are no interactions between units, and there is only one version of each treatment, that is, if two different units i and j have the same value for their treatment variable, then they receive the same treatment.
Assumption 3.3 (Ignorability). The potential outcome (T' t) is independent of the treatment assignment given all covariates. Formally, ((T= = 1 X
Assumption 3.4 (Positivity). Every unit should have non-zero probabilities to be assigned in each treatment group. Formally, P(T = t|X x) # 0, Vt E T, E
Assumption 3.5 (Generation of Covariates). Given a set of covariates, denoted X, we assume X follows the joint distribution of instrumental factors I, confounder factors A, and adjustment factors T. Formally, P(X) = P(T, ATT,A,T). T).
3.2 DTRNet Model
DTRNet is designed according to the causal graph in Fig. Tt first learns disentangled representations of covariates and subsequently estimates T and Y based the corresponding representations. Finally, it corrects for selection bias by utilizing relevant representations instead of the entire set. In particular, three contracted feedforward neural networks are utilized to obtain disentangled representations of three factors (T(:;), (x;), T(x:)) defined in Section 3.1 Then r(x;) and A(ra) are concatenated to predict the distribution of 4, using conditional density estimator P()[TTax;, ((x))). A(x;) and Tx:)) are used to predict the final outcome through a varying coefficient network go(t. T(x)), while Txx;) attempts to encode little information about treatment. Typically, re- re-weighting function is responsible for precisely adjusting for selection bias. The DTRNet framework is shown in Fig. [b), with the objective function defined as follows:
J(X,T,Y) E ( ((m)), + alr + LLind + >Lrre)) + BL
(3)
where [(t,, A (zx)) denotes the re-weighting function to mitigate selection bias; Ly and LT are the prediction losses for outcome Y and treatment. T, respectively. L.iise quantifies discrepancies between latent representations (T(z;), A(x;), Txx));; Lind promotes the independence between Txx;) and treatment 4;; while Lreg serves as regularization term against overfitting. a, 8, 7, and À are hyperparameters balancing the different terms in the objective function. In the following, we present the details of each term.
Factual Loss. Factual loss is used to force A(x;) and (x;)) to extract more predictive information from the covariates. As three arrows pointing to Y in Fig. [ we aim to estimate the factual outcome y; from three input, A(z;), T(x;), and assigned treatment & for unit i. To simultaneously preserve the influence of treatment and maintain the continuity of the dose-response curve, we adopt a varying coefficient neural network De(t.) to estimate outcomes (Nie et al. 2021). Factual loss computed by comparing the ground truth yr with our estimated value:
Ly= L (yr, Gou ((xs),(aa)))
where I. represents the loss function and we use the mean square error in our paper. The varying coefficient structure utilizes a function Go(t) with varying parameters ((t) instead of fixed parameters to predict outcomes. By leveraging this structure, continuous treatment effect can be estimated by incorporating continuous in outcome estimation. Especially, B-spline of degree p with q knots, resulting in k = p ++ q + 1 basis, is used to model ((t). Let B = [bi, b2,,.,,, by] E R denote the spline basis for the treatment T E 1 For single-layer leedforward network with m inputs and n outputs, the function is given by Saco = A (XW)) where We E Roxnnk is the optimizable weight.
We incorporate a treatment loss in our model to enhance the encoding of r(x;) and (xx;) with respect to f, Prior studies have usually predicted the probability of treatment using the full covariate representations (Shi et al., 2019 Nie et al., (202)). This approach can inadvertently leverage irrelevant information, such adjustment factors, from the representations. To mitigate this, we estimate
Treatment Loss,
5

---

Published in Transactions on Machine Learning Research (11/2024)
the probability of t, from the concatenation of T(24) and A(x;) using conditional density estimator T (a)) (xx)), yielding more accurate treat predictions. In this paper we adopt a naive density estimator (Nee et al. 2021), which approximates the conditional density by dividing & E [0, 1] equally into B intervals, and estimating the conditional density π (t;[( (xx)) on the B B+1 grid points using a simple neural network x N (T),),, A(x;)) = softmar W. W.cconcal([ (T,,), ((x))) E RB+1, where W, is the parameter for the network, and the densities for other values of are derived via linear interpolation. Performance is measured using negative logarithmic likelihood loss.
Lr [o][P(())( (s.)))
Discrepancy Loss. The discrepancy loss is employed to enhance the separation of three distinct repre- sentations, namely T(x), A(r), and Yxx). This ensures that each representation encodes only its specific information, thereby strengthening disentanglement in the latent space. The formula for this loss is presented as follows:
Liisc Lo(T(z), + Lo I(x))
(6)
where
1 N D LD(T()), = N Σ Σ Axx;) log
A(ra) I(x,)
N D N Σ Σ (o;)()((AA - log(T
(7)
where D is the dimension of the latent embedding, (T((x,), A(x)) denotes the average divergence loss between ((x) and (xx) across the units inspired by the KT. divergence. Tr ((x)) and A(x;) are identical for all unites, the result will be O. On the contrary, if ((;)) and (xx,) are distinctly different for all units, the value will be higher. Thercfore, through the collaboration of other modules, Laise encourages all disentanglod representations to encode only the relevant information. The definition of Ln and additional details can be found in Appendix D
Independent, loss is specifically applied to balance the T(xx) instead of the entire representations, ensuring that the learned factors Izx; do not contain any information about t, and that all information rolated to f, is encoded in ((x;) and (xx;). Previous studics (Shalit ct al. 2017 Johansson let al., 2016 Hassanpour & Greiner 2019a b. Acharki et al. 2023 Guo et al. 2023) have emphasized the need to balance adjustment representations for the treatment group and control group (binary treatment. setting), most of them done by inducing the desired independence using a discrepancy loss. However, within the framework of continuous treatments, the endeavor to achieve balance adjustment representations for every value 4 becomes infeasible. Hence, we intend to push T(x,) to embed little information about the treatment by forcing the performance of the treatment probability estimation from adjustment representation to be poor, which motivates us to minimize the following "positive" log-likelihood loss:
Independent Loss.
Lind = log log(P(4
(8)
In particular, this independent loss allows us to encode all information t, in I(x;) and A(x;) instead of T(x;), which facilitates the precisely adjustment for selection bias through the re- weighting function (discussed in next paragraph). This one of the key contributions of our paper. Furthermore, compared to previous studies that balance the entire representations of covariates, our approach does not balance confounder factors since they contain valuable information about outcome prediction (Hassanpour & Greiner 2019b). Moreover, we also exclude instrumental factors since they are related to the treatment assignment and independent of (Wu et al. 2022). Hence, they should not be balanced.
all confounders
Re-weighting Function. Rocall one of our objectives is to prociscly climinate selection bias. Inspired by [Imbens] 2000 Kloek Van Dijk 1978), we derive "propensity score"
((x;)) from I(a.) and A (Ix)) and use the inverse of to re-weight the prediction loss of outcomes as follows:
1
I(t,, = ((T,,
(9)

---

Published in Transactions on Machine Learning Research (11/2024)
Table 1: Performance comparison between DTRNet and baselines. Numbers reported are MISE/ +
standard deviation) on Synthetic, IHDP, and News with 50 runs.
Syuthetic Data
News
IHDP
Method
MISE
AMSE
MISE
AMSE
MISE
AMSE
Dragonet
0.1854 + 0.0232
0.0415 = 0.0081
1.3241 = 0.1617
0.0635++ 0.0053
4.7034
0.5860
0.9549 + 0.3005
Dragonet TR
0.1720 + 0.0219
0.0281 = 0.0095
1.3147 = 0.1594
0.0401 = 0.0062
4.2877
0.1226
0.6190 + 0.1660
DRNet
0.1849 0.0232
0.0409
1.3248
0.1616
0.0542
0.9581
DRNet TR
0.1752 + 0.0334
0,0315 0.0235
1.3148
0.1601
0.0403
0.0060
4.1313
0.6320
0.6140 0.1954
VCNet.
0.1545 + 0.0248
0.0173 0.0093
2.3372 = 0.1808
0.0384
0.0367
3.6651
0.6409
0.6735: 0.4873
VCNet_ TR
0,1118 + 0.0299
0,0112
0.0072
2.3289 = 0.2009
0.0378 = 0.0101
3.7935
1.3625
1.2302# 1.2198
Tra
0.2033 + 0.0978
0.0552
0.0884
1.2849 + 0.1587
0.0153 + 0.0066
4.1562
0.8053
1.8529 + 1.1155
CIRCE
0,6854 + 0.4050
0.5406
0.4040
1.7607 = 0.1995
0.6030 = 0.0291
10.4413 + 5.7308
8.218 : 5.4327
DTRNEs (Ours)
0.1414 + 0.0256
0.0131
0.0072
1.7846 + 0.2202
0.0104 + 0.0044
3.5376
0.5285
0.4254 + 0.3710
where (x;; (zx)) is the direct output of the conditional density estimator for treatment. Hence, it does not require additional computation prior knowledge about the treatment distribution as in (Hassanpour & Greiner 2019b). Furthermore, we can preciscly remove bias attributable to the confounder and instrumental factors instead of the unrelated part (x:))) A detailed proof follows.
3.3 Theoretical Proof Bias Elimination
In this section, we outline the theoretical derivation of our ro-weighting function, inspired by the importance sampling theory (Hassanpour Greiner 2019a Klock debiasing ability.
Van Dijk 1978). We then provide proofs for its
Definition 3.6. Let A, + R be the representation functions for confounder factors and adjustment y be the prediction function defined in the representation space R x R. The expected loss for the unit and treatment pair (2r, is:
factors, respectively. Let Ga(e) R R
I, (I, I, L(Y (0Y( (())))
(10)
Definition 3.7. The expected unbiased IDRF loss across all treatment t E T is:
= Ex / (0,
(x,t)df) = - La Ioco (z, ())((t)))ddt
(11)
Lemma 3.8. (Importance Sampling [Kloek & Van Dijk, 1978)) Let p(x) be probability density for a random variable X defined on Rdd, then for any probability density q(z) E IRd that salisfies g(x) > 0 whenever I f(x)p(x) * O, we have:
- (f p(x) g(x)
E Ex-p(>)) =
The Lemma suggests that importance sampling facilitates the computation of the expectation a target function f(x) under an unknown distribution p(z) by weighting the function with p(z) under known distribution g(x)-
((x)
Theorem 3.9. Let p(x, t'), () denote the counterfactual and factual probability density for unit x, respectively. Under the conditions of the lemma 3.8 the expected loss functions for factual and counterfactual oulcomes are:
CF = Enttsnttn (2, (x,t)).
(12)
CCF = - B2,t Us g (a,, 1)] R (x. II goce (I, (c,()P(a p(z,, 1'), p t)
(13)
where t! represents all counterfactual treatments of x, especially 1' = (1'l1' ET, #1).

---

Published in Transactions on Machine Learning Research (11/2024)
Table 2: Results of ablation study. The top row shows the average MISE and AMSE of DTRNet across 50 runs for the three datascts. Subsequent rows present performance when the respective component is disabled.
A values in pesentheses represent the percentage change in AMSE and MISE relative to the top row results.
Syuthetlc Data
News
IHDP
Method
MISE
AMSE
MISE
AMSE
MISE
AMSE
original
0.1414
0.0131
1.7846
0,0104
3.5376
0.4254
alpha
0.1413(A 0.0% t)
0.0132(A 0.8% f)
2.6400(A
47.9% 7) 0.0291 (A 179.8% f)
3.606
1.9% f]
0.5115 A 39.2% f1
beta
0,1414fA 0.0% f)
0.0137(
1.6% f)
2,6517(A
48.8% 7)
0.0319( A 206.7% ?)
3.7856(A 7.0% 1)
0.6310(4 A 48.3% f)
gamma
0.1411 (A 0.2% 4)
0.0131(A
0.0% f)
2.6871
50.5% ))
0.0297 A 185.5% +)
3.5951 (A 2.0% f)
0.4955 A 16.5% f)
re-weighting
0,1640(AA 16.0% f)
0.0212[A 61.8% 1)
2.2790(
27.7% +)
0.0136 30.7% ft
3.6804(A 4.0% 1)
0.9521
123.8% 1)
The theorem implics that the expected loss function for counterfactual outcomes can be derived by ro-weighting the factual loss. Hence, we can cooperate the counterfactual loss with the factual loss and optimise them
together through designed weight; w = -1+ ((a.t) Appendix A
P(a.')
P()T))AA))))) For further details, please refer to
Theorem 3.10. (Bias Removal with Weighted Loss) Under the theorem 3.9 and all assumptions, we have:
E Bu,tow((s II goc1) (2, gocc
This theorem states that the weighted loss is an unbiased estimation of IDRF loss, indicating our ro-weighting function is advanced in eliminating selection bias. See Appendix A for a detailed proof.
3.4 Discussions
Our proposed DTRNet is built upon existing works such as (Hassanpour & Greiner 2019b) and (Nie et al. 2021j. However, our approach extends the existing works in several significant ways. (1) The disentangled method in (Hassanpour & Greiner, 2019b) can only provide a causal effect under binary treatments, while our method facilitates a mixod type treatments, including continuous treatments. (2) We introduce a novel design for the independent loss. Instead of simply and brutally minimizing the discrepancy between the entire treatment representations and the control representations, that can be extracted from adjustment factors, which facilitates the precise adjustment for selection bias in re-weighting function. (3) We use the direct output of the conditional density estimator for re- roiicighting, which helps precisely eliminate bias and does not require additional compulation or prior knowledge about the treatment distribution. Furthermore, in the ablation study, we demonstrate the importance of this component for model performance. (4) We provide rigorous thcorctical proofs substantiating the debiasing ability of the re-weighting function. Finally, our DTRNet employs discrepancy loss between different deep representations, which ensures that each representation only encodes the relevant information. Overall, DTRNett offers extensive ability to generate disentangled representations on which selection bias can be precisely adjusted.
minimize the amount of treatment information
4 Experiments
In this sontion, we present extensive experimental results on three datasets demonstrate the effectiveness of DTRNet and address the following three research questions: Q1: How effective is DTRNet in estimating IDRF and adjusting for selection bias compared to the state-of-the-art methods? Q2: What is the individual contribution of each component in model, including the treatment loss, discrepancy loss, independent loss, and re-weighting function? Q3: Can deep disentangled representations be effective in identifying the three underlying factors?
4.1 Datasets and Baselines
As the ground truth of treatment effects is often unknown in practice, especially for continuous treatments, existing studies mainly rely synthetic datasets and semi-synthetic datasets that manually construct.
8

---

Published in Transactions on Machine Learning Research (11/2024)
3.5
3.0
3.0
25
Figure 2: MISE sensitivity analysis with different values of a, 857, À and different proportions of re-weighting value on 50 repeats of IHDP datascts. The standard deviation band is also plotted.
140
1.0
20.
375
8
7wdgs
Figure 3: AMSE sensitivity analysis with different values of a, 3,7, and different proportions of re-weighting value on 50 repeats of IHDP datascts. The standard deviation band is also plotted.
treatments and outcomes given real-world features (Curth et al. 2021 Wang al., 2022 Bellot al. 2022
Nie et al. 2021). We follow this convention build one synthetic dataset and semi-synthetie datasets: News and THDP (Hill 2011 for evaluation. To evaluate the effectiveness of selection bias adjust.ment, we intentionally design all training sets to contain selection bias, while test sets are unbiased. Hence, if model is trained on a biased training set and performs well on the corrosponding test set, it provides evidence of the model's ability to eliminate bias. We generate the three datasets following (Nie et al., 2021 Details can be found in Appendix G We evaluate the performance of our proposed D'TRNet against several state-of-the-art methods for IDRF estimation, including Dragonet, Dragonet TR (Shi et al. 2019 DRNet, DRNet TR (Schwab et al. 2020),, VCNet, VCNet_ TR (Nic al. 2021). and TransEE ZZhng where TR refers to targeted regularization. The details of baselines as follows.
et al.
Dragonet: (Shi et al. 2019 used three-headed architecture predict the propensity score and conditional outcome from covariates and treatment information. The model was later improved NyiNie et al. 2021 by using separate heads for treatment in different intervals to adjust for continuous
treatments.
DRNet: (Schwab et al. 2020 proposed to divide continuous treatments into several intervals and assign one head to cach interval to generate the dosc-respond curve. Following (Nie et al. 2021), DRNet was improved by adding a conditional density estimation head for treatment estimation.
Vcnet: Nie al. (2021) introduced a varying coefficient structure to allow the prediction head parameters be functions of continuous treatments.
TransTEE: [Zhang al.) adopted the transformer backbones to estimate treatment effect.
CIRCE: (Pogodin et al. 2022 introduces CIRCE regularizer in settings where X is used to predict Y. Specifically, the regularizer aims to learn neural features s(X) of the data X such that ((X) is independent of some metadata Z given Y.
We also include non-neural network based models: Causal Forest (Wager & Athey 2018 Bart (Chipman et al. 2010 and GPS Imbens 2000), due to the space limitation, we show the performance in the Appendix D
hhttps:/ http///rcccccccccco wwt..cd//m///daaaae///t./oo http:///aaaaiiie....cccc////////

---

Published in Transactions on Machine Learning Research (11/2024)
4.2 Implementation Details
All the neural network-based methods are trained for 800 epochs with the SGD optimizer (momentum 0.9). To mitigate the risk of overfitting or underfitting, we apply an early stop technique. For the three deep reprosentation networks in our model, we implement them as fully connected networks with two hidden layers, and each layer has 50 hidden units using ReLU activation. We also use wwo-hideen-laer settings (each with 50 hidden units) for the Y prediction network. The choice of the model structure aligns with previous approaches (Zhang al. Nie et al. 2021). We use grid search tuning to tune the following hyperparameters: a, S {0.1, 0.2, 0.4, 0.6}. The choice of this range is based the common practice of weighting loss terms between 0 and 1. Specifically, in our paper, these lyperparameters refer to the weighting of each loss term, except for the factual loss (the loss between Vi and ū.). Since the factual loss represents the prediction y, which is a crucial component of the total loss, assigning a high weight to other terms could diminish the impact of the factual loss. Therefore, we chose the range {0.1, 0.2, 0.4, 0.6), also considering computational and time costs. Typically, we choose the learning rate Ire {0.0001, 0.00005, 0.00001). For other hyperparameters, e.g., the number of knots and the degree of B-spline, we follow the setting of (Nie al. 2021) which is also tuned the same configurations of datasets. For each dataset, we generate 50 runs for training and 20 runs for tuning the aforementioned hyperparameterrs. The best hyperparameter settings for our method as well as the bascline can be found in Appendix F
4.3 Results and Analysis
To answer Q1, we follow previous works (Zhang et al. Nie et al. 2021 to report the mean integrated squared error (MISE)
Bellot et al.
2022
Zhu et al.
2024a
/11,,0000 v (1))
MISE
(14)
and the mean squared error (AMSE). More details can be found in Appendix B
AMSE / - (a())) (())))
(15)
to evaluate the performance of the models in estimating the individual level and the population level dose response curve, respectively. Since the integral of is intractable, we apply all values existing in the current datasets on each unit approximate the MISE and AMSE, that is,
MÎSE 1 11 1 Σ
Σ
(0i(t) vx(t))²
(16)
n TT| i= TET
AMSF = Σ LET
(1) - (()),
(17)
where [T] is the number of different treatment values. To ensure fair and reliable comparisons, we evaluate the performance of all models on 50 repetitions of three different datasets and report the mean and standard deviation of the MISE and AMSE. As presented in Table I DTRNet consistently outperforms the majority of baselines across all datasets, achieving satisfactory MISE and the lowest AMSE while demonstrating commendable stable performance. Specifically, the less satisfactory MISE performance on the News dataset is attributable to the distinct data generation process, where all features act as confounders. This does not align with real-world scenarios or our study assumptions, which why the performance is slightly lower compared to some of the basclines. These results demonstrate not only the effectiveness of DTRNet in IDRF estimation but also its ability to adjust for selection bias.
4.4 Ablation Study
To answer Q2, we conduct several ablation studies to evaluate the three major components of our model, including the treatment. loss Lr, discrepancy loss Tudssey independent loss Lind; and the re-weighting function
10

---

Published in Transactions on Machine Learning Research (11/2024)
Instrumental Variable
Confounder Variable
Adjustment Variable
Deep Hepresentators traindd mode 1:ME pou ecnes
Desp Nopessetated tranndd medel 15ME orgertion
Deep hepressntan ers trained metel 15ME praedtie
E4
E4
fdaa
tola
-206
-20 153 -108
153 -200
103
Figure 4: t-SNE plots of the three deep representations with respect to the instrumental variable, confounder variable, and adjustment variable. Point shapes represents the three types of deep representation. Point colors represents the value of the corresponding variable. Typically, arrows in the figures are manually added to provide clear visual indication of the directions in which the color changes in the corresponding representations.
Without Laise KL-D: 3.48
With Laise
KL-D: : 8.08
100
-100
115
12990
Figure I-SNE plots the deep representations with/without Ldiser KL-D represents KI-Divergence between three representations.
w(t,, ((t;,, ((z)). We demonstrate their roles by setting the corresponding hyperparameter to while keeping the other hyperparameters fixed at the best-tuned values. Especially, to evaluate the re-weighting function, we set the value of re-weighting to 1. As shown in Table 2 all components contribute to the model performance, as evidenced by the varying extent of performance drop in most scenarios when part is removed. Moreover, we find that the re-weighting function and the discrepancy loss are two critical components our model due to the significant increase of test error when they are disabled. In other words, adjusting for selection bias and forcing all representations to encode the corresponding information independently are predominant for the model performance. Additionally, to give a visual demonstration of the contribution of the discrepancy loss, show the t-SNE plots of the model trained with and without Laise synthetic dataset in Fig. 5 After incorporating the discrepancy loss, the distinct representations become more separate, lcading to larger KL-divergence (c.g. 8.08 vs 3.48). However, the ro-weighting function contributes less lo the News dataset than the other dataset due to the distinct data generation process of News as discussed above. In particular, all features are associated with treatment assignment and outcome generation, which mcans that all features act as confounders. Consequently, presents challenges for I(æ;) and T=x;) in accurately learning instrumental and adjustment factors. Inaccuracies in I(x;;) and M(x)) can also affect re-weighting function estimation, leading to limited contribution.
11

---

Published in Transactions on Machine Learning Research (11/2024)
Sensitivity Analysis. We have also investigated the effects of different values of a, 8, 7, X, and different proportions of the re-weighting values on the model performance. The results shown in Fig. 2 suggest that the re-weighting function and the discrepancy loss (3) have more substantial influence on the model's performance, which consists with the findings from the previous ablation study. Typically, the figure shows that a relatively large beta value results belter performance. Moreover, the results indicate that DTRNet can learn an accurate re-weighting value to improve performance, as evidenced by the fact that the current proportion (1.0) of reighhig values yields the lowest error in IDRF estimation. Similarly, the results in terms of AMSE are shown in Fig. 3 and the trend mirrors that of the MISE.
4.5
Disentanglement Performance
To answer Q3, we attempt to explore whether the representations capture the corresponding factors by utilizing t-SNE to visualize the three deep representations in a synthetic dataset with respect to different. types of variables, as shown in Fig. Since we have knowledge of the data generation process of the synthetic dataset, we know the ground truth of which features are instrumental, confounder, and adjustment variables in the dataset. Hence, for each type of variable, we choose one feature to show the relationship between it and the three deep representations. In Fig. 4. the color shade indicates the value of the corresponding variable, while the shape of the points denotes the lype of representations. For example, the first plot, we aim verify if the representations of instrumental factors embed the information of instrumental variables in the covariates. The plot shows that the color shading changes with the direction of the corresponding arrow in the instrumental representation (dot points), indicating that instrumental representations encode information about the instrumental variable. Specifically, this means that the instrumental representations learn information about the instrumental variables and can distinguish between large and small values of the corresponding variable. However, in other representations, such as A and (cross points and squared points, respectively); the color shade does not change regularly, indicating that the confounder and adjustment. representations do not encode information about instrumental variables. Therefore, our DTRNet can decently disentangle the three factors.
5 Conclusion
In this paper, we introduce Disentangled Representation Network (DTRNet), a novel model designed to estimate the individualized dose-response function (IDRF) with high precision while accounting for selection bias through disentangled representations of instrumental, confounder, and adjustment factors. Our experiments on synthetic and semi-synthetic datasets show the exceptional disentanglement capabilities of DTRNet and highlight its impressive performance on estimating IDRF, surpassing current SOTA methods.
5.1 Limitations and Future Work
Although we propose an advanced model for estimating the unbiased IDRF, there are still areas for improve- ment. First, our method currently employs a basic density estimator to estimate the treatment effect. This choice was made for proof of concept. Hence, exploring more sophisticated methods could further enhance performance. Similarly, our use of independent loss may not be optimal. Alternative approaches, such as the Hilbert-Schmidt Independence Critcrion (HSIC), could potcntially improve results. Hence, future work could focus on exploring these alternatives to improve model performance.
Moreover, our study utilizes only synthetic or semi-syuthetic datasets, which allows to generate the ground- truth IDRF using predefined equations. However, in real-world scenarios, it is challenging to obtain the true DDF and to differentiate between instrumental, confounder, and adjustment features. This complicates the evaluation of the model and its ability to correctly disentangle these features in practical settings. An expert in the relevant domain may be necessary to properly assess the model performance on real-world datasets. Therefore, key direction for future research is to investigate the deployment of our model in worldworld situations.
12

---

Published in Transactions on Machine Learning Research (11/2024)
References
Alberto Abadie, David Drukker, Jane Leber Herr, and Guido W Imbens. Implementing matching estimators for average treatment effects in stata. The stata journal, 4(3):290-311, 2004.
Naoufal Acharki, Ramiro Lugo, Antoine Bertoncello, and Josselin Garnier. Comparison of meta-learners for estimating multi-valued treatment. heterogeneous effects. In International Conference on Machine Learning, pp. 91 132. PMLR, 2023.
Ahmed M Alaa and Mihaela. van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes, arXiv preprint arXiv:1704.00880 2017.
Alexis Bellot, Anish Dhir, and Giulia Prando. Generalization bounds and algorithms for estimating conditional average treatment effect of dosage. arXiv preprint arXiv:2205. 14692, 2022.
Ioana Bica, James Jordon, and Mihaela van der Schaar. Estimating the elfects of interventions using generative adversarial networks. Advonces in Neural Information Processing Systems, 33:16434-16445, 2020.
valued
Vinod K Chauhan, Soheila Molaei, Marzia Hoque Tania, Anshul Thakur, Tingting Zhu, and David A Clifton. Adversarial de-confounding in individualised treatment effects estimation. In International Conference on Artificial Intelligence and Statistics, pp. 337-849 PMLR, 2023.
Hugh A Chipman, Edward George, and Robert E McCulloch. Bart: Bayesian additive regression trees. The Annals of Applied Statistics, 4(1):266 2000-9988222 2010.
Zhixuan Chu and Sheng Li. Causal effect estimation: Recent progress, challenges, and opportunities. Machine Learning for Causal Inference, pp. 79-100, 2023.
Zhixuan Chu, Stephen L Rathbun, and Sheng Li. Matching in selective and balanced representation space for treatment effects estimation. In Proceedings of the 29th ACM International Conference on Information 8 Knowledge Management, pp. 205-212 214, 2020.
Zhixuan Chu, Stephen L Rathbun, and Sheng Li. Graph infomax adversarial learning for treatment effect estimation with networked obsorvational data. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery & Data Mining, pp. 176-184, 2021.
Zhixuan Chu, Stephen L Rathbun, and Sheng Li. Learning infomax and domain-inndependent representations for causal effect inference with woal-world data, In Proceedings of the 2022 SIAM International Conference on Data Mining (SDM), pp. 433 441. STAM, 2022a.
Zhixuan Chu, Stephen L Rathbun, and Sheng Li. Multi-task adversarial learning for treatment effect estimation in basket trials. In Conference on health, inference, and learning, pp. 79 91. PMLR, 2022b.
Zhixuan Chu, Mechelle Claridy, Jose Cordero, Sheng Li, and Stephen L Rathbun. Estimating propensity scores with deep adaptive variable selection. In Proceedings of the 2023 SIAM International Conference on Data Mining (SDM), pp. 730-738. SIAM, 2023a.
Zhixuan Chu, Ruopeng Li, Stephen Rathbun, and Sheng Li. Continual causal inference with incremental obscrvational data. In 2023 IEEE 39th International Conference on Data Engineering (ICDE), pp. 3430 3439. IEEE, 2023b.
Zhixuan Chu, Mengxuan Hu, Qing Cui, Longfei Li, and Sheng Li. Task-driven causal feature distillation: Towards trustworthy risk prediction. In Proceedings of the AAAT Conference on Artificial Intelligence, volume 38, pp. 11642 11650, 2024.
Alicia Curth and Mihacla van der Schaar. Nonparametric estimation of heterogencous treatment effects: From theory to learning algorithms. In International Conference on Artificial Intelligence and Statistics, pp. PMLRR 2021.
13

---

Published in Transactions on Machine Learning Research (11/2024)
Alicia Curth, David Svensson, Jim Weatherall, and Mihaela van der Schaar. Really doing great at estimating cate? critical look at ml benchmarking practices in treatment effect estimation. In Thirly-fifth conference on neural information processing systems datasets and benchmarks track (round 2), 2021.
Erdun Gao, Howard Bondell, Wei Huang, and Mingming Gong. A variational framework for estimating continuous treatment effects with measurement error. In The Twelfth International Conference Learning Representations, 2024.
Sander Greenland. Invited commentary: variable selection versus shrinkage in the control of multiple confounders. American journal of epidemiology, 167(5):523 -2992995,99 2008.
Xingzhuo Guo, Yuchen Zhang, Jianmin Wang, and Mingsheng Long. Estimating heterogeneous mrealment effects: mutual information bounds and learning algorithms. In International Conference on Machine Learning, pp. 12108-12121 PMLR, 2023.
Ben B Hansen. The prognostic analogue of the propensity score. Biometrika, 95(2):481-48) 2008
Negar Hassanpour and Russell Greiner. Counterfactual regression with importance sampling weights. In IJCAI, pp. 5880-5887, 2019a.
Negar Hassanpour and Russell Greiner. Learning disentangled representations for counterfactual regression. International Conference on Learning Representations, 2019b.
Jennifer L Hill. Bayesian nonparametric modcling for causal inference. Journal of Computational and
Graphical Statistics, 20(1):217-240, 2011.
Wei Huang and Zheng Zhang, Nonparametric estimation of the continuous treatment effect with measurement error. Journal of the Royal Statistical Society Series B: Statistical Methodology, 85(2):474-495 2023.
Yiyan Huang, WANG Siyi, Cheuk Hang Leung, WU Qi, WANG Dongdong, and Zhixiang Huang. Dignet:
Learning decomposed patterns representation balancing for treatment effect estimation. Transactions on Machine Learning Research, 2024.
Guido W Imbens. The role of the propensity score in estimating doso-response functions. Biometrika, 87(3): 706-710, 2000.
Andrew Jesson, Alyson Douglas, Peter Manshausen, Maëlys Solal, Nicolai Meinshausen, Philip Stier, Yarin Gal, and Uri Shalit. Scalable sensitivity and uncertainty analyses for cansal-effect estimates of contimuons-valued interventions. Advances in Neural Information Processing Systems, 35:13892 139.138999797,,,, 2022.
Ziyang Jiang, Zhuoran Hou, Yiling Liu, Yiman Ren, Keyu Li, and David Carlson. Estimating causal effects using multi-task deep ensemble. In International Conference Mochine Learning, pp. 15023 15040. PMLR, 2023.
Fredrik Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In International conference on machine learning, pp. 3020 3029. PMLR, 2016.
Toun Klock and Herman K Van Dijk. Bayesian estimates of equation system parametors: an application of integration by monte carlo. Econometrica: Journal of the Econometric Society, pp. 1-19, 1978.
Kun Kuang, Peng Cui, Bo Li, Meng Jiang, Shiqiang Yang, and Fei Wang. Treatment effect estimation with data-driven variable decomposition. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
Christos Louizos, Uri Shalit, Joris Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. arXiv preprint ar xrXiv:1705.08822 2017.
Jessica Myers, Jeremy A Rassen, Joshua J Gagne, Krista F Huybrechts, Scbastian Schnocweiss, Kenneth J Rothman, Marshall M Joffe, and Robert Glynn. Effects of adjusting for instrumental variables on bias and precision of effect estimates. American journal of epidemsology, 174(11):1213-12)) 2011.
14

---

Published in Transactions on Machine Learning Research (11/2024)
Lokesh Nagalapatti, Akshay Iyer, Abir De, and Sunita Sarawagi. Contimuous treatment effect estimation using gradient interpolation and kernel smoothing. arXiv preprinl or xxxx:2/01.15447, 2024.
Lizhen Nie, Mao Ye, Qiang Liu, and Dan Nicolae. Venet and functional targeted regularization for learning causal effects of continuous treatments. arXiv preprint arXiv:2103.0788 2021.
Judea Pearl. Causal inference in statistics: An overview. Statistics surveys, 3:96 1466-16,600 2009.
Roman Pogodin, Namrata Deka, Yazhe Li, Danica J Sutherland, Victor Veitch, and Arthur Gretton. Elficient. conditionally invariant representation learning. arXiv preprint arXiv:2212.05645.. 2022.
Donald B Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of educational Psychology, 66(5):688, 1974.
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 AA Conference on Artificial Intelligence, volume 34, pp. 5612 5619, 2020.
Uri Shalit, Fredrik D Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In International Conference on Machine Learning, pp. 3076-3085. 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.
Yongduo Sui, Caizhi Tang, Zhixuan Chu, Junfeng Fang. Yuan Gao, Qing Cui, Longfei Li, JUN ZHOU, and Xiang Wang. Invariant graph learning for treatment effect estimation from networked observational data. In The Web Conference 2024.
Stefan Wager and Susan Athey. Estimation and inference of heterogencous treatment effects using random forests. Journal of the American Statistical Association, (13()23): 1228 1242, 2018.
Xin Wang, Shengfei Lyu, Xingyu Wu, Tianhao Wu, and Huanhuan Chen. Generalization bounds for estimating causal effects of continuous treatments. Advances in Neural Information Processing Systems, 1488005-6611 2022.
Anpeng Wu, Kun Kuang, Junkun Yuan, Bo Li, Runze Wu, Qiang Zhu, Yueting Zhuang, and Fei Wu. Learning decomposed representation for counterfactual inference. arXiv preprint arXiv:2006.00 07040, 2020.
Anpeng Wu, Kun Kuang, Bo Li, and Fei Wu. Instrumental variable regression with confounder balancing. In International Conference on Machine Learning, pp. 24056-24077 PMLR, 2022.
Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learning for treatment effect estimation from observational data. Advances in Neural Information Processing Systems, 31, 2018.
Liuyi Yao, Zhixuan Chu, Sheng Li, Yaliang Li, Jing Gao, and Aidong Zhang. A survey 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.
Lijun Zhang, Tianbao Yang, and Rong Jin. Empirical risk minimization for stochastic convex optimization: o(l/n)-and o(1/n?)-type of risk bounds. In Conference on Learning Theory, pp. 1954-1979. PMLR, 2017.
YiFan Zhang, Hanlin Zhang, Zachary Chase Lipton, Erran Li, and Fric Xing. Exploring transformer backbones for heterogeneous treatment effect estimation. Transactions on Machine Learning Research.
Minqin Zhu, Anpeng Wu, Haoxuan Li, Ruoxuan Xiong, Bo Li, Xiaoqing Yang, Xuan Qin, Peng Zhen, Jiecheng Guo, Fei Wu, et al. Contrastive balancing representation learning for heterogencous dosc-response curves estimation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 17175-17183, 2024a.
15

---

Published in Transactions on Machine Learning Research (11/2024)
Yaochen Zhu, Yinhan He, Jing Ma, Mengxuan Hu, Sheng Li, and Jundong Li. Causal inference with latent variables: Recent advances and future prospectives. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 6677 6687, 2024b.
Yuchen Zhu, Limor Gultchin, Arthur Gretton, Matt J Kusner, and Ricardo Silva. Causal inference with treatment mcasurement error: nonparametric instrumental variable approach. In Uncertainty in Artificial Intelligence, pp. 2414 2424. PMLR, 2022.
16

---

Published in Transactions on Machine Learning Research (11/2024)
A Proofs
Lemma A.1. (Importance Sampling) Let p(x) be probability density for a random variable X defined on Rd, then for any probobility density q(x)) Rd that satisfies q(x) > whenever f(z)p(z) * we have:
[/)))=)][(=)]] Ex (f p(x) g(x)
Proof.
(emp(=))[(tt)) J (x)p(x)do
S f())p(z) ((z) (())d
(18)
= =E (f) (f(x)) p(z) ((x)
Theorem A.2. Let p(x,() pla,1) 1) denole the counterfaclual and Jaclual probabilily densily for unil. x, respectively. Under the conditions of the lemma A.1, the expected loss functions for factual and counterfactual outcomes ore:
CCF -E, -Bg,t-p(a.t) l yous (x,
= -Bp, (a) (z, t) p(x,tt) t') p(x,t)
(19)
= E tmpla la aace (x, 1) p(a, t)
p(a, t)
(20)
teplla a. goce) (x, 1) 1),
where I' represents all counterfaclual trealments of x. defined as l' (l'll' TT,' The union of i and t constitutes the universal set of treatments. Therefore, given the probability density function f(x t) we have,
[,C,, Jrie Max1))d = 1 -(x,1)d
= P(x),
(21)
The theorem implies that the expected loss function for counterfactual outcomes be derived by reweighting the factual loss. Hence, cooperate the counterfactual loss with factual loss and optimise them together through designed weight (Hassanpour & Greiner 2019a
F(x,t) P(x) P(c, t) P(x,t)
(22)
P({)) P(ITT()) (()))
The second equality is by Equation 21 the third equality is by the law of conditional probability, the forth equality is by the Assumption 4 and the independency between adjustment factors and treatment.
Thcorem A.3. (Bias Removal with Weighted Loss) Under the theorem and all assumptions, we have:
E (
17

---

Published in Transactions on Machine Learning Research (11/2024)
Proof.
p(a,t')
E popooppt =E(1
) a..r(( (x,
p(a,t) t)
p(r, p(2,1))
-E pa ra (2, + l
p(r, )
-Bntepla,t) la (x, t)] - E M () (2,)
= -I laT(R, ()((r,, l)drdl
(23)
xxxT
fxx (s,r(x, t' dxxdt
IxxT
xxx (P)
This theorem states that the weighted loss is an unbiased estimation of the IDRF loss, which indicates the re-weighting function in our can precisely eliminate selection bias. Notably, each term in our loss function contributes meaningfully the theoretical proof. Specifically, some components of the loss function serve to support the assumptions underpinning functions are introduced as follows. In Assumption 4, we posit that. X should be decomposed into three different factors, where discrepency loss Lais is used to ensure this. Lind enforces us to encode all information of f; in T(z;) and (ax;) instcad of (x;), thoreby facilitating a more precise and accurate estimation of the weight w. Moreover, the weight w before factual loss the key of the proof, hence we use the LT loss to get the accurate weight.
proof, as outlined in Assumptions 1-4 in our paper. Their
B Reason to choose the MISE and AMSE metric
MISE and AMSE are widely used metrics for IDRF evaluation (Zhang al. Nie al. 2021 Bellot al., 2022 Zhu et al.) 2024ap. Therefore, we follow these previous works in adopting these metrics as evaluation metrics. Typically, MISE and AMSE assess the performance of the models by computing the difference between the estimd curve and the ground-truth curve at both the individual and population levels. Given the specific design of our dataset, if model is trained on the biased dataset and shows low MISE and AMSE on the unbiased evaluation dataset, it indicates that the model capable of adjusting for selection bias.
C Baseline Methods
We compare our model DTRNet with several state-of-thh-art methods on continuous treatment effect estimation, including Dragonet, Dragonet TR (Shi et al. 2019), DRNet, DRNet_ TR (Schwab et al. 2020). VCNet, VCNet TR (Nie et al. 2021), and TransEE (Zhang et al.). TR refers to targeted regularization, technique used to improve accuracy. The details of baselines are as follows.
Dragonet: Shi et al. 2019 used a three-headed architecture to predict the propensity score and conditional outcome from covariates and treatment information. The model was later improved by(Nie et al. 2021 by using separate heads for treatment in different intervals to adjust for continuous
treatments.
DRNett (Schwab et al. 2020 proposed to divide continuous treatments into several intervals and assign one head each interval to generate the dose-respond curve. Following (Nie al. 2021). DRNet was improved by adding conditional density estimation head for treatment estimation.
18

---

Published in Transactions on Machine Learning Research (11/2024)
Table Performance comparision between DTRNet and non-neural network baselines
Method
Synthetic Data
IHDP
News
Causal Forest
0.043 0.0021
0.97 + 0.034
0.211 0.003
BART
0.040 0.0013
0.33 = 0.005
0.066 0.003
GPS*
0.028 0.0016
0.67 = 0.025
0.022 0.001
Ours
0.013 + 0.0070
0.37 + 0.330
0.010 + 0.007
Vcnet: Nie et al. (2021 introduced a varying coefficient structure to allow the prediction head paramcters be functions of continuous trcatments.
TransTEE: Zhang et al.) adopted the transformer backbones to estimate treatment effect.
D Performance comparison between DTRNet and non-deep neural network model
We conducted our experiments under identical settings with (Nie al. 2021 including the same data generation process, dataset, and problem setting. Therefore, we compared our methods against non-deep learning approaches such causal forest, BART, and GPS (reported by (Nic et al. 2021)) these three datasets. The results from our experiments consistently indicate that our methods exhibit superior, at the very least, comparable performance in comparison to these alternative approaches.
E Details about the divergence loss
Typically, do not include the divergence loss between the instrumental represontations [(x;) and the adjustment representations T(x:). In our model, instrumental [(x;)) and confounder representations A(r;) are used to estimate the probability of treatment, while confounder A(x;) and adjustment representations T(x;) are used to predict the outcome. Hence, instrumental and confounder representations are more similar, as are confounder and adjustment representations, since they are optimized for similar goals. Therefore, we only penalize these. To avoid making our model cumbersome, we choose not to include the divergence loss between instrumental I'(x;) and adjustment representations Tx;)). In our work, we specifically use the implementation in PyTorch to compute Lp: KLDivLoss.
F Implementation Details and Hyperparameter Tuning
All the neural network-based methods are trained for 800 epochs with the SGD optimizer (momentum 0.9). To mitigate the risk of overfitting or underfitting, we apply an early stop technique. For the three deep representation networks in our model, we implement them fully connected networks with two hidden layers, and each layer has 50 hidden units using ReLU activation. We also use two-hidden- layer settings (each with 50 hidden units) for the Y prediction network. We used grid search tuning to tune the following hyperparameters: a, R,B,7 E (0.1,0.2, 0.4, 0.6) and the learning rate (Ir) {0.0001, 0.00005, 0.00001}. }. For other hyperparameters, e.g., the number knots and the degree of B-spline, we follow the setting of (Nie et al.) 2021) that is also tuned on the same configurations of datasets. For each dataset, we generate 50 runs for training and 20 runs for tuning the aforcmentioned hyperparamoters. The best hyperparameter settings are as follows:
For the baselines, used the best hyperparameters for each dataset reported by the authors Zhang ct al. (2017); Nic et al. (2021). The bascline models were also tuned on the same datasets with 20 runs using grid search. The tuning ranges for Dragonet, DRNet, and VCNet were as follows: learning rate (Ir) E {0.05, 0.005,000000. 0.001, 0.0005, 0.0001} and a E {1, 0.5). For the TR versions of these models, the learning rate for c(t) was Ir {0.001, 0.0001), and 3 E (20, 10, 5) x nå. For TransTEE, it was also tuned on these three datasets, but the authors only provided the final best hyperparameters used in their paper without specifying the range. Below are the best hyperparameters that we adopted.
19

---

Published in Transactions on Machine Learning Research (11/2024)
Table 4: The best hyperparameters in DTRNet.
Synthetic
News
IHDP
alpha
0.6
0.6
0.6
beta
0.2
0.2
0.6
gamma
0.6
0.6
0.1
Ir
0.00001
0.0001
0.00005
Table 5: The best embodding dimension, Lr. S indicates the schoduler of the learning rate (Cos is the cosine anncaling Learning
of TransTEE. Bsz indicates the batch size, # Emb indicates the
rate).
Dataset
Bsz
# Emb
# Layers
# Heads
Lr
Lr. S
Simu
500
10
1
2
0.01
Cos
THDP
128
10
1
2
0.0005
Cos
News 256
10
1
2
0.01
Cos
G
Dataset Generation
Our generation process is in line with Nie et al. 2021 The generation process of the three datasets are as follows:
Semi-synthetic Data. We construct 50 synthetic datasets for training and testing our method and baselines. Each dalaset contains 500 training samples and 200 lest samples. Furthermore, we simulate another datasets to tune the hyperparameters. In this synthetic dataset all features æ E Rº follow the i.i.d. distribution,
π 10 (21, 23)) + max (3,, F4, Ns) + sin(0.5zz) + exp(æ1 0.5xz))
1 (x1 + xz)²
+ +23 + 2sin(())) + 2rs - 6.5 + (.,05
(24)
u/ cos(22(( (t max 2x
sin(ra)) ((0,,0.25))
+ 2n
where t = (1 + exp(-t)) According to the above equation, I2, are instrumental variables, 31, R2, 24 are confounder variables, and 26 is the adjustment variable.
IHDP. Infant Health and Development Program (IHDP), an RCT dataset, originally compiled to estimate the causal effect of binary treatment (home visits of specialists) on future cognitive test scores. The original dataset contains 747 samples with 25 features, in order to estimate continuous causal effect, we generate the treatment and outcome following (Nie et al. 2021). We randomly split. the dataset into training set (67%) and testing set (33%).:
2r1
2 (
E EiESan (x: - 02)
+
2 tanh(5
4+ NN 0.25),
(1 + x2)
0.2 min(x3, W6, 26)
Sdis
(25)
sin (3Tt)
Eiesan (xt
C1)
exp(0.2((1) - 26))
y/r, (
1.2 (Lanh(5
(.0,0..)))
Sdiss
0.5 + 5 min [22, D3. 45)
where = -11+ exp(-t)) Soon = 2,,,23,5, is the index set of contimuous features, S
{4. 7,8,9,10,111 12, 13, 14, 14), Saha,2 = {16, 17, 18, 19, 20, 21, 22, 23, 24, 25} and Sais.1 U Saie.
Σ.e Σnesu
Σ.esuu 3
[25] Sco. Here c₁
C2 E
18tt,
20

---

Published in Transactions on Machine Learning Research (11/2024)
Simulation
News
IHDP
Deep resentations model prgection
upiton
300
30-
Without Beta
-008
-13
-20
-200
-250
200
-350
-36
-000
Goop Appresentaatums trained T&NE
Appresentations 1
100
With Beta
-100
-50
-20
-200
-156
-000
-132
-102
-130
150
Figure 6: t-SNE plots of the deep represetations with/without Liite across three datasets. For each dataset, we choose one repeat for visualisation due to the space limit.
The News benchmark is composed of news items from the NY Times corpus, and each item is represented by its word counts. Previous work used it to evaluate the causal effect of devices (binary ment) on readers' opinions. The original News datasel contains 3000 item samples. Since the dataset is used to estimate the effect of binary treatment, we follow (Nie et al. 2021 to generate continuous treatments as well as the corresponding outcome as follows, we first generate vi, v/2 and " D), where y vgg from (0, 1), and compute vù v///|o//22- Treatment is generate from Beta 1 is generated by:
News.
U.7 I 0.3), 017 x
y' 11, t exp(
(26)
5' |x, 1. = 2m2xxx22
+ 200Tx)) (4(t 0.5)² sin( 0.5))sin({ + NO(0,
According the generation algorithm, X and T are highly associated, meaning that the dataset contains confounding bias. We then split this biased dataset into training and testing sets. However, we do not use this biased test set directly for evaluation,
goal is to estimate the IDRF for each æ, rather than estimating y for single t corresponding to Furthermore, we follow the approach described in (Nic et al.) 2021 Zhang et al.).. Specifically, to construct the IDRF for each individual, all existing values in the test set are used to generate the IDRF using y'|x,, along with its x value. This process ensures that there no association between r and t. The average all y are used as the ground truth for evaluation. This approach ensures that there no correlation between and x, as the generation of no longer relies on f/z, thereby climinating selection bias and allowing us to test on an unbiased dataset.
measure the debiasing capability of the proposed methods,
21

---

Published in Transactions on Machine Learning Research (11/2024)
H Discrepancy loss on other dataset
In addition, to give a visual demonstration of the contribution of the discrepancy loss, we show the t-SNE plots of models trained with and without Laise on each dataset in Fig. 6 Points belonging to different. factors (i.e., {T()), A( (x:)) are further away under the model trained with the discrepancy loss, indicating that the discrepancy loss indeed helps learn separable representations. Moreover, except for the News dataset, points belonging to the same factor are closer together with Ldiec: suggesting that they extract similar innormation. However, T(2;) in the News dataset becomes more separale, and Y(x;) gathers less close compared to other datasets, indicating that it cannot well extract the corresponding information, which is consistent with the explanation of re-weighting above.
An instantiation of causal graph
In this section, provide example of the instantiation of our causal graph. Similar to WW et al. 2020 In the context of medical health record, we might collect extensive historical data from each patient, including the patients" features X (e.g,, age, gender, living environment, doctor-in-charge) treatment of patients T (taking a particular mediciine or not), and the final outcome Y (cured or not). Among these features, age and gender simultaneously affect the treatment (as physician would consider these factors when choosing a treatment) and the outcome (since they can also affect the paticnt's recovery ratc); therefore, they are confounding factors A. In contrast, the doclor-in-charge would influence only the treatment decision, without affecting the outcome, thus being instrumental factor r. The environment (e.g. hygiene situation), which only affects the outcome but not the treatment, falls into the category of adjustment factors Y.
J Sensitivity Analysis based on AMSE
In this section, present sensitivity analysis based on AMSE, shown in Figure The trend observed is similar to that of the MISE, which is presented in the main text. This analysis demonstrates that the re-weighting function and the discrepancy loss B have more substantial influence on the model's performance. These findings are consistent with those from our previous ablation study.
22


Reviews:
[{"review_id": "kGD6sRU2zi", "reviewer": "Reviewer_s1mR", "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:\nNovelty: First to tackle selection bias for continuous treatments via disentangled representations.\nComprehensive approach: Rigorous theoretical grounding with clear contributions on disentangling confounders from instrumental factors.\nSolid experiments: Extensive tests with synthetic/semi-synthetic datasets showing superiority over competing models.\n\nWeaknesses:\nReal-world applicability: Lack of testing on real-world datasets limits the model’s practical validation.\nBasic density estimator: The use of a simple density estimator may hinder performance.\nIndependent loss design: Current approach may not be optimal; more advanced methods like HSIC should be explored.", "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": "The paper proposes DTRNet, a novel model for estimating the individualized dose-response function (IDRF) under continuous treatment settings. DTRNet introduces disentangled representations of instrumental, confounder, and adjustment factors to address selection bias.\nTheoretical proofs support the model's ability to eliminate bias. Experiments on synthetic and semi-synthetic data show that DTRNet outperforms existing state-of-the-art models.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "None", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "gK0pyL6i7Y", "reviewer": "Reviewer_YgC7", "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:\n1. It is solid to include theoretical proofs for the effectiveness of Disentangled Representation Network for disentangled representations learning. The proofs are easy to follow. \n2. The experiments cover three important questions analyzing the different perspectives about the adjustment for selection bias. \n\nWeakness:\n1. The baseline methods do not contain many state-of-art casual learning methods also via neural networks. For example, like methods in this paper “Efficient Conditionally Invariant Representation Learning”. Also not sure how the biased training dataset and unbiased testing dataset is divided and whether it’s fair.", "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 way to adjust for selection bias under continuous settings. Disentangled Representation Network(DTRNet) learns disentangled representations and adjusts for selection bias, by estimating the individualized dose-response function (IDRF). The paper also includes some theory proofs for the effectiveness of DTRNet. By experiments in this paper, the authors evaluate their methods in three ways: 1. The effectiveness of DTRNet in estimating IDRF and adjusting for selection bias compared to the state-of-art methods. 2. The individual contributions of each complement in the model. 3. The effectiveness of the deep entangled representations in identifying the three underlying factors.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. I have this question and hope the authors can answer: One assumption for the theoretical proof is that there are no interactions between units and there is only one version of each treatment. But how realistic is it in the case of continuous settings, especially not by dividing the continuous treatments into bins? \n2. Like mentioned above, it would be more convincing to add comparison with other state-of-art casual learning methods also via neural networks, like the “circe” method in “Efficient Conditionally Invariant Representation Learning”.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XfuAnT18zC", "reviewer": "Reviewer_iCdT", "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\n•\tCausal effect estimation with continuous treatment is a relatively underexplored area of research and the problem being addressed by the paper is important.\n\n•\tThe background and methodology parts of the paper are well written and easy to understand in general.\n\n•\tThe experimental results demonstrate that the model effectively captures the IDRF, outperforming other state-of-the-art methods.\n\n•\tThe ablation study is valuable in highlighting the necessity of the different terms in the loss function.\n\n**Weaknesses**:\n\n•\tBased on Equation 3 and the experimental setup described in Section 4.2, there are several hyperparameters that require tuning, including the weighting of each loss term and the architecture of the neural networks used for extracting deep representations. While the authors state that they employed grid search to identify the optimal combination of these hyperparameters, it is unclear how the parameter ranges were determined. Further clarification on this point would be beneficial. Additionally, it would be useful to know whether the same fine-tuning strategy was applied to the benchmark methods when generating the results presented in Table 1.\n\n•\tIn Section 4.3, it makes sense to use the MISE and AMSE metrics for evaluating how well the model captures the IDRF. However, I don’t think these metrics directly demonstrate the model’s capability in adjusting for the selection bias. Since the experiments mainly rely on synthetic and semi-synthetic datasets, is it possible to calculate the bias, i.e., the difference between the ground truth expected IDRF and the re-weighted expected loss?\n\n•\tFigure 4 seems hard to understand. In particular, why does the fact that the shade of the color varies along a certain direction indicates disentanglement? I think a more intuitive visualization is to show the distribution of instrumental, confounder, and adjustment factors with and without the discrepancy loss $L_{ind}$.", "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": "The paper proposes Disentangled Representation Network (DTRNet) to precisely estimate the individual dose-response function (IDRF) with continuous treatments and adjust for selection bias. Specifically, the authors use 3 contracted neural networks to learn separate representations for instrumental variables, confounders, and adjustment variables. These representations are then combined in different ways to predict the treatment $T$ and the outcome $Y$, and to encourage independence. In addition, the authors provide a theoretical proof on how the re-weighted expected loss eliminates selection bias using the basic concept of importance sampling. Finally, DTRNet is tested on synthetic and semi-synthetic datasets to evaluate its accuracy in estimating IDRF and reducing selection bias. Ablation studies are also performed to demonstrate the effectiveness of each component in the loss function, especially the re-weighting and the discrepancy loss terms.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "•\tThe imbalanced regularization term $L_{imb}$ in Figure 1(b) is never used in the actual loss function in Eq. 3. Should it be $L_{ind}$ instead?\n\n•\tShould it be $Y(T = t) \\perp T | X$ instead of $Y(T = t) \\perp t | X$ in Assumption 3.3?\n\n•\tWhat is the factual loss for comparing the ground truth with the estimated targets in Eq. 4, mean-squared error?\n\n•\tWhat is the input $w$ in $\\pi^{NN}(\\Gamma(x_i), \\Delta(x_i)) = softmax(w, \\Gamma(x_i), \\Delta(x_i))$?\n\n•\tThere’s a typo in Eq. 12. It should be $\\epsilon_F = \\mathbb{E}\\_{x,t \\sim p(x,t)}$.\n\n•\tIn Section 2.2, additional citations can be included for estimating ITE with machine learning [1] and deep representation learning [2-4] methods.\n\n•\tIn Section 3.2 – Independent loss, there should be relevant citations when the authors argue that instrumental factors should not be balanced according to the causal theory.\n\nReferences:\n\n[1]. Alaa, A. M., & Van Der Schaar, M. (2017). Bayesian inference of individualized treatment effects using multi-task gaussian processes. Advances in neural information processing systems, 30.\n\n[2]. Louizos, C., Shalit, U., Mooij, J. M., Sontag, D., Zemel, R., & Welling, M. (2017). Causal effect inference with deep latent-variable models. Advances in neural information processing systems, 30.\n\n[3]. Yoon, J., Jordon, J., & Van Der Schaar, M. (2018, February). GANITE: Estimation of individualized treatment effects using generative adversarial nets. In International conference on learning representations.\n\n[4]. Jiang, Z., Hou, Z., Liu, Y., Ren, Y., Li, K., & Carlson, D. (2023, July). Estimating causal effects using a multi-task deep ensemble. In International Conference on Machine Learning (pp. 15023-15040). PMLR.", "broader_impact_concerns": "I do not foresee any ethical concerns of this work, and a Broader Impact statement section may not be necessary at this point.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper clearly meets the publication criteria of TMLR. It makes a concrete contribution for the literature and it otherwise conducted with good scientific rigour, and despite somewhat limited novelty compared to previous works it makes a clear contribution for the literature on this topic. The manuscript was clearly improved during the review process, for instance expanding the empirical evaluation and providing additional information and clarifications regarding the method. There are no open issues that would need to be resolved.
Decision: Accept (poster)