scores, Regression.

Causal inference is a cornerstone of modern statistical and machine learning research, with applications spanning economics, social sciences, and biomedical studies. The primary goal of causal inference is to estimate the effect of interventions or treatments on outcomes of interest, often in the presence of confounding variables that can bias results. In observational studies, where treatments are not randomly assigned, addressing confounding is particularly challenging. Propensity score methods, introduced by [1], have become a fundamental tool for reducing bias by balancing covariates between treated and control groups. These methods have been extensively studied and refined, with Austin [2] providing a comprehensive overview of their applications and limitations. Recent advancements in machine learning have significantly enhanced the toolkit for causal inference. Techniques such as Bayesian Additive Regression Trees (BART) [3], doubly robust estimation [4], and kernel-based methods [5] have enabled more flexible and accurate estimation of treatment effects. For instance, BART [3] offers a nonparametric Bayesian approach to modeling treatment effects, while doubly robust estimators [4] combine the strengths of outcome regression and propensity score weighting to achieve consistent estimates even under model misspecification. Kernel-based methods, such as the kernel two-sample test by [5], provide powerful tools for comparing distributions and ensuring covariate balance.

In this paper, we propose a novel framework for causal inference that utilizes the strengths of traditional methods with modern machine learning techniques. Our approach, which we refer to as Subgroup-Aware Holistic Ridge (SAHR), is based on a holistic ridge regression model that incorporates an Auto-supervised balancing mechanism to account for confounding variables. The model leverages Gaussian Mixture Models (GMM) [6] for subgroup identification, Nyström kernel approximation [7] for scalability, and a weighted ridge regression framework [8] to balance predictive accuracy and covariate balancing. This approach is inspired by recent work on patient representation learning [9] and individualized treatment effect estimation [10], as well as by advancements in distributional metrics such as the Wasserstein distance [11] and the Maximum Mean Discrepancy (MMD) [5]. By combining these elements, the SAHR model aims to provide a robust and scalable framework for estimating treatment effects in observational studies. Its importance lies in the following key contributions:

I. Scalability meets precision: SAHR uncovers hidden variations in treatment effects across subgroups using a Gaussian Mixture Model, delivering precise, personalized causal insights, and leverages the Nyström method to efficiently handle high-dimensional data while maintaining computational scalability, ensuring fast, accurate

II. Adaptive balance, robust results: SAHR introduces a holistic Bayesian ridge regression framework that adaptively balances covariates derived from a BNN, ensuring robust inference by aligning subgroup-specific patterns with the overall population distribution without deviation.

Causal inference, a field at the intersection of statistics, machine learning, and social sciences, seeks to estimate the effects of interventions or treatments on outcomes of interest. A central challenge in this domain is addressing confounding, particularly in observational studies where treatments are not randomly assigned. Propensity score methods, introduced by Rosenbaum and Rubin [1], have emerged as a foundational tool for mitigating confounding by balancing covariates between treated and control groups. These methods estimate the probability of treatment assignment given observed covariates, enabling researchers to create comparable groups through matching, stratification, or weighting. Austin [2] provides a comprehensive review of propensity score methods, highlighting their versatility and limitations in various applications. Building on this, Imai and Ratkovic [12] introduced covariate-balancing propensity scores, which optimize covariate balance directly during estimation, thereby further enhancing the robustness of these methods. While propensity score methods are powerful, they rely heavily on the correct specification of the propensity score model. To address this limitation, doubly robust estimators, introduced by [4], combine outcome regression and propensity score weighting. These estimators are consistent if either the propensity score model or the outcome model is correctly specified, offering a safeguard against model misspecification. [13] Further explored the performance of doubly robust estimators under high variability in inverse probability weights, providing practical insights into their application. This dual robustness property has made doubly robust methods a cornerstone of modern causal inference.

The integration of machine learning into causal inference has further expanded the toolkit for estimating treatment effects. BART, introduced by [3], provides a flexible, nonparametric Bayesian framework for modeling treatment effects. BART's ability to capture complex, nonlinear relationships has made it a popular choice for causal inference tasks. Similarly, Wage and Athey [14] developed random forests for heterogeneous treatment effect estimation, enabling the identification of subgroups with varying treatment effects. These machine learning approaches are complemented by neural network-based methods, such as those proposed by [10] and [15], which leverage deep learning to model individualized treatment effects. These advancements highlight the growing synergy between machine learning and causal inference.

Subgroup analysis is another critical area of research, as treatment effects often vary across subpopulations. Su et al. [16] introduced recursive partitioning for subgroup analysis, providing a systematic approach to identifying subgroups with differing treatment effects. In [14], this approach was extended to random forests, enabling more flexible and accurate subgroup identification. These methods are particularly valuable in personalized medicine and policy evaluation, where understanding heterogeneous treatment effects is essential. Ensuring covariate balance between treated and control groups is a fundamental requirement for valid causal inference. Distributional metrics, such as MMD [5] and Wasserstein distance [11], provide powerful tools for assessing and achieving this balance. Gretton et al. [5] proposed MMD as a kernel-based method for comparing distributions, offering a rigorous framework for evaluating covariate balance. Villani [11] provided a theoretical foundation for optimal transport and Wasserstein distance, which have become widely used in machine learning and causal inference for measuring distributional differences. These metrics are particularly useful in high-dimensional settings, where traditional methods may struggle.

Regularization techniques, such as ridge regression [8], play a crucial role in addressing multicollinearity and improving the stability of causal inference models. Imai and Ratkovic [8] introduced ridge regression as a method for biased estimation in nonorthogonal problems, providing a robust framework for handling high-dimensional data. Further advanced this field by proposing the Nyström method for scalable kernel approximation, enabling the application of kernel-based methods to large datasets. These techniques are essential for ensuring the scalability and efficiency of modern causal inference methods [7]. Evaluating the performance of causal models is a critical step in ensuring their validity and reliability. LaLonde [17] highlighted the challenges of evaluating causal models using experimental data and provided insights into the limitations of regression adjustment. Schochet [18] discussed the limitations of regression adjustment in the Neyman model for causal inference, emphasizing the importance of robust evaluation methods. These studies underscore the need for rigorous evaluation frameworks in causal inference. Covariate balancing remains a critical component of causal inference, as it ensures that treated and control groups are comparable. Belthangady et al. [19] proposed BCAUS, a method for automatically balancing covariates in massive multi-arm observational studies, providing a robust framework for handling high-dimensional data. This approach aligns with the broader trend of leveraging machine learning for covariate balancing, as seen in [12].

Generative models, such as GANITE [15], have emerged as powerful tools for estimating individualized treatment effects. Yoon et al. [15] introduced GANITE, a generative adversarial network for estimating individualized treatment effects, demonstrating the potential of deep learning for causal inference tasks. These models are particularly valuable in settings where counterfactual outcomes need to be estimated, such as personalized medicine and policy evaluation. Theoretical foundations, such as those provided by [20], [21], and Heckman and Vytlacil [22], are essential for understanding the principles of causal inference. Pearl [20] introduced the concept of Directed Acyclic Graphs (DAGs) for causal reasoning, providing a powerful framework for modeling causal relationships. Imbens and Rubin [21] and Heckman and Vytlacil [22] contributed to the econometric evaluation of social programs, emphasizing the importance of structural models and policy evaluation. These theoretical advancements have shaped the development of modern causal inference methods. The field of causal inference has evolved significantly, driven by advancements in propensity score methods, machine learning, and distributional metrics. These developments have enabled more accurate and flexible estimation of treatment effects, paving the way for applications in personalized medicine, policy evaluation, and beyond.

Despite significant advancements in causal inference methodologies, several challenges and gaps remain in the field, particularly in observational studies.

I. Handling heterogeneity in treatment effects: Traditional causal inference methods often assume homogeneous treatment effects across the population, which is rarely the case in real-world scenarios. While methods like recursive partitioning [16] and random forests [14] have been developed to address heterogeneity, they often lack a robust mechanism for balancing covariates while simultaneously modeling subgroup-specific treatment effects. This gap limits their ability to provide accurate and interpretable estimates of treatment effects across diverse subpopulations.

II. Scalability in high-dimensional settings: Many existing methods, such as propensity score matching [1] and BART [3], struggle with scalability when applied to high-dimensional datasets. Kernel-based methods, such as those proposed by [5], offer a solution but often require significant computational resources. The Nyström method [7] addresses this issue to some extent, but its integration with causal inference frameworks remains underexplored.

III. Balancing predictive accuracy and covariate balance: A key challenge in causal inference is achieving a balance between these two. While doubly robust estimators [4] and covariate balancing propensity scores [12] provide robust frameworks, they often rely on strong assumptions about model specification. There is a need for methods that can adaptively balance covariates while maintaining high predictive accuracy, especially in settings with complex, nonlinear relationships.

IV. Integration of subgroup analysis and distributional metrics: Subgroup analysis is critical for understanding heterogeneous treatment effects, but existing methods often lack a systematic approach to integrating subgroup identification with distributional metrics such as MMD [5] and Wasserstein distance [11]. This integration is essential to ensuring that subgroups are both interpretable and balanced with respect to covariate distributions.

V. Leveraging modern machine learning techniques: While machine learning methods like neural networks [10] and Generative Adversarial Networks (GANs) [15] have shown promise in causal inference, their application often lacks a principled framework for ensuring covariate balance and interpretability. There is a need for models that can leverage the flexibility of machine learning while adhering to the rigorous requirements of causal inference.

The SAHR model fills critical gaps in the causal inference literature by providing a robust, scalable, and interpretable framework for estimating treatment effects in observational studies. Its integration of subgroup analysis, kernel approximation, and adaptive covariate balancing addresses key challenges in the field, making it a significant advancement in causal inference methodology. By bridging traditional methods and modern machine learning techniques, the SAHR model has the potential to drive innovation and improve decision-making across healthcare, economics, and the social sciences.

Let X ∈ R denote a vector of d covariates for each individual, T = {0,1} denote a binary treatment assignment (where T = 1 indicates treatment and T = 0 indicates control), and Y denote the observed outcome. Let Y0 and Y1 represent the potential outcomes under control and treatment, respectively. The observed outcome Y can be expressed as Y = T. Y1 + (1 − T)Y0 to estimate causal effect from observational data, we rely on the following fundamental assumptions.

Assumption 1 (Stable Unit Treatment Value Assumption (SUTVA)). The potential outcomes for any individual are unaffected by the treatment assignments of other individuals. Formally, for each individual i, Yi = Yi(Ti), where Yi(Ti) is the potential outcome for individual i under treatment Ti. This assumption ensures that the treatment of one individual does not influence the outcomes of others and the observed outcome corresponds to the potential outcome under the assigned treatment.

Assumption 2 (ignorability (unconfoundedness)). The treatment assignment T is conditionally independent of the potential outcomes Y0 and Y1, given the covariates X. Formally, T ⊥ (Y0, Y1)|X. This assumption implies that all confounding factors are captured by the covariates X, ensuring that treatment assignment is as good as random conditional on X.

Assumption 3 (positivity (overlap)). Every individual has a non-zero probability of receiving either treatment or control, given their covariates X. Formally, for all i, 0 < P(T = 1|X = xi) < 1 this assumption ensures sufficient overlap in the covariate distributions between treated and control groups, making causal comparisons valid.

We evaluated our SAHR framework on the Jobs dataset. We compared its performance against several state-of-the-art methods, including OLS/LR-1, OLS/LR-2, BLR, k-NN, BART, Random Forest, Causal Forest, AIPW, BNN, TARNet, CFR Wass, GANITE, and BCAUSS. The results are summarized in Table 1, which presents the in-sample (ϵtr) and out-of-sample (ϵte) errors for the ATT.

Our SAHR framework achieved an in-sample error (ϵtr) of 0.036± 0.019 and an out-of-sample error (ϵte) of 0.047 ± 0.036 for the ATT. These results are competitive with the state-of-the-art methods, particularly BCAUSS, which achieved an out-of-sample error of 0.05 ± 0.02. Our method outperformed several traditional methods, such as k-NN, Random Forest, and Causal Forest, and performed comparably to more advanced methods, such as GANITE and BCAUSS.

Table 2 compares the performance of Dragonnet, BCAUSS, and SAHR in terms of distributional metrics for covariate balance between treated and untreated groups. SAHR demonstrates the best performance overall, with the lowest KS (p-value < 1%) (0.7% train, 1% test), indicating the least frequent rejection of the null hypothesis that the treated and untreated groups are from the same distribution. Additionally, SAHR achieves the lowest Wasserstein Distance (0.033 train, 0.081 test) and MMD measures across all kernels (Linear, RBF, Poly), highlighting its superior ability to maintain covariate balance. While BCAUSS performs better than Dragonnet in most metrics, it falls short of SAHR's performance. Dragonnet lags significantly behind, with higher KS (p-value < 1%), Wasserstein Distance, and MMD measures, indicating poorer covariate balance. Overall, SAHR emerges as the best-performing model, demonstrating excellent generalization and a balanced performance across both the training and test sets.

Fig. 2 provides a detailed, formal visualization of the SAHR model’s efficacy in balancing covariates between the treated and control groups on the Jobs dataset. Each subplot represents a distinct scenario: the training set before adjustment, the training set after adjustment, the test set before adjustment, and the test set after adjustment. These visualizations reveal a pronounced imbalance in the original data, with substantial differences in the distributions of covariates between the treated and control groups. The SAHR model effectively mitigates this imbalance through its self-supervised balancing mechanism and subgroup analysis, as evidenced by the improved alignment of distributions in the adjusted plots. Annotations within each subplot provide critical statistical metrics, such as the mean, variance, and Wasserstein distance, enabling a quantitative evaluation of distributional discrepancies and the model's enhancements. For instance, the Wasserstein distance decreases significantly from 0.0306 (training set before adjustment) to 0.0016 (training set after adjustment), and from 0.0593 (test set before adjustment) to 0.0027 (test set after adjustment), highlighting the model's success in reducing distributional gaps. Similarly, the variances for treated and control groups converge closely after adjustment, further demonstrating the model's balancing capability. The side-by-side comparison of before-and-after scenarios underscores the model's robustness in maintaining covariate balance, a crucial requirement for accurate causal inference.

Our SAHR framework represents a significant advancement in causal inference, demonstrating exceptional performance in estimating treatment effects from observational data. By integrating a self-supervised balancing mechanism and utilizing GMM for subgroup analysis, SAHR effectively reduces confounding bias and ensures robust covariate balance between treated and control groups. The use of kernel approximation techniques further enhances scalability, making SAHR suitable for high-dimensional datasets while maintaining computational efficiency. These features make SAHR a powerful tool for addressing the challenges of causal inference in real-world scenarios. One of SAHR's key strengths is its superior ability to maintain covariate balance. The framework achieves the lowest KS (p-value < 1%) (0.7% train, 1% test), Wasserstein Distance (0.033 train, 0.081 test), and MMD measures across all kernels (Linear, RBF, Poly), outperforming state-of-the-art methods like Dragonnet and BCAUSS. It demonstrates its effectiveness in ensuring that treated and control groups are well-balanced, which is critical for accurate causal inference. Additionally, SAHR achieves competitive estimates of the treatment effect, with an in-sample error (ϵtr) of 0.036 ± 0.019 and an out-of-sample error (ϵte) of 0.047 ± 0.036 for the ATT on the Jobs dataset. These results are comparable to advanced methods like GANITE and BCAUSS, while outperforming traditional approaches such as k-NN, Random Forest, and Causal Forest.

Despite its strengths, SAHR has some limitations that need to be addressed. Currently, the framework is designed to handle only binary treatments, which restricts its applicability to problems involving multiple or continuous treatments. Extending SAHR to support these scenarios is an important direction for future work. Additionally, while SAHR performs exceptionally well on the Jobs dataset, further validation on a broader range of real-world datasets is necessary to assess its generalizability across different domains and data structures. The use of GMM for subgroup analysis, while effective, also introduces additional complexity in model training and hyperparameter tuning. Simplifying this process or developing automated tuning methods could improve usability. Finally, like most causal inference methods, SAHR relies on the assumption of unconfoundedness, and violations of this assumption can lead to biased estimates. Future work could explore robustness to unmeasured confounding.

Looking ahead, several promising directions exist for extending and improving SAHR. Expanding the framework to handle multiple treatments and continuous treatment variables would broaden its applicability to a wider range of causal inference problems. Developing methods to assess and mitigate unmeasured confounding would enhance the robustness of SAHR in real-world settings. Incorporating automated hyperparameter tuning for the GMM and kernel approximation components would make SAHR more user-friendly and accessible to non-experts. Benchmarking SAHR on a diverse set of real-world datasets, including those with complex data structures and varying levels of confounding, would provide deeper insights into its generalizability and performance. Finally, integrating SAHR's balancing mechanism with deep learning architectures could further enhance its ability to model complex, high-dimensional relationships in observational data.

The SAHR framework provides a robust and flexible approach to causal inference, combining the interpretability of traditional methods with the scalability and performance of modern machine learning techniques. Its state-of-the-art performance on the Jobs dataset, particularly in terms of covariate balance and treatment effect estimation, demonstrates its potential to drive meaningful progress in the field. While there are limitations to address, the modular design and computational efficiency of SAHR provide a strong foundation for future advancements. By addressing these challenges and expanding its capabilities, SAHR could become a widely adopted tool for causal inference across research and practice.

The authors sincerely appreciate all those who contributed to the completion of this research.

This research was conducted without financial support from any public, private, or non-profit organization.

The data used in this study are available from the corresponding author upon reasonable request.