# Heterogeneous Treatment Effect with Trained Kernels of the Nadaraya-Watson Regression
## 4 The TNW-CATE Description  
It has been mentioned that the main idea behind TNW-CATE is to replace the Nadaraya-Watson regression with the neural network of a specific form. The whole network consists of two main parts. The first part implements the Nadaraya-Watson regression for training the control function \(g_{0}(\mathbf{x})\). In turn, it consists of \(n\) identical subnetworks such that each subnetwork implements the attention weight \(\mathbf{\alpha}(\mathbf{x},\mathbf{x}_{i})\) or the kernel \(K(\mathbf{x},\mathbf{x}_{i})\) of the Nadaraya-Watson regression. Therefore, the input of each subnetwork is two vectors \(\mathbf{x}\) and \(\mathbf{x}_{i}\), i.e., two vectors \(\mathbf{x}\) and \(\mathbf{x}_{i}\) are fed to each subnetwork. The whole network consisting of \(n\) identical subnetworks and implementing the Nadaraya-Watson regression for training the control function will be called the control network. In order to train the control network, for every vector \(\mathbf{x}_{i}\), \(i=1,\ldots,c\), \(N\) subsets of size \(n\) are randomly selected from the control set \(\mathcal{C}\) without example \((\mathbf{x}_{i},y_{i})\). The subsets can be regarded as \(N\) examples for training the network. Hence, the control network is trained on \(N\cdot c\) examples of size \(2d\cdot n\). If we have a feature vector \(\mathbf{x}\) for estimating \(\tilde{y}\), i.e., for estimating function \(g_{0}(\mathbf{x})\), then it is fed to each trained subnetworks jointly with each \(\mathbf{x}_{i}\), \(i=1,\ldots,c\), from the training set. In this case, the trained subnetworks or kernels of the Nadaraya-Watson regression are used to estimate \(\tilde{y}\). The number of subnetworks for testing is equal to the number of training examples \(n\) in every subset. Since the trained subnetworks are identical and have the same weights (parameters), their number can be arbitrary. In fact, a single subnetwork can be used in practice, but its output depends on the pair of vectors \(\mathbf{x}\) and \(\mathbf{x}_{i}\). The same network called the treatment network is constructed for the treatment group. In contrast to the control network, it consists of \(m\) subnetworks with inputs in the form of pairs \((\mathbf{z},\mathbf{z}_{i})\). In the same way, \(M\) subsets of size \(m\) are randomly selected from the training set of treatments without the example \((\mathbf{z}_{j},h_{j})\), \(j=1,\ldots,t\). The treatment network is trained on \(M\cdot t\) examples of size \(2d\cdot m\). After training, the treatment network allows us to estimate \(\tilde{h}\) as function \(g_{1}(\mathbf{z})\). If \(\mathbf{z}=\mathbf{x}\), then we obtain an estimate of CATE \(\tau(\mathbf{x})\) or \(\tau(\mathbf{z})\) as the difference between estimates \(\tilde{h}\) and \(\tilde{y}\) obtained by using the treatment and control neural networks. It is important to point out that the control and treatment networks are jointly trained by using the loss function defined below. Let us formally describe TNW-CATE in detail. Consider the control group of patients \(\mathcal{C}=\{(\mathbf{x}_{1},y_{1}),\ldots,(\mathbf{x}_{c},y_{c})\}\). For every \(i\) from set \(\{1,\ldots,c\}\), we define \(N\) subsets \(\mathcal{C}_{i,r}\), \(r=1,\ldots,N\), consisting of \(n\) examples randomly selected from \(\mathcal{C}\backslash(\mathbf{x}_{i},y_{i})\) as:  
\[\mathcal{C}_{i,r} =\{(\mathbf{x}_{1}^{(r)},y_{1}^{(r)}),\ldots,(\mathbf{x}_{n}^{(r )},y_{n}^{(r)})\},\] \[r =1,\ldots,N, \tag{10}\]  
where \(\mathbf{x}_{j}^{(r)}\) is a randomly selected vector of features from the set \(\{\mathbf{x}_{1},\ldots,\mathbf{x}_{c}\}\backslash\mathbf{x}_{i}\), which forms \(\mathcal{C}_{i,r}\); \(y_{j}^{(r)}\) is the corresponding outcome.