Heterogeneous Treatment Effect with Trained Kernels of the Nadaraya–Watson Regression

Abstract

A new method for estimating the conditional average treatment effect is proposed in this paper. It is called TNW-CATE (the Trainable Nadaraya–Watson regression for CATE) and based on the assumption that the number of controls is rather large and the number of treatments is small. TNW-CATE uses the Nadaraya–Watson regression for predicting outcomes of patients from control and treatment groups. The main idea behind TNW-CATE is to train kernels of the Nadaraya–Watson regression by using a weight sharing neural network of a specific form. The network is trained on controls, and it replaces standard kernels with a set of neural subnetworks with shared parameters such that every subnetwork implements the trainable kernel, but the whole network implements the Nadaraya–Watson estimator. The network memorizes how the feature vectors are located in the feature space. The proposed approach is similar to transfer learning when domains of source and target data are similar, but the tasks are different. Various numerical simulation experiments illustrate TNW-CATE and compare it with the well-known T-learner, S-learner, and X-learner for several types of control and treatment outcome functions. The code of proposed algorithms implementing TNW-CATE is publicly available.

BENK: The Beran Estimator with Neural Kernels for Estimating the Heterogeneous Treatment Effect

“…A new method called BENK for solving the CATE problem under the condition of censored data has been presented. It extends the idea behind TNW-CATE proposed in [16] to the case of censored data. In spite of many similar parts of TNW-CATE and BENK, they are different because BENK is based on using the Beran estimator for training and can be successfully applied to survival analysis of controls and treatments.…”

“…At the same time, a small amount of training data may lead to overfitting and unsatisfactory results. Therefore, the problem of overcoming this possible limitation motivated researchers to introduce a neural network of a special architecture, which implements the trainable kernels in the Nadaraya-Watson regression [16].…”

“…This is another important peculiarity of the TNW-CATE. Results provided in [16] illustrated outperformance of the TNW-CATE in comparison with other methods when the treatment group was very small and the feature vectors had complex structure. Following the ideas behind the TNW-CATE, we propose the CATE estimation method, called BENK (the Beran Estimator with Neural Kernels), dealing with censored time-toevent data in the framework of survival analysis.…”

“…One of the approaches for constructing regression models for controls and treatments is the application of the Nadaraya-Watson kernel regression [14,15], which uses standard kernel functions, for instance, the Gaussian, uniform or Epanechnikov kernels. In order to avoid selecting a standard kernel, Konstantinov et al [16] proposed to implement kernels and the whole Nadaraya-Watson kernel regression by using a set of identical neural subnetworks with shared parameters, with a specific way of the network training. The corresponding method called TNW-CATE (Trainable Nadaraya-Watson regression for CATE) is based on an important assumption that domains of the feature vectors from the treatment and control groups are similar.…”