Cluster-aware continuous treatment effect estimation via Hellinger-based balanced representation learning

Signal Processing, vol. 251, pp. 110866

Abstract

Estimating individualized treatment effects under continuous treatments (e.g., medication dosage) is challenging because treatment-response heterogeneity is often entangled with treatment-assignment bias in observational data. This paper proposes H-AdvCE, a Hellinger distance-based representation learning framework for cluster-aware continuous treatment effect estimation. The key idea is to learn a latent representation that preserves outcome-relevant heterogeneity while reducing dependence between the representation and treatment assignment. Theoretically, we derive a new counterfactual generalization error bound based on the Hellinger distance and show that it is tighter than commonly used Kullback–Leibler (KL) divergence-based alternatives. To optimize this objective, we introduce a variational Hellinger regularizer and formulate the learning problem as a min–max adversarial game. Empirically, H-AdvCE achieves strong performance on IHDP, News, and TCGA benchmarks for continuous dose-response estimation. Beyond prediction accuracy, the learned latent space reveals response-aware subgroups with distinct dose-response patterns, demonstrating its potential for clustering-driven interpretation in continuous treatment effect analysis. Our code is available at https://github.com/syzhao117/AdvCE .

Authors 5

1. Siyuan Zhao

   William & Mary

   Affiliation as printed

   Department of Computer Science, The College of William & Mary, Williamsburg, USA

2. Yiran Zhuo

   University of International Relations

   Affiliation as printed

   School of National Security, University of International Relations, Beijing, China

3. Hongchang Liu

   University of International Relations

   Affiliation as printed

   School of National Security, University of International Relations, Beijing, China

4. Lincen Yang corresponding Aachen Leiden Institute of Advanced Computer Science

   Leiden University

   Affiliation as printed

   Leiden Institute of Advanced Computer Science, Leiden University, Leiden, The Netherlands

5. Shujian Yu corresponding

   Vrije Universiteit Amsterdam

   Affiliation as printed

   Department of Artificial Intelligence, Vrije Universiteit Amsterdam, Amsterdam, The Netherlands