Space efficient data structure to store precomputed All Nearest Neighbors in high dimensions

Question

Seeking an indexing data structure that is smaller than quadratic in space.

As part of an NLP algorithm using word embeddings of 300-dimensions, I am trying to improve the speed of Word Mover's Distance (WMD). (Isn't everybody?) Computing the cost matrix is expensive, as it is size MxN (where M and N are the number of distinct non-stop words in documents one and two, respectively). I figured that if I only compute the distances to the k-nearest neighbors of each word and set all others to the mean distance between word pairs, I can perform a high quality approximate value for WMD. (I am not the first to consider this flavor of speedup. See "Speeding up Word Mover’s Distance and its variants via properties of distances between embeddings" by Matheus Werner and Eduardo Laber, 2020 on arXiv.) The trick is knowing what those nearest neighbors are a priori.

The corpus size is S = 250,000 words.

The dimensionality is 300 features per word embedding.

What if you preprocess the corpus of words, performing all nearest neighbors search for every word and store it in a file? That would allow you filter the k-nearest words for each word present in a given document comparison from the index, saving the computation of the distance function for most word pairs and the sorting time of (M + N) * k log k.

It is trivial to store a matrix of size S * S with the neighbors of each word sorted from nearest to farthest.

Is there a more space efficient structure than quadratic storage, like a variant on a Trie? I do not need it to store the actual distances, just the ordering. I can recalculate the distances for the pairs I care about. The time complexity of the preprocessing is immaterial.

Clarification of K:

Say that you want the k-nearest neighbors of the word "happiness". You might think the index storage requirement is k*S. However, what I want is NOT the K-nearest neighbors of happiness in the whole corpus of S words. I want the k-nearest neighbors drawn from the set of words in the opposite document. The 10th nearest word to happiness in the document could be the 10,000th nearest word in the whole corpus! That is the problem. To guarantee that I have all k-nearest words for ANY document that is produced for comparison, I need to store ALL neighbors for happiness from nearest to most distant.

As an exception, I will allow that once the distance to a word reaches the average distance between word pairs in S, it can be dropped from the index and the distance assumed to equal that mean distance. This automatically reduces the index size by half.

Research:

I have just learned about LSH, HNSW and IVF indices. You have control over accuracy versus speed, but the memory usage is substantial.

I have toyed with using Hilbert indices with lists of line segments sorted to give a loose ordering of neighbors that will need to be sorted later. That reduces number of computations of distances but does not do away with the K log K sort at the end.

Answer 1

After pondering this for days, I have arrived at a possible data structure. It may be some time before I get to testing whether it performs as well as hoped.

The index will be derived from the Hilbert Curve ordering of the points. The index construction time will be O(N-cubed).

1. Normalize and Quantize coordinates of the word embedding vectors to form integer valued points. 16-bits per coordinate should suffice, and likely fewer bits.
2. Obtain Hilbert index of each point.
3. Sort points by Hilbert index.
4. Once sorted, each point will be associate with its position in Hilbert order. Instead of being a huge Hilbert index of 300x16 bits, it will be an integer in [0,S) where S is the number of words in the corpus.
5. For each point, find all nearest neighbors exhaustively, comparing distance from every point to every other point and sorting.
6. For each point, create an index of neighbors that is a list of Hilbert Curve line segments.
7. Store the index to a file.

The sixth step is the tricky part. A line segment will specify:

• start Hilbert position
• length of segment in word vectors
• Highest position P such that there are no gaps of missing near neighbors with a Hilbert position less than or equal to P in the set of all points included in the union of this segment with all prior segments.

Each constructed line segment will be longer than the previous line segment. The exponential growth factor will be the square root of two. (I will experiment with different growth factors.) By using exponentially growing line segments, we can guarantee that the storage requirement for the index is S log S. The first segment will be length ten. Thus the segment lengths will be 10, 14, 20, 28, 40, 56, 80, etc.

Segments for each point will be defined as follows:

1. The first segment for a point will begin at the position of that point’s nearest neighbor.
2. The segment will extend for ten positions in increasing Hilbert position.
3. A Boolean array called FOUND of length S will record whether the point at that Neighbor index has been included in the current or any prior segment. Found will refer to points in nearest neighbor order for that point, not Hilbert order.
4. Iterate through all points in the segment and mark off in FOUND that they have been captured by a segment.
5. Iterate over FOUND to find the position M of the lowest point in neighbor order that is missing. Record in the segment record the value M-1 as the highest neighbor index without gaps in the segment and all prior segments.
6. Since the segment lengths will be growing and ordered haphazardly, new segments may overlap prior segments. Truncate as much overlap from the end of the segment as possible if it overlaps one or more other segments. Take care not to truncate gaps between the overlapping segments. For example, if a new segment C ranges from 1000-2000 and overlaps segments A from 500-600 and B from 1700-2200, only truncate the new segment C to 1000-1699, otherwise the gap of 601-1699 between A and B will be lost.
7. Each subsequent segment will begin with the nearest neighbor of the target point that has not yet been captured by the index.

Quality of Approximation. Because of the locality preserving nature of the Hilbert curve, each of the early segments will contain mostly near neighbors of the target point, but will miss some and include more distant neighbors, thus the approximate nature of the index. Research using the Hilbert curve for clustering shows that the Hilbert curve tends to divide data into 1.5 to 3 times as many clusters as are really present in the data. A Python experiment to explore the viability of this exponential segmentation idea shows that when the cumulative group of segments reaches a place where it guarantees that it has captured the K nearest neighbors, it will have done so by capturing a set of 1.5 to 2K neighbors, and on rare instances up to 2.5K. This is more than acceptable.

Thus my expectation (that must be proven with more experiments upon a fully constructed index) is that when you ask for the K-nearest neighbors of a point, you will typically be given 2K neighbors, but will not miss any of the desired near neighbors.

In the WMD algorithm, we need K neighbors and optimally desire to compute ½ K-squared distances for our cost matrix. Since we get 2K neighbors instead, we instead need to compute 2K-squared distances. This is vastly superior to S-squared!

Cost Matrix. In the end, to compose the WMD cost matrix:

1. Load from the index only the lists of segments for the M+N words in your documents.
2. Walk the segments for each word.
3. Toss all words not in the documents.
4. Continue until the set has gathered at least K words from your documents.
5. Perform distance calculations between all retained words.
6. Sort by distance.
7. Keep the K-nearest.

Note:

The index written to the file Contains:

• List of Word Vector ids sorted in Hilbert index order
• A list for each Word Vector of Hilbert segments as defined above