Using representation balancing to learn conditional-average dose responses from clustered data
Authors: Christopher Bockel-Rickermann
Toon Vanderschueren
Jeroen Berrevoets
Tim Verdonck
Wouter Verbeke
Venue: Accepted by TMLR ()

Abstract: Estimating the response to an intervention with an associated dose conditional on a unit's covariates, the "conditional-average dose response" (CADR), is a relevant task in a variety of domains, from healthcare to business, economics, and beyond. Estimating such a response is challenging for several reasons: Firstly, it typically needs to be estimated from observational data, which can be confounded and negatively affect the performance of intervention response estimators used for counterfactual inference. Secondly, the continuity of the dose prevents the adoption of approaches used to estimate responses to binary-valued interventions. That is why the machine learning (ML) community has proposed several tailored CADR estimators. Yet, the proposal of most of these methods requires strong assumptions on the distribution of data and the assignment of interventions, which go beyond the standard assumptions in causal inference.  Whereas previous works have so far focused on smooth shifts in covariate distributions across doses, in this work, we will study estimating CADR from clustered data and where different doses are assigned to different segments of a population. On a novel benchmarking dataset, we show the impacts of clustered data on model performance. Additionally, we propose an estimator, CBRNet, that enables the application of representation balancing for CADR estimation through clustering the covariate space and a novel loss function. CBRNet learns cluster-agnostic and hence dose-agnostic covariate representations through representation balancing for unbiased CADR inference. We run extensive experiments to illustrate the workings of our method and compare it with the state of the art in ML for CADR estimation.

Published in Transactions on Machine Learning Research (02/2025)
Using representation balancing to learn conditional-average dose responses from clustered data
Christopher Bockel-Rickermann
Brrisoopher-iceemmmmm@@hheeeee
KU Leuven
Toon Vanderschueren
ton..uundrsscherreekkuuuuaarr
KU Lewven
University of Antwerp
Jeroen Berrevoets
University of Cambridge
Tim Verdonck
tim, Sim..serroonckkkuuueueeeeer.
University of Antwerp imec
KU Leuven
Wouter Verbeke
wwwter..erbeeke/Bauueeee be
KU Leuven
Reviewed on OpenReview: https: openreview. net/ forum? id= USEMindyg4
Abstract
Estimnttig the response lo an intervention with an associated dose conditional on a unit's covariates, the "conditional-average dose response" (CADR), is relevant task in variety of domains, from healthcare to business, economics, and beyond. Estimating such response is challenging for several reasons: Firstly, it typically needs to be estimated from observa- tional data, which can be confounded and negatively affect the performance of intervention response estimators used for counterfactual inference. Secondly, the continuity of the dose prevents the adoption of approaches used to estimate responses to binary-valued interven- tions. That is why the machine learning (ML) community has proposed several tailored CADR estimators. Yet, the proposal of most of these methods requires strong assump- tions on the distribution of data and the assignment of interventions, which go beyond the standard assumptions in causal inference. Whereas previous works have far focused on smooth shifts in covariate distributions across doses, in this work, we will study estimating CADR from clustered data and where different doses are assigned to different segments of a population. On a novel benchmarking dataset, we show the impacts of clustered data on model performance. Additionally, we propose an estimator, CBRNet, that enables the ap- plication of representation balancing for CADR estimation through clustering the covariate space and novel loss function. CBRNet learns cluster-agnostic and hence dose-agnostic covariate representations for unbiased CADR inference. We extensive experiments to illustrate the workings of our method and compare it with the state of the art in ML for CADR estimation.
Introduction
Predicting responses to interventions conditional on unit's covariates, so-called "conditional-averag" inter- vention responses" has become popular field of machine learning (MT.) research (Shalit et. al., 2017; Shi

---

Published in Transactions on Machine Learning Research (02/2025)
P(s)
y(s)
X1
(a) Input space
(b) Dose distributions
(c) Dose responses
Figure 1: Illustration of confounding by cluster (CBC). (a) We assume that there are clusters of similar units in the data, visualized by the green and blue point cloud in exemplary two-dimensional covariate space. (b) The dose assignment mechanism is function of the cluster so similar units are assigned similar doses. This results in different distributions of dose assignment per cluster. (c) As dose responses (v(ss) are heterogeneous and depend on a unit's covariates (x1, X2), the dose response curves of units in different clusters (green and blue curves) are different, resulting in confounded data.
et al., 2019).". While a large share of studies focuses on estimating responses to categorical and especially binary-valued interventions (Shalit et al., 2017; Yoou et al., 2018; Shi et al., 2019; Johansson et al., 2022), less focus has been put on outimuous-valuee interventions, so interventions with an associated intensity or dose (Schwab et al., 2019). Yet, estimating the conditional-averago "dose response" (CADR) is of great interest in many domains of applications, such as marketing (Farrelly et al., 2005), medicine (Calabrese, 2016), or education (Turk, 2019), especially when too much or too little intervention can have negative effects (Frei & Canellos, 1980), or when decksion-makers must optimize resource usage (Vanderschueren et al., 2023).
Such responses must often be estimated from observational data. Hence, next to high-dimensional covariates and non-linear dose responses, an important challenge is for estimator to adjust for confounding (Bica et al., 2020), so the presence of a covariate vector that influences the dose assignment, as well as the response (Hancuse, 2016). Not accounting for confounding might hinder methods from learning unbiasod estimates of CADR by overfitting the dose assignment mechanism (Shalit et al., 2017). Many approaches to estimating responses to binary-valued interventions under observed confounding. however, cannot casily be transferred to the dose response setting (cf. Section 2). This is why the ML community has proposed several tailored methods to learn CADR from observational data (Schwab et al., 2019; Bica et al., 2020; Nie et al., 2021; Wang et al., 2022). Next to the standard assumptions necessary for causal inference (Stone, 1993), most of these estimators rely on further assumptions on the underlying data, such as on the distribution of covariates and the assignment. of interventions and doses (Nie et al., 2021; Wang et al., 2022). These assumptions drive algorithm development, yet are often not motivated by investigations of real-life scenarios and their accompanying data-generating processes (DGPs). This potentially explains why ML estimators for intervention response estimation have seen little adoption in practice (Curth et al., 2021). One such assumption is smooth shift in covariate distributions across different doses, required by Wang et al. (2022) and satisfied in the synthetic benchmarking datasets by Bica et al. (2020) and Nie et al. (2021) which are frequently used in the literature.
In this paper, we investigate CADR estimation from data in which such assumptions are violated. Especially, we will investigate scenarios in which units are clustered and where different doses are assigned to different segments of a population, leading to "confounding by cluster" (CBC; cf. Figure 1). Such data prevalent in many domains, such as loan and insurance pricing, where groups of similar customers receive similar prices (Phillips, 2013), epidemiology, where the impacts of clustered data modeling have been studied widely (Berlin et al., 1999; Seaman et al., 2014). In ML research, however, we find that little interest has been paid to understanding the impact of clustered data on intervention and dose response estimates. In this paper, we take a step towards closing this gap in the literature. By analyzing the real-world applications of CADR estimation in healthcare, public policy, and business, we abstract DGPs for clustered data and create
Some literature also refers the conditional-vverae response as the "individual" response, For a discussion on why the "conditional--e avverage is preferred over the term "Individual", we refer the reader Vegetabile (2021).
2

---

Published in Transactions on Machine Learning Research (02/2025)
novel semi-syuthetic benchmarking dataset for dose response estimation from clustered data with varying levels of confounding. We evaluate several MT. methods, including traditional supervised learning methods and tailored ML methods for CADR estimation. In addition, we propose a new method for estimating CADR from clustered observational data with cluster-based dose assignment. Our method is called CBRNet (Cluster-robust dose response estimation through Balanced Representation learning with neural Networks, cf. Figure 3) and is motivated by precoding studies on estimating intervention responses under confounding through the balancing of covariate distributions (Shalit et al., 2017; Schwab et al., 2018).
Contributions. Our paper makes three important contributions to the ML. literature on dose response estimation: (1) To the best of our knowledge, we are the first to investigate the impacts learning from clustered data on the performance of ML estimators for CADR; (2) we propose a novel semi-synthetic benchmark that simulates CBC in observational data and that allows for testing its impacts on estimators; (3) we introduce CBRNet, method for CADR cstimation from clustered data. By introducing novel loss function and by clustering the covariale space, we present an intuilive and effective approach to use representation balancing for CADR estimation. We perform extensive experiments to test the performance of CBRNet, both on our newly presented benchmark and on previously established datasets.
Outline. Section discusses related work ML methods for CADR estimation and preceding bench- marking practices. Our problem formulation follows in Section 3. We present our method, CBRNet, in Section 4, discussing motivation and architecture. Our experimental setup is presented in Section 5, with results following in Section 6. We conclude with Section 7, discussing the impact of our work and future research directions.
2 Context
Background. The majority of ML literature estimating intervention responses has been concerned with "treatment effects", so the differences in outcome between applying a specific intervention, the "treatment", or not (Johansson et al., 2016; Shalit et al., 2017; Louizos et al., 2017; Nie & Wager, 2021; Wager & Athey, 2018; Shi et al., 2019). Less aention has been paid to settings with interventions that are continuous- valued, so that have an associated intensity or dose. Nevertheless, estimating contimuuus valuee intervention responscs, or simply "dose responses", is relevant in a plcthora of domains of application (Farrolly et al., 2005; Calabrese, 2016; Turk, 2019; Vanderschueren al., 2023). Compared to estimating treatment effects, dose response estimation comes with unique challenges (Schwab et al., 2019), which make this interesting and relevant field of research.
For estimating treatment elfects, controlled experiments, like randomized controlled (rials (RCTs; Deaton & Cartwright, 2018), are often considered the gold standard. Yet, due to the practically infinite amount of possible doses, their application is signficantly complicated for dose response estimation (Holland-Lete & Kopp-Schneider, 2014). Even when suitable number of units is available, such experiments might be prohibitively costly.
Alternatively, dose responses can be estimated using observational data. Yet, with treatment effect estima- tion, this comes with several challenges, most importantly: a) The impossibility of observing counterfactual outcomes (Holland, 1986) and b) the presence of a dose assignment mechanism, which may have led to con- founding in the data (Varadhan & Sceger, 2013). These challenges may prevent the adoption of traditional supervised learning methods (Schwab el al., 2019).
ML dose response estimators. Phenomena, like observed confounding in data, are well described and widely researched for treatment effect estimation (Curth et al., 2021). However, many methods presented for estimating treatment effects do not translate to the dose response setting. Methods such as matching (Ho et al., 2007), the propensity score (Roscnbaum & Rubin, 1983), representation balancing (Johansson et al., 2016), representation learning (Louizos et al., 2017) and tailored tree-based estimators (Hill, 2011; Wager & Athey, 2018) rely on the presence of two distinct groups, such treated units and control units.
3

---

Published in Transactions on Machine Learning Research (02/2025)
Hence, several tailored ML estimators for dose responses were proposed. Hirano & Imbens (2004) propose the Hirano-Tmbent estimator (HTE); built on the idea of the generalized propensity score (GPS), extending the propensity score to contimuous-valued interventions. Schwab et al. (2019) propose DRNet, a neural ar- chitocture to learn conditional-average responses to multiple intervention options with an associated dose. Their method can be combined with several balancing techniques to tackle the confounding of intervention assignment, yet not the confounding of assigned doses. Alternatively, Bica et al. (2020) propose SCIGAN, which leverages generative adversarial networks (GANs; Goodfellow al., 2020), a method that learns how to generate counterfactual responses. Generating counterfactual responses data, enabling learning the dose-response via supervised learning methods. Nie et al. (2021) propose VCNet, a generalization of DRNet that accounts for the continuity of dose responses, by training a varying coeffi- cient network (Hastie & Tibshirani, 1993). Nie et al. (2021) sludy the estimation of average dose responses, by combining estimates of the conditional-average response through functional targeted regularization, a method to build doubly-robust average dose response estimators inspired by van der Laan & Rubin (2006). Wang et al. (2022) also lackle estimating the average dose response. Their estimator, ADMIT, minimizes the maximal distributional difference between units of discrete dose intervals, assuming smooth shift in the probability distribution of covariates for units of different doses. Zhang et al. (2023) propose TransTEE, method using the transformer architecture (Vaswani et al., 2017) to build intervention effect estimators. They claim that the attention mechanisms in these are superior in modeling observed confounders. Bellot et al. (2022) propose an architecture similar to DRNet, which uses multiple prediction heads for differ- ent interventions with a dose, and use representation balancing to overcome intervention confounding for conditional-average response estimation. Dose confounding is not addressed in their work.
remove confounding in the
Note that, even though methods such as VCNet and ADMIT were proposed for average dosc response estimation, their methodologies can be used to estimate conditionnllaveragee responses. Specifically, any preceding method for average response estimation first trains a conditional-averaee response estimator. A simple estimator of the average response can be obtained by the law of total expectation and averaging the conditionnalavvragee estimates over a population of interest (Abrevaya et al., 2015). For a detailed description of our problem setup, see Section 3.
Benchmarking practices in ML for dose response estimation. As counterfactual outcomes cannot be observed, ML estimators for intervention responses cannot easily be evaluated using observational data (Schwab et al., 2019; Bica et al., 2020). This is different from, for example, supervised learning (Hastie & Tibshirani, 1993). Alternatively, many researchers derive theoretical guarantees for their estimators' perfor- mance based on assumptions or use semi- or fully synthetic datascts for an empirical evaluation. Neither of these two approaches ensures good real-life performance. Datasets are usually constructed according to a synthetic data-generating process (DGP), which defines causal relationships between covariates, interven- tions, and outcomes. Prominent examples of such datasets for dose response estimation include the ones presented by Bica et al. (2020) and Nie el al. (2021).
Curth et al. (2021) show that most research lacks analysis of the aspects of a synthetic dataset, and that challenges in a dataset are often incompletely understood. In contrast to the preceding works, we argue that the creation of dose response estimators should be use-case driven and that we need a better understanding of the effects of different DGPs on model performance. Hence, we base the methodology of this work on abstractions of real-life use cases of CADR estimation, motivating the creation of our new method.
3 Problem formulation
We leverage the Neyman-Rubin potential outcomes framework (Splawa-Neyman et al., 1990; Rubin, 1974) and assume that we observe units defined by a pre-intervention covariate vector X E X C Rm in some m- m-dimensional feature space. Every unit is assigned a positive dose S e S C Rt. For the remainder of our paper and without loss of generality, we set = [o, 1] to be the unit interval. Every unit has potential responses, or "outcomes" Yss) E y C R, which are the responses had received the dose s. Following Schwab et al. (2019), we call (s) the "dose response". The expected response for a dose given a unit's
4

---

Published in Transactions on Machine Learning Research (02/2025)
X
Y(8)
S
Figure 2: Single world intervention graph (SWIG) illustrating causal dependencies between variables in the training data. Our goal is to estimate the response of unit dose s.
covariates, the "conditional-average dose response" (CADR), is subsequently given by
M|s x) ([
(1)
Our goal is to estimate the CADR from observational data. We expect have access to data in the form of Dn ((x,, 8t, ==)] where is the total number of observed units. X;; 81, and yi are the covariates, the assigned dose, and the observed response of the i-th unit respectively. We also refer to Y1 the "factual response". We operate under the slandard fundamental problem of causal inference (Holland, 1986), responses of unit i to doses different from S;. so-called "counterfactual responses", are never observed.
One challenge in CADR estimation from obsorvational data lies in the relationships between variables X, S, and Y. In observational dala, unlike in controlled environments, doses were likely assigned according lo some (partially) unknown assignment mechanism based on units' covariates. In the case of heterogenuous dose responses, such mechanism might introduce confounding, in which the covariate vector influences both dose assignment and the respective dose response. We illustrate the relationships between the different variables in the data in single-world intervention graph (SWTG; Richardson & Robins, 2013) in Figure 2. To learn an unbiased estimate of the CADR, a method must hence adjust for any potential confounding (Shalit et al., 2017).
Finding unbiased estimates of intervention responses relies on a set of untestable assumptions (Stone, 1993). In line with preceding works, we require the following standard assumptions hold:
Assumption 1. (Consistency) The observed outcome Y; for a unit that was assigned dose 8; is the potential outcome (s:)..
Assumption 2. (No hidden confounders) The assigned dose S is conditionally independent of the potential outcome (s) given the covariates X, so (Y(s) E S} 1 S|X
Assumption 3. (Overlap) Every unit has a greater-than-zero probability of receiving any dose, so 4s E Vx E X with F(x) > 0: < (Ps))) <1
Note: As indicated in Section 2, estimating the CADR is also relevant for estimating average dose responses (ADR, ((s))- Methods for ADR estimation are typically first training CADR estimator fudd, x) and are subsequently deriving an estimate of the ADR through the law of total expectation by calculating ((s) = 1 Em x;).
4 CBRNet: A method to tackle confounding by cluster in CADR estimation
Motivation. It is widely understood that use cases and applications should drive method selection (Luo, 2016), and that algorithm architecture plays major role in estimator performance (Alaa & van der Schaar, 2018). This is especially true for intervention response estimators, whose performance cannot be studied on real data. We hence argue that algorithm development should be aligned more closely with the study of real- world applications. By now, the MI community has relied mostly on small set of benchmarking datascts (Curth et al., 2021), such the IHDP dataset presented by Hill (2011) for treatment effect estimation, or
5

---

Published in Transactions on Machine Learning Research (02/2025)
the datasets by Bica et al. (2020) and Nie et al. (2021) for dose response estimation. While pos- ing distinct challenges to estimators and relying on covariates acquired from real-world studies, these datasets are created using synthetic DGPs specily- ing intervention assignment and potential outcomes. These mechanisms are often created without further motivation, raising questions about their realistic- ness.
MSE
L
IPM
A
For this manuscript, we take an alternative ap- proach and will focus on CADR estimation from clustered data with cluster-based dose assignment. Whereas in precoding works intervention and dose assignment are lypically specific to unit's covari- ates, we assume that similar units form clusters and that dosc assignment is conditional on the clus- ter (cf. Figure 1). So, every unit is assigned a cluster & according to some clustering function g(x): : X ++C. C. that assumption holds, we can re- formulate "No hidden confounders" (cf. Section 3) (Y(s)|s S) 1. S/C. While not previously stud- ied in the ML litcrature dose response cstimation, such mechanisms occur in several domains of application. Assuming that dose responses are heterogeneous, such custer-based assignment can introduce confounding in the data, often referred to as "confounding by cluster" (Localio et al., 2002; Scaman et al., 2014; Zetterqvist et al., 2016; Berlin et al., 1999). We identified several real-life scenarios in which such mechanisms apply:
Figure 3: Architecture of CBRNet. The network
consists of three parts. representation learner +; an inference network 1, and a clustering function A. To overcome confounding by cluster, & is trained to learn dose-agnostic representation by minimizing a tailored integral probability metric (IPM) over the re- sponse space given the clusters identified by A. Subse- quently, the inference network I is trained to learn the response of a unit to a dose by minimizing a standard mean squared error loss (MSE). For a full description of our method see Section 4.
In the pricing of loan and insurance products, companies might want to identify the effect of price on a customer's willingness to buy or renew a product, yet in accordance with established business processes similar customers were assigned to different price tiers (Yeo et al., 2002; Phillips, 2013);
in marketing, a company might want to predict the effect of exposure to marketing campaign, yet customers were grouped by loyalty programs and have hence received distinct amounts of discounts (Jonker al., 2004; Ho et al., 2012);
in medicine, practitioners might want to understand the effect of varying amounts of drug health outcomes, yet patients were grouped by biomarkers to receive certain doses (Schacht et al., 2014; Verhaart et al., 2014);
in education, one might want to predict the elfectiveness of certain learning programs students' educational success, yet students were divided based on ability, leaving them exposed to similar interventions (Betts & Shkolnik, 2000):
in public policy, policymakers might want to estimate the effects of subsidies on research organi- zations, yet organizations often form clusters with similar characteristics, which apply for grants jointly (Broekel et al., 2015).
In all of the above scenarios, the data will likely be confounded. To learn an unbiasod estimate of the CADR, we must find an estimator robust to such observed confounding. As doses are driven by the clus- ter, supervised estimators that do not address confounding or treatment assignment bias might overfit the assignment mechanism and learn biased estimates of dose responses. Feed-forward neural networks, for ex- ample, tend to memorize the relation between the dosc and the outcome as observed in the available data (Schwab et al., 2019). However, due to the non-random treatment assignment mechanism, some clusters will be over- or underrepresented at a certain dose, compared to their proportion in the overall population. These "traditional" supervised methods do not account for confounding, resulting from the dose assignment mechanism (cf. below). Therefore, the resulting models may not generalize well to units of different clusters, with different. characteristics and different. responses to doses. As illustrated in Figure 1e, the effect. of
6

---

Published in Transactions on Machine Learning Research (02/2025)
particular dose on the outcome is different for observations from different cluster. We propose to solve this challenge by building model that learns cluster-agnostic, and hence dose-agnostic representations of a unit's covariates, for unbiased CADR estimation.
Architecture. The architecture CBRNet is shown in Figure CBRNet consists of three parts Φ, A, and 1, based on our motivation in the paragraph above. 4 : X + R is a representation learning function mapping the covariates into the representation space R C R" learn a cluster-agnostic representation. We use a standard feed-forward neural nerk for Φ where both the number of layers and the number of hidden nodes per layer are hyperparameters. A A S + {1...., k} is clustering function mapping unit to one of & clusters by taking as input the covariates and doses of unit. We assume that dose assignment is conditional on the cluster, hence adding the dose as input to the clustering function is expected to improve performance in separating clusters with heterogeneous doses. For an overview of possible clustering algorithms, see Madhulatha (2012). We propose to use k-means (Lloyd, 1982) as a clustering function, minimizing Euclidian distances between units of a certain cluster. It is widely adopted in business and beyond (Wu, 2012) and has previously been used in treatment effect estimation (Berrevoets et al., 2020), We train A on all available training data and do not alter it during the training of the remaining network components (see paragraph below). I: R x S + I is an inference function laking as input the covariates in representation space R and a dose, to learn the CADR. As for 4, we propose I to be a feed-forward neural network with flexible hyperparametors, as similarly adopted by Shalit et al. (2017), Schwab et al. (2019), and Bica et al. (2020).
Model sclection and training. CBRNet is trained using gradient descent over the training datasct D, {(xi, 86, y minimizing loss L defined as
((x,8,9) (r( + X Lo L0,
(2)
L is standard mean squared error (MSF) loss
L MSE(y.))
1 N
Σ (yt 01)
(3)
where yi is the true outcome of unit i and 9 is the estimated outcome as calculated by the inference function T. Lpp is a loss that ensures that & learns a cluster-agnostic representation. To achieve this, we use integral probability metrics (IPMs), which have been previously used to improve ML models for treatment effect estimation by minimizing distances between treated and untreated observations in the representation space (Shalit et al., 2017). Metrics, such as the maximum mean discrepancy (MMD; Gretton et al., 2006) or the Wasserstein distance (Kantorovich, 1960; Vaserstein, 1969), are defined for strictly two distributions. To enable calculating distances between the distributions of different clusters in the representation space R, we calculate Lg as
( k- Σ IPM
(4)
Intuitively, we choose one base cluster & and take the average over the pair-wise distances between the remaining clusters and the base cluster as regularization loss. This approach is inspired by the generalization of the PEHE metric (Precision in Estimating Hetcrogencous Effects, Hill, 2011) multiple treatments as proposed by Schwab et al. (2018). is hyperparameter to balance the two losses, L, and Le. Our proposition of the regularization loss allows for flexible choice of IPM and cluster number k, such as linear or kernel MMDs or the Wasserstoin distance. We will test the impact of different choices in our empirical evaluation in the following sections. In practice, we choose as base k the cluster with the most observed units, in order to ensure accurate empirical approximation of the IPMs.
As discussod in previous works on ML for CADR estimators, model selection and hyperparameter tuning are inherently difficult due the unavailability of counterfactual outcomes (Schwab et al., 2019; Bica et al., 2020). For overview, Curth & Van Der Schaar (2023) discuss several approaches to model selection.
7

---

Published in Transactions on Machine Learning Research (02/2025)
Dose distribution
Dose responяea
Dose distribution
Dose reeponses
Deasity
Octoome
Density
Onlocae
11
12
Dose
Doee
Dose
Dose
(a) No confounding
(b) High level of confounding (a 3, )))
(a 3, 0)
Figure Visualization of level on confounding dry bean dataset. Per subligure, the left plot visualizes the dose distribution across the population and provides color legend. The right plot visualizes the dose response space. Per dose response, the line color corresponds to the assigned dose, which is marked with dot (-). In the unconfounded case (a), doses are homogeneously distributed across units. In the confounded case (b), doses are assigned conditional on clusters, as evident by the color separation of different dose responses.
Choosing procedure over another can significantly impact response estimation and model bias. For our empirical evaluation and to not confuse the performance of CADR estimators with the performance of model selection procedures, we propose to use a standard mean squared error over observed units in a validation set to optimize hyperparameters of CBRNet. For further detail, we refer to Section and our technical implementation Appendix E. We propose to use a validation set to optimize hyperparameters and to choose the best model in terms of the validation set MSE.
5 Experimental Evaluation
We evaluate CBRNet empirically by comparing it several benchmarking methods on novel semi-synthetic dataset, the "Dry bean-DR data". The following paragraphs will discuss the creation of this data, the benchmarking methods, and the metrics used for evaluation. The dataset is available publically and has been provided with the code for this manuscript. We want to mention that we do not claim the Dry bean- DR dataset to be superior dataset for evaluating CADR estimators. Instead, our proposal of the dataset is meant to complement the set of available benchmarking datascts. The key novelty compared to other established datasets is the cluster-based dose assignment.
Data generation. To simulate clustered data and cluster-based intervention assignments, synthetic setup. While previous works primarily motivate the realisticness of their data by the use of covariates sampled from real-life experiments, such the TCGA data (Cancer Genome Atlas Research Network et al., 2013), a dataset on gene sequences of cancer cells, we put strong emphasis on the intervention assignment mechanisms as a prime factor driving confounding. Our data generation starts by taking the covariates of the dry bean dataset. (Koklu & Ozkan, 2020). The dataset contains 13,611 samples of seven different types of beans, listing 16 unique fcatures per sample derived from computer vision analysis. We relate this data to intervention response estimation by considering situation in which we want to understand the effect of different irrigation levels crop yields (Goldstein et al., 2017; Dehghanisanij et al., 2022). Understanding the conditionallaverage responsos to irrigation could lead to improved yields, yet we have to estimate these responses from observational data, stemming from established farming practices. Hence, the data likely is confounded. The data generation now proceeds in three steps:
semi-
Step (Clustering): We take as clusters the types of beans the original data and randomly aggregate them form three clusters. For unit i, we denote the corresponding cluster as e E (1,2, 3). During model training and application, expect the cluster affiliation to be unknown.
8

---

Published in Transactions on Machine Learning Research (02/2025)
Step 2 (Dose assignment): Different from preceding datasets (Schwab et al., 2019; Bica et al., 2020; Nie et al., 2021), we assign doses based on clusters which will lead to a more rapid shift in covariate distributions across different dose levels in the data. We refer back to Section 4 for real-world examples of such mechanisms. Every cluster j is assigned a different "modal dosc" my, which is randomly drawn from { [_ 1+8 without replacement. The parameter B E [0, 1] determines the variability of doses across clusters. For 3 = the dose assignment homogenous across clusters. For B = 1, dose assignment is maximally heterogencous. We set B 2 for our main experiments, and show the impact of different levels of 3 in Appendix C. As the cluster of unit is determined by its location in the covariate space, the dose is implicitly dependent on the covariates.
'2 2
To add variability of doses within clusters, per unit i, we sample the individual dose 8, from a Beta- distribution
S ~ Betala, a
mc. + (1-aa) a)):
(5)
as motivated by Bica et al. (2020), ensuring that the mode of S ma. The parameter a determines the variance of the Beta-distribution. For a 0, the dose of a unit is sampled from a uniform distribution, leading to no confounding, as the dose is unconditional on unit's covariates. For - co, the dose assignment is fully deterministic, and every unit in cluster will be assigned the model dose. We can think of a as the confounding strength. We visualize the impacts of paramcters and B in Appendix A.
Slep 3 (Response calculation): We define a ground truth dose response per unit, which allows for the calculation of counterfactual outcomes. The conditional-averaee dose response for a covariate vector x, and a dose defined as
3 w
x,) = 10 w|x + 12s 8
(6)
4 w.Xx
where Wt E 1 for E (1, is a weight vector. We randomly sample half of the weights from a uniform dist ron 4(0, 1) and set the remaining weights to zero to add noise covariates (Shi el al., 2019). We test the effect of different degrees of sparseness in the weight vector in Appendix D. The individual dose response per unit is finally calculated as y: uus;, x;) + where c <~ (0, 1) is random error term. The final data for different values of a and / is visualized in Figure 4 illustrating the confounding.
Benchmarks. We compare CBRNet against several relevant baselines. First, we compare it against su- pervised learning algorithms, namely linear regression, regression tree, xgboost, as state-of-the-art imple- mentation of gradient-boosted trees and a feedforward multi-layer perceptron (MIP).. Second, we compare it against DRNet (Schwab et al., 2019) and VCNet (Nic ct al., 2021), two popular and competitive dosc response estimators. For a discussion of these methods, see Section 2. To understand the impacts of different IPMs on model performance, we train CBRNet using the linear MMD (MMD)m)). a kernel MMD (MMD,bf) using the radial basis function kernel, and the Wasserstein distance.
Performance metrics. We evaluate the performance of each method using the MISE metric introduced by Schwab et al. (2019), which measures the ability of method to estimate the CADR over all units in the test data and all dose levels. For a test dataset with units, the MISE is defined as
n MISE Σ
1 n i=1
/ (u(s,x)) (18, xi)) ds.
(7)
6 Empirical Results
Performance on the dry bean dataset. For a total of 17 different combinations of a and B, we generate 10 random instances of the Dry bean-DR data. We use 70% of the data for training, 10% as a validation set for hyperparameter tuning, and 20% as a test set for calculating performance metrics.
Fixing 3 1. we report the performance of all nine estimators for different levels of a in Table 1. Performance under different. levels of B is stated in our extended results section in Appendix C. CBRNet outperforms all
9

---

Published in Transactions on Machine Learning Research (02/2025)
Table 1: MISE per method on dry bean dataset for 3 2 and varying levels of a. We compare CBRNet against several benchmarking methods over varying levels of confounding by cluster. The best results are printed in bold, second-best results are in italics, Results for different values of can be found in Appendix C.
8
Model
1.0
2.0
3.0
4.0
Linear Regression
3.32 = 0.08
3.40 = 0.06
3.30 - 0.02
3.32 = 0.03
CART
1.31 0.06
1.40
0.09
1.56 0.15
1.53
0.13
xgboost
0.91 i 0.07
1.16 + 0.09
1.15 1 0.07
1.28 1 0.04
MLP
0.54 + 0.04
0.67 + 0.08
0.83
1.06++0000
DRNet
0.62 0.08
0.78 + 0.13
0.82 0.07
0.98 + 0.06
VCNet
0.75 + 0.11
+ 1.07
+
1.35 + 0.07
(MRNet(MMD(a))
0.44 + 0.09
0.47 0.12
0.44 + 0.07
0.63 + 0.15
(BRNe((MMD.or)
0.45 = 0.10
0.53 = 0.14
0.50 0.10
0.71 + 0.17
CBRNet (Wass)
0.43 + 0.09
0,46 0.10
0.45 0.08
0.55 = 0.11
benchmark methods across different levels of confounding, indicating that proposed model architecture is well-suited to learn CADR from clustered data.
CBRNet shows higher robustness to confounding by cluster than its benchmarks when the confounding strength a increases. Irrespective of the chosen IPM for regularization, the performance CBRNet is stable for low to moderate levels of a, only increasing for the highest value of a = 4. For all other models, performance is strictly decrcasing for increases in a.
Similar behavior can be observed for different values of B.. A notable case is given for 0, hence when there is variability of doses between clusters (ef. Table 2 in Appendix C). For 3 = 0, the data is unconfounded, as doses are sampled from the same distribution for all units and all clusters. So, CBRNet also outperforms the benchmark methods when learning CADR from clustered, but unconfounded data.
Hyperparameter scnsitivity We further investigate the sources of gain of our proposed architecture by testing the impact of two critical hyperparameters: (a) The strength of the regularization X, and (b) the number of clusters identified by the k-means clustering. The results of those experiments are summarized in Figure 5. For both experiments, we fix the hyperparameters of CBRNet and iterate over varying levels of À and k. To stabilize results, for every parameter configuration, we generate 10 random instances of the Dry bean-DR data with = 3 and 3.
In Appendix B we visualize the impact of the IMP regularization on the representation space. We find a strictly positive impact on model performance for low to moderate levels of regularization (A < 0.1). These results support the positive impact of regularizing for differences in cluster distributions the representa- tion space. These results are also independent of the chosen IPM, with MMDs and Wasserstein distance performing comparably. Only for an increased level of regularization, does the predictive performance of CBRNet decrease. This indicates that overregularization is possible
The performance of CBRNet appears to be robust to the specified number of clusters k, with performance converging on the Dry bean-DR data for larger k. We hence expect our method to have favorable asymptotical properties for unclustered data, as long as the dose assignment can be approximated sufficiently well by a large enough number of clusters and if sufficient data available to approximate the IPMs. In applications, next beforo-mentioned model selection procedures, the number of clusters can be selected together with domain experts, or by visual analysis.
10

---

Published in Transactions on Machine Learning Research (02/2025)
1.4
1.0
1.2
0.9
MISE
MISE
1.0
0.5
CBRNNet MM)(lin)
CBRNet MMD((rb))
0.8
0.7
CBRNet+ Wass
0.6
0.6
10
10
10"
10
10P
10
Reg. strength
Number of Clusters
(a) Impact of reg. strength
(b) Impact of num. clusters
Figure 5: Hyperparameter robustness. (a) Moderate levels of regularization improve model performance.
Regularizing by the MMD with radial basis function kernel (rbf) appears to overregularize earlier than for linear MMD and Wasserstein distance.
(b) We see robustness to the number of clusters. Performance
converges for all IPMs.
Performance on established datasets To judge the performance of CBRNet on unclustered data, we run experiments on previously established benchmarking datasets proposed by Bica ct al. (2020)² and Nic et al. (2021). CBRNet performs competitively, beating all of its benchmark methods on three of the four datasets. However, further decomposition is needed to understand the precise challenges these datasets attribute performance either the general architecture of CBRNet, or the IPM regularization. We provide the full results these datasets in Table 5 in Appendix C. While good performance benchmarking dataset is guarantee for performance in real-world application, we consider these results to indicate the effectiveness of our newly developed approach.
7 Conclusion
We studied the estimation of conditionaa-average dose responses (CADR) from clustered observational data. Practitioners face clustered data in several real-life sccnarios, covering applications in, e.g., business, cco- nomics, and healthcare (cf. Section 4). In such scenarios, dose assignment is oflen cluster-specifie, potentially leading to confounding. This confounding mechanism is distinct from scenarios previously studied in the ML literature on intervention response estimation and has impacts on the validity of assumptions that are often made beyond standard ones necessary for applying causal inference (cf. Section 2).
To enable thorough testing of estimators under the presence of such data, created novel dataset, the "Dry bean-DR data", based the covariates initially presented by Koklu & Ozkan (2020). Experiments this dataset reveal that traditional supervised estimators and ML estimators tailored to CADR estimation suffer from confounding by cluster. To work towards solving this problem, we proposed a new MI. estimator, CBRNet. Using a novel loss function that averages integral probability metrics over any number of clusters the covariates space, we can leverage representation balancing for CADR estimatiio and prevent. confounding biases.
In comparison with its benchmarks, CBRNet shows improved performance, indicating approach's ap- propriateness. Further experiments using established datasets revealed that CBRNet can similarly achieve competitive performance on unclustered data.
Our research has shown that understanding the context and unique characteristics of CADR estimation problem imperative. Variations the underlying data-generating process observational data can have adverse effects on model performance. To ensure good performance, especially given the inherent difficulty of model selection in causal inference, practitioners must check the validity of any method-specific assumptions, for example, by involving domain experts (Pearl, 2022).
2The DGP proposed by Bica et (2020) incorporate up three different. interventions, each with associated dose. comply with the architecture CBRNet, we generate responses only single intervention.
11

---

Published in Transactions on Machine Learning Research (02/2025)
Our implementation of CBRNet is available online for practitioners and fellow researchers to build upon (cf. Appendix F). Our implementation uses a basic scikit-learn syntax (Pedregosa et al., 201 1), enabling efficient model setup, training and inference:
from src. methods neural import CBRNet model CBRNet
3 model (ft_e.ia Y_train, S_train)
4 model prdeictr (pedict(X_tessee teste S_test)
All relevant hyperparameters can be specified. The architecture of our method was motivated by the analysis of real-world scenarios, leveraging and extending the concept of representation balancing, first proposed in the literature on ML for causal inference by Shalit et al. (2017). We highlight three potential future research directions: (1) A theoretical analysis of CBRNet and asymptotical guarantees of model performance in clustered and unclustered data. Especially, an analysis of the necessary amount of data for effective regularization could benefit the adoption of our model, as the IPMs for regularization are approximatod empirically. (2) Further research on model selection and hyperparameter tuning. While using validation set MSE can be practical first guess (Curth & Van Der Schaar, 2023), alternative approaches could further add performance guarantees when training models real-life data. This research is not specific to CBRNet, but is relevant to the wider related literature. (3) Finally, our research is specific to interventions with an associated dose, hence extending our experiments to binary-valued intervent ions could be a promising and insightful direction, especially given the limited mumber of benchmarking datasets for treatment effect estimation (Curth al., 2021).
Acknowledgments
We would like to thank the anonymous reviewers for their feedback during the revision process.
This work was supported by the Research Foundation - Flanders (FWO research projects G015020N and 1117322N).
References
Jason Abrevaya, Yu-Chin Hsu, and Robert P. Lieli. Estimating conditional average treatment effects. Journal of Business and Economic Stalislics, 33(4):485 505, 2015. doi: 10.1080/073550153 5735500152004.975555
Ahmed Alaa and Mihaela der Schaar. Limits of estimating heterogeneous treatment effects: Guidelines for practical algorithm design. In International Conference Machine Learning, pp. 129-138. PMLR,
2018.
Alexis Bellot, Anish Dhir, and Giulia Prando. Generalization bounds and algorithms for estimating condi- tional average treatment effect dosage. arXiv preprint 3rXiv:2205.14699 2022.
Jesse A. Berlin, Stephen E. Kimmel, Thomas R. Ten Have, and Mary D. Sammel. An Empirical Comparison of Several Clustcred Data Approaches Under Confounding Duc to Cluster Effects in the Analysis of Complications of Coronary Angioplasty. Biometrics, 55(2):470 476, 1999. doi: 10.111/j..000-3411119999 00470.x.
Jeroen Berrevoets, James Jordon, Ioana Bica, and Mihaela der Schaar. OrganITE: Optimal transplant donor organ offering using an individual treatment effect. NIPS, 33:20037-20050, 2020.
Julian R. Betts and Jamie L. Shkolnik. The effects of ability grouping on student achievement and resource allocation in secondary schools. Economics of Education Review, 19(1):1-15, 2000. doi: 10.1016/s0272-7771 (0/1)//22227277779999004442
Toana Bica, James Jordon, and Mihacla van der Schaar. Estimating the effects of contimuous-valued inter- ventions using generative adversarial networks. In Advances in Neural Information Processing Systems, volume 33, pp. (8434-16447
Curran Associates, Inc., 2020.
12

---

Published in Transactions on Machine Learning Research (02/2025)
Tom Brockel, Dirk Fornahl, and Andrea Morrison. Another cluster premium: Innovation subsidies and R&& collaboration networks. Reseorch Policy, 44(8):143111444, 2015. doi: rep00///.
Edward Calabrese, The emergence of the dose response concept in biology and medicine. International
Journal of Molecular Sciences, 17(12):2034, 2016. doi: 10.3390/jjms11122033
Cancer Genome Atlas Research Network, John N. Weinstein, Eric A. Collisson, Gordon B. Mills, Kenna R. Mills Shaw, Brad A. Ozenberger, Kyle Ellrott, Πya Shmulevich, Chris Sander, and Joshua M. Stuart. The Cancer Genome Atlas Pan-Cancer analysis project. Nature Genetics, 45(10):1113 1120, 2013. doi: 10.1038/ng /t.///g..222 2764.
Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Dala Mining, KDD '16. ACM, 2016. doi: 10.1145 11444//993672.999999
Alicia Curth and Mihaela Van Der Schaar. In search of insights, not magic bullets: towards demystification of the model selection dilemma in heterogeneous treatment effect estimation. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
Alicia Curth, David Svensson, Jim Weatherall, and Mihaela van der Schaar. Really doing great at estimating CATE? a critical look at ML benchmarking practices in treatment effect estimation. In Thirty-fifth Conference on Neural Information Processing Systems Dalasets and Benchmarks Track (Round 2), 2021.
Angus Deaton and Nancy Cartwright. Understanding and misunderstanding randomized controlled trials.
Social Science & Medicine, 210:2 21, 2018.
Hossein Dehghanisanij. Hojjat Emami, Somayeh Emami, and Vahid Rezaverdinejad. A hybrid machine learning approach for estimating the water-use efficiency and yield in agriculture. Scientific Reports, 12 (1), 2022. doi:
William Falcon. Pytorch lightning. /ittt///ii/////////////////////a
2020.
Matthew C. Farrelly, Kevin C. Davis, M. Lyndon Haviland, Peter Messeri, and Cheryl G. Healton. Evidence of a dose response relationship between "truth" antismoking ads and youth smoking prevalence. American doi: 10.2105 2spph..004..499922
Journal of Public Health, (543) 431, 2005.
Emil Frei and George P. Cancllos. Dose: critical factor in cancer chemotherapy. The American Journal of Medicine, 69(4):585- 594, 1980. doi:
Anat Goldstein, Lior Fink, Amit Meitin, Shiran Bohadana, Oscar Lutenberg, and Gilad Ravid. Applying ma- chine learning on sensor data for irrigation recommendations: revealing the agronomist's 's tacit knowledge. Precision Agriculture, 19(3):421 444, 2017. doi:
Ian Goodfellow, Jean Pouget-Abadis Mehdi Mirza, Bing Xu, David Marde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11): 139-144, 2020. doi:
Arthur Gretton, Karsten Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alex Smola. A kernel method for the two-sample-problem. NIPS, 19, 2006.
Sebastien Haneuse. Distinguishing Selection Bias and Confounding Bias in Comparative Effectiveness Re- search. Medical Care, 54(4):e23 29, 2016. doi: 001097/ALL..0000000000000
Trevor Hastic and Robert Tibshirani. Varying-cocfficent models. Journal of the Royal Statistical Society: Series B (Methodologicaal) 55(4):757 779, 1993. doi: 10.111/::511-61611198811
Jennifer L. Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and
Graphical Stalistics, 20(1):217 240, 2011. doi: 10.1198//egs.2010088133
13

---

Published in Transactions on Machine Learning Research (02/2025)
Keisuke Hirano and Guido W. Imbens. The Propensity Score with Continuous Treatments, chapter 7, pp. 73 84. John Wiley & Sons, Ltd, 2004. ISBN 9780470090459. doi:
Daniel Ho, Kosuke Imai, Gary King, and Elizabeth A. Stuart. Matching as nonparametric preprocessing for reducing model dependence in paramotric causal inference. Political Analysis, 15(3):199-22) -229 2007. doi: 10.1093/pam 1000//an/mpp11. /mp013.
G. T. S. Ho, W. H. Ip, C. K. M. Lee, and W. L. Mou. Customer grouping for better resources allocation using GA based clustering technique. Expert Systems with Applications, 39(2):1979-1987, 2012. doi: 10.1016/j.cswa.201 wtw.s/a.....0.044....
Paul W. Holland. Statistics and Causal Inference. Journal of the Americon Statistical Association, 81 (396): 945 960, 1986. doi: 10.1080/0162
10478354.
Tim Holland-Lotz and Annette Kopp-Schneider. Optimal experimental designs for dose studies with continuous endpoints. Archives of Toricology, 89(11):2059 2068,, 2014. doi: 10.1007/s002004-0444 1335-2..
studics
Fredrík Johansson, Uri Shalit, and David Sontag. Learning representations for counterfactual inference. In Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 3020-3029, New York, New York, USA, 2016. PMLR.
Fredrik D. Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and repre- sentation learning for estimation of potential outcomes and causal effects. Journal of Machine Learning Research, 23(166):1 50, 2022.
Jedid-Jah Jonker, Nanda Piersma, and Dirk Van den Poel. Joint optimization of customer segmentation and marketing policy to maximize long-term profitability. Expert Systems with Applications, 27 2004. doi:
L. V. Kantorovich. Mathematical methods of organizing and planning production. Management Science, 6 (4):366 422, 1960. doi: 10.1287 10.1287/mnsc. w.m.nc.....36..
Murat Koklu and mkker Ali Ozkan. Multiclass classification of dry beans using computer vision and machine learning techniques. Computers and Electronics in Agriculture, 174:105501 2020. doi: 10.1016/j.compag, 2020.105507.
Stuart Lloyd. Least squares quantization in PCM. IEEE Trans Inf Theory, 28(2) 1982.
A. Russell Localio, Jesse A. Berlin, and Thomas R. Ten Have. Confounding Due to Cluster in Multicenter Studies Causes and Cures. Health Services and Outcomes Research Methodology, 3(3):195-210 2002. doi: (10223/A:102889881111
Cuadies-Causes
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. Advances in neural information processing systems, 30, 2017.
Gang Luo. A review of automatic selection methods for machine learning algorithms and hyper-parameter values. Network Modeling Analysis in Health Informatics and Bioinformatics, 5(1):18, 2016. doi: 10.1007/ s13721-016-0125-6
T. Soni Madhulatha. An overview on clustering methods. orXiv preprint 3rXiv:1205.11117 2012.
Lizhen Nie, Mao Ye, Qiang Liu, and Dan Nicolae. VCNet and functional targeted regularization for learning causal effects of continuous treatments. In International Conference Learning Representations, 2021.
Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108 (2):299 319999999,9000 2021.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, et al. Auto- matic differentiation in Pororc. In 31st Conference on Neural Information Processing Systems, 2017.
14

---

Published in Transactions on Machine Learning Research (02/2025)
Judea Pearl. Causality. Cambridge University Press, Cambridge, second edition, reprinted with corrections edition, 2022. ISBN 9780521 9780521895606. 895606. Hier auch später erschienene, unveränderte Nachdrucke.
F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine Learning in Python. Journal of Machine Learning Reseorch, 12:2825 2830, 2011.
Robert Phillips. Optimizing prices for consumer credit. Journal of Revenaee and Pricing Management, 12: 360 377, 2013.
Thomas Richardson and James M. Robins. Single world intervention graphs (SWIGs): A unification of the counterfactual and graphical approaches to causalityy Center for the Stotistics and the Social Sciences, University of Washington Series. Working Paper, 128(30):2013, 2013.
Paul Rosenbaum and Donald Rubin. The central role of the propensity score in obsorvational studics for causal effects. Biometrika, 70(1):41 55, 1983. doi: 10.1093 10.1093/biomet/ biooet/770.....
Donald B. Rubin. Estimating causal effects of treatments in randomized and nonrandomized studies. Journal of Educalional Psychology, 66(5):688 701, 1974. doi: 10.1037 111037/h0033500
Alexander Schacht, Philip Gorwood, Philip Boyce, Ayal Schaffer, and Hernan Picard. Depression symptom clusters and their predictive value for treatment outcomes: Results from an individual patient data meta- analysis of duloxetine trials. Journal of Psychiatric Research, 53:54 61, 2014. doi: 10.y0i6/j..ssychires. 2014.02.00
Patrick Schwab, Lorenz Linhardt, and Walter Karlen. Perfect Match: A Simple Method for Learning Representations For Counterfactual Inference With Neural Networks. arXiv preprint Xrx: 10..1810000000 0.000066, 2018.
Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M. Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. In Proceedings of the AAAI Conference on Artificial Intelligence, 2019. doi: 10.1609 1609///aiii/ aaaaaaiii v34.04..40446
Skipper Seabold and Josef Perktold. statsmodels: Econometric and statistical modeling with python. In 9th Python in Science Conference, 2010.
Shaun Seaman, Menelaos Pavlou, and Andrew Copas. Review of methods for handling confounding by cluster and informative cluster size in clustered data. Statistics in Medicine, 33(3)):5371 5387, 2014. doi: 10.1002/sim.6277.
Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: generaliza- tion bounds and algorithms. In Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, 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, volume 32. Curran Associates, Inc., 2019.
Jerzy Splawa-Neyman, D. M. Dabrowska, and T. P. Speed. On the application of probability theory to agricultural experiments. essay on principles. section 9. Stalistical Science, 5(4):465 1474247777, 1990.
Richard Stone. The Assumptions on Which Causal Inferences Rest. Journal of the Royal Stalistical Sociely: Series B (Methodological), 55(2):455-4660 1993. doi: 10.11111.257--1111199... 199...1119.9....
Jonathan M. Turk. Estimating the impact of developmental education on associate degree completion: A dosc-response approach. Research in Higher Education, 6098)11000-1122, 2019. doi: 10.1007/
-11162-011--99549999
Mark van der Laan and Daniel Rubin. Targeted maximum likelihood learning. The International Journal of Biostatistács, 2(1), 2006. doi: 10.2202/1657-4679.1
15

---

Published in Transactions on Machine Learning Research (02/2025)
Guido Van Rossum, Fred L. Drake, al. Python reference manual. Centrum voor Wiskunde en Informatica Amsterdam, 1995.
Toon Vanderschueren, Robert Boute, Tim Verdonck, Bart Baesens, and Wouter Verbeke. Optimizing the preventive maintenance frequency with causal machine learning. International Journal of Production Economics, 258:108798, 2023. doi: 1..016//.jpp.2222..187789
Ravi Varadhan and John D. Seeger. Estimation and Reporting of Heterogeneity Treatment Effects. In Developing Protocol for Observational Companative Effectiveness Research: User's Guide. Agency for Healthcare Research and Quality (US), 2013.
Leonid Nisonovich Vaserstein. Markov processes over denumerable products of spaces, describing large systems automata. Problemy Peredači Informacii, 5(3):64-77) 72, 1969.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Brian G Vegetabile. On the distinction between' conditional average treatment effects' effects'(cate) and' individual treatment elfects' elTects"(ite) under ignorability assumptions. arXiv preprinl or Xiv:220008 04939, 2021.
René Verhaart, Valerio Fortunati, Gerda M. Verduijn, Theo van Walsum, Jifke F. Veenland, and Mar- garethus M. Paulides. CT-bbaed patient modeling for head and neck hyperthermia treatment planming: Manual versus automatic normal-tisnneseegnnntaaoo Radiotherapy and Oncology. 111(1):558-163 2014. doi: 10.1006//.raadon::114....1.1
Stefan Wager and Susan Athey. Estimation and inference of heterogeneous treatment effects using random forests. Journal of the Americon Statistical Association, 113(5233:1228--11 1242,, 2018. doi: 10.1080/01621459.
2017.1319839.
Xin Wang, Shengfei Lyu, Xingyu Wu, Tianhao Wu, and Huanhuan Chen. Generalization bounds for esti- mating cansal effects of continuous treatments. In Advances in Neural Information Processing Systems, volume 35, pp. 8605-8617. Curran Associates, Inc., 2022.
Junjie Wu. Advances in K-means clustering: a data mining thinking. Springer Science & Business Media, 2012.
A C Yeo, K A Smith, R J Willis, and M. Brooks. A mathematical programming approach to optimisc insurance premium pricing within a data mining framework. Journal of the Operulional Research Sociely, 53(11):1197-122) -203, 2002. doi:
26014113.
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.
Johan Zetterqvist, Stijn Vansteelandt, Yudi Pawitan, and Arvid Sjölander. Doubly robust methods for handling confounding by cluster. Biostatistics, 17(2):264-215 2016. doi: 10.1093/biose httss/biossatiitisskkvkkk
YiFan Zhang, Hanlin Zhang, Zachary Chase Lipton, Li Erran Li, and Eric Xing. Exploring transformer backbones for heterogeneous treatment effect estimation. Transactions on Machine Learning Research, 2023.
16

---

Published in Transactions on Machine Learning Research (02/2025)
A
Visualization of confounding mechanism
Pob
Prob
Prob.
Prob.
Dosage
Dosage
Dosage
Dosage 3 a
(a) B 0, a=
(b)
0, a = 3
(c) B
(d)
Figure 6: Visualization of dose distributions per cluster. Subfigures visualize the dose distributions of
different clusters for varying levels of confounding. Parameter 3 determines the difference in doses bet ween individual clusters. Parameter a determines the variability of doses within a cluster.
B Effect of regularization on covariate representation
Dinecssion 2
Domeaty
Dimerrsion
Dimersioo
Doee
Dimension
DEmension 1
Dimension 1
(a) Dose distribution
(b) Input space
(c) Unregul. rep.
(d) Regul. rep.
Figure 7: Regularizing impact of integral probability metric (IPM). We visualize the regularizing impact of the TPM. Figure (a) visualizes the density of different doses the test data and provides color legend. The preceding figures are two-dimensional t-SNE plots of (b) the input space in X, (c) the learned representation of CBRNet without regularization, and (d) the learned representation with regularization by the MMD with radial basis function kernel. The regularization effectively removes the clustering by dose level.
17

---

Published in Transactions on Machine Learning Research (02/2025)
C Extended results
Table 2: MISE per method on dry bean dataset for B = 0 and varying levels of a.
8 0
Model
0.0
1.0
2.0
3.0
4.0
Linear Regression
2.83
3.27
3.80
4.19
4.44
CART
0.86 + 0.03
1.04 + 0.03
1.44 + 0.07
1.93 + 0.10
2.18 0.11
xgboost
0.55
0.69 0.02
1.11
1.58
1.88
MLP
2.83 0.01
0.54 + 0.08
1.17 0.09
1.98 = 0.05
2.52 0.09
DRNet
0,41 = 0.03
0.60 + 0.06
0.99 0.12
1.74 - 0.19
1.92 0.23
VCNet
0.34 = 0.04
0.69
141 0.09
1.69 0.17
(MMNec(MMD
0.25 0.01
0.35 = 0.05
0.41 + 0.10
0.64 + 0.08
0.88 0.14
CBRNet (MMD)r)
0.26 0.04
0.34 0.05
0.42 1 0.07
0.61 0.08
0.92 0.09
CBRNet (Wass)
0.28 = 0.05
0.34 0.05
0.42 + 0.10
0.62 0.08
0.86 1 0.16
Table 3: MISE per method on dry bean dataset for = and varying levels a.
B=
Model
1.0
2.0
3.0
4.0
Linear Regression
3.43 + 0.03
3.81 0.04
4.06 + 0.03
4.35 0.04
CART
1.17 = 0.09
1.47 + 0.11
1.74 = 0.08
2.04 + 0.21
xgboost
0.77 0.07
1.04 + 0.07
1.37
1.66 + 0.13
MLP
0.57 + 0.06
1.01
1.46
2.14
DRNet
0.72 + 0.07
0.96 1 0.07
1.30 0.16
1.54 4: 0.20
VCNet
0.68
1.16
1,48 0.16
1.83 0.12
(BRNet(MMDDaa)
0.41 0.09
0.57 + 0.17
0.71 + 0.17
0.79
(MMNet(
0.42 0.11
0.57 + 0.15
0.75 + 0.26
0.87 + 0.15
CBRNet (Wass)
0.44 = 0.08
0.57 = 0.15
0.69 0.28
0.74 0.17
18

---

Published in Transactions on Machine Learning Research (02/2025)
Table 4: MISE per method dry bean dataset for 3= 2 and varying levels of a.
Model
1.0
2.0
3.0
4.0
Linear Regression
2.91 + 0.01
2.92 + 0.02
3.01 + 0.08
3.01 + 0.02
CART
1.24
1.31
1.32 0.10
1.43 0.1.1
xgboost
1.04 = 0.05
1.04 0.08
1.00 = 0.06
1.09 0.06
MLP
1.19 0.09
1.19 0.1
0.15
DRNet
= 0.07
0.59 0.06
0.71 = 0.05
0.70 + 0.08
VCNet
0.55 + 0.10
0.62 1 0.07
0.60 + 0.05
0.66 0.00
CBRNe((MMD(a))
0.35 0.07
0.39 1 0.11
0.39 +: 0.09
0.45 0.06
(MMNer)
+ 0.11
0.49 0.18
0.42 + 0.06
0.51 0.13
CBRNet(Wass)
0.34 0.06
0.38 + 0.10
0.42 0.08
0.45 0.07
Table 5: MISE on established benchmark datasets. CBRNet performs competitively against bench- mark methods across established datasets, achieving state-of-the-art performance three out of four datascts, and improving over standard MLP all datasets.
Dataset
Model
TCGA-2
IHDP-1
News-2
Synth-1
MLP
2.01 1 0.1
2.79 = 0.09
1,22 1 0.12
1.60 1 1.22
DRNet
0.35 0.0
2.54
0.08
0.81
0.07
1.52 + 1.66
VCNet
0.68 + 0.30
1.42
0.19
0.78 + 0.06
0.90 + 0.55
(MRNet(MMDin))
0.24 + 0.07
1.60 + 0.33
1.12 + 0.06
0.99 + 0.66
(BRN,t(MMMD,w))
0.26 0.09
1.56 0.17
1.10 + 0.05
0.99 + 0.61
CBRNet(Wass)
0.24 + 0.08
1.40 + 0.25
1.10 + 0.05
0.89 + 0.56
19

---

Published in Transactions on Machine Learning Research (02/2025)
D Dataset ablation
In this section we test the effect of different levels of sparseness in the weight vectors determining the dose response in the Dry bean-DR data (ef. Section 5). For four different levels of sparseness, a = and 8 = 1 we generate 10 instances of the Dry bean-DR data. We then train all methods on these instances of the data and report results in Table 6. The results indicate that increasing sparsity does not immediately lead to a decrease in model performance. Across models, performance increases for spareness levels up 50%. Only for sparseness of 75% model performance decreases. While CBRNet. is similarly affected by high levels of sparsity, the decrease in model performance is less severe when compared with the benchmark methods.
Table 6: Model performance (MISE) per weight vector sparseness. We are iterating over different levels of spareness in the weight vectors determining the dose response in the Dry bean-DR data (ef. Sec- tion 5). Increasing sparsity does not immediately lead to decrease in model performance. Across models, performance increases for spareness levels up to 50%. Only for sparseness of 75% model performance decreases.
a= B 8
Sparseness
0%
25%
50%
75%
Linear Regression
4.88 + 0.05
4.39 = 0.04
3.78 + 0.04
18.91 + 0.20
CART
2.45 + 0.26
1.70 = 0.08
1.37 + 0.13
11.33 + 1.22
xgboost
2.25 + u.16
1.55 =
1.23 + 0.08
6.25
MLP
3.09 + 0.21
0.98 = 0.04
1.31 + 0.28
5.41 - 0.31
DRNet
1.18 + 0.05
0.81 0.06
0.66
+ 0.06
5.85 0.30
VCNet
1.24 1.24 + 0.27
0.73 0.21
0.58 0.13
5.70 0.40
(MMNec
0.88 0.13
0.67 +
0.50 +
2.51 +
CBRNet(MMDrss)
1.10 + 0.30
1.05 - 0.28
0.70 + 0.35
2.71 = 1.40
CBRNet (Wass)
0.90 0.18
0.71 - 0.22
0.51 + 0.25
2.51 + 0.87
20

---

Published in Transactions on Machine Learning Research (02/2025)
E Implementation, reproducibility, and hyperparameter optimization
Our experiments are written in Python 3.10 (Van Rossum et al., 1995) and were exccuted on an Apple M2 Pro SoC with 10 CPU cores, 16 GPU cores, and 16 GB of shared memory. The system needs approximately one day for the iterative execution of all experiments. We aim for maximal reproducibility. All methods considered in our manuscript are consistently implemented in an sklearn style, The code to reproduce all experiments, results, and figures paper can be found online via https:/ gi thub. com/ christopher-brr CBRNet.
For VCNet, build on the original implementation provided by Nie et al. (2021) (uttps://githn) com/ lushleaf/ evaryiii coefficient were implemented in Pylorrhh (Paszke et al., 2017) using Lightning (Falcon, 2020). Xgboost is imple- mented using the xgboost library (Chen & Guestrin, 2016). CART models were implemented using the Scikit-Learn library (Pedregosa al., 2011). Linear regression was implemented using the statsmodels library (Seabold & Perktold, 2010).
All remaining neural network architectures
For the TCGA-2 dataset, linear regression models were trained using the first 50 principal components of the covariate matrix lo reduce compulational complexity.
Hyperparameter optimization. Results are not to be compared to the original papers, as the optimiza-
tion scheme and parameter search ranges differ from the original records. If not specified differently, the remaining hyperparameters are set to match the specifications of the original authors.
Table 7: Hyperparameter search range for Linear Regression:
Parameter
Values
Penalty
(Elastic net, Lasso, None}
Table 8: Hyperparameter search range for CART:
Parameter
Values
Max depth
(5, 15, None]
Min sample split
{2, 5, 20}
Min samples per leaf
{1,5, 10}
Max features per split
(None, V p(x))
Splitting criterion
{Gini]
Table Hyperparameter search range for xgboost:
Parameter
Values
Learning rale
(0.01,0.1, 0.2)
Max depth
(3,5,7,9)
Subsample
(0.5,0.7, 0.7,1.0) 1.0)
Min child weight
{1,3,5}
Gamma
00.,,0.1.023
Columns sampled per tree
(0.3,0.5.0.7)
21

---

Published in Transactions on Machine Learning Research (02/2025)
Table 10: Hyperparameter search range for MLP:
Parameter
Values
Learning rate
{0.0001,
1.2 regularization
(0.0,0.1)
Batch size
{64, 128}
Hidden size
(82,48)
Num steps
[5000}
Num layers
{2)
Optimizer
{Adam)
Table 11: Hyperparameter search range for DRNet:
Parameter
Values
Learning rate
{0.0001, 0.001)
1.2 regularization
(0.0,0.1}
Batch size
{64, 128)
Hidden size
{32.48)
Num dose strata
{10}
Num steps
(5000)
Num layers
{2}
Optimizer
{Adam)
Table 12: Hyperparameter search range for VCNet:
Parameter
Values
Learning rate
{0.001, 0.01}
Batch size
{128, 256}
Hidden size
332}
Num steps
(5000}
Optimizer
[Adam]
Table 13: Hyperparameter search range for CBRNet:
Parameter
Values
Learning rate
(0.001, 0.013
Batch size
{128,
Hidden size
(32)
Num steps
(5000)
Optimizer
{Adam}
IPM Regularization
{0.001, 0.01, 0.13
22


Reviews:
[{"review_id": "5bOIlPdlVf", "reviewer": "Reviewer_Ytfk", "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 .This paper has clear writing and is easy to follow. \n\nWeakness:\n1. The innovation in the paper is not enough in my opinion. Causal invariant learning is studied in many papers. Even though studying it in the case of CADR estimation, it is not by nature very different. The method proposed in this paper is also similar to many used in causal invariant learning.\n2. The experiment parts are not convincing enough. It proposed a new benchmarks to evaluate the CADR\testimation. However it is not clear the superiority of this benchmark compared to other common CADR estimation benchmarks.\n3. The result table 5 doesn’t show a clear signal of the advantage for CADR estimation method proposed in this paper, as the gap between CADR and the best baseline estimation method is very small.", "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 provides a method(CBRNet) estimating CADR from clustered data and where different doses are assigned to different segments of a population. It studies the impacts of learning from clustered data on the performance of ML estimators for CADR and proposes a benchmark to evaluate.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Can the authors explain more about the motivation of using this benchmark and the superiority of it?\n2. Can the authors explain how their method is different from other literatures from the causal invariant learning fields?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "V0lg16gt2s", "reviewer": "Reviewer_AqkL", "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": "Strength: On a novel benchmarking dataset, the authors show the impacts of clustered data on model performance and propose an estimator, CBRNet, that learns cluster-agnostic and hence dose-agnostic covariate representations through representation balancing for unbiased CADR inference; they run extensive experiments to illustrate the workings of our method and compare it with the state of the art in ML for CADR estimation.\n\nWeaknesses: NA (updated).", "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 authors study estimating conditional average dose response from clustered data where different doses are assigned to different segments of a population.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to Weakness.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "rOJTERESh1", "reviewer": "Reviewer_vAMe", "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": "Overall, the paper would benefit from a more rigorous explanation of the proposed method, along with theoretical guarantees and justifications. Please refer to the detailed requests below.\n\n\nI may have misunderstood certain aspects of the paper. If the authors could address the requested changes and clarify any misunderstandings on my part, I would be happy to reconsider increasing the score.", "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 to use cluster-based ML methods to estimate the conditional average dose response.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "The paper is somewhat unclear and requires clarification on several points:\n\n1. The term \"conditional-average dose response\" is confusing. Typically, \"average treatment effect\" refers to an expectation taken over certain covariates. In your case, you're conditioning on the covariate, so the natural interpretation would involve taking the expectation over the dosage. However, as defined in Equation (1), this isn't the case—you are estimating the expected response under a fixed dosage and fixed covariate. It might be more accurate to rename this term as \"conditional dose response.\"\n\n2. Assumption 2 states that the covariates X are pretreatment, which is challenging to guarantee in practice. This is a strong assumption and should be explicitly highlighted, in addition to the no hidden confounding assumption, which only implicitly suggests this condition.\n\n3. The paper assumes that \"similar units form clusters and that dose assignment is conditional on the cluster.\" However, in practice, covariate assignment decisions are often based on a subset of covariates. In Section 6, the experiments are shown under a setup where half of the coefficients were set to zero. It would be helpful to vary the sparsity of the weights and report how this affects the performance of the algorithms.\n\n4. The architecture of the estimation procedure in Equation (2) lacks theoretical justification. For instance, you mention that ϕ learns a cluster-agnostic representation, but there is no explanation of which part of the causal quantity this representation is capturing.\nSimilarly, there is no formal mathematical definition or justification for the statement that \"similar units form clusters and the dose assignment is conditional on the cluster.\" The absence of theoretical guarantees for the proposed estimation procedure weakens the rigor of the paper.\n\n5. It's unclear whether the authors are attempting to learn the number of clusters k in their procedure. If so, the problem is closely related to work on the optimal adjustment set, and baselines from that literature should be included. If not, then the problem loses relevance—practically, one would not know k in advance.\n\n6. The experimental setup is somewhat confusing. It's unclear how S depends on x. Specifically, how does $m_{c_i}$ relate to x as shown in Figure 2?", "broader_impact_concerns": "na", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The recommendations from the three reviewers were one vote for "leaning accept" and two for "leaning reject".

However, this submission went through a little unusual process of review because of the extension of the discussion period after the reviewers posted recommendations. After the discussions, most of the concerns seem to have addressed, and I recommend acceptance with minor revisions, according to the exchanges between the authors and the reviews including myself.

There is a few very small concerns about a citation and an equation format, which should be easily addressed in preparing the final version.
Decision: Accept (poster)