abstract: $t$-SNE and hierarchical clustering are popular methods of exploratory data analysis, particularly in biology. Building on recent advances in speeding up $t$-SNE and obtaining finer-grained structure, we combine the two to create tree-SNE, a hierarchical clustering and visualization algorithm based on stacked one-dimensional $t$-SNE embeddings. We also introduce alpha-clustering, which recommends the optimal cluster assignment, without foreknowledge of the number of clusters, based off of the cluster stability across multiple scales. We demonstrate the effectiveness of tree-SNE and alpha-clustering on images of handwritten digits, mass cytometry (CyTOF) data from blood cells, and single-cell RNA-sequencing (scRNA-seq) data from retinal cells. Furthermore, to demonstrate the validity of the visualization, we use alpha-clustering to obtain unsupervised clustering results competitive with the state of the art on several image data sets. Software is available at <https://github.com/isaacrob/treesne>.
author: Isaac Robinson
Yale University
New Haven, CT 06511
`email@example.com`
Emma Pierce-Hoffman
Yale University
New Haven, CT 06511
`firstname.lastname@example.com`
title: Tree-SNE: Hierarchical Clustering
  and Visualization Using $t$-SNE

# Introduction

$t$-SNE (van der Maaten and Hinton 2008) is a widely-used method of visualizing high-dimensional data in low dimensions. It is motivated by minimizing the Kullback-Leibler divergence between the distributions of pairwise affinities among observations in the high-dimensional and low-dimensional spaces. Its predecessor, SNE (Hinton and Roweis 2002), uses a Gaussian kernel to transform the low-dimensional distances into affinities, while $t$-SNE uses a heavier-tailed $t$-distribution with one degree of freedom. As noted by van der Maaten and Hinton (2008), the heavier tails of the $t$-distribution compared to the Gaussian distribution help to alleviate the "crowding problem" so that distinct blobs appear in the low-dimensional embedding.

Since $t$-SNE gives better-separated clusters than SNE, Kobak et al. (2019) used the Fourier transform (FFT)-accelerated interpolation-based $t$-SNE (FIt-SNE) approximation from Linderman et al. (2019) to implement a fast version of $t$-SNE for smaller, fractional degrees of freedom. They used a scaled $t$-distribution kernel defined as $$k(d) = \frac{1}{(1+ d^2/\!\alpha)^\alpha}$$ where the degrees of freedom are $v = 2\alpha - 1$. Since the degrees of freedom are positive ($v > 0$), $\alpha$ is also positive ($\alpha>0$). Kobak et al. found that $\alpha < 1$ (which corresponds to $v < 1$) produces $t$-SNE plots with tighter, smaller blobs, capable of capturing finer-grained structures than standard $t$-SNE. For instance, standard $t$-SNE with $\alpha = 1$ separates the images of handwritten digits from the MNIST data set by digit, whereas Kobak et al. found that two-dimensional $t$-SNE with $\alpha = 0.5$ results in further separation into blobs representing different handwriting styles of each digit (see Figure ).

In a lecture at Yale University in fall 2019, Stefan Steinerberger argued that the $t$-SNE optimization problem can be interpreted as a dynamic systems problem as opposed to a probability distribution matching problem, thereby justifying using $0<\alpha<0.5$, for which the resulting kernel is not a valid probability distribution (Steinerberger 2019). He argued that this allows for an even more fine-grained look at the high dimensional data. $t$-SNE has also been shown to cluster well-separated data reliably in any embedding dimension (Linderman and Steinerberger 2019). Therefore, the current theory suggests that for well-clustered data, one-dimensional $t$-SNE can be leveraged to convey the same information as two-dimensional $t$-SNE in a more compact manner (Linderman et al. 2019).

Based on those results, and motivated by the popularity of $t$-SNE and hierarchical clustering individually, we present tree-SNE, a hierarchical clustering method based on one-dimensional $t$-SNE with decreasing values of $\alpha$ and perplexity at each level. Tree-SNE allows for visualization and elucidation of high-dimensional hierarchical structures by creating $t$-SNE embeddings with increasingly heavy tails to reveal increasingly fine-grained structure, and then stacking these embeddings to create a tree-like structure. We then run spectral clustering on each one-dimensional embedding, computationally determining the number of distinct clusters in the embedding. The number of clusters will increase as $\alpha$ decreases. We define the alpha-clustering of the data to be the cluster assignment that is stable across the largest range of $\alpha$ values, and then we demonstrate that alpha-clustering is competitive with the state of the art in unsupervised clustering algorithms on several data sets, including MNIST (LeCun et al. 1998) and COIL-20 (Nene et al. 1996a). Figure shows example visualizations produced by tree-SNE on the USPS (Hull 1994) handwritten digits data set.

# Methods

In this section, we provide a description of our algorithm and motivations for various design decisions. Section 2.1 discusses tree-SNE and section 2.2 discusses alpha-clustering.

## Overview of tree-SNE

Tree-SNE reveals hierarchical structures in high dimensional data by creating many $t$-SNE embeddings with changing perplexity and degrees of freedom of the $t$-distribution. As both perplexity and degrees of freedom decrease, smaller, more granular clusters are produced in the embedding. In this paper, we restrict ourselves to one-dimensional $t$-SNE embeddings for computational efficiency and because it makes clustering easier. In the future, two-dimensional embeddings could also be used to create 3-D tree structures that are potentially more informative. However, we will demonstrate that these one-dimensional embeddings are sufficiently informative to provide potentially valuable insight into hierarchical structures of data.

### Generating tree-SNE embeddings

We begin with a standard one-dimensional $t$-SNE ($\alpha = 1$) embedding with high perplexity, by default the square root of the number of data points. A high starting perplexity increases the effective number of neighbors used by $t$-SNE (van der Maaten and Hinton 2008), which means that larger clusters will tend to form, capturing more global structures in the data (Kobak and Berens 2019). By using a large perplexity at the beginning, tree-SNE can display the entire spectrum of data organization, from global structures at the base of the tree to very fine-grained structures at the top. The default initial perplexity value of $\sqrt{N}$ was inspired by analysis of perplexity in $t$-SNE by Oskolkov (2019) and gives favorable results in practice. The exact starting value of perplexity did not appear to be important across many trials, as long as it is high, as $t$-SNE is fairly robust to small changes in perplexity (van der Maaten and Hinton 2008), and most of the interesting features of the data emerge in further embeddings from the adjustments of perplexity and $\alpha$.

Many one-dimensional $t$-SNE embeddings (typically in the range of 30 to 100) are stacked on top of each other to create the tree-SNE visualization, with the first layer at the bottom of the plot, and $\alpha$ and perplexity decreasing in each subsequent layer. To generate each sequential level, the $\alpha$ used in the previous level is multiplied by a constant factor $0<r<1$, which can be expressed as $\alpha_{n+1} = r\alpha_{n}$, where $\alpha_n$ refers to the value of $\alpha$ used in the $n$th embedding layer. With $p_n$ representing the the perplexity on level $n$, we let $p_{n+1} = p_{n}^r$. In practice, $r$ is very close to $1$. In this way, $\alpha$ approaches $0$ and perplexity approaches $1$ as the number of levels $n$ gets large, for a constant $r$. Perplexity of $1$ corresponds to entropy of $0$ in the embedding, and $\alpha$ of $0$ means our kernel function $k(x, y) = 1$ for any points $x$, $y$ in the low-dimensional embedding. Both of these parameters can be thought of as optimizing for clusters containing only a single data point each, as there is no pressure for points to form clusters under these settings. Accordingly, we observe that clusters become smaller and more numerous moving upward through the layers.

When creating each layer, the $t$-SNE embedding is initialized with the embedding from the previous level, rather than a standard random initialization. This is done so that each layer is a refinement of the clustering found on the previous layer, with larger clusters breaking into smaller clusters on subsequent levels of the tree. As a result, the path of a single observation or a group of observations can be traced vertically through the resulting tree visualization.

While the FIt-SNE paper (Linderman et al. 2019) suggests that high initial exaggeration and learning rates improve the embedding quality, we use smaller learning rates and no early exaggeration after the first embedding to ensure relative stability from layer to layer. With smaller initial exaggeration and learning rates, each embedding does not change drastically from its initialization embedding, which is necessary for an interpretable tree structure to emerge. We do, however, use a constant exaggeration factor (default 12), as was done by Kobak and Berens (2019), to increase cluster separation, based on the finding by Linderman et al. (preprint 2017) that late stage exaggeration increases differentiation of clusters.

### Differentiating factors of tree-SNE

Visualizing data at multiple levels of granularity in this way can be useful both to uncover inherently hierarchical structures within the data and to explore new data, because it can be difficult to determine the optimal level of granularity on which to view or cluster unlabeled data. In standard $t$-SNE it is possible to attempt to evaluate the data across multiple scales by using different perplexity and/or degrees of freedom settings. However, the shattering of global structure is a common complaint of $t$-SNE (Kobak and Berens 2019), so it can be difficult to relate a $t$-SNE plot at one scale to another of a different scale, because the global organization of the data is not necessarily preserved. A hierarchical visualization approach such as this is necessary to enable comparison across scales. Visualizing the data with tree-SNE as opposed to $t$-SNE allows a researcher (or our alpha-clustering algorithm) to decide which scale is best for their data by showing them a more in-depth view of the hierarchical organization than is possible with standard $t$-SNE.

## Description of alpha-clustering

Alpha-clustering extends and validates the tree-SNE algorithm by assigning cluster labels based on the $t$-SNE embeddings and then automatically selecting the best level of granularity for clustering the data, the result of which can then be compared to other unsupervised clustering algorithms to show that the structure uncovered by tree-SNE indeed captures meaningful aspects of the data. A researcher does not of course have to use alpha-clustering to use tree-SNE. However, since alpha-clustering runs on top of an expressive visualization tool, it provides a powerful opportunity for explainable unsupervised clustering which is difficult to achieve in other approaches, many of which rely on deep neural networks. It is easy to understand why alpha-clustering makes its decisions, and cluster labels can be examined on multiple levels of granularity so a researcher can choose the level that fits their needs, or the alpha-clustering algorithm will recommend an optimal level.

Figure provides further insight into the alpha-clustering process. Cluster labels are assigned on each layer of the tree, and in the top image in Figure , each layer is colored by the cluster labels for that individual layer. Then, the most stable clustering out of all of the layers is defined as the alpha-clustering. In the second image in Figure , all layers of the tree are colored by the alpha-clustering labels.

### Outline

Alpha-clustering works by searching for a cluster assignment for the data that is maximally invariant across multiple levels of clustering granularity. As one-dimensional $t$-SNE tends to generate easily distinguishable clumps of data, we identify and count the clumps on each layer of the tree-SNE tree, then we define the alpha-clustering as the clustering for which the number of clumps does not change across the largest range of $\alpha$ values.

In other words, define a clustering algorithm $c(n)$ that accepts a layer numbered $n$ and produces a clustering. If we observe that the clustering algorithm $c$ produces the same clustering on all layers between layers $n$ and $m$ (where $m > n$) such that $c(n) = c(m)$, then we call clustering $c(n)$ a stable clustering. Each level's embedding is generated by a unique value of $\alpha$, so each stable clustering has a nonzero $\alpha$ range equal to $a_m-a_n$ where $n$ and $m$ are the first and last layers that produce the stable clustering. We define the alpha-clustering to be the stable clustering for which the $\alpha$ range is the largest. Recall that the relationship between the value of $\alpha$ used to generate embeddings $n$ and $n+1$ is $\alpha_{n+1} = r\alpha_n$ for some $0<r<1$. This means that $\alpha$ varies the most in the first few layers, and as we add more layers the difference in magnitude between $\alpha_n$ and $\alpha_{n+1}$ rapidly decreases. Therefore, since alpha-clustering determines clustering stability based on the magnitude of the range of $\alpha$ values that it spans, it will tend to select for clusterings that appear earlier in the tree structure. This helps prevent alpha-clustering from over-fitting to noise in the data, as could happen if we selected a clustering that persists across multiple levels higher up the tree.

### Spectral clustering and determination of $k$

Alpha-clustering uses spectral clustering on each layer of the tree. To perform spectral clustering, we first generate a graph as described in Section . Then we compute the graph Laplacian, for which the number of zero-valued eigenvalues gives the number of connected components in the graph (von Luxburg 2006). Therefore, as one-dimensional $t$-SNE embeddings tend to give clearly distinct clusters, if we can build a graph that captures that connectivity, then the graph Laplacian will tell us the number of clusters ($k$) that are present. This allows us to avoid the need for the user to specify a number of clusters, which can be a drawback of many traditional forms of clustering when the structure of the data is unknown. It is important to note that if we were to implement a version of tree-SNE based on two-dimensional $t$-SNE embeddings, this approach would not be as successful, as it would be more difficult to build a clearly meaningfully disconnected graph on the more complex shapes that can appear in two-dimensional embeddings.

### Shared nearest neighbors

For each one-dimensional embedding, we generate a graph via shared nearest neighbors (SNN) (Ertöz 2002). SNN tends to generate more disconnected graphs than k-Nearest Neighbors (k-NN), because in order for two vertices to be connected via an edge, they need to both be nearest neighbors of each other, which reduces the probability of two vertices on the edges of clusters being randomly connected. In order to reduce the probability that two distinct clusters are connected by an edge in the graph, we want to minimize the number of neighbors used; however, the number of neighbors should not be too small, because that could result in a cluster breaking into multiple disconnected components. To balance these requirements, we choose the number of neighbors to be $\beta \log{N}$, where $N$ is the number of observations, and $\beta = 2$ by default. This choice was inspired by the disconnection criterion of random graphs (Erdős and Rényi 1960) and was seen to perform well in practice.

### Subsampling

To improve the efficiency of spectral clustering on large data sets, since it is run on every layer of the tree, we only perform spectral clustering on a random subsample of 2000 observations. Then a k-nearest neighbors classifier is used to cluster the remaining observations. This greatly speeds up the implementation and does not significantly impact the performance of spectral clustering as long as we sample a large enough number of points. Using a k-nearest neighbors classifier also helps to stop individual points that appear separate from clusters from appearing to form their own clusters, by connecting them to the closest large cluster. For intuition as to why we can subsample in this way without ruining performance, see the Nyström method of kernel approximation (Williams and Seeger 2000). The Nyström method could also be used in potential future implementations of alpha-clustering to circumvent the need for a k-nearest neighbors classifier.

### Differentiating factors of alpha-clustering

An important characteristic of alpha-clustering is that it does not require the user to provide $k$, the number of clusters, and instead determines the optimal level of clustering granularity from the tree-SNE embedding. In addition to the optimal alpha-clustering labels, the clustering labels at every level of the tree-SNE embedding are also returned, so the user can explore patterns on each level of granularity.

Unlike most common clustering algorithms, alpha-clustering is easily interpretable because of the accompanying visualization. The user can see the separation between clusters on each individual level, which provides more insight into the data organization than the black box of a neural network and circumvents the need for an external data visualization method to accompany other clustering algorithms like Louvain (Blondel et al. 2008) or K-Means. Not only that, but the user can see the full depth of organization of the data through the hierarchical visualization, which itself provides more flexibility to color by different features or labels than a typical dendrogram. This enables the user to explore the reasons behind specific splits in the tree and track the progression of points and clusters through the hierarchy. This helps address the increasing need for explainable machine learning methods.

# Demonstration on MNIST

## Algorithm choices

To demonstrate the reasoning behind some of our algorithmic decisions, we run variations on tree-SNE with various aspects missing in Figure .

In the top visualization in Figure , we run tree-SNE without decreasing the degrees of freedom at each level, meaning we only decrease perplexity. Notice how the data is organized and for the most part captures the main structure of MNIST; the visualization is split into different colors showing the grouping of different MNIST digits, and except for fours and nines they are cleanly divided. However, there is no cluster separation, and the visualization does not divide the structure into any substructures.

The middle visualization in Figure does not decrease perplexity, and only decreases the degrees of freedom ($\alpha$). The clusters become tighter as $\alpha$ decreases, but they do not split into smaller clusters. We observe in Figure that in two-dimensional $t$-SNE, decreasing $\alpha$ is sufficient to reveal cluster substructure, and decreasing perplexity in tandem is not necessary, as was previously shown in Kobak et al. (2019). However, we see in the middle plot of Figure that decreasing $\alpha$ alone is not enough to produce finer-grained clusters in one-dimensional $t$-SNE embeddings.

The final visualization in Figure is our technique, tree-SNE, which decreases $\alpha$ and perplexity in tandem through the layers. Notice how clusters on earlier levels break into smaller subclusters on later levels, unlike the two other visualizations. We observe that decreasing both $\alpha$ and perplexity appears to be necessary to reveal cluster substructures in one-dimensional $t$-SNE, and tree-SNE therefore takes this approach, balancing the spreading and contracting forces of decreasing perplexity and the degrees of freedom to build a tree structure containing meaningful hierarchy.

## Showing hierarchy

In order to illustrate how tree-SNE displays hierarchical organization of data, we visualized the MNIST handwritten digits data set, then we examined the averaged digit images for each cluster on 17 of the levels (every fifth level between zero and 80, inclusive). The results are shown in two different formats in Figure .

The top plot in Figure shows the tree-SNE visualization of MNIST data that was used for this analysis of hierarchical organization. It was run with a different random seed than the visualization in Figure , but it is highly similar in appearance. This plot is colored by true digit labels, and we observe that the digits separate cleanly for the most part by level 21, the level chosen by alpha-clustering, with two notable exceptions: four is grouped with nine, and three is grouped with five. In both cases, these digits split farther up the tree. It makes sense that these pairs of numbers remain un-separated longer than other digits, considering their similar appearances.

The middle plot in Figure shows within-cluster average digits for each cluster identified by spectral clustering on each of 17 levels of the tree-SNE visualization from the top image. Each averaged image is plotted in the location corresponding to the middle of the cluster from which it was generated, so that the images can be viewed in the tree-SNE tree structure. We observe that the clusters are refined moving up the tree, as $\alpha$ and perplexity decrease. Clusters representing multiple digits separate by digit, and clusters representing a single digit split into smaller clusters representing different handwriting variations for that digit. For instance, ones split into vertical and tilted versions around level 35, and twos split into bubbly and straight sub-groups at level 55. There is a clear hierarchy of digit types, with distinctions becoming finer-grained at the top of the plot. From this visualization, we see that tree-SNE's branching structure conveys meaningful information about the hierarchical organization of the data.

Some overlap between images was unavoidable while preserving the tree structure, so the bottom visualization in the same figure (Figure ) is provided for larger views of the within-cluster average digits. Within-cluster average digits are shown for ten of the 17 layers shown in the middle plot, stacked vertically in the same order as they appear in the tree-SNE visualization. The second row from the bottom represents the stable clustering identified by alpha-clustering, in which only fours and nines and threes and fives have yet to separate cleanly. We draw attention to the top row, where we observe that the rightmost cluster of fours that is last to split from nines is triangular in shape, so in fact is more similar in appearance to nines than to the other fours. Based on these visualizations of within-cluster averaged digits, we see that tree-SNE produces meaningful clusters on each embedding, with hierarchical organization from layer to layer.

## Alpha-clustering

We apply alpha-clustering to MNIST. Recall that alpha-clustering determines the level with the clustering that best fits the overall structure of the data. Figure shows the tree-SNE embedding of MNIST from Figure in the top image, and the same embedding with data points on all layers colored by alpha-clustering labels in the bottom plot. As discussed in Section , alpha-clustering correctly groups six of the ten MNIST classes, and combines the remaining four into two different clusters of two each. It combines four and nine, and three and five. Looking closer at the tree structure in the top plot of Figure , we see that the split that separates three from five happens at roughly the same level that one splits into two different sub-classes. Since alpha-clustering selects the optimal cluster assignment from a single layer, choosing a clustering level that separates three and five would mean also subdividing the ones into multiple clusters.

In order to determine the quality of a clustering assignment, we consider the normalized mutual information (NMI) (Strehl and Ghosh 2002) of a clustering. The NMI is defined as $$\frac{I(X; Y)}{\sqrt{H(X)H(Y)}}$$ where $X$ and $Y$ are two different cluster assignments for the data, $I$ is the mutual information, and $H$ is the entropy. Values of NMI fall between $0$ and $1$, with $1$ representing a complete match between $X$ and $Y$. We can assess the quality of alpha-clustering by computing its NMI relative to the ground truth labels: $X$ = alpha-clustering labels and $Y$ = true labels. On MNIST (10,000 samples), alpha-clustering has an NMI score of 0.84. Importantly, there are is no clustering assignment on any individual layer with a higher NMI. This shows that alpha-clustering is selecting the best possible clustering from all levels of the tree.

# Experiments

We demonstrate tree-SNE on several popular clustering data sets and two biological data sets. The purpose of running tree-SNE on the clustering data sets is to show that it can generate a competitive and interpretable clustering, which means that one-dimensional $t$-SNE embeddings created in this way are indeed able to express meaningful organization of high-dimensional data. Applying tree-SNE to biological data sets allows us to demonstrate its performance on diverse types of data and compare it to standard visualization techniques in biology.

## Experimental Setup

Tree-SNE has three major parameters: the ratio $r$, the initial perplexity, and the number of layers. For the following experiments, we set the starting perplexity to be the square root of the number of observations in the data set ($\sqrt{N}$), as described in Section . Tree-SNE embeddings tend to be exhibit highly branched structures corresponding to very local structures in the data for $\alpha<0.01$, so by default we set the lower bound on $\alpha$ to 0.01. Given this fixed lower bound on $\alpha$, we can solve for the ratio $r$ required to decrease $\alpha$ from 1 to 0.01 in a given number of layers $n$: $$r = e^{\frac{\ln{0.01}}{ n}}$$ Using these sensible defaults, $r$ and perplexity are automatically determined for a given data set and a desired number of layers for the tree. For testing alpha-clustering without visualization, we use 30 layers, and for visualization, we use 100 layers. Using fewer layers has minimal effect on the alpha-clustering and is more computationally efficient for repeated experimentation, while a larger number of layers is ideal for visualization because it produces a higher-resolution plot. For example, for 30 layers, $r$ is $0.858$, and for 100 layers, $r$ is $0.955$. Alpha-clustering has one major parameter, $\beta$, where $\beta \log{n}$ is the number of nearest neighbors used to generated the shared nearest neighbors graph for spectral clustering. We report our results with the default of $\beta = 2$, unless otherwise noted.

## Results of alpha-clustering

We apply alpha-clustering to four standard clustering data sets and compare the NMI to that of several state of the art clustering methods. The results are shown in Table , which is based on Table 4 from Gultepe and Makrehchi (2008). Alpha-clustering achieves a competitive NMI on all of the data sets.

| Method               | COIL-20   | COIL-100      | USPS      | MNIST     |
|:---------------------|:----------|:--------------|:----------|:----------|
| Deep Learning        |           |               |           |           |
| 1-1 JULE\* (2016)    | **1.000** | **0.985**     | **0.913** | 0.913     |
| DBC\* (2017)         | 0.895     | 0.905         | 0.724     | **0.917** |
| IEC\* (2016)         | \-        | 0.787         | 0.641     | 0.542     |
| DEPICT\* (2017)      | \-        | \-            | **0.927** | **0.917** |
| AEC\* (2013)         | \-        | \-            | 0.651     | 0.669     |
| NMF-D\* (2014)       | 0.692     | 0.719         | 0.287     | 0.152     |
| TSC-D\* (2016)       | 0.928     | \-            | \-        | 0.651     |
| DCEC\* (2017)        | \-        | \-            | 0.826     | 0.885     |
| SpectralNet\* (2018) | \-        | \-            | \-        | **0.924** |
| 1-1                  |           |               |           |           |
| 1-1 ICA BSS\* (2018) | **0.965** | **0.962**     | 0.868     | 0.824     |
| AC-PIC\* (2013)      | 0.855     | 0.840         | 0.840     | 0.017     |
| K-means\*            | 0.774     | 0.775         | 0.613     | 0.490     |
| DBSCAN               | 0.892     | 0.705         | 0.299     | \-        |
| 1-1 Ours             | **0.958** | **0.926\*\*** | **0.885** | 0.864     |

Table 1. The NMI scores of various clustering algorithms on different data sets, with top three in bold. For DBSCAN, $\varepsilon$ was tuned until the correct number of clusters was obtained. We performed DBSCAN, K-Means, and tree-SNE benchmarking; the rest of the values are reported from the original papers. \* means the clustering algorithm requires knowing the number of clusters beforehand. \*\* this was run with $\beta=1.3$, whereas the rest were run with the default $\beta=2$. The methods in the table are: Joint Unsupervised Learning (JULE, Yang et al. 2016), Discriminatively Boosted Clustering (DBC, Li et al. 2018), Infinite Ensemble Clustering (IEC, Liu et al. 2016), Autoencoder-based Clustering (AEC, Song et al. 2013), NMF with Deep learning model (NMF-D, Trigeorgis et al. 2014), Task-specific Deep Architecture for Clustering (TSC-D, Wang et al. 2016), Deep Convolutional Embedded Clustering (DCEC, Guo et al. 2017), SpectralNet (Shaham et al. 2018), Independent Component Analysis Blind Source Separation (ICA BSS, Gultepe and Makrehchi 2018), Agglomerative Clustering via Path Integral (AC-PIC, Zhang et al. 2013), K-Means, Density-Based Spatial Clustering of Applications with Noise (DBSCAN, Ester et al. 1996). Benchmarking data sets: COIL-20 (Nene et al. 1996a), COIL-100 (Nene et al. 1996b), USPS (Hull 1994), MNIST 60,000 samples (LeCun et al. 1998).

Note that while most of the other methods featured in Table require specifying a number of clusters, alpha-clustering determines the best number of clusters from the data. In biomedical applications, the optimal number of clusters is often not known *a priori*, so algorithms such as tree-SNE that do not require the number of clusters as input may be particularly useful.

If the number of clusters is known beforehand, the $\beta$ parameter can be tuned as needed such that alpha-clustering yields closer to the correct number of clusters. This is particularly useful when the data has a large number of clusters, as some clusters may be artificially merged in the SNN graph, and a lower $\beta$ value will decrease the number of neighbors used to construct the graph and reduce this merging behavior. We find that decreasing $\beta$ from our default of $2$ to $1.3$ increases the performance of of alpha-clustering on COIL-100 from $0.899$ to $0.926$ and increases the number of clusters found from 76 to 101. Note that changing $\beta$ to give the correct number of clusters does not always increase the NMI, as the clusters that split apart with lower $\beta$ do not necessarily correlate with true label clusters. For example, tuning $\beta$ on COIL-20 to give the correct number of clusters results in a decreased NMI of 0.951.

It should be noted that all of the benchmarking data sets used are image data sets: the two Columbia Object Image Library data sets, COIL-20 (Nene et al. 1996a) and COIL-100 (Nene et al. 1996b); as well as two handwritten digit data sets, USPS (Hull 1994) and MNIST (LeCun et al. 1998). These data sets were chosen due to readily available comparison metrics from other methods. However, convolutional neural network-based approaches have an advantage on these data sets because they are designed to process image data.

## Results of tree-SNE

We demonstrate running tree-SNE on several different data sets, including mass cytometry (CyTOF) data from bone marrow cells (Bendall, et al. 2011) and single-cell RNA-sequencing (scRNA-seq) data from retinal cells (Shekhar et al. 2016). We compare tree-SNE with $t$-SNE (van der Maaten and Hinton 2008) and PHATE (Moon et al. 2019), standard approaches in biomedical data visualization.

### CyTOF

Figure shows a tree-SNE embedding of bone marrow single-cell mass cytometry (CyTOF) data (Bendall et al. 2011), colored by the expression levels of specific markers in each cell. The top visualization is colored by intensity of CD8, which is a marker for cytotoxic T cells, and the bottom is colored by intensity of CD45ra, which is a marker for naive T cells (Golubovskaya and Wu 2016). Notice how in the top visualization most of the cytotoxic T cells are in the same bright arm, which subdivides at higher embedding levels. This indicates that tree-SNE separates cytotoxic T cells on lower levels of the embedding, and then splits them into subgroups higher up as $\alpha$ decreases. Importantly, the CD8 branch splits in two about halfway up the plot. The CD45ra visualization (bottom of Figure ) reveals that the split in the cytotoxic T cell line is mediated by the presence of CD45ra: CD45ra is present in one of the branches but not the other. This indicates that tree-SNE is splitting cytotoxic T cells ($\textrm{CD8}^+$) into naive ($\textrm{CD45ra}^+$) and mature cells ($\textrm{CD45ra}^-$). In this way, tree-SNE very clearly reveals the hierarchical organization of the CyTOF data and facilitates identification of the features mediating the branching.

For comparison with established visualization methods for biological data, we show the same data embedded via $t$-SNE (van der Maaten and Hinton 2008) in Figure and via PHATE (Moon et al. 2019) in Figure .

### scRNA-seq

We demonstrate tree-SNE on single-cell RNA-sequencing (scRNA-seq) data from mouse retinal cells and compare the results with the findings from the original paper by Shekhar et al. (2016). Shekhar et al. used Louvain community detection (Blondel et al. 2008) to detect clusters, two-dimensional $t$-SNE to visualize the clusters, and hierarchical clustering with Euclidean distances and average linkage to produce a hierarchical depiction of the relationships between different mouse retinal bipolar cells (BCs). We run tree-SNE on the same data and label it according to the labels assigned by Shekhar et al. from clustering and morphological analysis. Some pre-processing, mediated by alpha-clustering, was performed to roughly isolate the cone bipolar cells, which encompass the cell subtypes evaluated by hierarchical clustering by Shekhar et al. The full process is described in Appendix A (Section ). Figure shows the tree-SNE embedding of the data subset.

Although the tree-SNE hierarchy and the dendrogram created by Shekhar et al. do not match entirely, there are strong similarities between the two. Notice how BC1A and BC1B split apart at a high level of the tree, and how they are more closely related to BC2, BC3A, BC3B, and BC4 than to the BC5 subtypes or BC7. Observe also how tree-SNE show the same organization of the four subtypes of BC5 as Shekhar et al. discovered: BC5A and BC5D are closely related, and BC5B and BC5C are closely related. Tree-SNE also shows further branching of some of the cell types, and these subdivisions may warrant further examination.

### Swiss roll

A common concern about $t$-SNE is that it does not fully capture some standard artificial data sets, including the famous Swiss roll (Linderman and Steinerberger 2019). In the case of the Swiss roll, $t$-SNE does not fully unwind the manifold, as seen in Figure . However, tree-SNE does for the most part unwind the Swiss roll, as seen in the top plot of Figure . Notice how on the bottom two levels of the embedding, the order of the points is wrong in the same way that it is wrong in the $t$-SNE embedding, with purple in the middle as opposed to at the end. However, tree-SNE's multi-scale approach fixes this issue and places the purple segment where it should be, correctly capturing the manifold. This phenomenon is highlighted in the bottom plot of Figure , which shows only the first four levels of the tree-SNE plot.

Additionally, notice that there is very little cluster separation in the tree-SNE embedding of the Swiss roll in Figure , and no well-defined branching. This shows that tree-SNE does not artificially introduce clusters or hierarchical structures when applied to data lacking this type of organization. We illustrate this point further in the next section (Section ). This suggests that the clusters that emerge in other datasets are unlikely to be artifacts of noise, as otherwise similar clusters would have emerged here.

Although we demonstrate tree-SNE on the Swiss roll data set because it is a common "toy" data set and the behavior of tree-SNE when applied to the Swiss roll is interesting, tree-SNE is by no means a dimensionality reduction method. It is designed to visualize hierarchical cluster structures and is limited in its ability to represent manifolds.

## Random noise

To demonstrate that tree-SNE does not artificially introduce clusters or hierarchical structure into data lacking those structures, we run tree-SNE on 5,000 random samples drawn from a 100-dimensional uniform distribution over the range \[0,1). The results are displayed in Figure . Observe how, unlike the tree-SNE visualizations on meaningfully organized data, the tree does not begin to branch until level 50, and the subdivisions thereafter are long, without further branching. This shows that tree-SNE does not tend to produce clear cluster separation or hierarchical structure by mere artifact. Further note that there are many points seemingly randomly scattered in the embedding, which does not tend to appear in well-organized data.

In Figure , we generate 5,000 samples of uniform noise in a two-dimensional plane in 100-dimensional space, then embed the samples with both tree-SNE and $t$-SNE. The tree-SNE plot does not have as long of an initial stem as it does in Figure , because the data lies along a plane rather than completely randomly throughout the 100 dimensions. The tree-SNE plot still has relatively uniformly distributed branches without a clear hierarchy, which does not tend to happen with organized data. The standard two-dimensional $t$-SNE embedding is provided for comparison. The $t$-SNE embedding gives the impression of clumping or clustering, despite the fact that the data is completely random, whereas tree-SNE does not show any clusters.

Together, these experiments on random noise illustrate that tree-SNE does not introduce artificial structure into unstructured data, and the tree-SNE plots of noise have a characteristic profile that is distinguishable from the plots of well-organized data.

# Conclusion

We present tree-SNE, a novel hierarchical visualization method for high-dimensional data, and a companion clustering method, alpha-clustering. Tree-SNE provides a way to explore the organization of data at multiple levels of granularity, which is not available in many other visualization techniques. Alpha-clustering is built on top of the tree-SNE embedding and provides quantitative validation that tree-SNE reveals meaningful structures in the data. It also achieves NMI scores on various data sets that are competitive with state-of-the-art clustering approaches, without ever needing to specify the number of clusters. Finally, we demonstrate the viability of this approach on several diverse real-world data sets. Our analysis suggests that tree-SNE offers promise in terms of yielding insights not currently revealed by other visualization approaches.

In the future, we intend to modify the tree-SNE algorithm to be able to leverage neural network embeddings, with the hope that this would help us be more competitive with neural network-based approaches, particularly on image or sequential data. We hope to create a means for out-of-sample extension of alpha-clustering such that the clustering approach can be used for more than exploratory data analysis, perhaps leveraging previous work on parametric $t$-SNE (van der Maaten 2009). We also want to explore the idea of a variation of alpha-clustering that varies scale ($\alpha$) differently for different branches, acknowledging that a global scale may not be optimal for all data sets. In addition, we want to investigate using alternative clustering approaches instead of subsampled SNN-based spectral clustering, perhaps leveraging the Nyström method. And finally, we want to extend tree-SNE to support stacked two-dimensional $t$-SNE plots.

# Acknowledgements

The authors thank Stefan Steinerberger for inspiration, support, and advice; George Linderman for enabling one-dimensional $t$-SNE with degrees of freedom \< 1 in the FIt-SNE package; Scott Gigante for data pre-processing and helpful discussions of visualizations and alpha-clustering; Smita Krishnaswamy for encouragement and feedback; and Ariel Jaffe for discussing the Nyström method and its relationship to subsampled spectral clustering.

# Appendix

## Appendix A: Methods for subsetting scRNA-seq data with alpha-clustering

In Section , we applied tree-SNE to a subset of a scRNA-seq data set from retinal cells (Shekhar et al. 2016). Now, we demonstrate using alpha-clustering to partition the full data set from Shekhar et al. to obtain the subset of interest used in Section . Prior to running tree-SNE, we use PCA to reduce the dimensions of the data set from 24,576 features to 100 dimensions to increase the efficiency of the algorithm. We start with the tree-SNE embedding of the full scRNA-seq data set, in Figure , colored by the labels from Shekhar et al. Notice that the Muller Glia (teal, leftmost branch) and Rod BC (dark green, rightmost branch) cell types take up a lot of space in the visualization, crowding the main group of cone bipolar cell (BC) subtypes of interest. We would like to remove the Müller glia and rod BCs to reduce this crowding, and we will use alpha-clustering labels to partition the data to do so.

Figure shows the same tree-SNE embedding of the full scRNA-seq data set, this time colored by alpha-clustering labels. Alpha-clustering automatically determined that the optimal number of clusters was three, and these three clusters correspond to (left to right) Müller glia, cone BC subtypes and photoreceptor (PR) cells, and rod BCs.

Since alpha-clustering has already partitioned the data, it is simple to remove clusters 0 and 2 and re-apply tree-SNE to only the observations within cluster 1, allowing for a visualization that has more space to preserve the hierarchical organization of the cells of interest. This new embedding is shown in Figure . Notice that on the far right, there is still a small group of rod PR cells present in this plot, which are not a cell type of interest.

Running alpha-clustering on the tree-SNE embedding from Figure yields an optimal cluster assignment of just two clusters; again, the number of clusters was automatically detected by the algorithm. Figure is colored by the alpha-clustering labels, showing that the two clusters discovered from this subset of the data correspond to the majority of the tree (left) and the rod PR cells (right).

The alpha-clustering labels make it trivial to remove cluster 1 and re-embed the remaining data from cluster 0. For this final round of tree-SNE, we also reduce the number of PCA components used to 37 to more closely match the methods from Shekhar et al. (2016) and enable more direct comparison. The final tree-SNE embedding is shown in Figure (which is the same as Figure shown previously), and is discussed in more detail in Section .

The ability to automatically detect, partition, and remove particular significantly different subsets of the data is, to the best of our knowledge, unique to tree-SNE and alpha-clustering. This allows a researcher to very easily select and hone in on subsets of interest in their data.


abstract: Despite the constant evolution of similarity searching research, it continues to face the same challenges stemming from the complexity of the data, such as the curse of dimensionality and computationally expensive distance functions. Various machine learning techniques have proven capable of replacing elaborate mathematical models with combinations of simple linear functions, often gaining speed and simplicity at the cost of formal guarantees of accuracy and correctness of querying.
  .
  The authors explore the potential of this research trend by presenting a lightweight solution for the complex problem of 3D protein structure search. The solution consists of three steps – (i) transformation of 3D protein structural information into very compact vectors, (ii) use of a probabilistic model to group these vectors and respond to queries by returning a given number of similar objects, and (iii) a final filtering step which applies basic vector distance functions to refine the result.
  .
  .
  [^1]
author: Jaroslav Olha; Terézia Slanináková; Martin Gendiar; Matej Antol ; Vlastislav Dohnal
bibliography: citations.bib
title: Learned Indexing in Proteins: Substituting Complex Distance Calculations with Embedding and Clustering Techniques

# Introduction

The methods and approaches developed by the similarity searching community are used by a wide range of scientific fields, both within computer science and beyond. While some applications require provable accuracy guarantees, verifiable algorithms, or support for complex similarity functions, many of the similarity searching problems emerging in various areas of research do not have such strict formal constraints.

Often, the representations of data or their similarity functions are inherently unreliable, so the idea of a perfect similarity search is unattainable regardless of the quality of the search method. Other times, the problems are too computationally complex, and instead of a perfect answer that would take too long to compute, a best-effort answer in a reasonable amount of time is preferable (perhaps to be refined later using more costly methods). Since the relevance of research on similarity searching is determined by its ability to solve current problems, it should strive to employ the widest possible array of problem-solving tools.

There are multiple similarity search indexes, falling under the umbrella of approximate searching, capable of adjusting to such use cases by lowering their accuracy thresholds or returning partial results . However, in recent years, an entirely new approach has begun to gain traction – along with most other areas of computer science, the area of data retrieval has started to incorporate various machine learning approaches. These approaches have the potential to not only expand the available toolbox of problem-solving instruments, but also to offer a new way of thinking about the problems themselves. Notably, in 2018, Kraska et al.  suggested that all conventional index structures could be viewed as models of data distributions, implying that machine and deep learning models could be used in their place. Even though the idea was originally proposed and tested on structured data, this reframing of the problem has already inspired similar work in the realm of unstructured datasets .

To investigate the potential of these approaches further, we have chosen to examine the problem of 3D protein structure similarity search. This is an important open problem in biochemical and medical research, which can be viewed as an instance of similarity searching in non-vector datasets, because similarity between a pair of protein structures is usually calculated using a series of non-trivial, computationally expensive operations. Additionally, the amount of 3D protein structure data is currently exploding due to a recent breakthrough in the field , and the demand for versatile similarity searching approaches is likely to grow in the near future.

In this paper, we demonstrate that even a relatively complex interdisciplinary problem such as 3D protein structure retrieval can be tackled with fast and lightweight solutions. We present a simple pipeline where protein structures are first transformed into short vectors and used to train multiple partitioning and classification models – these are linked together to form a learned index structure. The index then answers queries by returning several candidate leaf nodes, and filtering the objects stored therein using basic vector (similarity) functions[^2].

This approach, while based on probabilities rather than mathematical guarantees, provides a reasonable quality of results at a fraction of the computational costs required by previous methods. In addition, its modularity allows us to change algorithms or their parameters for various trade-offs between complexity and accuracy, depending on the particular use case.

# Related work

*Learned indexing* was first introduced in  with the core idea of learning a cumulative distribution function (CDF) to map a key to a record position. This proposition challenged the long-standing paradigm of building indexes solely with classic data structures such as B-trees and allowed for reduction in searching costs compared to the traditional methods. To allow for indexing of large data collections, the authors introduced *Recursive model index (RMI)* – a hierarchical tree index structure of simple interconnected machine learning models, each learning the mapping on a subset of the dataset. RMI is, however, limited to sorted datasets of structured data, and cannot accommodate multi-dimensional data.

The generalization of the learned indexing concept to spatial and low-dimensional data was explored primarily by the *Z-order model* , which makes use of the space-filling curve encoding to produce 1D representation of the original data, and *ML-index*  which achieves the same with the use of *iDistance* . *RSMI*  introduced a recursive partitioning strategy, building on the Z-order model. Furthermore, *LISA*  and *Flood*  both partition the data into grid cells based on data distribution, improving the range and kNN performance of prior approaches. These approaches were recently directly compared to (and surpassed by) *LIMS* , which generalizes to metric spaces and establishes a new state-of-the-art performance on datasets of up to 64 dimensions.

Indexing solutions for approximate, rather than precise queries were explored by Chiu et al.  introducing a probability-based ranking model able to learn the neighborhood relationships within data, which can be integrated with existing *product quantization (PQ)* methods. This method was tested on 1-billion datasets with approximately 100 dimensions and has been shown to boost the performance of traditional PQ methods.

Following the architectural design of RMI, we proposed the *Learned metric index (LMI)* , which can use a series of arbitrary machine learning models to solve the classification problem by learning a pre-defined partitioning scheme. This was later extended to a fully unsupervised (data-driven) version introduced in , which is utilized in this work.

Finally, the ideas of creating an indexing solution with machine learning have found their use in many different domains, e.g., in trajectory similarity search  or information retrieval, using a single transformer model .

## Protein representation

To enable computational approaches to the problem of uncovering functional properties of proteins, a great amount of research attention has been directed to creating representative (numerical) embeddings of protein structures. There are two distinct categories of embeddings based on the input data – those that operate with *sequences* and those working with *3D structures*. Although they serve a similar purpose, these two categories are completely distinct in terms of the technical approaches they employ. Specifically, sequence embeddings use techniques such as hidden Markov models  or various natural language processing techniques  to derive meaning from protein sequences, treating them as encoded sequences of characters – this approach is not applicable to our research.

Embeddings representing protein 3D structures are generally less elaborate, since the information content is more robust to begin with. The most common encoding is a protein distance map, which produces a symmetric 2D matrix of pairwise distances between either atoms, groups of atoms or amino acid residues. This distance map can be transformed into a protein contact map, which is a binary image where each pixel indicates whether two residues are within a certain distance from one another or not. Contact maps have been used in conjunction with machine learning techniques for prediction of protein structure ; another studied problem is the reconstruction of 3D structure based on information in contact maps . While these techniques are related to our own approach, we produce embeddings that are considerably more compact and reflective of our similarity searching use case, as will be shown in the following sections.

# Data domain

We have chosen to test our approach on 3D protein structures for several reasons. First, while protein structure data is very widely used, and the study of this data is vital for almost every area of biochemical research, the issue of efficient search and comparison of protein structures is still unresolved to some extent, with many databases still relying on time-consuming brute-force linear search . This data is also publicly available in a single database, called the Protein Data Bank (PDB), which is used by the majority of protein researchers and widely agreed upon as the standard. Even though this database is large enough to require efficient search methods, its size ( structures as of 2022) still makes it possible to download the entire database and scan it exhaustively for ground truth answers if necessary.

Just as importantly, it is clear that the issue of efficient search within this data will only become more crucial and challenging in the next few years. Firstly, the common dataset of empirically solved protein structures continues to grow exponentially  – its current contents constitute a mere fraction of all the protein structures in nature, and the complex laboratory procedures needed to obtain these structures are being refined every year. Secondly, the computational prediction of protein structures from their sequences has recently seen rapid improvement with the release of AlphaFold 2 in 2021  – this has already resulted in hundreds of thousands of new reliable 3D structures that are of great interest to researchers, with tens of millions more to be added in the coming months .

Interestingly, the question of efficient protein structure search is not only important for storage and retrieval purposes – for instance, researchers often discern the function of unknown proteins by comparing them to other, better-known proteins. Since the function of a protein is entirely dictated by its 3D structure, this is equivalent to searching a database for the most similar protein structures. However, the specific needs of this type of research can vary – while some researchers are looking for extremely specific deviations among a group of very similar proteins (e.g., when studying mutations or conformational changes), others might be looking for much broader patterns of similarity between distant protein families.

Structurally, a protein chain consists of a linear sequence of interconnected building blocks called amino acids. Within the right biochemical environment, this linear sequence folds in on itself to form a complex 3D structure. Protein structures are sometimes cited as a typical example of complex unstructured data, since they cannot be meaningfully ordered according to any objective criteria (any search method needs to rely on pairwise similarity), and the similarity of two protein structures often cannot be determined by a single vector operation. Typically, protein molecule data are stored using the three-dimensional coordinates of each of their atoms, with the protein randomly oriented in space. In order to compare a pair of protein structures, they first need to be properly spatially aligned in terms of translation and rotation, and a subset of atoms must be selected for alignment. This typically involves gradual optimization of a spatial distance metric (such as the root-mean-square deviation of all the atom coordinates), which is a computationally expensive process that cannot be directly mapped to a simple vector operation.

One commonly-used measure of protein similarity is the $Q_{score}$ , which is calculated by dividing the number of aligned amino acids in both protein chains by the spatial deviation of this alignment and the total number of amino acids in both chains. Even though this measure is imperfect and not appropriate for all use cases, it is used in several prominent applications, including the PDB's own search engine.

Note that two identical structures have a $Q_{score}$ of 1, and completely different structures have a Q-score of 0: as a result, the score needs to be inverted in order to be used as a distance metric ($d(x,y) = 1 - Q_{score}(x,y)$). In the following sections, we will refer to this inverted value as $Q_{distance}$.

While there are a few types of protein structure embeddings that are invariant to the spatial alignment of the molecules (see the Related Work section), these were not developed for the purpose of fast data retrieval. As a result, they tend to be too detailed and cumbersome to serve as simple data descriptors. By contrast, the embedding presented in this work has been designed specifically to contain the optimal amount of information for efficient and accurate similarity searching, as will be shown in the following sections.

# Fast searching in proteins

We present a pipeline (visualized in Figure ) consisting of three separate components: (i) a simple embedding technique for protein data in the PDB format, (ii) the use of a machine-learning-based index – Learned Metric Index (LMI) – to locate a candidate set of similar protein structures, and (iii) fast filtering to produce the final query answer.

The embedding we propose divides the protein sequence into $N$ consecutive sections – the positions of the atoms within each section are averaged, and the section is subsequently treated as a single point in space. We then calculate distances between each pair of these sections, creating an incidence matrix. In this matrix, we prune all the values exceeding a cutoff, and normalize the rest. The matrix is symmetrical and all the diagonal values are 0. The half of this matrix (omitting the main diagonal) is then reduced into a single row in a $M \times (\frac{N^2-N}{2})$ matrix, where $M$ is the number of proteins in the entire dataset (see Figure ).

This produces a very compact embedding for all the proteins, and the entire dataset can be represented by a file that is up to two orders of magnitude smaller than the original database – see Table .

| Embedding Size | File Size | Index Build Time (256-64) | Index Build Time (128-128) |
|:--:|---:|---:|---:|
| $5\times5$ | $16$ MB | 246s | 184s |
| $10\times10$ | $51$ MB | 350s | 270s |
| $30\times30$ | $456$ MB | 927s | 655s |
| $50\times50$ | $1275$ MB | 2391s | 1814s |

File size of the protein dataset (518,576 protein chains) stored using protein embeddings, and build times of two different LMI architectures. Note that the size of the original database is 8.2 GB.

To reduce the search space to a small number of candidate protein structures, we used the Learned Metric Index (LMI), a tree index structure where each internal node is a learned model trained on a sub-section of data assigned to it by its parent node . Specifically, we used the data-driven version of LMI, where the partitioning is determined in an unsupervised manner. We explored different architectural setups – both in terms of the number of nodes at each level (index breadth), as well as the number of levels (index depth). As the learned models, we explored K-Means, Gaussian Mixture Models, and K-Means in combination with Logistic regression (see  for details regarding the model setups). For the sake of compactness, in the experimental evaluation we only present the results achieved with the best-performing setup – a two-level LMI structure with arity of 256 on level 1 and 64 on level 2 (i.e., 256 root descendants, each of them with 64 child nodes), with K-Means chosen as the partitioning algorithm. After LMI is built, we search within it using a query protein structure and return target candidate sets; the size of the candidate sets is determined by a pre-selected stop condition (for instance, a stop condition of 1% of the dataset corresponds to $\sim5,000$ candidate answers per query).

In the final step, we filter the candidate set according to a particular distance function. In our experiments, we have examined filtering based on the Euclidean distance as well as the cosine distance of the vector embeddings, but the filtering step could theoretically be performed using any distance metric, or even the original $Q_{score}$ similarity of the full protein structures. The filtering step returns a subset of the candidate set based on the specified criteria (i.e., kNN or range).

# Experimental evaluation

We evaluated our approach using range queries, with 512 randomly chosen protein chains from the dataset used as query objects. In order to compare our results against the ground truth, we needed to know the $Q_{distances}$ (based on $Q_{score}$) between the 512 protein chains and all the other chains in the database – these distances were kindly provided by the researchers behind , where the same 512 objects were used as the pivots for their search engine. The objects were chosen uniformly randomly with respect to protein chain length, which ensures that even very long proteins are represented among our queries (despite constituting a relatively small portion of the dataset).

We expected the performance of our method to deteriorate as the range of the queries expands, since a wider search range would require the method to correctly identify more objects which are less similar. To examine this effect, we have chosen three representative query ranges of 0.1, 0.3 and 0.5 – in a real use case, the range would be chosen by a domain expert based on the particular use case. As a rule of thumb, a range of 0.1 represents a high degree of similarity, while a range of 0.5 represents low (but still biologically significant) similarity; the biological relevance of answers drops sharply beyond this range .

First, we evaluated the performance of the LMI, before the filtering step. The recall shown in Figure pertains to the entire candidate set of objects (i.e. how much of the ground truth answer is contained in the 1%/5%/10% of the dataset returned by the LMI for further filtering)[^3].

This figure presents us with two important pieces of information – firstly, it is clear that LMI can reach very high recall even when trained on the smaller 10x10 embedding – this makes the embedding a natural choice for further evaluation, since it is efficient while significantly reducing the memory and CPU costs of training compared with the larger embeddings. It can also be seen that, especially in the lower query ranges which are of most interest to us, the 1% stop condition represents a sensible trade-off between recall and search time, returning relatively few candidate objects () while minimizing the amount of false negatives.

In addition to recall, we also need to investigate the distribution of objects in the index, to ensure that the occupancy of leaf nodes (i.e. data buckets) is not overly imbalanced – an extremely imbalanced index would achieve high recall, but it would also be more likely to return overly large candidate sets, which would be detrimental to the filtering step. The size distribution of the buckets is shown in Figure , and it confirms that the distribution is not overly skewed towards large buckets, even when the embeddings are quite small, as is the case with N=10. However, the embeddings can only be condensed up to a point – the 5x5 embedding (which transforms each protein structure into a vector of only 10 values) causes a large portion of the objects to concentrate in a single bucket, as the LMI can no longer distinguish among groups of objects.

During the filtering step, recall naturally decreases over time (since the method occasionally filters out relevant answers), while precision should improve as the portion of relevant objects in the candidate set increases. This effect will be different depending on the distance metric used for filtering, as well as the dataset – we show the effects for two distance metrics (Euclidean distance and cosine distance) in Figure . By analyzing the correlation between $Q_{distance}$ and each of the two distance metrics, we have determined that the Euclidean distance is the better filtering function for this dataset.[^4]

Table  shows the final results of the range queries with the best-performing configuration of parameters: embedding size N=10, the K-means clustering model, 256-64 LMI architecture, and filtering after the 1% stop condition using the Euclidean distance metric. The results, especially in the lower query ranges, are very encouraging, although the filtering stage seems to introduce a surprisingly large amount of false negatives by filtering out parts of the correct answer. It is likely that the filtering metric we have chosen was slightly too naïve, and the filtering step could have benefited from a different distance function, or at least a different weighting of the vectors before calculating their Euclidean distance. In the future, this presents a natural point of focus to improve our results even further.

|  |  |  |  |
|:---|---:|---:|---:|
|  | Range 0.1 | Range 0.3 | Range 0.5 |
|  | Mean \# of objects: 83 | Mean \# of objects: 236 | Mean \# of objects: 519 |
| LMI Recall | 0.973 (1.000) | 0.895 (0.999) | 0.755 (0.867) |
| Recall after filtering | 0.742 (0.878) | 0.649 (0.711) | 0.530 (0.637) |
| F1 after filtering | 0.712 (0.855) | 0.669 (0.766) | 0.592 (0.673) |

Overall evaluation of protein range queries: the average values, as well as the median values (in parentheses) are shown.

Figure  further illustrates the strengths and weaknesses of our approach by showing the relationship between the $Q_{distance}$ metric and the Euclidean distances of vector embeddings. There is clear correlation between these two metrics, making a strong case for the simpler one, and LMI is much more successful at finding objects that are more similar to the query (left side of the graph) than it is at finding less similar objects (right side of the graph). It should be restated that even though we use the $Q_{distance}$ metric as the "ground truth", it is merely a subjective similarity metric, and the relevance of results that are close to the edge of the query range is debatable (i.e., an object with a $Q_{distance}$ of 0.499 is not necessarily more relevant than an object with a $Q_{distance}$ of 0.501).

Distance computations on long proteins pose a considerable challenge for similarity searching methods, typically requiring a disproportionate amount of computing time . This does not apply to our embedding approach, which transforms all the proteins into fixed-length vectors using the same algorithm, regardless of their length. It would therefore intuitively follow that such an approach should achieve lower recall on the longer protein queries, since more information is lost in the embedding. However, in practice, that is not the case, as can be seen in Figure .

This is probably due to the simple fact that the distribution of chain lengths in protein databases is not uniform – in fact, the long chains are much less numerous. As a result, losing a lot of information about the long chains is not a problem, since they are still relatively easy to find. This gives the fixed-length embedding approach a significant performance advantage – in a dataset where there are relatively few long protein chains, it would be a waste of resources to calculate and store an excessive amount of data about them, the way many other methods do.

As has been mentioned in the previous section, in addition to variable chain lengths, the proteins also have a variable number of neighbors in any given range. This can, again, call into question the reliability of the results based on recall, since recall is a relative measure, and the size of the error will be different based on the size of the actual answer. For instance, if all range(0.1) query answers only consisted of a single easy-to-find object, evaluating recall on its own would give us a biased idea of the searching efficiency. Figure shows how many neighbors each protein structure has according to the ground truth, and how many of these neighbors have been found using our approach – while there are significantly more proteins with fewer than 10 neighbors in the lower query ranges, these are not the majority, and the errors are distributed evenly relative to query answer size.

Finally, to provide broader context for the pipeline's performance, we have evaluated it against a more conventional approach. While there is no analogous method for searching 3D protein structures using range queries, there is a similar, recently-published method for nearest neighbor retrieval in the same database of 3D protein structures . This method uses a three-stage search engine which compares bit-strings in the Hamming space ("sketches") to approximate the distance of protein chains.

To allow for a relatively fair comparison of these two methods, we modified our search parameters to more closely match the search results presented in . Specifically, since their similarity queries were mainly 30NN queries limited by the range 0.5, we have performed range queries in the same range (even though our method's performance is substandard in this range), and in the filtering step, we filtered out all objects beyond the 30 best-ranking answers. Since the sketch-based method was originally compared with the linear search of the PDB database, we present this benchmark as well – naturally, this is always the slowest method by far, but it requires no index and always finds the exact answer.

All of these results can be found in Table  – while our method clearly does not match the high accuracy of the sketch-based method in this experimental setup, it is faster by at least an order of magnitude, occasionally exceeding 4 orders of magnitude since it does not suffer from an extreme "tail" of worst-case search times caused by evaluation of long proteins.

|                   |  LMI + Filtering | Sketch-based method | PDB Engine |
|:------------------|-----------------:|--------------------:|-----------:|
| Accuracy (median) |          $0.660$ |               $1.0$ |      $1.0$ |
| Accuracy (mean)   |          $0.626$ |             $0.937$ |      $1.0$ |
| Time (median)     |        $0.094$ s |             $2.5$ s |    $183$ s |
| Time (max)        |        $0.145$ s |            $6109$ s |  $14321$ s |
| Index size        |          $87$ MB |            $178$ MB |        N/A |

The accuracy, search times, and memory requirements of 30NN protein search queries with a maximum distance radius of 0.5.

# Summary and Conclusions

In an effort to investigate the potential of new data retrieval techniques in the field of similarity searching, we have developed and evaluated a novel approach to the problem of protein structure search, resulting in a short pipeline consisting of a concise vector embedding, learned indexing, and distance-based answer filtering. By successfully applying this approach on a well-established database of 3D protein structures, we have shown that even in a domain that may, at first, seem poorly suited to simple vector-based transformations, a surprising amount of information can be discerned by learned models.

One advantage of our modular approach is that every part of the pipeline can be evaluated separately, allowing experts to identify the weakest spots and alter them based on the current use case and dataset. The experiments presented in this paper serve as a good example – after evaluating each part of our own pipeline, it is clear that we have chosen an overly simplistic filtering method for our data. In the future, we plan to investigate more sophisticated options for vector-based filtering, as well as a completely different approach to reducing the size of query answers.

While it is difficult to compare our work with the state of the art (since there are no direct analogues to our method in the chosen data domain), we have made an effort to modify our method for the fairest possible comparison with a recent, more conventional similarity searching approach in the same domain. In this comparison, our solution, although coming up short in terms of accuracy, is consistently faster by multiple orders of magnitude, and maintains much lower memory requirements.

Our work aims to make a case for less conventional solutions to similarity searching problems – ones that rely on learned approaches, rather than traditional indexing models, to discern the natural distribution of data. This is by no means an argument for universal adoption of machine-learning-based techniques in all similarity searching applications. The approach in this paper is problem-specific, and has several potential drawbacks that need to be weighed against its strengths. Nevertheless, we provide some insight as to how the trend of machine-learned pattern recognition, which has already caused minor and major revolutions in most fields of computer science, could be applied in practical solutions to current similarity searching problems.

Based on these results, we remain convinced that learned indexing approaches (such as the Learned Metric Index used in this work) will play an integral part in shaping the future of similarity searching.

[^1]: The publication of this paper and the follow-up research was supported by the ERDF "CyberSecurity, CyberCrime and Critical Information Infrastructures Center of Excellence" (No.CZ.02.1.01/0.0/0.0/16_019/0000822). Part of this work was carried out with the support of ELIXIR CZ Research Infrastructure (ID LM2018131, MEYS CR). Computational resources were supplied by the project "e-Infrastruktura CZ" (e-INFRA CZ LM2018140 ) supported by the Ministry of Education, Youth and Sports of the Czech Republic.

[^2]: The entire functionality is publicly available as a search engine prototype at <https://disa.fi.muni.cz/demos/lmi-app-protein/>

[^3]: Precision is not evaluated in this step – at this point in the pipeline, it is very low and not particularly relevant.

[^4]: Note that the actual $Q_{distance}$ function scales differently from the Euclidean distance used to filter answers (see Figure ) – this requires simple re-scaling (e.g., to find range=0.5 queries, we set the Euclidean distance cutoff at 0.75).