pith. machine review for the scientific record. sign in

arxiv: 2603.19186 · v2 · submitted 2026-03-19 · 💻 cs.LG

Recognition: 2 theorem links

· Lean Theorem

Improving RCT-Based CATE Estimation Under Covariate Mismatch via Calibrated Alignment

Amir Asiaee , Samhita Pal
Authors on Pith no claims yet

Pith reviewed 2026-05-15 07:59 UTC · model grok-4.3

classification 💻 cs.LG
keywords CATE estimationcovariate mismatchembedding alignmentcalibrationRCTobservational studiesheterogeneous treatment effectscausal inference
0
0 comments X

The pith

CALM learns embeddings to align mismatched covariates between RCTs and observational studies, transferring and calibrating outcome models to improve CATE estimates without imputation.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper introduces CALM to handle the problem of covariate mismatch when combining randomized controlled trials with observational studies for estimating conditional average treatment effects. Instead of imputing missing covariates, it learns embeddings that project features from each data source into one shared representation space. Outcome models trained on the observational study are then moved into the RCT's embedding space and adjusted using the randomized trial data. This keeps the causal guarantees from randomization intact. The approach comes with risk bounds that break down the error into alignment, model complexity, and calibration parts, and simulations show gains especially when effects are nonlinear.

Core claim

CALM bypasses imputation by learning 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, identifying when embedding alignment outperforms imputation. Under the calibration-based linear variant, the framework provides protection against negative transfer; the neural variant can be vulnerable under severe distributional shift. Under sparse linear models, the embedding approach strictly generalizes imputati

What carries the argument

Learned embeddings that project mismatched covariates into a shared space, followed by transfer of observational outcome models and calibration on RCT data.

If this is right

  • Finite-sample risk bounds identify when embedding alignment outperforms imputation by decomposing total error into alignment, complexity, and calibration terms.
  • The linear calibration variant protects against negative transfer from the observational data.
  • Under sparse linear models the embedding method strictly generalizes imputation.
  • The neural embedding version outperforms imputation in all simulated nonlinear settings.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same alignment-plus-calibration pattern could be tested on other causal tasks that combine randomized and observational sources with partial covariate overlap.
  • If embedding quality can be monitored in practice, the method might reduce the need for complete covariate overlap when pooling data sources.
  • Real-world datasets with known external validation of CATE would provide a direct check on whether simulation advantages hold outside controlled settings.

Load-bearing premise

The mapping to a shared feature space must preserve the true conditional treatment effects, and calibration with randomized data must fully correct any remaining differences between sources.

What would settle it

Apply CALM and a standard imputation baseline to a large RCT with artificially masked covariates; the true CATE is known from the full RCT, so check whether CALM's estimates stay within the derived risk bounds while imputation does not.

Figures

Figures reproduced from arXiv: 2603.19186 by Amir Asiaee, Samhita Pal.

Figure 1
Figure 1. Figure 1: How OS data inform pseudo-outcome construction in [PITH_FULL_IMAGE:figures/full_fig_p004_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: RMSE of CATE estimation across three experimental sweeps. Mean over 20 replicates. In all panels, the blue band groups [PITH_FULL_IMAGE:figures/full_fig_p008_2.png] view at source ↗
Figure 3
Figure 3. Figure 3: RMSE of CATE estimation as a function of intrinsic dimension [PITH_FULL_IMAGE:figures/full_fig_p013_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: RMSE of CATE estimation as a function of RCT sample size [PITH_FULL_IMAGE:figures/full_fig_p014_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: RMSE across outcome model types (linear, quadratic, sinusoidal), with [PITH_FULL_IMAGE:figures/full_fig_p014_5.png] view at source ↗
Figure 6
Figure 6. Figure 6: RMSE as a function of shared covariate proportion [PITH_FULL_IMAGE:figures/full_fig_p015_6.png] view at source ↗
Figure 7
Figure 7. Figure 7: RMSE under increasing outcome shift between OS and RCT, with [PITH_FULL_IMAGE:figures/full_fig_p016_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: RMSE as a function of shared-covariate signal weight [PITH_FULL_IMAGE:figures/full_fig_p017_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: RMSE as a function of latent coupling strength [PITH_FULL_IMAGE:figures/full_fig_p018_9.png] view at source ↗
Figure 10
Figure 10. Figure 10: RMSE across three nonlinear CATE functional forms (sinusoidal, quadratic, absolute value) in the nonlinear-CATE DGP [PITH_FULL_IMAGE:figures/full_fig_p018_10.png] view at source ↗
Figure 11
Figure 11. Figure 11: Semi-synthetic benchmark using IHDP covariates: mean RMSE of CATE estimation over 50 replicates (error bars: [PITH_FULL_IMAGE:figures/full_fig_p019_11.png] view at source ↗
read the original abstract

Randomized controlled trials (RCTs) are the gold standard for estimating heterogeneous 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 bypasses imputation by learning 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, identifying when embedding alignment outperforms imputation. Under the calibration-based linear variant, the framework provides protection against negative transfer; the neural variant can be vulnerable under severe distributional shift. Under sparse linear models, the embedding approach strictly generalizes imputation. Simulations across 51 settings confirm that (i) calibration-based methods are equivalent for linear CATEs, and (ii) the neural embedding variant wins all 22 nonlinear-regime settings with large margins.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 1 minor

Summary. The manuscript proposes CALM (Calibrated ALignment under covariate Mismatch) to improve CATE estimation from underpowered RCTs by aligning embeddings from observational studies with partially overlapping covariates. OS outcome models are transferred into the RCT embedding space and calibrated on trial data, with finite-sample risk bounds decomposing into alignment error, outcome-model complexity, and calibration complexity. The linear variant claims protection against negative transfer, while simulations across 51 settings show equivalence for linear CATEs and large gains for the neural variant in all 22 nonlinear regimes.

Significance. If the embedding alignment preserves the conditional treatment effect mapping, the method offers a practical alternative to imputation for combining RCT and OS data. The explicit risk decomposition and the linear variant's negative-transfer protection are notable strengths, as are the broad simulation results. The work could influence causal ML practice if the CATE-preservation property is established more rigorously.

major comments (2)
  1. [Finite-sample risk bounds] The finite-sample risk bounds (abstract) decompose risk additively into alignment error + outcome complexity + calibration complexity without interaction terms between alignment and the treatment-by-covariate surface. When partially overlapping covariates contain source-specific treatment interactions, any marginal or joint distribution-matching embedding can rotate or collapse those interactions; the subsequent RCT calibration then lacks the lost information, undermining the claim that randomization supplies the correct conditional expectation after alignment.
  2. [Abstract] The abstract states that under sparse linear models the embedding approach strictly generalizes imputation, yet no derivation or explicit comparison of the alignment objective versus imputation is provided to show how this generalization holds without distorting the CATE surface.
minor comments (1)
  1. [Abstract] The abstract reports 51 simulation settings with wins in nonlinear regimes but does not list the specific ranges of covariate overlap, shift severity, or interaction strength used, making it difficult to assess coverage of the regime where alignment may distort CATE.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive and detailed comments. We address each major point below and outline the revisions we will make to strengthen the manuscript.

read point-by-point responses

  1. Referee: [Finite-sample risk bounds] The finite-sample risk bounds (abstract) decompose risk additively into alignment error + outcome complexity + calibration complexity without interaction terms between alignment and the treatment-by-covariate surface. When partially overlapping covariates contain source-specific treatment interactions, any marginal or joint distribution-matching embedding can rotate or collapse those interactions; the subsequent RCT calibration then lacks the lost information, undermining the claim that randomization supplies the correct conditional expectation after alignment.

    Authors: We appreciate the referee highlighting this subtlety in the risk decomposition. Our bounds treat alignment error as the primary term capturing any distortion of the conditional treatment effect surface, including loss of source-specific interactions; under the stated assumptions, large alignment error would dominate the bound and correctly signal that the transferred model cannot be reliably calibrated. We agree, however, that an explicit interaction term is absent and that the current presentation does not fully address the case of severe source-specific interactions. We will revise the relevant section and appendix to (i) discuss this scenario explicitly, (ii) clarify that randomization in the RCT guarantees unbiasedness only conditional on the aligned representation, and (iii) add a remark on how such interactions would appear as elevated alignment error. A full extension of the bound with higher-order terms is left for future work but will be noted as a limitation. revision: partial

  2. Referee: [Abstract] The abstract states that under sparse linear models the embedding approach strictly generalizes imputation, yet no derivation or explicit comparison of the alignment objective versus imputation is provided to show how this generalization holds without distorting the CATE surface.

    Authors: We agree that the abstract claim would benefit from an explicit derivation. In the full manuscript (Section 4 and Appendix B), we show that under sparse linear models the alignment objective recovers the imputation estimator as a feasible solution while permitting lower-variance alignments that remain CATE-preserving; the proof proceeds by showing that the population alignment loss is minimized by any embedding that preserves the linear span of the observed covariates, with imputation corresponding to the coordinate-wise completion. To make this transparent to readers, we will add a concise derivation and side-by-side comparison of the two objectives in the revised main text, together with a short proof that the CATE surface is unchanged under the sparsity assumption. revision: yes

Circularity Check

0 steps flagged

No significant circularity in derivation chain

full rationale

The paper introduces CALM for CATE estimation under covariate mismatch via learned embeddings for alignment followed by calibration on RCT data. Finite-sample risk bounds are decomposed into alignment error, outcome-model complexity, and calibration complexity; this is a standard additive risk decomposition rather than a reduction of the target quantity to its inputs by construction. No equations or steps are shown where a prediction equals a fitted parameter, where an embedding is defined circularly in terms of the CATE it is meant to preserve, or where uniqueness is imported solely via self-citation. The protection against negative transfer in the linear variant follows from the explicit calibration step on randomized data, which is an independent modeling choice rather than a tautology. The derivation therefore remains self-contained against external benchmarks and does not meet the criteria for any enumerated circularity pattern.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 0 invented entities

The framework rests on standard causal identification from randomization plus the assumption that embeddings can be learned to preserve relevant conditional distributions. No new invented entities; free parameters appear in the embedding and calibration steps but are not enumerated in the abstract.

axioms (1)
  • domain assumption Randomization in the RCT identifies the CATE conditional on the observed covariates in the RCT space.
    Invoked to preserve causal identification after transfer and calibration.

pith-pipeline@v0.9.0 · 5501 in / 1253 out tokens · 75637 ms · 2026-05-15T07:59:14.501870+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Reference graph

Works this paper leans on

16 extracted references · 16 canonical work pages

  1. [1]

    Amir Asiaee, Chiara Di Gravio, Cole Beck, Yuting Mei, Samhita Pal, and Jared D. Huling. Improving precision of RCT-based CATE estimation using data borrowing with double calibration.arXiv preprint arXiv:2306.17478,

    work page arXiv

  2. [2]

    Adaptive combination of randomized and observational data.arXiv preprint arXiv:2111.15012,

    David Cheng and Tianxi Cai. Adaptive combination of randomized and observational data.arXiv preprint arXiv:2111.15012,

    work page arXiv

  3. [3]

    Bénédicte Colnet, Imke Mayer, Guanhua Chen, Awa Dieng, Ruohan Li, Gaël Varoquaux, Jean-Philippe Vert, Julie Josse, and Shu Yang

    doi: 10.1515/jci-2021-0059. Bénédicte Colnet, Imke Mayer, Guanhua Chen, Awa Dieng, Ruohan Li, Gaël Varoquaux, Jean-Philippe Vert, Julie Josse, and Shu Yang. Causal inference methods for com- bining randomized trials and observational studies: a re- view.Statistical Science, 39(1):165–191,

    work page doi:10.1515/jci-2021-0059 2021

  4. [4]

    Dahabreh, Sarah E

    Issa J. Dahabreh, Sarah E. Robertson, Jon A. Steingrimsson, Elizabeth A. Stuart, and Miguel A. Hernán. Extending inferences from a randomized trial to a new target popu- lation.Statistics in Medicine, 39(14):1999–2014,

    work page 1999

  5. [5]

    Irina Degtiar and Sherri Rose

    doi: 10.1002/sim.8426. Irina Degtiar and Sherri Rose. A review of generalizability and transportability.Annual Review of Statistics and Its Application, 10:501–524,

    work page doi:10.1002/sim.8426

  6. [6]

    Arthur Gretton, Karsten M

    doi: 10.5705/ss.202018.0416. Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test.Journal of Machine Learning Research, 13:723–773,

    work page doi:10.5705/ss.202018.0416

  7. [7]

    Heerman, Russell L

    William J. Heerman, Russell L. Rothman, Lee M. Sanders, Jonathan S. Schildcrout, Kori B. Flower, Alan M. De- lamater, Melissa C. Kay, Charles T. Wood, Rachel S. Gross, Aihua Bian, Laura E. Adams, Evan C. Sommer, H. Shonna Yin, and Eliana M. Perrin. A digital health behavior intervention to prevent childhood obesity: The Greenlight Plus randomized clinical...

    work page 2068

  8. [8]

    2010.08162

    doi: 10.1198/jcgs. 2010.08162. Brian P. Hobbs, Bradley P. Carlin, Sumithra J. Mandrekar, and Daniel J. Sargent. Hierarchical commensurate and power prior models for adaptive incorporation of his- torical information in clinical trials.Biometrics, 67(3): 1047–1056,

    work page doi:10.1198/jcgs 2010

  9. [9]

    Da- habreh, and Jesse H

    Rickard Karlsson, Piersilvio De Bartolomeis, Issa J. Da- habreh, and Jesse H. Krijthe. Robust estimation of hetero- geneous treatment effects in randomized trials leveraging external data.arXiv preprint arXiv:2507.03681,

    work page arXiv

  10. [10]

    Understand- ing the risks and rewards of combining unbiased and possibly biased estimators, with applications to causal inference.arXiv preprint arXiv:2205.10467,

    Michael Oberst, Alexander D’Amour, Minmin Chen, Yuyan Wang, David Sontag, and Steve Yadlowsky. Understand- ing the risks and rewards of combining unbiased and possibly biased estimators, with applications to causal inference.arXiv preprint arXiv:2205.10467,

    work page arXiv

  11. [11]

    Huling, and Amir Asiaee

    Samhita Pal, Jared D. Huling, and Amir Asiaee. Improving RCT-based CATE estimation under covariate mismatch via double calibration.arXiv preprint arXiv:2603.17066,

    work page arXiv

  12. [12]

    Rosenstein, Zvika Marx, Leslie Pack Kaelbling, and Thomas G

    Michael T. Rosenstein, Zvika Marx, Leslie Pack Kaelbling, and Thomas G. Dietterich. To transfer or not to transfer. InNIPS 2005 Workshop on Transfer Learning,

    work page 2005

  13. [13]

    Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang

    doi: 10.1186/s40537-016-0043-6. Liuyi Yao, Sheng Li, Yaliang Li, Mengdi Huai, Jing Gao, and Aidong Zhang. Representation learning for treatment effect estimation from observational data. InAdvances in Neural Information Processing Systems, volume 31,

    work page doi:10.1186/s40537-016-0043-6

  14. [14]

    Step 1 (CATE calibration layer).By the pseudo-outcome regression framework (Asiaee et al

    We decompose the error in three steps: the CATE calibration layer, the augmentation error, and the assembly. Step 1 (CATE calibration layer).By the pseudo-outcome regression framework (Asiaee et al. [2023], Theorem 6), the CATE estimation error satisfies ∆2 2(ˆτCALM , τ r)≤∆ 2 2(F, τ r) +C 1 1 +P a ∆2 2,r ˆµcal a ,¯µr a R2 nr(F) +C 2 log(1/γ) nr ,(17) whe...

    work page 2023

  15. [15]

    The first term is the approximation error of F; the second captures how the quality of the CMO estimate amplifies the CATE estimation error. Cross-fitting ensures that the nuisance estimates ˆµcal a are independent of the pseudo-outcome regression sample, allowing us to bound the stochastic term via Rademacher complexity and a concentration inequality. St...

    work page 2006

  16. [16]

    Mean over 20 replicates

    The DGP uses pz = 30, pu = 10, pv = 20, linear outcomes, and default parameters nr = 500, no = 10,000, σ2 V = 1.0, dtrue = 5, shift 2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0 Intrinsic Dimension (d_true) 1.0 1.5 2.0 2.5 3.0RMSE of CATE Calibration group (RACER, SR/MR-OSCAR, CALM-Lin) Naive CALM-NN HTCE-T HTCE-DR Figure 3: RMSE of CATE estimation as a function o...

    work page 2011