Model Averaging for Heterogeneous Treatment Effects via Proximity Matching

Abstract

Accurate estimation of heterogeneous treatment effects (HTEs) serves as a cornerstone of personalized decision-making, especially in observational studies where treatment assignment is not randomized. However, the presence of confounding and complex covariate structures poses significant challenges to reliable inference. In this study, we develop an innovative model averaging framework, which leverages proximity-based matching to enhance the accuracy of HTE estimation. The method constructs pseudo-outcomes via proximity score matching and subsequently applies an optimal model averaging procedure to these matched samples. We demonstrate that the proposed estimator achieves asymptotic optimality when the standard regularity conditions are met. Simulation studies, adapted from benchmark settings for evaluating HTE model averaging, confirm its superior finite-sample performance. Compared to standard HTE estimation approaches, the proposed method achieves consistently lower estimation errors and reduced variability. The method is further validated on a clinical dataset from the CPCRA trial, demonstrating its practical value for individualized causal inference.

1. Introduction

Identifying how binary interventions influence outcomes represents a central challenge in both statistics and econometrics, with extensive applications across fields such as economics, healthcare, marketing, and policy analysis. Classical examples include the evaluation of job training programs, college education, and medical interventions [1,2,3]. In empirical research, it is often essential to move beyond estimating the average treatment effect (ATE) and explore how treatment responses vary among individuals—referred to as heterogeneous treatment effects (HTEs). Understanding HTE is critical for personalized interventions, policy targeting, and efficient resource allocation [4,5].

Traditional approaches to causal inference, such as linear regression or difference-in-means estimators, typically assume constant treatment effects and rely on strict parametric conditions. These limitations have motivated the development of flexible, data-driven methods for HTE estimation. Representative examples include Bayesian Additive Regression Trees (BART) [6,7], causal trees [8], and causal forests [5]. While these methods offer greater flexibility and improved detection of effect heterogeneity, they also introduce new challenges, such as model uncertainty, sensitivity to model specification, and potential overfitting in finite samples.

Averaging across multiple candidate models has become a well-established approach to mitigate the risks of model uncertainty, as opposed to depending on a single model choice. Early studies in this area emphasized Bayesian Model Averaging (BMA), which utilizes Bayes’ theorem to compute posterior weights for candidate models and enables statistically consistent inference when appropriate priors are specified [9,10,11]. In addition to Bayesian methods, numerous model averaging techniques have been proposed from a frequentist perspective. A number of these approaches incorporate smoothed likelihood-based criteria, including SAIC, SBIC, and SFIC [12,13,14].

To further improve predictive performance and robustness, adaptive model averaging methods have also been proposed. Notable examples include Adaptive Regression by Mixing (ARM) and its L1-penalized variant (L1-ARM), which estimate model weights directly from data with theoretical guarantees [15,16,17]. Other methods aim to minimize prediction risk, including Mallows Model Averaging (MMA) [18,19], Jackknife Model Averaging (JMA) [20], KL-divergence–based methods [21,22], and Parsimonious Model Averaging (PMA) [23].

Although much of the existing research emphasizes the estimation of overall treatment effects, recent developments have adapted model averaging techniques to the study of heterogeneous effects. In observational settings, Seng and Li [24] proposed a model averaging method that incorporates instruments into a two-stage estimation routine for treatment effect inference. A method for model averaging that seeks to minimize the approximate mean squared error for semiparametric estimators was developed by [25]. More directly related to HTE, Rolling et al. [26] constructed weights via nearest neighbor matching and derived finite-sample risk upper bounds, whereas Zhao et al. [27] also utilized nearest neighbor matching but focused on achieving asymptotic optimality of the CATE estimator. These developments highlight the growing interest in integrating model averaging with HTE estimation. In essence, model averaging provides a principled way to integrate the uncertainty across multiple HTE estimators, thereby yielding more robust and accurate treatment effect estimates in observational studies.

An essential stream of causal inference research involves matching techniques, designed to lessen confounding through the alignment of treated and untreated observations exhibiting analogous covariate patterns. Traditional approaches include exact matching, Mahalanobis distance matching [28,29], and the widely used propensity score matching [30,31]. These methods construct matched samples that approximate randomized experiments by balancing covariates between groups. Subsequent developments such as prognostic score matching [32] and full matching [33,34] further improve finite-sample performance and covariate balance.

Recent research has shifted toward data-adaptive and machine learning–enhanced matching strategies. For example, Gao et al. [35] proposed proximity score matching based on random forests to measure unit similarity. Their combo (Proximity score distance with FACT matching)method solves an optimal assignment problem using proximity distances, yielding better covariate balance and improved estimation efficiency compared to classical matching algorithms. These innovations pave the way for adaptive matching methods in the context of HTE estimation.

Despite these advances, the integration of matching strategies and model averaging techniques remains theoretically underdeveloped. In particular, few studies have rigorously examined the optimality properties of model averaging estimators based on matched samples. Bridging this gap is crucial for constructing robust procedures that leverage both model-based and design-based inference principles.

In this study, a new methodological framework is introduced to estimate HTEs through the integration of proximity matching and model averaging. We adopt the proximity matching approach of [35] to construct high-quality pseudo-observations and integrate it with a penalized model averaging procedure tailored to the matched sample. We construct an optimal weight selection rule aimed at reducing an asymptotically unbiased estimate of the expected prediction error, and prove that the resulting estimator achieves asymptotic optimality under standard regularity conditions. Our contributions extend classical model averaging theory to matched-design settings and establish a principled foundation for combining these two powerful paradigms in HTE estimation. Compared to the partition matching strategy used in [27], our method employs proximity score matching, which enables more adaptive unit pairing and demonstrates superior empirical performance in simulation studies.

To evaluate the effectiveness of our approach, we carry out comprehensive simulation experiments adapted from [35], specifically designed to evaluate model averaging estimators. The proposed framework demonstrates consistent advantages over standard alternatives, including selection-based strategies and model averaging methods without proximity integration, in terms of estimation precision and predictive reliability. These findings underscore the practical utility and theoretical soundness of our approach in modern causal inference settings. While prior literature has explored model averaging and matching separately for treatment effect estimation, the integration of these two approaches remains underdeveloped. Our study addresses this gap by proposing a unified framework that leverages proximity-based matching to construct pseudo-observations and applies model averaging to mitigate estimator instability. This joint approach improves robustness to model misspecification and enhances accuracy in heterogeneous treatment effect estimation, especially in observational studies with complex covariate structures.

The contributions of this work are summarized below.

• We introduce a novel HTE estimation framework that combines proximity matching with model averaging, addressing both model uncertainty and matching-induced bias.
• The resulting estimator, constructed under the combo matching framework, is shown to achieve asymptotic optimality through a derived optimal weighting criterion.
• We validate the proposed method through simulation studies, demonstrating superior performance over existing approaches.

The subsequent sections are organized as outlined below. We begin in Section 2 by outlining the proposed approach, which integrates proximity matching with model averaging. Section 3 develops the theoretical results and establishes asymptotic optimality. Section 4 presents the simulation results, while Section 5 illustrates the method with a real-world dataset. Section 6 concludes with a discussion of future directions. Proofs of the theoretical results appear in Appendix A. Appendix B includes the PM-OPT pseudocode along with summaries of key notations and abbreviations for ease of reference.

2. Methodology

This section outlines our approach to estimating HTEs through the combination of proximity matching and model averaging. We first present the structural assumptions and model setup, including the potential outcomes framework and the definition of CATE. Next, we describe the construction of pseudo-observations based on proximity matching, followed by the formulation of multiple candidate models. Finally, we introduce the model averaging procedure and a penalized criterion for weight selection that ensures asymptotic optimality. Together, these components form a unified and robust estimation strategy for HTEs.

2.1. Model Framework and Assumptions

We begin by introducing the general framework and notations. Denote by $(X_i,T_i,Y_i)$ the observed covariates, treatment, and outcome for unit i, with $i=1,\dots ,n$. These units are assumed to be drawn i.i.d. from the population. For each unit, we observe a covariate vector $X_i\in {\mathbb{R}}^p$, a binary treatment indicator $T_i\in \{0,1\}$ (where $T_i\in \{0,1\}$ indicates treatment assignment, with 1 representing treated units and 0 indicating controls), and an outcome variable $Y_i\in \mathbb{R}$.

Under the potential outcomes framework [30], each unit i is associated with a pair of potential responses: $Y_i\left(1\right)$ if treated, and $Y_i\left(0\right)$ if untreated. The observed response is

$$Y_i=T_i{Y}_i\left(1\right)+(1-T_i)Y_i\left(0\right),$$

which implicitly relies on the Consistency assumption.

Assumption 1

(Consistency). For each unit i, the observed outcome equals the potential outcome under the received treatment: $Y_i=Y_i\left(T_i\right)$.

This assumption presumes no interference between units and a well-defined treatment assignment, commonly referred to as the Stable Unit Treatment Value Assumption (SUTVA).

The primary estimand of interest is the CATE

$$\tau \left(x\right)=\mathbb{E}\left[Y\right(1)-Y(0)\mid X=x],$$

which quantifies how the causal effect varies with covariates. We estimate $f_1\left(x\right)$ and $f_0\left(x\right)$ separately on the treated and control samples, respectively, using linear regression models without including the treatment indicator or interaction terms in the design matrices. The HTE function $\tau \left(x\right)$ is defined as the difference between the two fitted functions.

To nonparametrically identify $\tau \left(x\right)$ from observational data, we impose the following standard assumptions.

Assumption 2

(Unconfoundedness).

$$\left(Y\right(1),Y(0\left)\right)\perp T\mid X.$$

Assumption 3

(Overlap). There exists a constant $\eta >0$ such that for all x in the support of X,

$$\eta <\mathbb{P}(T=1\mid X=x)<1-\eta .$$

Assumption 2 rules out unmeasured confounding, while Assumption 3 ensures that both treatment and control groups are sufficiently represented at each covariate value. Together, Assumptions 1–3 guarantee the identifiability of $\tau \left(x\right)$.

In this paper, we further assume a structural model for the potential outcomes of the form

$$Y_i\left(t\right)=f_t\left(X_i\right)+{\epsilon }_i\left(t\right),t\in \{0,1\},$$

(1)

where $f_0\left(x\right)$ and $f_1\left(x\right)$ are unknown regression functions, and ${\epsilon }_i\left(t\right)$ is a mean-zero error term with finite variance.

Assumption 4

(Error Structure).

(i)

$\mathbb{E}[{\epsilon }_i\left(t\right)\mid X_i]=0$ for $t\in \{0,1\}$;

(ii)

$\mathit{Var}({\epsilon }_i\left(t\right)\mid X_i)={\sigma }_t^2<\infty$;

(iii)

The error terms ${\epsilon }_i\left(0\right)$ and ${\epsilon }_j\left(1\right)$ are independent for $i\ne j$.

Assumption 4 specifies a standard error structure often adopted in causal inference literature. Part (i) assumes that the error terms have zero conditional mean given the covariates, ensuring that the systematic component $f_t\left(X_i\right)$ captures all predictable variation in potential outcomes. Part (ii) imposes that the error variances are finite and may differ between treatment groups, which allows for heteroskedasticity but ensures well-behaved inference. Part (iii) assumes that the error terms are independent across units, meaning that ${\epsilon }_i\left(0\right)$ and ${\epsilon }_j\left(1\right)$ are independent for $i\ne j$. This cross-unit independence is typical in cross-sectional settings and supports valid asymptotic analysis. We do not require ${\epsilon }_i\left(0\right)$ and ${\epsilon }_i\left(1\right)$ to be independent, as only one potential outcome is observed per unit. This pruning strategy ensures that the resulting matched graph consists of small connected components, each containing at most one treated (or control) unit. As a result, the matched residuals are unlikely to exhibit cross-unit dependence, supporting the plausibility of the independence assumption in Assumption 4 (iii); see [35] for related discussion.

Given this structure, the observed outcome model becomes

$$Y_i=T_i{f}_1\left(X_i\right)+(1-T_i)f_0\left(X_i\right)+{\epsilon }_i,$$

where ${\epsilon }_i=T_i{\epsilon }_i\left(1\right)+(1-T_i){\epsilon }_i\left(0\right)$.

Under this formulation, the CATE is equivalently expressed as

$$\tau \left(x\right)=f_1\left(x\right)-f_0\left(x\right).$$

This setup forms the foundation for the model averaging estimator developed in the following sections. Our goal is to estimate $\tau \left(x\right)$ by combining multiple candidate models fitted on proximity-matched samples, and to establish theoretical guarantees for its asymptotic performance.

2.2. Candidate Models and Model Averaging Formulation

To estimate the CATE function $\tau \left(x\right)$, we employ a model averaging strategy based on linear candidate models. Each candidate model approximates the potential outcome functions $f_t\left(x\right)=\mathbb{E}[Y\left(t\right)\mid X=x]$ for $t\in \{0,1\}$ using a different subset of covariates. This construction explicitly accounts for model uncertainty in variable selection.

Let there be K candidate models indexed by $k=1,\dots ,K$. Denote by $S_k\subseteq \{1,\dots ,p\}$ the index set of covariates used in model $M_k$, and let $x^{\left(k\right)}\in {\mathbb{R}}^{l_k}$ denote the subvector of $x\in {\mathbb{R}}^p$ corresponding to the covariates in $S_k$, where $l_k=\left|S_k\right|$. Then, model $M_k$ approximates the outcome model as

$$Y_i=\left\{\begin{array}{cc}{f}_1^{\left(k\right)}\left(X_i\right)+{\epsilon }_{i1}^{\left(k\right)}=x_i^{\left(k\right)\top }{\beta }_1^{\left(k\right)}+{\epsilon }_{i1}^{\left(k\right)},\hfill & \mathrm{if}{T}_i=1,\hfill \\ f_0^{\left(k\right)}\left(X_i\right)+{\epsilon }_{i0}^{\left(k\right)}=x_i^{\left(k\right)\top }{\beta }_0^{\left(k\right)}+{\epsilon }_{i0}^{\left(k\right)},\hfill & \mathrm{if}{T}_i=0,\hfill \end{array}\right.$$

where ${\beta }_t^{\left(k\right)}\in {\mathbb{R}}^{l_k}$ are model-specific coefficients for treatment group $t\in \{0,1\}$, and ${\epsilon }_{it}^{\left(k\right)}$ are mean-zero errors. We emphasize that the true functions $f_0\left(x\right)$ and $f_1\left(x\right)$ are not assumed to be linear; instead, we approximate them via linear least squares within each candidate model, following standard practice in model averaging under potential misspecification.

For model $M_k$, we estimate ${\beta }_1^{\left(k\right)}$ and ${\beta }_0^{\left(k\right)}$ using least squares on the treated and control samples, respectively. Let $X_1^{\left(k\right)}\in {\mathbb{R}}^{n_1\times l_k}$ and $X_0^{\left(k\right)}\in {\mathbb{R}}^{n_0\times l_k}$ denote the design matrices of treated and control units under model $M_k$, and ${\mathit{Y}}_1$, ${\mathit{Y}}_0$ the corresponding outcomes. Then the estimators are given by

$$\begin{array}{cc}\hfill {\widehat{\beta }}_1^{\left(k\right)}& ={\left(X_1^{\left(k\right)\top }{X}_1^{\left(k\right)}\right)}^{-1}{X}_1^{\left(k\right)\top }{\mathit{Y}}_1,\hfill \\ \hfill {\widehat{\beta }}_0^{\left(k\right)}& ={\left(X_0^{\left(k\right)\top }{X}_0^{\left(k\right)}\right)}^{-1}{X}_0^{\left(k\right)\top }{\mathit{Y}}_0.\hfill \end{array}$$

At covariate value x, the CATE estimated by model $M_k$ is

$${\widehat{\tau }}^{\left(k\right)}\left(x\right)=x^{\left(k\right)\top }{\widehat{\beta }}_1^{\left(k\right)}-x^{\left(k\right)\top }{\widehat{\beta }}_0^{\left(k\right)}.$$

Let $\mathit{w}={(w_1,\dots ,w_K)}^{\top }$ be model weights satisfying

$$H_n=\left\{\mathit{w}\in {[0,1]}^K|\sum _{k=1}^K{w}_k=1\right\}.$$

We define the model averaging estimator of $\tau \left(x\right)$ as a convex combination of the K candidate estimators

$$\widehat{\tau }\left(\mathit{w}\right)=\sum _{k=1}^K{w}_k{\widehat{\tau }}^{\left(k\right)}\left(x\right)=\sum _{k=1}^K{w}_k\left(x^{\left(k\right)\top }{\widehat{\beta }}_1^{\left(k\right)}-x^{\left(k\right)\top }{\widehat{\beta }}_0^{\left(k\right)}\right).$$

Since $\mathit{w}\in H_n$, the model averaging estimator $\widehat{\tau }\left(\mathit{w}\right)$ represents a convex combination of the candidate estimators ${\widehat{\tau }}^{\left(k\right)}\left(x\right)$, ensuring the final estimator lies within their convex hull. This estimator allows for improved robustness to model misspecification by integrating information across diverse model structures. The subsequent sections introduce our proximity-matching-based strategy for constructing pseudo-observations and describe how we select the weight vector $\mathit{w}$ by minimizing an empirical approximation of the conditional prediction risk.

3. Weight Choice and Asymptotic Optimality

In this section, we present the methodology for determining optimal model averaging weights in the context of proximity-matched pseudo-observations. We begin by describing the construction of matched pairs via the combo matching approach. Building on this, we formulate a risk-based criterion for weight selection tailored to treatment effect estimation. We then rigorously analyze the asymptotic behavior of the resulting estimator and establish its optimality under a set of regularity conditions.

The proposed PM-OPT framework extends conventional model averaging approaches by integrating proximity-based matching to construct pseudo-outcomes, enabling more accurate estimation of heterogeneous treatment effects. Compared to existing methods such as OPT or TEEM, PM-OPT leverages local balance through matching and aggregates predictions across multiple models using a risk-based weighting scheme. This integration not only improves robustness against model misspecification and sampling variability, but also offers a practical pathway for applying model averaging in causal inference settings where high-quality matching is difficult to achieve.

3.1. Proximity Matching and Pseudo-Observations

We adopt the combo matching strategy introduced in Gao et al. [35], which combines full matching constraints with average distance minimization. Compared to propensity score matching and nearest-neighbor matching, combo matching simultaneously achieves global covariate balance and local similarity through a constrained optimization framework based on proximity score distances. The proximity score distance is a model-based measure of similarity between units, computed from an auxiliary predictive model such as random forests or gradient boosting trees. It reflects how often two units fall into the same terminal nodes across trees, capturing nonlinear relationships and interactions among covariates. This allows the matching procedure to account for complex patterns in the covariate space beyond what traditional distances can capture. To enforce balance and reduce estimation variance, the algorithm constructs a matching graph with a path length constraint of at most three. This structure prevents the formation of large connected components and ensures that each unit appears in at most one matched pair. A pruning procedure is applied to iteratively remove redundant or suboptimal edges from the matching graph, ensuring that each unit appears in at most one matched pair while preserving high-quality matches. This design allows matched units to remain in the same validation fold during sample splitting, thereby avoiding the breakdown of pseudo-outcomes and ensuring low-variance estimates. In the context of PM-OPT, combo matching provides stable and informative matched pseudo-observations that are critical for accurately estimating HTEs.

Let L denote the total number of matched treatment-control pairs constructed using the combo algorithm [35]. For each matched pair indexed by $l=1,\dots ,L$, let $(1_l,0_l)$ denote the indices of the treated and control units, respectively. The observed covariates and outcomes are denoted as $(X_{1_l},Y_{1_l})$ and $(X_{0_l},Y_{0_l})$, with $T_{1_l}=1$ and $T_{0_l}=0$.

Under the potential outcomes framework, a pseudo-treatment effect is defined for every matched pair of units

$${\tilde{Y}}_l=Y_{1_l}-Y_{0_l}=\tau \left(X_{1_l}\right)+b_l+{ϵ}_l,$$

where the decomposition includes

• $b_l=f_1\left(X_{1_l}\right)-f_1\left(X_{0_l}\right)$ representing the bias due to imperfect covariate matching;
• ${ϵ}_l={\epsilon }_{1_l}-{\epsilon }_{0_l}$ representing the composite noise from both units.

To ensure the validity of the pseudo-outcomes, it is important to control the bias term $b_l$. Under the assumption that the regression functions $f_t\left(x\right)$ satisfy a smooth continuity condition, and the covariate space is compact, the combo algorithm ensures high-quality matches in the covariate space.

Lemma 1.

Suppose the function $f_0\left(x\right)$ is smooth over a compact domain in the sense that its local variation over matched covariate pairs is bounded by a constant multiple of $\parallel X_{1_l}-X_{0_l}{\parallel }^{\alpha }$ for some $\alpha >0$. Then, under the proximity matching procedure, the matching bias is uniformly bounded as

$$\underset{1\le l\le L}{sup}\left|b_l\right|={\mathcal{O}}_p\left(n^{-\alpha /p}\right),$$

(2)

where p is the covariate dimension and the result holds for sufficiently well-matched pairs.

Proof. 

Under the assumption that $f_0\left(x\right)$ is smooth continuous with exponent $\alpha$, we have

$$|f_0\left(X_{1_l}\right)-f_0\left(X_{0_l}\right)|\le C\parallel X_{1_l}-X_{0_l}{\parallel }^{\alpha }.$$

According to Gao et al. [35], the combo matching algorithm minimizes the total matching cost based on proximity scores from a random forest, which effectively controls the pairwise distance between matched treated and control units. In high-dimensional settings, it is well established (e.g., [3]) that proximity-based matching yields a rate $\parallel X_{1_l}-X_{0_l}\parallel ={\mathcal{O}}_p\left(n^{-1/p}\right)$. Combining this with the smooth condition yields

$$|b_l|=|f_0\left(X_{1_l}\right)-f_0\left(X_{0_l}\right)|={\mathcal{O}}_p\left(n^{-\alpha /p}\right),$$

uniformly over all matched pairs $l=1,\dots ,L$.    □

This result implies that, under smoothness conditions and sufficiently accurate matching, the deviation of the pseudo-outcome ${\tilde{Y}}_l$ from the true conditional average treatment effect $\tau \left(X_{i_l}\right)$ decays at a polynomial rate as the sample size n increases.

Remark 1.

Although prior empirical evidence ([35]) indicates that combo matching yields low estimation bias, our result formalizes this observation by linking the bias to smoothness conditions on the outcome function and deriving an explicit convergence rate.

3.2. Weight Choice Criterion

Given the matched treatment-control pairs ${\left\{(1_l,0_l)\right\}}_{l=1}^L$ constructed using the combo matching algorithm, define the pseudo-outcome for each pair as

$${\tilde{Y}}_l=Y_{1_l}-Y_{0_l},l=1,\dots ,L,$$

and let $\tilde{\mathit{Y}}={({\tilde{Y}}_1,\dots ,{\tilde{Y}}_L)}^{\top }\in {\mathbb{R}}^L$ be the vector of all pseudo-treatment effects.

For each candidate model $M_k$, define the vector of predicted treatment effects evaluated at the treated units $\left\{1_l\right\}$ as

$${\widehat{\mathit{\tau }}}_k={\left({\widehat{\tau }}_k\left(X_{1_1}\right),\dots ,{\widehat{\tau }}_k\left(X_{1_L}\right)\right)}^{\top }.$$

Let $X_1^{\left(k\right)}\in {\mathbb{R}}^{n_1\times l_k}$ and $X_0^{\left(k\right)}\in {\mathbb{R}}^{n_0\times l_k}$ denote the full design matrices of the treated and control units under model $M_k$, respectively. Let ${\tilde{X}}_1^{\left(k\right)}\in {\mathbb{R}}^{L\times l_k}$ and ${\tilde{X}}_0^{\left(k\right)}\in {\mathbb{R}}^{L\times l_k}$ be the corresponding submatrices for the matched pairs, where the rows correspond to the units $\left\{1_l\right\}$ and $\left\{0_l\right\}$, respectively.

We define the individual projection matrices as

$$P_1^{\left(k\right)}={\tilde{X}}_1^{\left(k\right)}{\left(X_1^{\left(k\right)\top }{X}_1^{\left(k\right)}\right)}^{-1}{X}_1^{\left(k\right)\top },P_0^{\left(k\right)}={\tilde{X}}_1^{\left(k\right)}{\left(X_0^{\left(k\right)\top }{X}_0^{\left(k\right)}\right)}^{-1}{X}_0^{\left(k\right)\top }.$$

The model-averaged projection matrices are defined as

$${\overline{P}}_1\left(\mathit{w}\right)=\sum _{k=1}^K{w}_k{P}_1^{\left(k\right)},{\overline{P}}_0\left(\mathit{w}\right)=\sum _{k=1}^K{w}_k{P}_0^{\left(k\right)},$$

where $\mathit{w}={(w_1,\dots ,w_K)}^{\top }$ denotes the model averaging weights.

To align dimensions with the full length L, we define the zero-padded projection matrices as

$$P_{1,0}^{\left(k\right)}={\left({\left(P_1^{\left(k\right)}\right)}^{\top },{\mathbf{0}}_{n_1\times (L-n_1)}\right)}^{\top },P_{0,0}^{\left(k\right)}={\left({\left(P_0^{\left(k\right)}\right)}^{\top },{\mathbf{0}}_{n_0\times (L-n_0)}\right)}^{\top },$$

and their corresponding model-averaged forms are

$$P_1\left(\mathit{w}\right)=\sum _{k=1}^K{w}_k{P}_{1,0}^{\left(k\right)},P_0\left(\mathit{w}\right)=\sum _{k=1}^K{w}_k{P}_{0,0}^{\left(k\right)}.$$

We define the penalized criterion as

$$C_n\left(\mathit{w}\right)={∥\widehat{\mathit{\tau }}\left(\mathit{w}\right)-\tilde{\mathit{Y}}∥}^2+2{\sigma }_1^2\mathrm{tr}\left(P_1\left(\mathit{w}\right)\right)+2{\sigma }_0^2\mathrm{tr}\left(P_0\left(\mathit{w}\right)\right),$$

(3)

where ${\sigma }_1^2$ and ${\sigma }_0^2$ are the noise variances for the treated and control groups, respectively.

The optimal weight is obtained by

$$\widehat{\mathit{w}}=arg\underset{\mathit{w}\in H_n}{min}{C}_n\left(\mathit{w}\right).$$

Let $\mathit{\tau }={\left(\tau \left(X_{i_1}\right),\dots ,\tau \left(X_{i_L}\right)\right)}^{\top }$ denote the corresponding vector of true CATEs, and define the squared error loss and conditional risk of $\widehat{\mathit{\tau }}\left(\mathit{w}\right)$ as

$${\mathcal{L}}_n\left(\mathit{w}\right)={∥\widehat{\mathit{\tau }}\left(\mathit{w}\right)-\mathit{\tau }∥}^2,R_n\left(\mathit{w}\right)=\mathbb{E}\left[{\mathcal{L}}_n\left(\mathit{w}\right)\mid \mathcal{X}\right].$$

where $\mathcal{X}$ denotes the observed covariates used in the projection matrices. We determine the weight vector $\mathit{w}$ by minimizing an empirical approximation of the risk function $R_n\left(\mathit{w}\right)$.

Theorem 1

(Asymptotic Unbiasedness). Provided that ${sup}_l\left|b_l\right|={\mathcal{O}}_p(1/logn)$. Then, the penalized criterion satisfies

$$\mathbb{E}(C_n\left(\mathit{w}\right)\mid \mathcal{X})=R_n\left(\mathit{w}\right)(1+o\left(1\right))+L({\sigma }_0^2+{\sigma }_1^2),a.s.$$

uniformly over $\mathit{w}\in H_n$.

Proof. 

See Appendix A.1.    □

Remark 2.

This result shows that the weight selection criterion $C_n\left(\mathit{w}\right)$ serves as an asymptotically unbiased estimator of the conditional risk $R_n\left(\mathit{w}\right)$ up to a constant. The combo-matching framework ensures that the approximation error $b_l$ is sufficiently small, and the projection-based penalty terms control the model complexity.

3.3. Asymptotic Optimality

In this section, we show that our proposed model averaging estimator under combo matching enjoys an asymptotic optimality property. Unless otherwise stated, all limiting processes discussed in this and subsequent sections are with respect to $n\to \infty$.

Now, write ${\mathit{f}}_1={\left(f_1\left(X_{1_1}\right),\dots ,f_1\left(X_{1_L}\right)\right)}^{\top },{\mathit{f}}_0={\left(f_0\left(X_{0_1}\right),\dots ,f_0\left(X_{0_L}\right)\right)}^{\top },{\widehat{\mathit{f}}}_1\left(\mathit{w}\right)={\overline{P}}_1\left(\mathit{w}\right){\mathit{Y}}_1$, and let ${\widehat{\mathit{f}}}_0\left(\mathit{w}\right)={\overline{P}}_0\left(\mathit{w}\right){\mathit{Y}}_0$ be the model-averaged predictions based on the treatment and control arms respectively, where ${\overline{P}}_1\left(\mathit{w}\right)={\sum }_{k=1}^K{w}_k{P}_1^{\left(k\right)}$ and ${\overline{P}}_0\left(\mathit{w}\right)={\sum }_{k=1}^K{w}_k{P}_0^{\left(k\right)}$.

We define the squared error and the conditional risk as

$${\mathcal{L}}_{n_1}\left(\mathit{w}\right)={∥{\widehat{\mathit{f}}}_1\left(\mathit{w}\right)-{\mathit{f}}_1∥}^2,R_{n_1}\left(\mathit{w}\right)=\mathbb{E}\left[{\mathcal{L}}_{n_1}\left(\mathit{w}\right)\mid \mathcal{X}\right],$$

$${\mathcal{L}}_{n_0}\left(\mathit{w}\right)={∥{\widehat{\mathit{f}}}_0\left(\mathit{w}\right)-{\mathit{f}}_0∥}^2,R_{n_0}\left(\mathit{w}\right)=\mathbb{E}\left[{\mathcal{L}}_{n_0}\left(\mathit{w}\right)\mid \mathcal{X}\right].$$

Now, write ${\xi }_n={inf}_{\mathit{w}\in H_n}{R}_n\left(\mathit{w}\right)$. The subsequent theoretical analysis will be based on these conditional risk functions.

Condition 1

(Moment condition). For all matched pairs $l=1,\dots ,L$, and some fixed integer $1\le G<\infty$,

$$\mathbb{E}\left({ϵ}_l^{4G}\mid \mathcal{X}\right)\le \kappa <\infty ,a.s.$$

Condition 1 imposes a uniform bound on the conditional moments of the pseudo-error terms across all matched pairs. This type of moment condition is standard in the matching literature and is similar to condition (7) in [19] and condition (17) in [27]. This condition is readily satisfied when the error terms follow Gaussian or sub-Gaussian distributions, which are common in practice.

Condition 2

(Risk concentration condition). Let ${\mathit{w}}_k^0\in {\mathbb{R}}^K$ denote the unit vector with 1 at the kth position. And,

$$K{\xi }_n^{-2G}\sum _{k=1}^K\left[{\left(R_{n_1}\left({\mathit{w}}_k^0\right)\right)}^G+{\left(R_{n_0}\left({\mathit{w}}_k^0\right)\right)}^G\right]\stackrel{a.s.}{\to }0.$$

Condition 2 resembles Condition (8) in [19] and is also related to the bias-variance trade-off condition discussed in [18,27]. This assumes that the true model is not exactly contained in the candidate model set, which is a standard and realistic assumption in the model averaging literature. In practice, all candidate models are usually simplified approximations of the data-generating process, especially when the model space is of moderate size and does not exhaust all possible higher-order interactions. This condition ensures that the approximation bias does not vanish entirely, aligning with the setup commonly adopted in nonparametric regression frameworks (see [18]).

Condition 3

(Minimum eigenvalue condition). For each $k=1,\dots ,K$, there exist constants $C_0>0$ such that

$$C_0\le {\lambda }_{min}\left({\displaystyle \frac{X_0^{\left(k\right)\top }{X}_0^{\left(k\right)}}{n_0}}\right),a.s.,$$

where ${\lambda }_{min}(·)$ denotes the smallest eigenvalue of a symmetric matrix.

Condition 3 is a minimum eigenvalue condition, ensuring sufficient identifiability of the covariate design matrix within each candidate model. This requirement aligns with condition C.4 of [22] and condition 1 of [36], and is necessary for ensuring the stability of least squares estimators and the convergence of projection matrices. This condition is generally easy to satisfy in practice, as it only requires that the covariates are not perfectly collinear and have reasonable variation, which is common in empirical data sets.

Condition 4

(Growth condition).

$$L{n}^{-\alpha /p}{\xi }_n^{-1}=o\left(1\right).$$

Condition 4 restricts the relative growth rates of the matched sample size L, the number of covariates p, and the bias decay rate parameter ${\zeta }_n$. Such a condition is commonly required in matched sampling frameworks to ensure that the approximation bias induced by matching does not dominate the overall estimation error. In our setting, it reflects a balance between the smoothness of the regression functions and the size of the matched sample. For instance, when the bias decay rate satisfies ${\xi }_n=\mathcal{O}\left(n^{1-r}\right)$ for some $0<r\le 1$, Condition 4 can be satisfied by choosing $L=\mathcal{O}\left(n^{\theta }\right)$ with sufficiently small $\theta$, especially when the regression functions are smooth (i.e., large $\alpha$).

Theorem 2

(Asymptotic Optimality). Suppose Assumptions 1–4 and Conditons 1–4 hold, and the number of matched pairs satisfies $L\to \infty$ as $n\to \infty$. Then the model averaging estimator $\widehat{\mathit{\tau }}\left(\widehat{\mathit{w}}\right)$ satisfies

$${\displaystyle \frac{{\mathcal{L}}_n\left(\widehat{\mathit{w}}\right)}{{inf}_{\mathit{w}\in H_n}{\mathcal{L}}_n\left(\mathit{w}\right)}}\stackrel{p}{\to }1,$$

i.e., the model averaging estimator achieves asymptotic optimality among all convex combinations of the K candidate models.

Proof. 

See Appendix A.2.    □

Remark 3.

When ${\sigma }_1^2$ and ${\sigma }_0^2$ are replaced by consistent estimators ${\widehat{\sigma }}_1^2$ and ${\widehat{\sigma }}_0^2$, Theorem 2 remains valid under mild regularity conditions. The literature provides similar insights, for instance in Hansen [18] and Zhao et al. [27].

4. Numerical Experiments

This section evaluates the effectiveness of Proximity Matching OPT estimator (“Proximity Matching OPT” (PM-OPT) denotes the proximity–matching criterion $C_n\left(W\right)$ with the OPT weight solution). against several competing model averaging methods under a comprehensive set of simulation scenarios. All simulations are implemented in R, and each setting is replicated 1000 times to ensure stable Monte Carlo error. The evaluation focuses on each method’s ability to estimate HTEs across multiple experimental designs, differing in treatment assignment mechanisms, functional forms of effect heterogeneity, signal-to-noise ratios (SNR), and sample sizes. The code used for the numerical results section is available on our GitHub repository at https://github.com/zhzhao07 (accessed on 4 August 2025).

4.1. Simulation Design

We conduct a series of simulation experiments to systematically compare the performance of various model selection and model averaging methods for estimating HTEs. Specifically, we evaluate classical criteria such as the Akaike Information Criterion (AIC) and Bayesian Information Criterion (BIC); their smoothed model averaging extensions, SAIC and SBIC [12]; treatment-effect-oriented criteria including the Treatment Effect Cross-Validation (TECV) criterion [37] and the Treatment Effect Estimation by Mixing (TEEM) method [26]; the asymptotically optimal averaging method OPT [27]; and the proposed Proximity-Matching-based Optimal Averaging method, PM–OPT. The simulation designs are crafted to reflect realistic challenges in observational data, including model misspecification, covariate imbalance, and treatment effect heterogeneity. In our study, the parameters for the Proximity matching and Prognostic score distance were determined by optimizing the out-of-bag (OOB) mean-squared error of the random forest. These settings follow the methodology outlined by [35]. Specifically, we used the same settings for the number of trees and the number of variables randomly sampled at each split in the random forest to compute the proximity scores. This ensures consistency with previous work and reliable comparisons. For the other model estimation processes, no additional tuning parameters were adjusted, as we focused on maintaining consistency and comparability across the methods.

The AIC and BIC are well-established model selection criteria that identify a single best model by trading off goodness-of-fit with model complexity. For the kth candidate model, we define the residual variance as

$${\tilde{\sigma }}_k^2={\displaystyle \frac{1}{n}}\sum _{i=1}^n{(Y_i-{\widehat{f}}_{T_i}^{\left(k\right)})}^2,$$

and compute the scores

$${\mathrm{AIC}}_k=nlog{\tilde{\sigma }}_k^2+2l_k,{\mathrm{BIC}}_k=nlog{\tilde{\sigma }}_k^2+(logn)l_k,$$

$l_k$ denotes the model complexity (i.e., number of parameters) in $M_k$.

To mitigate model selection uncertainty, ref.  [12] proposed SAIC and SBIC estimators, which form model averages rather than choosing a single model. The weight assigned to model k under SAIC is given by

$$w_k={\displaystyle \frac{exp(-\frac{1}{2}{\mathrm{AIC}}_k)}{{\sum }_{k=1}^{K}exp(-\frac{1}{2}{\mathrm{AIC}}_k)}},$$

with an analogous formulation for SBIC based on BIC scores.

Recognizing the limitations of prediction-based criteria in causal inference settings, ref.  [37] introduced the TECV criterion, which selects models based on their accuracy in estimating treatment effects. Building on this, ref. [26] developed the TEEM method, a computationally intensive approach that adaptively averages models to minimize conditional treatment effect estimation error.

We also include the OPT method proposed by [27], which determines averaging weights by minimizing an unbiased estimator of the mean squared error of the CATE estimator. This method enjoys strong asymptotic optimality guarantees under fixed-parameter asymptotics.

Finally, we introduce and evaluate our PM–OPT estimator, which extends the OPT approach by incorporating proximity-score-based matching [35] to construct pseudo-outcomes for both treatment and control groups. This matching step reduces bias from covariate imbalance, and the subsequent model averaging enhances robustness and finite-sample efficiency. Together, this two-stage strategy is designed to combine the strengths of flexible matching and asymptotically optimal model averaging for HTE estimation.

To form the pool of candidate models for averaging, we consider every possible combination of covariates along with their interaction terms with the treatment indicator. Given a dataset consisting of an outcome variable Y, a binary treatment indicator T, and a set of p covariates $(X_1,X_2,\dots ,X_p)$, we rename the covariates as $(u_1,u_2,\dots ,u_p)$ for consistency. The model space is formed by taking the power set of these covariates, which yields $2^p$ subsets. For each subset $\mathcal{S}\subseteq \{u_1,\dots ,u_p\}$, we construct a linear regression model that includes the treatment main effect T, the main effects of the covariates in $\mathcal{S}$, and the interaction terms between T and each covariate in $\mathcal{S}$. Accordingly, each model has the form

$$Y={\beta }_0+{\beta }_{T}T+\sum _{u_j\in \mathcal{S}}{\beta }_j{u}_j+\sum _{u_j\in \mathcal{S}}{\gamma }_j(T·u_j)+\epsilon ,$$

where $\epsilon$ denotes a error term. All candidate models are estimated via ordinary least squares. This full enumeration strategy ensures that the model averaging procedure operates over a rich and structured set of specifications, ranging from simple models with no covariates to complex ones that incorporate both covariate main effects and treatment–covariate interactions. The resulting model space captures a broad spectrum of potential treatment effect heterogeneity structures and facilitates a robust assessment of model averaging performance.

For each replication, we generate five covariates $X={(X_1,\dots ,X_6)}^{\top }$ independently from a uniform distribution on the interval $[-1,1]$. Two treatment assignment mechanisms are considered. In the first case, treatment is assigned completely at random, so that the propensity score is fixed at $e\left(X\right)=0.5$. In the second case, the treatment assignment depends on the first covariate, with the propensity score modeled as $e\left(X\right)=exp\left(X_1\right)/(1+exp\left(X_1\right))$.

To represent different forms of heterogeneity in treatment effects, we consider two types of true CATE. The linear CATE is specified as ${\tau }_{\mathrm{lin}}\left(X\right)=2X_1-1.5X_2+X_3$, while the nonlinear CATE takes the form ${\tau }_{\mathrm{non}}\left(X\right)=X_1^2+sin\left(\pi X_2\right)-0.5X_3^2$.

We define the potential outcomes by $Y\left(0\right)=f_0\left(X\right)+\epsilon$, with the treatment outcome given by $Y\left(1\right)=Y\left(0\right)+\tau \left(X\right)$, where $\epsilon \sim N(0,1)$ is an independent noise term and the baseline outcome function is given by $f_0\left(X\right)=1+0.5X_1-X_2+0.3X_4$.

To control the level of treatment effect signal relative to outcome noise, we follow [27] and define the SNR as the design-level quantity $R^2=\mathrm{Var}(Y-\epsilon )/\mathrm{Var}\left(Y\right).$ We vary $R^2$ over the grid $\{0.1,0.3,0.5,0.7,0.9\}$ to represent different levels of treatment signal strength. In each case, the function $\tau \left(X\right)$ is scaled appropriately to ensure the empirical $R^2$ matches the desired target value.

Finally, we examine the performance of each method under three sample size regimes: $n=100$ (small), $n=200$ (medium), and $n=400$ (moderately large).

Each simulation scenario is repeated for $J=1000$ replications. Estimator accuracy is evaluated based on their associated risk, quantified by the average squared error (ASE)

$$\mathrm{Risk}={\displaystyle \frac{1}{J}}\sum _{j=1}^J{\mathrm{ASE}}_j.$$

To enable fair comparison, we scale the risk by the corresponding value from the optimal least squares estimator (assumed to be infeasible) within each replication and design.

Table 1. Simulation settings.

Design	CATE Form	Propensity e(X)	${\mathit{R}}^2$ Grid
A	Linear	$0.5$	$\{0.1,0.3,0.5,0.7,0.9\}$
B	Linear	$exp\left(X_1\right)/(1+exp{X}_1)$	same
C	Nonlinear	$0.5$	same
D	Nonlinear	$exp\left(X_1\right)/(1+exp{X}_1)$	same

compiles the four designs. Design A adopts the linear ${\tau }_{\mathrm{lin}}$ with random assignment; Design B keeps the linear effect but introduces covariate–dependent assignment; Design C employs the nonlinear ${\tau }_{\mathrm{non}}$ under random assignment; Design D combines ${\tau }_{\mathrm{non}}$ with the logistic propensity score.

Remark 4.

Although we adopt full enumeration of candidate models in the current low-dimensional setting, we acknowledge that this strategy does not scale well when the number of covariates increases. Several practical extensions are available to address this limitation. For moderately high-dimensional settings, nested models based on variable importance rankings or greedy model selection algorithms such as forward or backward selection can be used to generate a manageable model space. In higher-dimensional applications, it is common to apply Lasso-based variable preselection to screen covariates prior to constructing model subsets. In addition, within the context of HTE estimation, the model space can be further simplified by restricting interaction terms. For example, one may consider only main effects and exclude higher-order interactions, which reduces computational burden while maintaining interpretability.

4.2. Results

Following the simulation design and proximity matching setup in [35], we adopt the same covariate structure and distance metrics. Since [35] provides a comprehensive evaluation of the matching quality using standardized mean differences and other diagnostics, we do not replicate those assessments here to avoid redundancy.

Building upon the simulation framework described above, we now turn to the empirical evaluation of various model averaging strategies. The results, summarized in Figure 1, Figure 2, Figure 3 and Figure 4, highlight how each method adapts to changes in the data environment, revealing both the advantages and limitations of classical, smoothed, and targeted model averaging approaches. In what follows, we describe the main findings from each design in turn.

Design A. This scenario features a linear treatment effect under randomized assignment. As shown in Figure 1, all methods perform similarly when the SNR is high and the sample size is moderate to large. However, under small sample sizes ($n=100$) or lower $R^2$, traditional criteria such as AIC and BIC exhibit increased risk relative to more tailored approaches. Smoothed variants (SAIC and SBIC) consistently outperform their unsmoothed counterparts, particularly in low signal settings. The TECV and TEEM estimators demonstrate competitive performance throughout, while the proposed OPT and PM-OPT estimators achieve the lowest risks across nearly all configurations. The benefits of PM-OPT are especially apparent in smaller samples, where matching helps stabilize estimation.

Design B. Introducing covariate-dependent treatment assignment while maintaining linear treatment effects, Design B presents a more realistic observational setting. As illustrated in Figure 2, performance trends mirror those in Design A, but the relative advantage of targeted methods becomes more pronounced. The gap between OPT/PM-OPT and classical methods widens under moderate $R^2$ values, suggesting that accounting for assignment bias via flexible modeling and matching can lead to more accurate estimates. Notably, BIC deteriorates more sharply in small samples, reflecting its tendency toward overly parsimonious models in complex settings.

Design C. This setting imposes a nonlinear treatment effect structure under randomized assignment. The results in Figure 3 show that the performances of classical model averaging approaches (AIC, BIC, SAIC, SBIC) suffers under model misspecification. In contrast, TECV, TEEM, and especially PM-OPT are more robust to this complexity. PM-OPT consistently yields the best performance across all sample sizes and noise levels, owing to its capacity to adaptively weight models while preserving local balance via proximity matching.

Design D. This setting represents the most challenging scenario, combining nonlinear effects with covariate-dependent treatment assignment. Figure 4 confirms that traditional methods fail to capture the underlying heterogeneity, leading to elevated risk. Smoothed methods show modest gains, while TECV and TEEM again demonstrate superior adaptability. OPT remains competitive, but PM-OPT clearly dominates in this setting, underscoring the importance of leveraging both model flexibility and data structure through matching.

Based on the observed performance of the estimators in both simulations and the empirical study, we recommend the following

• Nonlinear treatment effects: The PM–OPT method is particularly effective when treatment effects exhibit significant heterogeneity or nonlinearity. As shown in the results, it outperforms traditional methods such as AIC and BIC in these complex settings.
• Model uncertainty: When there is uncertainty about the correct model specification or variable inclusion, PM–OPT provides more robust and reliable estimates compared to classical criteria, which may underfit or overfit the data.
• Moderate sample sizes: In studies with moderate sample sizes, PM–OPT offers improved accuracy and stability over traditional model selection methods, which may not perform as well under these conditions.
• Observational studies with covariate imbalance: In observational studies, where covariate imbalance is often an issue, PM–OPT is a good choice, as it is designed to handle such imbalances more effectively than traditional model selection methods.

We also note that TECV and TEEM are viable alternatives when the goal is to reduce variability, particularly in simpler models. However, for more accurate and stable estimates, PM–OPT is recommended. Overall, the simulations confirm that while classical model averaging techniques remain serviceable in simple or large-sample settings, their performance degrades in the presence of complex heterogeneity or treatment selection. Smoothed and targeted methods substantially improve estimation accuracy, with PM-OPT providing the most reliable performance under various realistic data-generating conditions.

5. An Empirical Example

To demonstrate the practical utility of our approach, we reanalyze a clinical trial dataset from the Community Programs for Clinical Research on AIDS (CPCRA), which has been examined in prior studies including MacArthur et al. [38], Rolling et al. [26], and Zhao et al. [27]. The trial, referred to as FIRST, compared two antiretroviral treatment strategies in terms of their efficacy in increasing CD4-cell counts among HIV-positive patients.

The dataset consists of $n=1191$ observations, with 799 individuals assigned to the less intensive therapy and 392 to the more intensive alternative. The primary outcome is defined as the change in the square root of CD4-cell counts ($\mathrm{CD}{4}_0$), computed by subtracting the square root of the baseline value from that of the average CD4 count recorded at or beyond 32 months after enrollment. Covariates include $\mathrm{CD}{4}_0$, the log-transformed baseline HIV RNA level (${\mathrm{RNA}}_0$), age, and a binary indicator T denoting treatment assignment.

Following the modeling strategy in [27], we construct candidate regression models that all include an intercept and $\mathrm{CD}{4}_0$ as core covariates, while allowing for flexible inclusion of the remaining covariates and their interactions with treatment T. This results in a diverse collection of models capturing various potential forms of effect heterogeneity.

To assess the performance of competing model averaging approaches in a realistic setting, we conduct a guided simulation study. Specifically, we select several representative models from the candidate set to serve as data-generating processes. For each chosen model, synthetic outcomes are generated by adding Gaussian noise with variance estimated from the model fit to the predicted values, conditional on observed covariates and treatment assignments. This preserves the empirical distribution of the data while allowing repeated evaluation.

We apply each model averaging method to estimate the individual treatment effect $\tau \left({\mathit{x}}_i\right)$ for every unit. The accuracy of each estimator is evaluated using the mean squared error

$${\displaystyle \frac{1}{n}}\sum _{i=1}^n{\left(\tau \left({\mathit{x}}_i\right)-\widehat{\tau }\left({\mathit{x}}_i\right)\right)}^2.$$

Each scenario is replicated 100 times to ensure stable conclusions. This empirical framework provides a rigorous and interpretable benchmark for comparing the effectiveness of model averaging strategies under realistic treatment effect heterogeneity.

Figure 5 reveals that PM-OPT consistently outperforms conventional information criterion-based approaches. In particular, the proposed PM–OPT and OPT methods achieve the lowest median errors and reduced variability, indicating both accuracy and stability. TECV and TEEM also perform competitively, outperforming conventional model selection criteria in most cases.

In contrast, AIC and BIC exhibit higher median risks and wider spreads, especially for BIC, which tends to underfit in complex settings. The smoothed variants SAIC and SBIC yield moderate improvements over their classical counterparts, particularly in reducing variability, but still fall short of the performance attained by the targeted ensemble methods.

Overall, the evidence confirms that incorporating both model flexibility and treatment–covariate interactions, as done by PM–OPT, provides more robust and accurate estimation of HTEs in realistic settings.

To investigate the heterogeneity of treatment effects across baseline covariates, we visualize the estimated CATEs against two clinically relevant variables: age and baseline CD4 count. These variables represent respectively the patient’s immunological status and biological aging, both of which are known to moderate antiretroviral therapy response in HIV patients. We create two scatter plots (Figure 6), where each dot shows one individual. Each dot is colored by treatment group and positioned according to the individual’s estimated CATE.

Figure 6 displays the estimated CATEs as a function of patient age and baseline CD4 count. These CATEs are not derived from synthetic or simulated outcomes, but are instead estimated using the PM-OPT method based on the observed data in the CPCRA cohort. Overall, the CATE estimates are centered around zero, indicating no substantial average benefit or harm from early ART initiation in the full sample. This observation aligns with previous evidence showing that dual- and triple-drug regimens have comparable effects on CD4 counts at the population level ([37,38]).

Nevertheless, subtle heterogeneity is evident. Patients with lower baseline CD4 counts and those of younger age tend to show slightly more positive CATE estimates, suggesting a greater benefit from early treatment in these subgroups. These trends align with clinical intuition that younger individuals and those with weaker immune systems may respond more favorably to intensive ART. While the effects are not sharply pronounced, the PM-OPT estimator appears to capture meaningful variation in treatment response across individuals, underscoring its practical value in real-world data settings where the ground-truth treatment effect is unobservable.

6. Conclusions

This paper introduces a novel framework for estimating HTEs through model averaging combined with proximity-based matching. The proposed estimator, referred to as PMOPT, incorporates proximity information to construct pseudo-outcomes and employs a data-driven weighting scheme across candidate models. This approach enhances predictive accuracy and robustness across diverse scenarios. Comprehensive simulation studies and a guided empirical application show that the proposed method consistently outperforms traditional model selection methods, as well as several recent model averaging techniques, especially in settings characterized by nonlinear effects or treatment selection bias. Based on the numerical findings, PM–OPT consistently outperforms traditional model selection criteria in terms of accuracy and stability. PM–OPT achieves the lowest median errors and reduced variability, particularly in complex settings where treatment effects are heterogeneous or nonlinear. In contrast, AIC and BIC exhibit higher median risks and wider spreads, with BIC tending to underfit in more complex scenarios. While the smoothed variants SAIC and SBIC show moderate improvements in reducing variability, they still fall short of the performance attained by PM–OPT and OPT. It is particularly effective in moderate sample sizes, where model uncertainty is substantial and treatment heterogeneity is expected. These findings underscore the advantages of PM–OPT in handling more complex models and highlight its superior robustness and accuracy when faced with model uncertainty and HTEs.

While the empirical study demonstrates the utility of the proposed method in a realistic setting, we acknowledge some general limitations. First, as noted in Gao et al. [35], matching procedures, particularly those that rely on proximity scores, can become computationally intensive and may perform less reliably when sample sizes are small or the covariate space is high-dimensional. Second, similar to other observational methods, our approach relies on the assumption of no unmeasured confounding, an assumption that cannot be empirically verified. These considerations point to promising directions for future research, such as scalable approximate matching, incorporating domain knowledge, and sensitivity diagnostics, some of which we elaborate on below.

Moreover, although our theoretical framework relies on technical assumptions such as smooth response functions, compact covariate support, and uniform matching quality, these conditions are often only approximately satisfied in real-world applications. Nevertheless, our simulation results demonstrate that the proposed method remains effective even under moderate violations of such assumptions—for example, in scenarios with covariate imbalance or mild non-smoothness in treatment response surfaces.

In addition, as suggested in Gao et al. [35], the matching quality can be improved through domain-informed adjustments. Specifically, one may incorporate exact or near-exact matching constraints on important covariates or modify the proximity metric by penalizing key mismatches. These strategies can help maintain robustness when proximity learning is challenged by small sample sizes or high-dimensional covariate spaces. We regard the integration of such techniques and the development of adaptive matching diagnostics as important directions for future research.

Despite these promising results, our current analysis is restricted to low-dimensional linear regression models with binary treatment assignment. Several important directions remain open for future investigation. First, this paper focuses on low-dimensional settings, where model averaging theory is more mature. Extending PM-OPT to high-dimensional regimes is nontrivial, as the theoretical guarantees and asymptotic properties can differ significantly. Nevertheless, the proposed framework could potentially accommodate dimensionality reduction techniques (e.g., PCA or variable screening) and regularized base learners (e.g., Lasso or Ridge) to improve performance in high-dimensional applications. We leave these extensions as promising directions for future research. Second, adapting the method to generalized linear models, such as logistic or Poisson regression, could allow for broader use in binary or count outcome settings. Third, relaxing the binary treatment assumption to accommodate continuous, multi-valued, or dosage-type treatments would further generalize the framework to a wider range of causal inference problems. These extensions would require new theoretical developments and algorithmic strategies, but they hold the potential to further advance the flexibility and utility of proximity-informed model averaging in HTE estimation.

Finally, while several limitations have been discussed previously, including computational scalability, reliance on the unconfoundedness assumption, and the restriction to low-dimensional linear models, it is also important to note that the current framework assumes fully observed covariates. Extending the methodology to handle missing data, for example by applying imputation techniques or developing robust proximity measures, represents a valuable direction for future research.