# Nearest Neighbours Graph Variational AutoEncoder
## 2 Materials and Methods
### Nearest Neighbour Graph VAE
#### 2.1.1 ReNN-Pool and Un-Pool  
In CNNs, pooling operations are widely used to decrease the dimensionality of feature maps while increasing the receptive fields of neurons when processing grid-like data. This is easily possible on images, where a standard pooling layer typically involves a downsampling operation such as the max or mean function applied to groups of nearby pixels. Conversely, applying this idea to graph structures is generally a challenging task. However, there are cases where, although data does not have a grid-like shape and can be better processed with GNNs, the graph structures are fixed and have a regular and known geometry. For example, some datasets may contain examples whose data are arranged in a cylindrical or spherical structure. For these cases, we developed a simple pooling operation (ReNN-Pool) that can sub-sample graph nodes in a regular way. The ReNN-Pool consists in a masking operation and a subsequent update of the adjacency matrix of the graph. First of all, nodes are sorted on the basis of their position in the graph. For example, if samples have a cylindrical structure, nodes can be hierarchically ordered on the basis of their positions along the \(z\), \(\theta\) and \(r\) axes. Then, the masking operation is carried out. Masking consists in dropping, i.e., removing, certain nodes from the graph. The process is performed in a recursive manner on the sorted node list. It begins with the first node, which is preserved, while all its nearest neighbours are dropped and removed from the node list. The process then continues with the next node and repeats, until all nodes in the list have been processed. After the masking, we rewire links between the \"survived\" nodes, connecting the ones that were 2 nd order neighbours before the masking. This is conducted substituting the adjacency matrix \(A\) with \(A^{2}\). If we call \(M\) the vector that contains all the indices of the \"survived\" nodes, \(X=\{x_{i}\}_{i=1}^{N}\) and \(A\), respectively, and the nodes' feature matrix and the adjacency matrix before the pooling, the application of ReNN-Pool gives in output:  
\[X^{\prime}=\{x_{i}^{\prime}\}_{i=1}^{N},\text{ where }x_{i}^{\prime}= \begin{cases}x_{ i},&\text{if }i\in\ M\\ 0,&\text{otherwise}\end{cases} \qquad\qquad A^{\prime}= \begin{cases}A_{ij}^{2},& \text{if }i,j\in\ M\\ 0,&\text{otherwise}\end{cases} \]  
The process is illustrated in Figure 1.  
In the case of regular geometrical graph structures, by construction, pooled and not-pooled nodes are evenly spread across the graph; thus, the choice of the starting node does not affect the performances of the model in which the ReNN-Pool is used. For irregular graph structures, the performance can depend on the choice of the first node. However, in principle this can become an hyperparameter to optimize using a validation set. Such a possibility will be explored in future works. Figure 1: **Pooling. ** The ReNN-Pool operation consists of two steps: masking and rewiring. In the first step, recursively on the whole graph, a node is selected and all its nearest neighbours are dropped. In the second step, nodes are linked to the ones that were their 2 nd order neighbours. On regular graphs, the masking operation allows one to define evenly spread clusters of nodes. In fact, the surviving nodes can be thought as centroids of the clusters made up by their nearest neighbours. These clusters can be used to perform a mean or max pooling on graphs in addition to the simple masking. A comparison of these methods is presented in Section 3. Due to the fact that the creation of masks and the adjacency matrices' augmentation in our pooling only depends on the graph structure of data, it is possible to compute and store them before the training. This has two main advantages. First, the pooling operation has no influence on the computational needs both in the training and the inference phase. Thus, it can be used inside lightweight models for resource-constrained devices and it is scalable for graphs of any size.