Using representation balancing to learn conditional-average dose responses from clustered data

TL;DR

This work tackles estimating the conditional-average dose response (CADR) from observational data when dose assignment is cluster-dependent, a setting termed confounding by cluster. It introduces CBRNet, a neural architecture that learns cluster-agnostic representations via integral probability metric regularization to enable unbiased CADR inference across doses in $[0,1]$. The authors provide a semi-synthetic Dry bean-DR benchmark to study cluster-based confounding, and show that CBRNet outperforms both standard supervised models and existing CADR estimators under varying confounding strengths, while remaining competitive on unclustered data. The results highlight the value of representation balancing for dealing with cluster-driven dose assignment and point to promising directions for theory, model selection, and broader applicability. Overall, the paper advances CADR estimation in realistic clustered environments and offers a practical tool for domains like finance, healthcare, and policy where cluster-based dosing is common.

Abstract

Estimating a unit's responses to interventions with an associated dose, the "conditional average dose response" (CADR), is relevant in a variety of domains, from healthcare to business, economics, and beyond. Such a response typically needs to be estimated from observational data, which introduces several challenges. 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 and propose an estimator, CBRNet, that 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.

Figures (7)

• Figure 1: Illustration of confounding by cluster (CBC). (a) We assume that there are clusters of similar units in the data. (b) The dose assignment mechanism is a function of the cluster so similar units are assigned similar doses. As dose responses are heterogeneous and depend on a unit's covariates (c), the resulting data is confounded.
• Figure 2: Single world intervention graph (SWIG) illustrating causal dependencies between variables in the training data. Our goal is to estimate the response of a unit to a dose $s$.
• Figure 3: Architecture of CBRNet. The network consists of three parts. A representation learner $\Phi$, an inference network $I$, and a clustering function $\Delta$. To overcome confounding by cluster, $\Phi$ is trained to learn a dose-agnostic representation by minimizing a tailored integral probability metric (IPM) over the response space given the clusters identified by $\Delta$. Subsequently, 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 \ref{['sec:CBRNet']}.
• Figure 4: Visualization of level on confounding in dry bean dataset. Per subfigure, the left plot visualizes the dose distribution across the population and provides a 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 a dot ($\cdot$). 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.
• Figure 5: Hyperparameter robustness. (a) Moderate levels of regularization improve model performance. Regularizing by the MMD with a 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.