# 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. t  
small training dataset, which does not allow us to apply many efficient machine learning methods directly. [PERSON] and [PERSON] [11] consider an example of the difficulties in collecting corresponding information about patients and to explore the side effects of drugs adopted on certain pediatric cancer treatments for several reasons, including the sensitivity of issues for the families, the rarity of the disease, and the effort required between hospitals and doctors to refer their patients. The authors propose to use a transfer learning approach, implemented by means of a neural network, to estimate treatment effects on small datasets. This is an efficient and interesting approach. However, it may encounter difficulties when training the neural network that models the regression functions for controls and treatments. The same problem has been pointed out in [4; 12]. Therefore, we need to develop a similar approach where the neural networks should be rather simple and play a secondary role in the modeling the regression functions. The second problem is that each patient cannot be simultaneously in the treatment and control groups, i.e., we either observe the patient outcome under treatment or control, but never both [13]. This is a fundamental problem of computing the causal effect. This problem is solved by the explicit or implicit construction of regression functions for the control and treatment patients. For instance, we can train neural networks [14], which predict outcomes for a new patient under condition that this patient belongs to one of the control or treatment groups. However, this method requires one to have large datasets to train the corresponding machine learning models. We thus return to the first problem of the treatment group size. In addition to the above two problems, there are many difficulties facing machine learning model development concerning noisy data, especially with the high dimension of the patients' health records, etc. [15]. One of the difficulties is a complex data structure. A lot of methods have some assumptions about the model parameters and the data structure. For instance, a Gaussian mixture for outcomes is proposed for use in [11]. In some cases, these assumptions are correct, and they correspond to the real structure of the treatment and control data. However, they can lead to errors in other cases where the data structure is rather complex. [PERSON] et al. [16] note that the analysis of observations, having a grouped structure, shows that the impact of the treatment exposure will vary across these groups. In this case, most current machine learning approaches ignore these varying effects. Therefore, it is necessary to develop models that take into account these problems.