Improving RCT-Based Treatment Effect Estimation Under Covariate Mismatch via Calibrated Alignment

Amir Asiaee, Samhita Pal
Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, PMLR 337:185-209, 2026.

Abstract

Randomized controlled trials (RCTs) are the gold standard for estimating treatment effects, yet they are often underpowered for detecting effect heterogeneity. Large observational studies (OS) can supplement RCTs for conditional average treatment effect (CATE) estimation, but a key barrier is covariate mismatch: the two sources measure different, only partially overlapping, covariates. We propose {CALM} (Calibrated ALignment under covariate Mismatch), which learns embeddings that map each source’s features into a common representation space. OS outcome models are transferred to the {RCT} embedding space and calibrated using trial data, preserving causal identification from randomization. Finite-sample risk bounds decompose into alignment error, outcome-model complexity, and calibration complexity terms, making explicit when the learned embedding is accurate enough to reduce variance. We instantiate {CALM} in two forms: a closed-form linear version, {CALM-Lin}, and a neural representation-learning version, {CALM-NN}. Across 51 simulation settings, calibration-based linear methods are effectively tied in linear-CATE regimes, while {CALM-NN} wins all 22 nonlinear-CATE settings by wide margins. Moreover, on two real-data studies {CALM-NN} delivers the largest gains over the trial-only baseline.

Cite this Paper

BibTeX
@InProceedings{pmlr-v337-asiaee26c, title = {Improving {RCT}-Based Treatment Effect Estimation Under Covariate Mismatch via Calibrated Alignment}, author = {Asiaee, Amir and Pal, Samhita}, booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence}, pages = {185--209}, year = {2026}, editor = {Perković, Emilija and Malinsky, Daniel}, volume = {337}, series = {Proceedings of Machine Learning Research}, month = {17--21 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v337/main/assets/asiaee26c/asiaee26c.pdf}, url = {https://proceedings.mlr.press/v337/asiaee26c.html}, abstract = {Randomized controlled trials (RCTs) are the gold standard for estimating treatment effects, yet they are often underpowered for detecting effect heterogeneity. Large observational studies (OS) can supplement RCTs for conditional average treatment effect (CATE) estimation, but a key barrier is covariate mismatch: the two sources measure different, only partially overlapping, covariates. We propose {CALM} (Calibrated ALignment under covariate Mismatch), which learns embeddings that map each source’s features into a common representation space. OS outcome models are transferred to the {RCT} embedding space and calibrated using trial data, preserving causal identification from randomization. Finite-sample risk bounds decompose into alignment error, outcome-model complexity, and calibration complexity terms, making explicit when the learned embedding is accurate enough to reduce variance. We instantiate {CALM} in two forms: a closed-form linear version, {CALM-Lin}, and a neural representation-learning version, {CALM-NN}. Across 51 simulation settings, calibration-based linear methods are effectively tied in linear-CATE regimes, while {CALM-NN} wins all 22 nonlinear-CATE settings by wide margins. Moreover, on two real-data studies {CALM-NN} delivers the largest gains over the trial-only baseline.} }
Endnote
%0 Conference Paper %T Improving RCT-Based Treatment Effect Estimation Under Covariate Mismatch via Calibrated Alignment %A Amir Asiaee %A Samhita Pal %B Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence %C Proceedings of Machine Learning Research %D 2026 %E Emilija Perković %E Daniel Malinsky %F pmlr-v337-asiaee26c %I PMLR %P 185--209 %U https://proceedings.mlr.press/v337/asiaee26c.html %V 337 %X Randomized controlled trials (RCTs) are the gold standard for estimating treatment effects, yet they are often underpowered for detecting effect heterogeneity. Large observational studies (OS) can supplement RCTs for conditional average treatment effect (CATE) estimation, but a key barrier is covariate mismatch: the two sources measure different, only partially overlapping, covariates. We propose {CALM} (Calibrated ALignment under covariate Mismatch), which learns embeddings that map each source’s features into a common representation space. OS outcome models are transferred to the {RCT} embedding space and calibrated using trial data, preserving causal identification from randomization. Finite-sample risk bounds decompose into alignment error, outcome-model complexity, and calibration complexity terms, making explicit when the learned embedding is accurate enough to reduce variance. We instantiate {CALM} in two forms: a closed-form linear version, {CALM-Lin}, and a neural representation-learning version, {CALM-NN}. Across 51 simulation settings, calibration-based linear methods are effectively tied in linear-CATE regimes, while {CALM-NN} wins all 22 nonlinear-CATE settings by wide margins. Moreover, on two real-data studies {CALM-NN} delivers the largest gains over the trial-only baseline.
APA
Asiaee, A. & Pal, S.. (2026). Improving RCT-Based Treatment Effect Estimation Under Covariate Mismatch via Calibrated Alignment. Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, in Proceedings of Machine Learning Research 337:185-209 Available from https://proceedings.mlr.press/v337/asiaee26c.html.

Related Material