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.
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.