# BENK: The Beran Estimator with Neural Kernels for Estimating the Heterogeneous Treatment Effect
## 6 Neural Network for Estimating CATE
The architecture of the neural network, consisting of \(n\) subnetworks that implement the neural kernels, is shown in Figure 2. Examples \(\mathbf{a}_{i}^{(r)}\) produced from the dataset of controls are fed to the whole neural network, such that each pair \((\mathbf{x}_{i},\mathbf{x}_{k}^{(r)})\), \(k\in\mathcal{I}^{(r)}\), is fed to each subnetwork, which implements the kernel function. The output of each subnetwork is kernel \(K(\mathbf{x}_{i},\mathbf{x}_{k}^{(r)})\). All subnetworks are identical and have shared weights. After normalizing the kernels, we obtain \(n\) weights \(W(\mathbf{x}_{i},\mathbf{x}_{k}^{(r)})\), which are used to estimate SFs by means of the Beran estimator in (14). The block of the whole neural network that implements the Beran estimator uses all weights \(W(\mathbf{x}_{i},\mathbf{x}_{k}^{(r)})\), \(k\in\mathcal{I}^{(r)}\), and the corresponding values \(\delta_{k}^{(r)}\) and \(f_{k}^{(r)}\), \(k\in\mathcal{I}^{(r)}\). As a result, we obtain SF \(\tilde{S}_{0}(t\mid\mathbf{x}_{i})\). In the same way, we compute SFs \(\tilde{S}_{0}(t\mid\mathbf{x}_{k})\) for all \(k=1,\ldots,c\). These functions are the basis for training. In fact, the normalization block and the block that implements the Beran estimator can be regarded as part of the neural network, and they are trained in an end-to-end manner. According to (13), expected lifetimes are used to compute the CATE \(\tau(\mathbf{z})\). Therefore, the whole network is trained by means of the following loss function:  
\[L=\frac{1}{c^{*}\cdot N}\sum_{i\in\mathcal{C}^{*}}\sum_{k=1}^{N}\Bigl{(} \tilde{E}_{k}^{(i)}-f_{k}^{(i)}\Bigr{)}^{2}. \tag{18}\]  
Here, \(\mathcal{C}^{*}\) is a subset of \(\mathcal{C}\), which contains only uncensored examples from \(\mathcal{C}\), \(c^{*}\) is the number of elements in \(\mathcal{C}^{*}\); \(f_{k}^{(i)}\) is the time to an event of the \(k\)-th example from the set \(\mathcal{C}^{*}\backslash(\mathbf{x}_{i},\delta_{i},f_{i})\) and \(\tilde{E}_{k}^{(i)}\) is the expected lifetime computed through SF \(\tilde{S}_{0}(t\mid\mathbf{x}_{k})\), obtained by integrating the SF:  
\[\tilde{E}_{k}^{(i)}=\sum_{j=1}^{n}(f_{j}^{(i)}-f_{j-1}^{(i)})\tilde{S}_{0}(f_{j} ^{(i)}\mid\mathbf{x}_{k}). \tag{19}\]  
The sum in (18) is taken over uncensored examples from \(\mathcal{C}\). However, the Beran estimator uses all the examples. One of the loss functions, which takes into account all data (censored and uncensored), is the C-index. However, our aim is not to estimate the SF or the CHF. We aim to estimate the difference between the predicted time to event and the expected time to event. Therefore, we use the standard mean squared error (MSE) loss function. But the censored times introduce bias into MSE and, therefore, they are not used. It is important to point out that our aim is to train subnetworks with shared training parameters, which are the neural kernels. By having the trained neural kernels, we can use them to compute kernels \(K(\mathbf{z},\mathbf{x}_{i})\) and \(K(\mathbf{z},\mathbf{y}_{i})\) and then to compute estimates of SFs \(\tilde{S}_{0}(t\mid\mathbf{z})\) and \(\tilde{S}_{1}(t\mid\mathbf{z})\) for controls and treatments, respectively, i.e., we realize the idea of transferring tasks from the control group to the treatment group. Let \(t_{1}^{(0)}<t_{2}^{(0)}<\ldots<t_{c}^{(0)}\) and \(t_{1}^{(1)}<t_{2}^{(1)}<\ldots<t_{s}^{(1)}\) be the ordered time moments corresponding to times \(f_{1},\ldots,f_{c}\) and \(h_{1},\ldots,h_{s}\), respectively.