Beyond Imputation: Embedding Alignment for Causal Inference Under Covariate Mismatch

The integration of randomized controlled trials with observational data represents one of the most pressing methodological challenges in modern biostatistics. While RCTs maintain their status as the definitive source for causal identification, they frequently lack the statistical power necessary to detect fine-grained heterogeneous treatment effects across patient subgroups. Conversely, large-scale observational studies, such as electronic health records and insurance claims databases, offer substantial sample sizes but introduce risks of confounding and, crucially, often measure entirely different sets of covariates than those collected in trials. This phenomenon, termed covariate mismatch, has historically forced researchers into an uncomfortable tradeoff between statistical efficiency and causal validity.

In their recent work, Improving RCT-Based Treatment Effect Estimation Under Covariate Mismatch via Calibrated Alignment, Amir Asiaee and Samhita Pal propose a framework called CALM that challenges the prevailing approach to this problem. Rather than attempting to impute missing covariates, which requires solving a high-dimensional reconstruction problem, CALM learns low-dimensional embeddings that align the heterogeneous feature spaces of RCTs and observational studies into a common representation. This approach preserves the causal identification afforded by randomization while leveraging observational data purely for variance reduction, offering both theoretical guarantees and substantial empirical improvements over existing methods.

The Architecture of Alignment

To understand the CALM framework, one must first appreciate the structural barrier that covariate mismatch creates. Consider the standard setup where covariates partition into three distinct blocks: Z, the shared variables measured in both sources; U, those observed only in the RCT (such as detailed behavioral assessments or survey instruments); and V, those available only in the observational study (such as comprehensive laboratory panels or longitudinal utilization metrics). Previous approaches, including the mismatch-aware extension of R-OSCAR (MR-OSCAR), addressed this tripartite structure by learning an imputation function that predicts V from Z using observational data, then filling in the missing V block for RCT units.

Asiaee and Pal identify a critical inefficiency in this imputation strategy. The goal of CATE estimation is not to reconstruct the full covariate vector V, but rather to obtain sufficient statistics for predicting outcomes and estimating discrepancy functions. When the outcome-relevant information in V concentrates on a low-dimensional manifold, attempting to recover the entire vector amounts to solving a problem harder than necessary, with error scaling in the dimensionality p_v and the complexity of P(V|Z).

CALM circumvents this by learning embedding functions ϕ_o: R^{p_o} → R^d and ϕ_r: R^{p_r} → R^d that map the distinct feature spaces of the observational study and RCT, respectively, into a shared d-dimensional representation space H. Outcome models are trained within the observational embedding space, then transferred to the RCT embedding space where they undergo a double-calibration procedure. This calibration step adjusts the observational predictions to respect the RCT distribution, ensuring that any systematic biases in the transferred model are corrected using the randomized trial data.

Crucially, this architecture maintains a strict separation of concerns: causal identification rests entirely upon the randomized contrast within the trial, while the observational data enter only through the quality of the counterfactual mean outcome estimation used to augment the pseudo-outcome regression. This preserves the internal validity of the RCT while potentially improving precision through external data integration.

Theoretical Guarantees and Empirical Performance

The paper derives finite-sample risk bounds that illuminate precisely when embedding alignment outperforms imputation. The bounds decompose into four interpretable terms: alignment error, outcome-model complexity, calibration complexity, and CATE-class complexity. When compared to the MR-OSCAR imputation bounds, which include an explicit imputation error term scaling as L²r²_im (where L represents a Lipschitz constant and r²_im the oracle imputation risk), the embedding approach can offer substantial improvements when d << p_v and the outcome-relevant structure is effectively low-rank.

Under sparse linear Gaussian models, the authors establish that CALM strictly generalizes imputation-based approaches, providing a theoretical foundation for the intuition that reconstructing unnecessary covariate dimensions introduces avoidable error. Perhaps more importantly, the linear variant of CALM incorporates protection against negative transfer; when the distributional shift between the observational study and RCT is severe, the framework can detect this misalignment and effectively downweight the external data to prevent degradation of the CATE estimate. The neural embedding variant, while more flexible and powerful in capturing nonlinear heterogeneity, lacks this automatic safeguard and can suffer performance degradation under extreme distributional shift.

The empirical validation spans 51 distinct simulation settings and an IHDP semi-synthetic benchmark, providing a comprehensive examination of the method's operating characteristics. In linear CATE regimes, the calibration-based linear methods perform equivalently, confirming the theoretical predictions. However, in the 22 nonlinear-regime settings tested, the neural embedding variant of CALM achieves superior performance in every instance, often by substantial margins. This uniform dominance across diverse nonlinear scenarios suggests that the representation learning approach successfully captures treatment effect heterogeneity that imputation-based methods miss, particularly when complex interactions between OS-specific covariates and treatment assignment drive the heterogeneity.

Critical Assessment and Broader Implications

While CALM represents a significant methodological advance, several limitations and open questions merit consideration. The effectiveness of the embedding approach hinges upon the assumption that the outcome-relevant information in the OS-only covariates V can be meaningfully compressed and aligned with the shared covariates Z via a low-dimensional mapping. In biomedical contexts where V comprises fundamentally distinct biological measurements (such as genomic markers absent from the RCT), this manifold assumption may fail; the learned embeddings might capture predictive correlations without preserving the causal structure necessary for valid transport of outcome models.

The dichotomy between the linear variant's robustness and the neural variant's vulnerability to severe distributional shift presents a practical dilemma for applied researchers. The very scenarios where neural networks offer the greatest advantage, namely complex nonlinear CATE surfaces, often coincide with settings where the external validity of observational data is most suspect. The lack of negative-transfer protection in the neural variant implies that practitioners must exercise caution when applying CALM to observational studies that differ substantially from their target trial population in terms of demographics, disease severity, or clinical practice patterns.

More broadly, the CALM framework contributes to a growing recognition in causal inference that identification and estimation can be modularized. By anchoring causal identification firmly in the RCT while using observational data purely for efficiency gains, the method sidesteps contentious debates about unconfoundedness assumptions in the observational study. This separation aligns with recent work on data fusion that treats external data as augmentation rather than substitution for randomized evidence.

Future Directions

The introduction of calibrated alignment raises several questions for future research. First, the optimal choice of embedding dimension d remains an open problem; while the theoretical results suggest benefits when d is small relative to the covariate dimensions, adaptive methods for selecting d based on the observed data would enhance practical utility. Second, the current framework assumes a static snapshot of covariates, yet electronic health records typically involve longitudinal measurements with irregular sampling. Extending CALM to handle temporal covariate mismatch, where the observational study captures time-series data while the RCT collects only baseline measurements, would significantly broaden its applicability.

Finally, the interpretability of the learned embeddings presents both a challenge and an opportunity. While the black-box nature of neural embeddings may suffice for prediction, clinical adoption requires understanding which aspects of the high-dimensional observational data drive the treatment effect estimates. Developing post-hoc interpretation methods or constrained embedding spaces that respect known biological pathways could bridge the gap between statistical efficiency and clinical actionable insight.

As precision medicine continues to demand ever more granular understanding of treatment heterogeneity, methods like CALM that navigate the tension between data abundance and causal validity will become increasingly essential. The shift from imputation to alignment reflects a deeper evolution in how we conceptualize the relationship between disparate data sources: rather than forcing heterogeneous studies into a common mold through reconstruction, we learn the transformations that make them commensurable while respecting their distinct structural properties.