# Heterogeneous Treatment Effect with Trained Kernels of the Nadaraya-Watson Regression
###### Abstract
Many methods for estimating CATE have been proposed and developed due to importance of the problem in medicine and other applied areas [12; 17; 18; 19; 20; 21; 22; 23; 24; 25; 26; 27; 28; 29]. This is only a small part of all publications which are devoted to the CATE estimation problem solution. Various approaches were used for solving the problem, including the support vector machine [30], tree-based models [15], neural networks [2; 13; 31; 32], and transformers [33; 34; 35; 36]. It should be noted that most approaches to estimating CATE are based on constructing regression models for handling the treatment and control groups. However, the problem of the small treatment group motivates us to develop various tricks that at least partially resolve the problem. We propose a method based on using the Nadaraya-Watson kernel regression [37; 38] which is widely applied to machine learning problems. The method is called TNW-CATE (the Trainable Nadaraya-Watson regression for CATE). The Nadaraya-Watson estimator can be seen as a weighted average of outcomes (patient responses) by means of a set of varying weights called attention weights. The attention weights in the Nadaraya-Watson regression are defined through kernels that measure distances between training feature vectors and the target feature vector, i.e., kernels in the Nadaraya-Watson regression conform with relevance of a training feature vector to a target feature vector. If we have a dataset \(\{(\mathbf{x}_{1},y_{1}),\ldots,(\mathbf{x}_{n},y_{n})\}\), where \(\mathbf{x}_{i}\in\mathbb{R}^{m}\) is a feature vector (key) and \(y_{i}\in\mathbb{R}\) is its target value or its label (value), then the Nadaraya-Watson estimator for a target feature vector \(\mathbf{z}\in\mathbb{R}^{m}\) (query) can be defined by using weights \(\alpha(\mathbf{z},\mathbf{x}_{i})\) as  
\[\hat{f}(\mathbf{z})=\sum_{i=1}^{n}\alpha(\mathbf{z},\mathbf{x}_{i})y_{i}, \tag{1}\]  
where  
\[\alpha(\mathbf{z},\mathbf{x}_{i})=\frac{K(\mathbf{z},\mathbf{x}_{i},\gamma)}{ \sum_{j=1}^{n}K(\mathbf{z},\mathbf{x}_{j},\gamma)}, \tag{2}\]  
\(K(\mathbf{z},\mathbf{x}_{i},\gamma)\) is a kernel and \(\gamma>0\) is a bandwidth parameter. Standard kernels widely used in practice are the Gaussian, uniform, or Epanechnikov kernels [39]. However, the choice of a kernel and its parameters significantly impact on results obtained from the Nadaraya-Watson regression usage. Moreover, the Nadaraya-Watson regression also requires a large number of training examples. Therefore, we propose a quite different way for implementing the Nadaraya-Watson regression. The method is based on the following assumptions and ideas. First, each kernel is represented as a part of a neural network implementing the Nadaraya-Watson regression. In other words, we do not use any predefined standard kernels such as the Gaussian one. Kernels are trained as the weight sharing neural subnetworks. The weight sharing is used to identically compute kernels under the condition that the pair \((\mathbf{z},\mathbf{x}_{i})\) of examples is fed into every subnetwork. The neural network kernels become more flexible and sensitive to a complex location structure of feature vectors. In fact, we propose to replace the definition of weights through the kernels with a set of neural subnetworks with shared parameters (the neural network weights) such that every subnetwork implements the trainable kernel, but the whole network implements the Nadaraya-Watson estimator. At that, the trainable parameters of the kernels are nothing else but the weights of each neural subnetwork. The above implementation of the Nadaraya-Watson regression by means of the neural network leads to an interesting result when the treatment examples are considered as a single example whose \"features\" are the whole treatment feature vectors. It should be noted that several authors [40; 41; 42; 43] proposed to use the Nadaraya-Watson kernel regression with standard kernels having the bandwidth parameter to construct the CATE estimator. However, they did not propose to learn the kernels.