Original: Generalization Bounds for Estimating Causal Effects of Continuous Treatments
Revised: Generalization Bounds for Estimating Causal Effects of Continuous Treatments

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: We focus on estimating causal effects of continuous treatments (e.g., dosage inmedicine), also known as dose-response function.
Revised: We focus on estimating causal effects of continuous treatments (e.g., dosage inmedicine), also known as dose-response function.

Original: Existing methods in causalinference for continuous treatments using neural networks are effective and tosome extent reduce selection bias, which is introduced by non-randomized treatments among individuals and might lead to covariate imbalance and thus unreliableinference.
Revised: Existing methods in causalinference for continuous treatments using neural networks are effective and tosome extent reduce selection bias, which is introduced by non-randomized treatments among individuals and might lead to covariate imbalance and thus unreliableinference.

Original: To theoretically support the alleviation of selection bias in the settingof continuous treatments, we exploit the re-weighting schema and the Integral
Revised: To theoretically support the alleviation of selection bias in the settingof continuous treatments, we exploit the re-weighting schema and the Integral

Original: Probability Metric (IPM) distance to derive an upper bound on the counterfactualloss when estimating the average dose-response function (ADRF).
Revised: Probability Metric (IPM) distance to derive an upper bound on the counterfactualloss when estimating the average dose-response function (ADRF).

Original: The IPM distance builds a bridge between a source (factual) and an inﬁnite number of target(counterfactual) domains.
Revised: The IPM distance builds a bridge between a source (factual) and an inﬁnite number of target(counterfactual) domains.

Original: We provide a discretized approximation of the IPMdistance with a theoretical guarantee in the practical implementation.
Revised: We provide a discretized approximation of the IPMdistance with a theoretical guarantee in the practical implementation.

Original: Based on thetheoretical analysis, we propose a novel algorithm, called A verage D ose-responseesti M at I on via re-weigh T ing schema (ADMIT).
Revised: Based on thetheoretical analysis, we propose a novel algorithm, called A verage D ose-responseesti M at I on via re-weigh T ing schema (ADMIT).

Original: ADMIT simultaneously learnsa re-weighting network, which aims to alleviate the selection bias, and an inference network, which makes factual and counterfactual estimations.
Revised: ADMIT simultaneously learnsa re-weighting network, which aims to alleviate the selection bias, and an inference network, which makes factual and counterfactual estimations.

Original: In addition,the effectiveness of ADMIT is empirically demonstrated in both synthetic andsemi-synthetic experiments by outperforming the existing benchmarks.
Revised: In addition,the effectiveness of ADMIT is empirically demonstrated in both synthetic andsemi-synthetic experiments by outperforming the existing benchmarks.

Original:  Introduction
Revised:  Introduction

Original: Causal inference, which estimates causal effects from observational data, is a fundamental problem inidentifying causal reactions between actions/treatments and effects.
Revised: Causal inference, which estimates causal effects from observational data, is a fundamental problem inidentifying causal reactions between actions/treatments and effects.

Original: It facilitates the decision-makingprocess in a wide variety of domains, such as economics [1], public policy [2], medicine [3, 4]and advertising [5, 6].
Revised: It facilitates the decision-makingprocess in a wide variety of domains, such as economics [1], public policy [2], medicine [3, 4]and advertising [5, 6].

Original: Speciﬁcally, causal inference aims to estimate the causal effects of takingdifferent treatments (e g., whether to take a medication or not) in the binary setting, where neuralnetwork-based methods have made a breakthrough [7, 8, 9, 10].
Revised: Speciﬁcally, causal inference aims to estimate the causal effects of takingdifferent treatments (e g., whether to take a medication or not) in the binary setting, where neuralnetwork-based methods have made a breakthrough [7, 8, 9, 10].

Original: In real-world applications, e.g.a medical experiment, treatments cannot be simpliﬁed as binary treatments since the dosage is acritical factor for the effect of the medication.
Revised: In real-world applications, e.g.a medical experiment, treatments cannot be simpliﬁed as binary treatments since the dosage is acritical factor for the effect of the medication.

Original: Compared with causal inference for binary treatments,causal inference for continuous treatments, also known as dose-response estimation, deals with aninﬁnite number of treatments, resulting in an inﬁnite number of counterfactuals, far more than onecounterfactual of the binary setting.
Revised: Compared with causal inference for binary treatments,causal inference for continuous treatments, also known as dose-response estimation, deals with aninﬁnite number of treatments, resulting in an inﬁnite number of counterfactuals, far more than onecounterfactual of the binary setting.

Original: The success of neural networks on causal inference for binary treatments [11, 12, 13] can be attributedin part to the alleviation of selection bias, which leads to unreliable counterfactual inference since thecovariates receiving different treatments are unbalanced due to non-randomized treatments amongthem.
Revised: The success of neural networks on causal inference for binary treatments [11, 12, 13] can be attributedin part to the alleviation of selection bias, which leads to unreliable counterfactual inference since thecovariates receiving different treatments are unbalanced due to non-randomized treatments amongthem.

Original: Specially, the two distributions of the treated and control groups are balanced using Integral
Revised: Specially, the two distributions of the treated and control groups are balanced using Integral

Original: Probability Metric (IPM) [11] or re-weighting schemes [12, 13].
Revised: Probability Metric (IPM) [11] or re-weighting schemes [12, 13].

Original: In the continuous setting, sincethe control groups can be inﬁnite, these standard methods for alleviating selection bias in binary
Revised: In the continuous setting, sincethe control groups can be inﬁnite, these standard methods for alleviating selection bias in binary

Original: [Figure]
Revised: [Figure]

Original: treatments cannot be easily extended to mitigate selection bias among a theoretically inﬁnite numberof subgroups.
Revised: treatments cannot be easily extended to mitigate selection bias among a theoretically inﬁnite numberof subgroups.

Original: Currently, the Dose Response network (DRNet) [14] is proposed that extends "two heads", which aretwo outcomes in a neural network for making factual and counterfactual estimations in the binarysetting, to "multiple heads" for a continuous dosage by discretizing the dosage into multiple intervals.
Revised: Currently, the Dose Response network (DRNet) [14] is proposed that extends "two heads", which aretwo outcomes in a neural network for making factual and counterfactual estimations in the binarysetting, to "multiple heads" for a continuous dosage by discretizing the dosage into multiple intervals.

Original: Several regularization schemes in DRNet, such as distribution matching [7], propensity dropout [15]and matching on balancing scores [16, 17], are exploited to handle selection bias.
Revised: Several regularization schemes in DRNet, such as distribution matching [7], propensity dropout [15]and matching on balancing scores [16, 17], are exploited to handle selection bias.

Original: However, theseapproaches are heuristic and do not provide theoretical evidence to back up their intuition.
Revised: However, theseapproaches are heuristic and do not provide theoretical evidence to back up their intuition.

Original: To theoretically support the alleviation of selection bias, this paper derives an upper bound onthe counterfactual loss during estimation of the average dose-response function (ADRF).
Revised: To theoretically support the alleviation of selection bias, this paper derives an upper bound onthe counterfactual loss during estimation of the average dose-response function (ADRF).

Original: The keypoints of the derivation lie in the re-weighted factual loss obtained from observational data andan Integral Probability Metric (IPM) that measures the distance between the factual (re-weighted)and counterfactual distributions, as illustrated in Figure 1.
Revised: The keypoints of the derivation lie in the re-weighted factual loss obtained from observational data andan Integral Probability Metric (IPM) that measures the distance between the factual (re-weighted)and counterfactual distributions, as illustrated in Figure 1.

Original: The inﬁnite number of counterfactualdistributions is the main challenge since the distance can not be calculated with ﬁnite samples inpractice.
Revised: The inﬁnite number of counterfactualdistributions is the main challenge since the distance can not be calculated with ﬁnite samples inpractice.

Original: In the practical implementation, we provide a discretized approximation of the IPM termand prove that the difference between its approximate and true values converges to zero in probabilityunder the assumption that the covariates’ distributions of subpopulations receiving different treatmentssmoothly shift.
Revised: In the practical implementation, we provide a discretized approximation of the IPM termand prove that the difference between its approximate and true values converges to zero in probabilityunder the assumption that the covariates’ distributions of subpopulations receiving different treatmentssmoothly shift.

Original: On the basis of the derived upper bound and its implementation with a theoreticalguarantee, we propose a novel algorithm called A verage D ose-response esti M at I on via re-weigh T ingschema (ADMIT).
Revised: On the basis of the derived upper bound and its implementation with a theoreticalguarantee, we propose a novel algorithm called A verage D ose-response esti M at I on via re-weigh T ingschema (ADMIT).

Original: ADMIT jointly learns an inference network and a re-weighting neural networkto reduce selection bias while making factual and counterfactual estimations.
Revised: ADMIT jointly learns an inference network and a re-weighting neural networkto reduce selection bias while making factual and counterfactual estimations.

Original: The effectiveness of
Revised: The effectiveness of

Original: ADMIT is empirically conﬁrmed by outperforming the existing benchmarks on both synthetic andsemi-synthetic experiments in the continuous treatment setting.
Revised: ADMIT is empirically conﬁrmed by outperforming the existing benchmarks on both synthetic andsemi-synthetic experiments in the continuous treatment setting.

Original: Our contributions in this paper are 4-fold: (1) to the best of our knowledge, this is the ﬁrst studythat provides a generalization bound for estimating ADRF, alleviating selection bias; (2) we providea discretized approximation of the IPM distance with a theoretical guarantee; (3) we proposean algorithm, ADMIT, capable of reducing selection bias while making factual and counterfactualestimations; (4) we conduct both synthetic and semi-synthetic experiments in the continuous treatmentsetting, empirically demonstrating the effectiveness of ADMIT with superiority over the existingbenchmarks.
Revised: Our contributions in this paper are 4-fold: (1) to the best of our knowledge, this is the ﬁrst studythat provides a generalization bound for estimating ADRF, alleviating selection bias; (2) we providea discretized approximation of the IPM distance with a theoretical guarantee; (3) we proposean algorithm, ADMIT, capable of reducing selection bias while making factual and counterfactualestimations; (4) we conduct both synthetic and semi-synthetic experiments in the continuous treatmentsetting, empirically demonstrating the effectiveness of ADMIT with superiority over the existingbenchmarks.

Original:  Related Work
Revised:  Related Work

Original: Causal inference for continuous treatments is broadly studied through statistical methods [18, 19,20, 21].
Revised: Causal inference for continuous treatments is broadly studied through statistical methods [18, 19,20, 21].

Original: These methods are often based upon the generalized propensity score (GPS) [22], orentropy balancing based method [23] that learns the unit weights by incorporating informationabout known sample moments to achieve covariate balance, e.g., entropy balancing for continuoustreatments (EBCT)
Revised: These methods are often based upon the generalized propensity score (GPS) [22], orentropy balancing based method [23] that learns the unit weights by incorporating informationabout known sample moments to achieve covariate balance, e.g., entropy balancing for continuoustreatments (EBCT)

Original: Adaptation of neural networks for causal inference in the continuoustreatments setting has attracted the attention of research community recently [14, 25, 26].
Revised: Adaptation of neural networks for causal inference in the continuoustreatments setting has attracted the attention of research community recently [14, 25, 26].

Original: Schwabet al.
Revised: Schwabet al.

Original: [14] present the Dose Response network (DRNet) to estimate counterfactuals for continuoustreatments, an extension of the network architecture proposed in [11].
Revised: [14] present the Dose Response network (DRNet) to estimate counterfactuals for continuoustreatments, an extension of the network architecture proposed in [11].

Original: Speciﬁcally, DRNet divides acontinuous treatment t into several successive intervals, and trains one separate head for each interval.
Revised: Speciﬁcally, DRNet divides acontinuous treatment t into several successive intervals, and trains one separate head for each interval.

Original: Bica et al.
Revised: Bica et al.

Original: [26] point out that DRNet might be inﬂexible since it could not determine the dividedintervals dynamically and propose SCIGAN that generates counterfactual outcomes for continuoustreatments based on a generative adversarial network (GAN) framework.
Revised: [26] point out that DRNet might be inﬂexible since it could not determine the dividedintervals dynamically and propose SCIGAN that generates counterfactual outcomes for continuoustreatments based on a generative adversarial network (GAN) framework.

Original: Meanwhile, for the issuethat DRNet produces discontinuous ADRF due to interval discretization of continuous treatments,
Revised: Meanwhile, for the issuethat DRNet produces discontinuous ADRF due to interval discretization of continuous treatments,

Original: Varying Coefﬁcient Network (VCNet) [25] lets the weights of the prediction head be continuousfunctions of the treatment to preserve the continuity of ADRF.
Revised: Varying Coefﬁcient Network (VCNet) [25] lets the weights of the prediction head be continuousfunctions of the treatment to preserve the continuity of ADRF.

Original: Moreover, Nie et al.
Revised: Moreover, Nie et al.

Original: [25] obtain adoubly robust estimator of ADRF by generalizing the targeted regularization proposed in [10], andprovide theoretical guarantees for its asymptotic correctness.
Revised: [25] obtain adoubly robust estimator of ADRF by generalizing the targeted regularization proposed in [10], andprovide theoretical guarantees for its asymptotic correctness.

Original: Unlike VCNet, our work is derivedfrom the generalization error to reduce the selection bias when estimating ADRF.
Revised: Unlike VCNet, our work is derivedfrom the generalization error to reduce the selection bias when estimating ADRF.

Original: To balance thecovariates among inﬁnite subpopulations, we learn re-sampling weights that reduce the IPM distancebetween observed and counterfactual groups.
Revised: To balance thecovariates among inﬁnite subpopulations, we learn re-sampling weights that reduce the IPM distancebetween observed and counterfactual groups.

Original: As such, we derive an upper bound on the estimatedcounterfactual error and demonstrate experimentally the proposed algorithm ADMIT based on thederived upper bound outperforms GPS, EBCT, DRNet, SCIGAN and VCNet.
Revised: As such, we derive an upper bound on the estimatedcounterfactual error and demonstrate experimentally the proposed algorithm ADMIT based on thederived upper bound outperforms GPS, EBCT, DRNet, SCIGAN and VCNet.

Original: For a discussion on research related to causal inference for binary or continuous treatments, see
Revised: For a discussion on research related to causal inference for binary or continuous treatments, see

Original: Appendix A. Appendix A also introduces the relationships between causal inference for continuoustreatments and domain adaptation [27, 28, 29].
Revised: Appendix A. Appendix A also introduces the relationships between causal inference for continuoustreatments and domain adaptation [27, 28, 29].

Original:  Problem Formulation
Revised:  Problem Formulation

Original: Let X denote the d -dimensional space of covariates and Y the outcome space.
Revised: Let X denote the d -dimensional space of covariates and Y the outcome space.

Original: We refer to T as acontinuous treatment random variable in an interval T =
Revised: We refer to T as acontinuous treatment random variable in an interval T =

Original: [ a, b ] .
Revised: [ a, b ] .

Original: Without loss of generality, we set T =[0 , 1] in this study.
Revised: Without loss of generality, we set T =[0 , 1] in this study.

Original: Suppose we have an i.i.d sample of units, indexed by i = 1 , 2 , · · · , n .
Revised: Suppose we have an i.i.d sample of units, indexed by i = 1 , 2 , · · · , n .

Original: Following
Revised: Following

Original: Rubin’s potential outcome framework [30], we postulate the existence of potential outcomes, Y t , fort ∈ T .
Revised: Rubin’s potential outcome framework [30], we postulate the existence of potential outcomes, Y t , fort ∈ T .

Original: For each unit i we observe the treatment s i , and the covariate x i .
Revised: For each unit i we observe the treatment s i , and the covariate x i .

Original: The observed outcome(denoted by Y i ) is deﬁned to be the potential outcome (denoted by Y s i i ) corresponding to the receivedtreatment:
Revised: The observed outcome(denoted by Y i ) is deﬁned to be the potential outcome (denoted by Y s i i ) corresponding to the receivedtreatment:

Original: Y i =
Revised: Y i =

Original: Y s i
Revised: Y s i

Original: i .
Revised: i .

Original: Given an input covariate x and a treatment t , we refer to µ ( · , · ) as the individual dose-responsefunction, which is deﬁned as:
Revised: Given an input covariate x and a treatment t , we refer to µ ( · , · ) as the individual dose-responsefunction, which is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: Our goal is to derive an estimator of the average dose-response function (ADRF):
Revised: Our goal is to derive an estimator of the average dose-response function (ADRF):

Original: [Equation]
Revised: [Equation]

Original: By the law of iterated expectation, µ ( t ) =
Revised: By the law of iterated expectation, µ ( t ) =

Original: E [ µ ( t, X )] .
Revised: E [ µ ( t, X )] .

Original: In this study, the estimator of ADRF µ ( t ) isbased on estimating the individual dose-response function µ ( t, x ) .
Revised: In this study, the estimator of ADRF µ ( t ) isbased on estimating the individual dose-response function µ ( t, x ) .

Original: The following assumptions have been made to ensure that ADRF is identiﬁable from observationaldata.
Revised: The following assumptions have been made to ensure that ADRF is identiﬁable from observationaldata.

Original: Assumption 1.
Revised: Assumption 1.

Original: (Unconfoundedness) Conditional on the covariates, X , the potential outcomes, Y t ,are independent of the treatment assignment, T ,
Revised: (Unconfoundedness) Conditional on the covariates, X , the potential outcomes, Y t ,are independent of the treatment assignment, T ,

Original: [Equation]
Revised: [Equation]

Original: This assumption is also known as no unmeasured confounding and ensures that the individual doseresponse function is identiﬁable from observational data, i.e., µ ( t, x ) =
Revised: This assumption is also known as no unmeasured confounding and ensures that the individual doseresponse function is identiﬁable from observational data, i.e., µ ( t, x ) =

Original: E [ Y t | X = x
Revised: E [ Y t | X = x

Original: ] = E [ Y | X
Revised: ] = E [ Y | X

Original: =x , T = t ] .
Revised: =x , T = t ] .

Original: Assumption 2.
Revised: Assumption 2.

Original: (Overlap)
Revised: (Overlap)

Original: There exists some constant c > 0 , ∀ x ∈ X and ∀ t ∈ T ,
Revised: There exists some constant c > 0 , ∀ x ∈ X and ∀ t ∈ T ,

Original: P T | X ( t | x ) ≥ c.
Revised: P T | X ( t | x ) ≥ c.

Original: This assumption is the condition that every point in the space of covariates has a non-zero conditionalprobability density of receiving any value of treatment.
Revised: This assumption is the condition that every point in the space of covariates has a non-zero conditionalprobability density of receiving any value of treatment.

Original:  Theoretical Analyses and Proposed Algorithm
Revised:  Theoretical Analyses and Proposed Algorithm

Original: This section ﬁrstly presents a theorem that the error of estimating ADRF is bounded by the marginalloss consisting of the factual and counterfactual losses.
Revised: This section ﬁrstly presents a theorem that the error of estimating ADRF is bounded by the marginalloss consisting of the factual and counterfactual losses.

Original: The counterfactual loss cannot be calculatedfrom observational data since counterfactual outcomes as the ground truth are unobserved.
Revised: The counterfactual loss cannot be calculatedfrom observational data since counterfactual outcomes as the ground truth are unobserved.

Original: Inthe meanwhile, the counterfactual inference might be unreliable due to selection bias.
Revised: Inthe meanwhile, the counterfactual inference might be unreliable due to selection bias.

Original: We thenderive theoretically optimal weights in alleviating selection bias from the perspective of importancesampling.
Revised: We thenderive theoretically optimal weights in alleviating selection bias from the perspective of importancesampling.

Original: Although the optimal weights are theoretically guaranteed, they are impractical to learndue to their instability.
Revised: Although the optimal weights are theoretically guaranteed, they are impractical to learndue to their instability.

Original: Under a similar re-weighing schema, a theoretical guarantee is provided toguide the learning of adaptive weights that mitigate the impact of selection bias.
Revised: Under a similar re-weighing schema, a theoretical guarantee is provided toguide the learning of adaptive weights that mitigate the impact of selection bias.

Original: We refer the readerto Appendix B for detailed proofs.
Revised: We refer the readerto Appendix B for detailed proofs.

Original: 4.1 Error in estimating ADRF
Revised: 4.1 Error in estimating ADRF

Original: Let ˆ µ ( t ) be an estimator for µ ( t ) , and we are interested in estimators with a small expected meansquared error ( EMSE ) in estimating ADRF,
Revised: Let ˆ µ ( t ) be an estimator for µ ( t ) , and we are interested in estimators with a small expected meansquared error ( EMSE ) in estimating ADRF,

Original: [Equation]
Revised: [Equation]

Original: Since µ ( t ) cannot be observed from observational data, this study uses hypotheses f : T ×X → Y toestimate the individual dose-response function µ ( t, x ) ﬁrst.
Revised: Since µ ( t ) cannot be observed from observational data, this study uses hypotheses f : T ×X → Y toestimate the individual dose-response function µ ( t, x ) ﬁrst.

Original: For ease of expression, we write f ( t, x )as f t ( x ) .
Revised: For ease of expression, we write f ( t, x )as f t ( x ) .

Original: After that, we combine these hypotheses to estimate µ ( t ) ,ˆ µ
Revised: After that, we combine these hypotheses to estimate µ ( t ) ,ˆ µ

Original: [Equation]
Revised: [Equation]

Original: t
Revised: t

Original: [Equation]
Revised: [Equation]

Original: E [ f t ( X )] .
Revised: E [ f t ( X )] .

Original: We bound the error EMSE( µ, ˆ µ ) by the loss of hypotheses f t , for t ∈ T , with respect to thecorresponding potential outcomes.
Revised: We bound the error EMSE( µ, ˆ µ ) by the loss of hypotheses f t , for t ∈ T , with respect to thecorresponding potential outcomes.

Original: Let L :
Revised: Let L :

Original: Y × Y → R + be a loss function.
Revised: Y × Y → R + be a loss function.

Original: The expected loss forthe treatment and covariate pair ( t, x ) is deﬁned as:
Revised: The expected loss forthe treatment and covariate pair ( t, x ) is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: Theorem 1.
Revised: Theorem 1.

Original: Let L be the squared loss function, i.e., L ( y, y (cid:48) )
Revised: Let L be the squared loss function, i.e., L ( y, y (cid:48) )

Original: = ( y − y (cid:48) ) 2 .
Revised: = ( y − y (cid:48) ) 2 .

Original: For hypotheses f tof individual dose-response function µ ( t, · ) with marginal loss (cid:15) ( f t ) =
Revised: For hypotheses f tof individual dose-response function µ ( t, · ) with marginal loss (cid:15) ( f t ) =

Original: E [ l f t ( X )] , there exists aconstant σ min ≥ 0 , such that,
Revised: E [ l f t ( X )] , there exists aconstant σ min ≥ 0 , such that,

Original: [Equation]
Revised: [Equation]

Original: Though µ ( t ) cannot be observed from observational data, Theorem 1 implies that its expected errorcan be bounded by the marginal loss (cid:15) ( f t ) , which are related to the prediction loss concerningpotential outcomes, i.e., l f t ( x ) .
Revised: Though µ ( t ) cannot be observed from observational data, Theorem 1 implies that its expected errorcan be bounded by the marginal loss (cid:15) ( f t ) , which are related to the prediction loss concerningpotential outcomes, i.e., l f t ( x ) .

Original: However, the model cannot obtain marginal loss due to the inabilityto observe the counterfactuals.
Revised: However, the model cannot obtain marginal loss due to the inabilityto observe the counterfactuals.

Original: A re-weighting schema to address this issue is presented next.
Revised: A re-weighting schema to address this issue is presented next.

Original: 4.2 Re-weighting via importance sampling
Revised: 4.2 Re-weighting via importance sampling

Original: The fundamental problem of estimating ADRF is that it is impossible to observe all potential outcomesfor any unit in data.
Revised: The fundamental problem of estimating ADRF is that it is impossible to observe all potential outcomesfor any unit in data.

Original: The potential outcome Y t could only be observed in the subpopulation thatreceived treatment t , i.e., T = t .
Revised: The potential outcome Y t could only be observed in the subpopulation thatreceived treatment t , i.e., T = t .

Original: We refer to p X | T ( x | t ) as the distribution of T = t subpopulation,and p X ( x ) as the distribution of population.
Revised: We refer to p X | T ( x | t ) as the distribution of T = t subpopulation,and p X ( x ) as the distribution of population.

Original: The two distributions, p X | T ( x | t ) and p X ( x ) , are notequal since X and T are not independent.
Revised: The two distributions, p X | T ( x | t ) and p X ( x ) , are notequal since X and T are not independent.

Original: Their difference leads to the margin loss (cid:15) ( f t ) not beingequal to the conditional loss of hypothesis f t , (cid:15) ( f t | T = s ) , which is deﬁned as:
Revised: Their difference leads to the margin loss (cid:15) ( f t ) not beingequal to the conditional loss of hypothesis f t , (cid:15) ( f t | T = s ) , which is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: The conditional loss (cid:15) ( f t | T = s ) is the expected loss with respect to the distribution of the subpopulation that received treatment s , i.e., T = s , and the marginal loss (cid:15) ( f t ) is the expected loss withrespect to the distribution of the population.
Revised: The conditional loss (cid:15) ( f t | T = s ) is the expected loss with respect to the distribution of the subpopulation that received treatment s , i.e., T = s , and the marginal loss (cid:15) ( f t ) is the expected loss withrespect to the distribution of the population.

Original: When the treatments t and s are equal, the conditional loss (cid:15) ( f t | T = s ) is called factual loss , otherwise it is called counterfactual loss .
Revised: When the treatments t and s are equal, the conditional loss (cid:15) ( f t | T = s ) is called factual loss , otherwise it is called counterfactual loss .

Original: The counterfactualloss is infeasible since the potential outcome Y t could only be observed in the subpopulation whoreceive the treatment t .
Revised: The counterfactualloss is infeasible since the potential outcome Y t could only be observed in the subpopulation whoreceive the treatment t .

Original: A minimizer f t of (cid:15) ( f t ) is required to minimize the error of estimating
Revised: A minimizer f t of (cid:15) ( f t ) is required to minimize the error of estimating

Original: ADRF according to Theorem 1, which may be different from a minimizer of (cid:15) ( f t | T = t ) sincep X | T ( x | t ) (cid:54) =
Revised: ADRF according to Theorem 1, which may be different from a minimizer of (cid:15) ( f t | T = t ) sincep X | T ( x | t ) (cid:54) =

Original: p X ( x ) .
Revised: p X ( x ) .

Original: Importance sampling is introduced to estimate (cid:15) ( f t ) in this work, which aimsto estimate E x ∼ g
Revised: Importance sampling is introduced to estimate (cid:15) ( f t ) in this work, which aimsto estimate E x ∼ g

Original: ( x ) [ h ( x )] of a particular distribution g ( x ) while only samples generated from adifferent distribution q ( x ) are available.
Revised: ( x ) [ h ( x )] of a particular distribution g ( x ) while only samples generated from adifferent distribution q ( x ) are available.

Original: The basic idea is to rewrite the expectation as:
Revised: The basic idea is to rewrite the expectation as:

Original: [Equation]
Revised: [Equation]

Original: where g ( x ) q ( x ) is called sampling weight.
Revised: where g ( x ) q ( x ) is called sampling weight.

Original: To estimate (cid:15) ( f t ) based on the factual loss (cid:15) ( f t | T = t ) , aweight w ( x ) is utilized for each unit to deﬁne a re-weighted factual loss :
Revised: To estimate (cid:15) ( f t ) based on the factual loss (cid:15) ( f t | T = t ) , aweight w ( x ) is utilized for each unit to deﬁne a re-weighted factual loss :

Original: [Equation]
Revised: [Equation]

Original: By Equation (8), the re-weighted factual loss is equal to the marginal loss when w ( x ) =
Revised: By Equation (8), the re-weighted factual loss is equal to the marginal loss when w ( x ) =

Original: w ∗ ( x ) ,where
Revised: w ∗ ( x ) ,where

Original: [Equation]
Revised: [Equation]

Original: The term p T | X ( t | x ) is the conditional density of receiving a treatment t given the covariate x , whichis also known as the Generalized Propensity Score (GPS) [22, 31].
Revised: The term p T | X ( t | x ) is the conditional density of receiving a treatment t given the covariate x , whichis also known as the Generalized Propensity Score (GPS) [22, 31].

Original: In binary treatment cases, thepropensity score could balance the distributions of covariates between the treated control groups [32].
Revised: In binary treatment cases, thepropensity score could balance the distributions of covariates between the treated control groups [32].

Original: Hirano and Imbens [31] stress that the GPS has a similar balancing property in continuous treatmentcases.
Revised: Hirano and Imbens [31] stress that the GPS has a similar balancing property in continuous treatmentcases.

Original: Fong et al.
Revised: Fong et al.

Original: [20] proposed the covariate balancing generalized propensity score methodologyby deﬁning a stabilized weight based on the GPS.
Revised: [20] proposed the covariate balancing generalized propensity score methodologyby deﬁning a stabilized weight based on the GPS.

Original: Interestingly, the weight deﬁned in [20] is exactlythe same as w ∗ ( x ) derived from the sample weight via importance sampling.
Revised: Interestingly, the weight deﬁned in [20] is exactlythe same as w ∗ ( x ) derived from the sample weight via importance sampling.

Original: The weight w ∗ ( x ) based on the GPS is theoretically optimal in expectation [22] but may be notpractical to obtain.
Revised: The weight w ∗ ( x ) based on the GPS is theoretically optimal in expectation [22] but may be notpractical to obtain.

Original: On the one hand, the approaches based on the propensity score may be unstablein the estimations [33, 34].
Revised: On the one hand, the approaches based on the propensity score may be unstablein the estimations [33, 34].

Original: What’s worse is that, this issue may be ampliﬁed in continuous treatmentcases since the conditional density is difﬁcult to estimate accurately [35].
Revised: What’s worse is that, this issue may be ampliﬁed in continuous treatmentcases since the conditional density is difﬁcult to estimate accurately [35].

Original: On the other hand, it couldbe proved that a large sampling weight g ( x ) q ( x ) may result in a high variance in the estimations, whichcould lead to a poor ﬁnite sample performance.
Revised: On the other hand, it couldbe proved that a large sampling weight g ( x ) q ( x ) may result in a high variance in the estimations, whichcould lead to a poor ﬁnite sample performance.

Original: Therefore, the following subsection focuses onlearning adaptive weights to balance observed covariates across subpopulations that received differenttreatments and giving a generalization bound of the marginal loss .
Revised: Therefore, the following subsection focuses onlearning adaptive weights to balance observed covariates across subpopulations that received differenttreatments and giving a generalization bound of the marginal loss .

Original: 4.3 Generalization bound of the marginal loss
Revised: 4.3 Generalization bound of the marginal loss

Original: The imbalance of covariates across subpopulations that received different treatments caused byselection bias is closely connected with covariate shift in domain adaptation.
Revised: The imbalance of covariates across subpopulations that received different treatments caused byselection bias is closely connected with covariate shift in domain adaptation.

Original: This study alleviatesthis imbalance using an IPM distance, which is similar to the distribution distance metrics in domainadaptation, and gives a generalization bound of the marginal loss when using adaptive weights in
Revised: This study alleviatesthis imbalance using an IPM distance, which is similar to the distribution distance metrics in domainadaptation, and gives a generalization bound of the marginal loss when using adaptive weights in

Original: Equation (9).
Revised: Equation (9).

Original: The Integral Probability Metric ( IPM ) is a class of metrics that measure the distance between twoprobability distributions, which is deﬁned as:
Revised: The Integral Probability Metric ( IPM ) is a class of metrics that measure the distance between twoprobability distributions, which is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: 11) where G is a function family of functions
Revised: 11) where G is a function family of functions

Original: g : X → R .
Revised: g : X → R .

Original: For simplicity, we abbreviate p X | T ( x | t ) andp X | T ( x | t ) w ( x ) as p t and p wt , respectively.
Revised: For simplicity, we abbreviate p X | T ( x | t ) andp X | T ( x | t ) w ( x ) as p t and p wt , respectively.

Original: We ﬁrst state a Lemma bounding the difference betweenthe re-weighted factual loss , (cid:15) w ( f t | T = t ) , and the counterfactual loss , (cid:15) ( f t | T = s ) , where s (cid:54) = t .
Revised: We ﬁrst state a Lemma bounding the difference betweenthe re-weighted factual loss , (cid:15) w ( f t | T = t ) , and the counterfactual loss , (cid:15) ( f t | T = s ) , where s (cid:54) = t .

Original: Lemma 1.
Revised: Lemma 1.

Original: Let G be a family of functions l : X → R .
Revised: Let G be a family of functions l : X → R .

Original: Assume the per-unit expected loss functionl f t ∈ G .
Revised: Assume the per-unit expected loss functionl f t ∈ G .

Original: Then for any s ∈ [0 , 1] and s (cid:54) =
Revised: Then for any s ∈ [0 , 1] and s (cid:54) =

Original: t , we have:
Revised: t , we have:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Based on Lemma 1, we could give a generalization bound of the marginal loss through the re-weightedfactual loss as follows.
Revised: Based on Lemma 1, we could give a generalization bound of the marginal loss through the re-weightedfactual loss as follows.

Original: Theorem 2.
Revised: Theorem 2.

Original: Let IPM max = max s ∈ [0 , 1] { IPM G ( p s , p wt ) } .
Revised: Let IPM max = max s ∈ [0 , 1] { IPM G ( p s , p wt ) } .

Original: The following holds under the conditionsof Lemma 1,
Revised: The following holds under the conditionsof Lemma 1,

Original: [Equation]
Revised: [Equation]

Original: Theorem 2 shows that the marginal loss is upper bound by the average of the re-weighted factual lossand an IPM distance.
Revised: Theorem 2 shows that the marginal loss is upper bound by the average of the re-weighted factual lossand an IPM distance.

Original: The ﬁrst term could be estimated with observed samples.
Revised: The ﬁrst term could be estimated with observed samples.

Original: The second term isdeﬁned based on the distance between the two probability distribution of the subpopulation receivingtreatment t and another subpopulation.
Revised: The second term isdeﬁned based on the distance between the two probability distribution of the subpopulation receivingtreatment t and another subpopulation.

Original: Choosing a family of norm-1 reproducing kernel Hilbert space(RKHS) functions leads to IPM being the Maximum Mean Discrepancy (MMD) metric [36].
Revised: Choosing a family of norm-1 reproducing kernel Hilbert space(RKHS) functions leads to IPM being the Maximum Mean Discrepancy (MMD) metric [36].

Original: Wheninvolving discrete domains, there may have samples { x p 1 , · · · , x pn } ∼ p and { x q 1 , · · · , x qm } ∼ q , anestimator of the MMD distance between p and q could be obtained as follows:
Revised: Wheninvolving discrete domains, there may have samples { x p 1 , · · · , x pn } ∼ p and { x q 1 , · · · , x qm } ∼ q , anestimator of the MMD distance between p and q could be obtained as follows:

Original: [Equation]
Revised: [Equation]

Original: where k denotes a differentiable kernel, e.g., the Gaussian RBF kernel.
Revised: where k denotes a differentiable kernel, e.g., the Gaussian RBF kernel.

Original: A certain number of samplesare needed to estimate the MMD as shown in Equation 14.
Revised: A certain number of samplesare needed to estimate the MMD as shown in Equation 14.

Original: When the treatment is discrete and ﬁnite,all the choices of the treatment can be observed in the samples and thus one can estimate the MMDfrom these samples.

Original: However, when the treatment is continuous, the samples that received sometreatment t may be unavailable since only a ﬁnite number of samples are observed while the choiceof t is inﬁnite.
Revised: When the treatment is continuous, thesamples that received some treatment t may be unavailable since only a ﬁnite number of samplesare observed while the choice of t is inﬁnite.

Original: In other words, the IPM term in Theorem 2involves continuously shifting domains(for each s ∈ [0 , 1] ), which cannot be estimated empirically.
Revised: In other words, the IPM term in Theorem 2 cannot beestimated empirically.

Original: To overcome this challenge, we boundthe difference between the IPM max and its discretization under the assumption that the probabilitydistributions of subpopulations that received different treatments shift smoothly.
Revised: To overcome this challenge, we bound the difference between the IPM maxand its discretization under the assumption that the probability distributions of subpopulations thatreceived different treatments shift smoothly.

Original: Assumption 3.
Revised: Assumption 3.

Original: Let p s 1 and p s 2 denote the conditional probability densities of subpopulations thatreceived treatment s 1 and s 2 , respectively.
Revised: Let p s 1 and p s 2 denote the conditional probability densities of subpopulations thatreceived treatment s 1 and s 2 , respectively.

Original: We assume that there is a constant α such that thefollowing inequality holds for all s 1 , s 2 ∈ [0 , 1] :
Revised: We assume that there is a constant α such that thefollowing inequality holds for all s 1 , s 2 ∈ [0 , 1] :

Original: [Equation]
Revised: [Equation]

Original: Intuitively, α indicates the rate of the probability distribution shift of subpopulations, and a reasonablysmall α means the covariates smoothly shift when receiving different treatments.
Revised: Intuitively, α indicates the rate of the probability distribution shift of subpopulations, and a reasonablysmall α means the covariates smoothly shift when receiving different treatments.

Original: Based on Assumption 3, the following Lemma states that the IPM max corresponding to continuouslyshifting domain (for each s ∈ [0 , 1] ) could be approximated as an IPM term based on discretedomains (for each s ∈ { s 1 , · · · , s n } ).
Revised: Based on Assumption 3, the following Lemma states that the IPM max corresponding to continuouslyshifting domain (for each s ∈ [0 , 1] ) could be approximated as an IPM term based on discretedomains (for each s ∈ { s 1 , · · · , s n } ).

Original: Lemma 2.
Revised: Lemma 2.

Original: Suppose we have n i.i.d sample of units, and the i -th unit received a treatment s i ∼ p ( t ) .
Revised: Suppose we have n i.i.d sample of units, and the i -th unit received a treatment s i ∼ p ( t ) .

Original: This study assumes Assumption 3 holds for a constant α .
Revised: This study assumes Assumption 3 holds for a constant α .

Original: Then the following holds,
Revised: Then the following holds,

Original: [Equation]
Revised: [Equation]

Original: Lemma 2 bounds the IPM max by a discrete set { IPM G ( p s i , p wt ) | i ∈ { 1 , · · · , n }} .
Revised: Lemma 2 bounds the IPM max by a discrete set { IPM G ( p s i , p wt ) | i ∈ { 1 , · · · , n }} .

Original: To empirically estimate IPM G ( p s i , p wt ) , the following lemma states that a point s i could be replaced by itsneighbourhood [ s i , s i + δ ] .
Revised: To empirically estimate IPM G ( p s i , p wt ) , the following lemma states that a point s i could be replaced by itsneighbourhood [ s i , s i + δ ] .

Original: Lemma 3.
Revised: Lemma 3.

Original: Let p ∆ s = P X | T ( x | t ∈ [ s, s + δ ]) ( 0 < δ < 1 ) denotes the conditional density ofcovariates when t ∈ [ s, s + δ ] .
Revised: Let p ∆ s = P X | T ( x | t ∈ [ s, s + δ ]) ( 0 < δ < 1 ) denotes the conditional density ofcovariates when t ∈ [ s, s + δ ] .

Original: Then the following holds under Assumption 3,
Revised: Then the following holds under Assumption 3,

Original: [Equation]
Revised: [Equation]

Original: Following Lemma 2 and Lemma 3, we bound the marginal loss as follows.
Revised: Following Lemma 2 and Lemma 3, we bound the marginal loss as follows.

Original: Theorem 3.
Revised: Theorem 3.

Original: Suppose we have n i.i.d sample of units, and the i -th unit received a treatment s i .
Revised: Suppose we have n i.i.d sample of units, and the i -th unit received a treatment s i .

Original: Let IPM ∆ max = max i ∈{ 1 , ··· ,n } { IPM G ( p ∆ s i , p t ) } .
Revised: Let IPM ∆ max = max i ∈{ 1 , ··· ,n } { IPM G ( p ∆ s i , p t ) } .

Original: This study assumes Assumption 3 holds for a 234 constant α .
Revised: This study assumes Assumption 3 holds for a 232 constant α .

Original: Then, for a neighbourhood size 0 < δ <
Revised: Then, for a neighbourhood size 0 < δ <

Original: we have,
Revised: we have,

Original: [Equation]
Revised: [Equation]

Original: Contrary to the second term in Theorem 2, IPM ∆ max could be estimated empirically, since we mayget samples that received treatment s ∈ [ s i , s i + δ ] by choosing an appropriate neighbourhood size δ .
Revised: Contrary to the second term in Theorem 2, IPM ∆ max could be estimated empirically, since we mayget samples that received treatment s ∈ [ s i , s i + δ ] by choosing an appropriate neighbourhood size δ .

Revised: Algorithm

Original: Input: Observed samples D = { ( x i , s i , Y i ) |
Revised: Input: Observed samples D = { ( x i , s i , Y i ) |

Original: i = 1 , · · · , n } , loss function L ( ·  · ) ,function family G for IPM , batch size n b , neighbourhood size δ .
Revised: i = 1 , · · · , n } , loss function L ( · , · ) ,function family G for IPM , batch size n b , neighbourhood size δ .

Original: Initialized parameters: θ =
Revised: Initialized parameters: θ =

Original: [ θ φ  θ h , θ ω ] 1: while θ not converged do 2:
Revised: [ θ φ , θ h , θ ω ] 1: while θ not converged do 2:

Original: Sample mini-batch { ( x i 1 , s i 1 , y i 1 ) , · · · , ( x i nb , s i nb , y
Revised: Sample mini-batch { ( x i 1 , s i 1 , y i 1 ) , · · · , ( x i nb , s i nb , y

Original: i nb ) } from D 3:
Revised: i nb ) } from D 3:

Original: Calculate the weight of each unit: w i j = ω ( s i j , φ ( x i j )) 4: for l = 1 to L = (cid:100) 1 /δ (cid:101) do 5: for k = 1 to L and k (cid:54)
Revised: Calculate the weight of each unit: w i j = ω ( s i j , φ ( x i j )) 4: for l = 1 to L = (cid:100) 1 /δ (cid:101) do 5: for k = 1 to L and k (cid:54)

Original: =
Revised: =

Original: l do 6:
Revised: l do 6:

Original: [Equation]
Revised: [Equation]

Original:  7: end for 8:
Revised: ) 7: end for 8:

Original: [Equation]
Revised: [Equation]

Original: 9: end for 10:
Revised: 9: end for 10:

Original: Calculate the gradients of the IPM term:
Revised: Calculate the gradients of the IPM term:

Original: [Equation]
Revised: [Equation]

Original: Calculate the gradients of the re-weighted factual loss :
Revised: Calculate the gradients of the re-weighted factual loss :

Original: [Equation]
Revised: [Equation]

Original: Update parameters:

Original: [Equation]
Revised: [Equation]

Original:  13: end while
Revised: b (cid:80) j w i j · L ( h ( s i j , φ ( x i j )) , y i j ) 12: Update parameters: [ θ φ , θ h , θ ω ] ← [ θ φ − η ( g 2 + g 4) , θ h − ηg 3 , θ ω − ηg 1 ] 13: end while

Original: Output: Learned representation function φ ( · ) , outcome function h ( · , · ) , re-weighting function ω ( · , · )
Revised: Output: Learned representation function φ ( · ) , outcome function h ( · , · ) , re-weighting function ω ( · , · )

Original: To minimize the error (cid:15) u in estimating the average dose-response function, Theorem 1 and Theorem3 suggest that we may minimize the upper bound in Equation (18) with respect to the re-weightedconditional loss and the IPM term, which leads us to Algorithm 1 detailed in the following subsection.
Revised: To minimize the error (cid:15) u in estimating the average dose-response function, Theorem 1 and Theorem3 suggest that we may minimize the upper bound in Equation (18) with respect to the re-weightedconditional loss and the IPM term, which leads us to Algorithm 1 detailed in the following subsection.

Original: 4.4 Proposed algorithm
Revised: 4.4 Proposed algorithm

Original: An algorithm, called average dose-response estimation via re-weighting schema (ADMIT) is proposedbased on the theoretical results derived in the previous subsections.
Revised: An algorithm, called average dose-response estimation via re-weighting schema (ADMIT) is proposedbased on the theoretical results derived in the previous subsections.

Original: As illustrated in Figure 2, we useneural networks to build the model, which is mainly composed of three components: a representationnetwork φ ( x ) , a re-weighting network ω ( t, φ ( x )) , and an inference network h ( t, φ ( x )) .
Revised: As illustrated in Figure 2, we useneural networks to build the model, which is mainly composed of three components: a representationnetwork φ ( x ) , a re-weighting network ω ( t, φ ( x )) , and an inference network h ( t, φ ( x )) .

Original: This work ﬁrstly learns a treatment-agnostic representation z
Revised: This work ﬁrstly learns a treatment-agnostic representation z

Original: = φ ( x ) of x using all data.
Revised: = φ ( x ) of x using all data.

Original: Theinﬂuence of treatment t is different from that of x [11], which might be lost during training if φ ( x ) ishigh-dimensional.
Revised: Theinﬂuence of treatment t is different from that of x [11], which might be lost during training if φ ( x ) ishigh-dimensional.

Original: To address this issue, the varying coefﬁcient neural network following [25] is usedto build h ( t, φ ( x )) and
Revised: To address this issue, the varying coefﬁcient neural network following [25] is usedto build h ( t, φ ( x )) and

Original: ω ( t, φ ( x )) .
Revised: ω ( t, φ ( x )) .

Original: Take the former as an example, h ( t, φ ( x )) is deﬁned as:
Revised: Take the former as an example, h ( t, φ ( x )) is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: ) where NN θ h ( t ) is a neural network with parameter θ h ( t ) .
Revised: (19) where NN θ h ( t ) is a neural network with parameter θ h ( t ) .

Original: Thus, the treatment t determines thenonlinear function in NN θ h ( t ) .
Revised: Thus, the treatment t determines thenonlinear function in NN θ h ( t ) .

Original: Based on the theoretical analysis aforementioned, the algorithm aims to seek a re-weighting functionω and outcome function h to minimize the error EMSE( µ, ˆ µ ) of estimating the average dose-responsefunction, using the following objective:
Revised: Based on the theoretical analysis aforementioned, the algorithm aims to seek a re-weighting functionω and outcome function h to minimize the error EMSE( µ, ˆ µ ) of estimating the average dose-responsefunction, using the following objective:

Original: [Equation]
Revised: [Equation]

Revised: where θ =

Revised: [ θ φ , θ ω , θ h ] .

Revised: To effectively calculate the IPM , the interval [0 , 1] is divided into (cid:100) 1 /δ (cid:101)intervals to conﬁrm the neighbourhood of the i -th unit, and details of the calculation are presented inlines 4 to 9 of the Algorithm 1.

Revised: The model is trained using stochastic gradient descent to minimizethe objective.

Revised: Finally, an estimator of ADRF is built as:

Revised: [Equation]

Original: X t
Revised: X t

Original: Z
Revised: Z

Original: Representation network ⋯ 0 % (*) 0 & (*) ⋯ ⋯
Revised: Representation network ⋯ 0 % (*) 0 & (*) ⋯ ⋯

Original: Inference network
Revised: Inference network

Original: Re-weighting network !" !
Revised: Re-weighting network !" !

Original: $ !"
Revised: $ !"

Original: ($% ! ,% ! )
Revised: ($% ! ,% ! )

Original: IPM( [Equation] ,+ !
Revised: IPM(+ " ,+ !# )

Original:  # ) where θ = [ θ φ , θ ω , θ h ] . To effectively calculate the IPM , the interval [0 , 1] is divided into (cid:100) 1 /δ (cid:101)intervals to conﬁrm the neighbourhood of the i -th unit, and details of the calculation are presented inlines 4 to 9 of the Algorithm 1. The model is trained using stochastic gradient descent to minimizethe objective.
Revised:  Experiments

Original: Finally, an estimator of ADRF is built as:

Original: [Equation]
Revised: [Figure]

Original: Experiments

Original: It is difﬁcult to get the ground truth treatment effect on real-world datasets due to the inability toobserve the counterfactuals.
Revised: It is difﬁcult to get the ground truth treatment effect on real-world datasets due to the inability toobserve the counterfactuals.

Original: To deal with this, existing literature [14, 25, 26] often use synthetic/semisynthetic data for meaningful evaluation.
Revised: To deal with this, existing literature [14, 25, 26] often use synthetic/semisynthetic data for meaningful evaluation.

Original: We use one synthetic dataset and two semi-syntheticdatasets, News [14, 7] and TCGA [37], to demonstrate the effectiveness of ADMIT.
Revised: We use one synthetic dataset and two semi-syntheticdatasets, News [14, 7] and TCGA [37], to demonstrate the effectiveness of ADMIT.

Original: Details ondatasets and implementation can be found in Appendix C.
Revised: Details ondatasets and implementation can be found in Appendix C.

Original: 5.1 Experimental setup
Revised: 5.1 Experimental setup

Original: Synthetic data generation.
Revised: Synthetic data generation.

Original: We generate covariates x ∼ Unif ( 0 , 1 ) ∈ R 6 , and the assigned treatments and their corresponding outcomes are generated as follows:
Revised: We generate covariates x ∼ Unif ( 0 , 1 ) ∈ R 6 , and the assigned treatments and their corresponding outcomes are generated as follows:

Original: • t =
Revised: • t =

Original: (1 + exp( − t (cid:48) )) − 1 , t (cid:48) ∼ N ( µ 1 , 0 . 5) where
Revised: (1 + exp( − t (cid:48) )) − 1 , t (cid:48) ∼ N ( µ 1 , 0 . 5) whereµ 1 = 10sin(max( x 1 , [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: ( x 4 ) + 2 x 5 − 6 . 
Revised: + x 23 + 2sin( x 4 ) + 2 x 5 − 6 . 5 .

Original: • y ∼ N ( µ 2 , 0 . 5)
Revised: • y ∼ N ( µ 2 , 0 . 5)

Original: where µ 2 = cos(2 π ( t − 0 . 5))
Revised: where µ 2 = cos(2 π ( t − 0 . 5))

Original: ( t 2 + 4 max( x 1 ,x 6 ) 3 1+2 xsin( x 4 )) .
Revised: ( t 2 + 4 max( x 1 ,x 6 ) 3 1+2 xsin( x 4 )) .

Original: Semi-synthetic data generation.
Revised: Semi-synthetic data generation.

Original: We obtain covariates from two real-world datasets, News [14, 7]and TCGA [37].
Revised: We obtain covariates from two real-world datasets, News [14, 7]and TCGA [37].

Original: To generate the assigned treatments and their corresponding outcomes, we generatea set of parameters, v i
Revised: To generate the assigned treatments and their corresponding outcomes, we generatea set of parameters, v i

Original: = u
Revised: = u

Original: i / || u i ||
Revised: i / || u i ||

Original: and i = 1 , 2 , 3 , where u i is sampled from a uniform distribution
Revised: and i = 1 , 2 , 3 , where u i is sampled from a uniform distribution

Original: Unif ( 0 , 1 ) , then
Revised: Unif ( 0 , 1 ) , then

Original: • t ∼ Beta ( α, β ) where β =
Revised: • t ∼ Beta ( α, β ) where β =

Original: α − 1 d ∗ + 2 − α and d ∗ = v T 3 x 2 v T 2 x , and• y ∼ N ( µ, 0 . 5) where µ =2(max( − 2 , exp(( v T 2 x v T 3 x − 0 .
Revised: α − 1 d ∗ + 2 − α and d ∗ = v T 3 x 2 v T 2 x , and • y ∼ N ( µ, 0 . 5) where µ =2(max( − 2 , exp(( v T 2 x v T 3 x − 0 .

Original: 3)))+10 v T 1 x ) × 4( t − 0 . 5) 2 × sin( π 2 t ) .
Revised: 3)))+10 v T 1 x ) × 4( t − 0 . 5) 2 × sin( π 2 t ) .

Original: As discussed in [26], a larger α will result in a higher selection bias.
Revised: As discussed in [26], a larger α will result in a higher selection bias.

Original: Without speciﬁc instruction, weset α = 2 in our experiments.
Revised: Without speciﬁc instruction, weset α = 2 in our experiments.

Original: Baselines and metrics.
Revised: Baselines and metrics.

Original: For neural network baselines, we compare ADMIT with DRNet [14], VCNet[25] and SCIGAN [26].
Revised: For neural network baselines, we compare ADMIT with DRNet [14], VCNet[25] and SCIGAN [26].

Original: With the same settings as in [25], we apply targeted regularization on both
Revised: With the same settings as in [25], we apply targeted regularization on both

Original: DRNet and VCNet, and add a conditional density estimation head for DRNet.
Revised: DRNet and VCNet, and add a conditional density estimation head for DRNet.

Original: For statistical baselines, [Table] we compare ADMIT with GPS [31] and entropy balancing for continuous treatments (EBCT)
Revised: For statistical baselines,we compare ADMIT with GPS [31] and entropy balancing for continuous treatments (EBCT)

Original: Itis worth mentioning that EBCT learns sample weights from data separately, rather than simultaneouslylearning sample weights and the inference model as in this paper.
Revised: Itis worth mentioning that EBCT learns sample weights from data separately, rather than simultaneouslylearning sample weights and the inference model as in this paper.

Original: When evaluating EBCT, we take
Revised: When evaluating EBCT, we take

Original: VCNet as its inference network to estimate the ADRF.
Revised: VCNet as its inference network to estimate the ADRF.

Original: For metrics, we use Expected Mean Square
Revised: For metrics, we use Expected Mean Square

Original: Error (EMSE) on the test dataset, where EMSE = 1 N (cid:80) [-Ni =
-] (ˆ µ ( t i ) − µ ( t i )) 2 .
Revised: Error (EMSE) on the test dataset, where EMSE = 1 N (cid:80) [-Ni =
-]N i =1 (ˆ µ ( t i ) − µ ( t i )) 2 .

Original: 5.2 Benchmarks comparison
Revised: 5.2 Benchmarks comparison

Original: As shown in Table 1, the proposed ADMIT is signiﬁcantly more accurate in estimating the averagedose-response in both synthetic and semi-synthetic datasets.
Revised: As shown in Table 1, the proposed ADMIT is signiﬁcantly more accurate in estimating the averagedose-response in both synthetic and semi-synthetic datasets.

Original: Considering the dimension of covariatesin TCGA and News is 4000 and around 3000 respectively.
Revised: Considering the dimension of covariates [Table] in TCGA and News is 4000 and around 3000 respectively.

Original: This result may indicate that ADMITis well adapted to the high-dimensional data.
Revised: This result may indicate that ADMITis well adapted to the high-dimensional data.

Original: Moreover, using sample weights obtained by preprocessing the data with EBCT also enhances the performance on the two semi-synthetic datasetscompared with VCNet.
Revised: Moreover, using sample weights obtained by preprocessing the data with EBCT also enhances the performance on the two semi-synthetic datasetscompared with VCNet.

Original: This result demonstrates the effectiveness of the re-weighting schema inenhancing the performance of estimating causal effects.
Revised: This result demonstrates the effectiveness of the re-weighting schema inenhancing the performance of estimating causal effects.

Original: Furthermore, the performance boost issmaller compared with ADMIT, which indicates the superiority of the proposed re-weighting strategyin this study.
Revised: Furthermore, the performance boost issmaller compared with ADMIT, which indicates the superiority of the proposed re-weighting strategyin this study.

Original: 5.3 Selection bias
Revised: 5.3 Selection bias

Original: [Figure]
Revised: [Figure]

Original: 8 selection bias 0. E M S E SCIGANDRNetVCNetADMIT As indicated in [26], the variance of the Betadistribution decreases with the increase ofα in the semi-synthetic data generation, resulting in higher selection bias.
Revised: As indicated in [26], the variance of the Betadistribution decreases with the increase ofα in the semi-synthetic data generation, resulting in higher selection bias.

Original: We comparethe performance of ADMIT with VCNet,
Revised: We comparethe performance of ADMIT with VCNet,

Original: SCIGAN and DRNet on TCGA when varying choices of α from 1 (no bias) to 8 (highbias) to assess the robustness of these methods in different levels of selection bias.
Revised: SCIGAN and DRNet on TCGA when varying choices of α from 1 (no bias) to 8 (highbias) to assess the robustness of these methods in different levels of selection bias.

Original: Onone hand, Figure 3 illustrates that the performance of the four models tends to decrease,i.e. increase √ EMSE , as selection bias increases, which is consistent with the theoretical analysis.
Revised: Onone hand, Figure 3 illustrates that the performance of the four models tends to decrease,i.e. increase √ EMSE , as selection bias increases, which is consistent with the theoretical analysis.

Original: On the other hand, we foundthat ADMIT shows consistent performanceand outperforms the three baselines acrossthe entire range of evaluated selection bias.
Revised: On the other hand, we foundthat ADMIT shows consistent performanceand outperforms the three baselines acrossthe entire range of evaluated selection bias.

Original: This result demonstrates the effectiveness of
Revised: This result demonstrates the effectiveness of

Original: ADMIT in empirically mitigating the impactof selection bias.
Revised: ADMIT in empirically mitigating the impactof selection bias.

Original:  Conclusion
Revised:  Conclusion

Original: In this paper we propose a novel re-weighting schema to mitigate the impact of selection bias incausal inference with continuous treatments.
Revised: In this paper, we propose a novel re-weighting schema to mitigate the impact of selection bias incausal inference with continuous treatments.

Original: Under this schema, this paper provides and proves ageneralization error bound on the estimation of ADRF based on an IPM distance between the observed(re-weighted) and counterfactual distributions, with theoretical evidence on learning sample weightsthat alleviate the unreliable counterfactual inference problem caused by selection bias.
Revised: Under this schema, this paper provides and proves ageneralization error bound on the estimation of ADRF based on an IPM distance between the observed(re-weighted) and counterfactual distributions, with theoretical evidence on learning sample weightsthat alleviate the unreliable counterfactual inference problem caused by selection bias.

Original: In the practicalimplementation, we provide a discretized approximation of the IPM distance with a theoreticalguarantee, leading to an algorithm ADMIT, that jointly learns re-weighting and inference neuralnetworks.
Revised: In the practicalimplementation, we provide a discretized approximation of the IPM distance with a theoreticalguarantee, leading to an algorithm ADMIT, that jointly learns re-weighting and inference neuralnetworks.

Original: The superiority of ADMIT is validated in both synthetic and semi-synthetic experiments.
Revised: The superiority of ADMIT is validated in both synthetic and semi-synthetic experiments.

Original: Experiments show that re-weighting schemas facilitate causal inference
Revised: Experiments show that re-weighting schemas facilitate causal inference.

Original: Checklist
Revised: Checklist

Original: For all authors...(
Revised: For all authors...(

Original: a
Revised: a)

Original: ) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope
Revised: Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?

Original: Yes
Revised: [Yes]

Original: (b ) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: No
Revised: [No]

Original: (c) Did you discuss any potential negative societal impacts of your work
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: No](
Revised: [No]

Original: d ) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] See AppendixB.
Revised: [Yes] See Appendix B.

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] See Appendix B.3.
Revised: [Yes] See Appendix B.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?

Original: [Yes] See AppendixC.
Revised: [Yes] See Appendix C.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes] See Appendix C.
Revised: [Yes] See Appendix C.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] See Section
Revised: [Yes] See Section

Original: (d) Did you include the total amount of compute and the type of resources used (e
Revised: (d) Did you include the total amount of compute and the type of resources used (e

Original: g., typeof GPUs, internal cluster, or cloud provider)?
Revised: g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Appendix C.
Revised: [Yes] See Appendix C.

Original: If you are using existing assets (e
Revised: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets...
Revised: g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [Yes] See Appendix C.
Revised: [Yes] See Appendix C.

Original: (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes]See Appendix C.
Revised: [Yes] See Appendix C.

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [N/A]
Revised: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable
Revised: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: [Table]

Original: information or offensive content?
Revised: information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if
Revised: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review
Revised: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?
Revised: Board (IRB) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount
Revised: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: [Table]

Original: spent on participant compensation?
Revised: spent on participant compensation?

Original: [N/A]
Revised: [N/A]