Distance-preserving set embeddings: theory and applications

Abstract

A common operation in machine learning is the aggregation of an unordered multiset of vectors into a single Euclidean vector. In the case of graph neural networks (GNNs), each node must aggregate the messages coming from each of its neighbors, while in the case of DNA sequence analysis, one commonly seeks to summarize a genome by aggregating the set of short sequences, called k-mers, of which it is composed (i.e. bases 1 through k, 2 through k+1, 3 through k+2, etc.). In the case of GNNs, one typically takes the elementwise sum or maximum of the incoming messages, possibly after applying a neural network to each message. A potential disadvantage of this approach, however, is that it is not bi-Lipschitz, namely, that distinct multisets can be mapped to arbitrarily close embeddings. In the case of DNA sequence comparison, one typically constructs a histogram of the k-mers present in a genome. A potential disadvantage of this approach is that it does not consider the content of the k-mers: For instance, it does not take into account the fact that AAAA and AAAC are closer to one another than AAAA and CGCG. A potential solution to these disadvantages is the Fourier Sliced Wasserstein (FSW) embedding, developed by Nadav Dym and Tal Amir at the Technion in Israel, which is a provably bi-Lipschitz multiset embedding with respect to the Wasserstein distance. However, prior to my work, the Lipschitz constants could not be precisely described or even bounded. My theoretical contribution in this thesis is the first theoretical bounds on the Lipschitz constants of a closely related operation. On the practical side, I attempted to improve graph neural networks' performances on tasks in the Open Graph Benchmark collection by replacing sum-pooling aggregations with FSW, and I sought to improve predictions for what bacterial host a given virus will infect and predictions for taxonomic tree distances among species by replacing the histogram embedding with FSW. In the domain of GNNs, inserting FSW had no effect on performance, while in the realm of DNA sequence comparison, the use of FSW degraded performance.