SLOSH: Set LOcality Sensitive Hashing
via Sliced-Wasserstein Embeddings

We use the Sliced-Wasserstein (SW) probability metric as a distance measure between the sets (viewed as empirical probability distributions). Later we will see that this choice allows us to effectively embed sets into a vector space such that the Euclidean distance between embedded sets is equal to the SW-distance between their corresponding empirical distributions. But first, let us briefly define the Wasserstein and Sliced-Wasserstein distances.

Let $\mu_{i}$ and $\mu_{j}$ be one-dimensional probability measures defined on $\mathbb{R}$. Then the p-Wasserstein distance between these measures can be written as:

where $F^{-1}_{\mu}$ is the inverse of the cumulative distribution function (c.d.f) $F_{\mu}$ of $\mu$, i.e., it is the quantile function. The one-dimensional p-Wasserstein distance for empirical distributions with $N$ and $M$ samples can be computed with $\mathcal{O}(Nlog(N)+Mlog(M))$, which is in stark difference from the generally cubic order for (d$>$1)-dimensional distributions. The one-dimensional case motivates the concept of Sliced-Wasserstein distances [52, 28]. For the rest of this paper we will consider only the case $p=2$, and for brevity we refer to 2-Wasserstein and 2-Sliced-Wasserstein distances as Wasserstein and Sliced-Wasserstein distances.

The main idea behind SW distances is to slice d-dimensional distributions into infinite sets of their one-dimensional slices/marginals and then calculate the expected Wasserstein distance between their slices. Let $g_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ be a parametric function with parameters $\theta\in\Omega_{\theta}\subseteq\mathbb{R}^{d_{\theta}}$, satisfying the regularity conditions in both inputs and parameters as presented in [28]. A common choice is $g_{\theta}(x)=\theta^{T}x$ where $\theta\in\mathbb{S}^{d-1}$ is a unit vector in $\mathbb{R}^{d}$, and $\mathbb{S}^{d-1}$ denotes the unit $d$-dimensional hypersphere. The slice of a probability measure, $\mu$, with respect to $g_{\theta}$ is the one-dimensional probability measure $g_{\theta\#}\mu$, with the density,

The generalized Sliced-Wasserstein distance is defined as

where $\sigma(\theta)$ is the uniform measure on $\Omega_{\theta}$, and once again for $g_{\theta}(x)=\theta^{T}x$ and $\Omega_{\theta}=\mathbb{S}^{d-1}$, the generalized Sliced-Wasserstein distance is simply the Sliced-Wasserstein distance. Equation (3) is the expected value of the Wasserstein distances between uniformly distributed slices (i.e., on $\Omega_{\theta}$) of distributions $\mu_{i}$ and $\mu_{j}$.

Locality Sensitive Hashing (LSH) [21, 17] strives for hashing near points in the high-dimensional input-space into the same hash bucket with a higher probability than distant ones, effectively solving the $(R,c)$-Near Neighbor problem. More precisely, let $H:=\{h:\mathbb{R}^{d}\rightarrow U\}$ denote a LSH function family with $U$ denoting the hash values. The LSH function family $H$, is called $(R,c,P_{1},P_{2})$-sensitive if for any two points $u,v\in\mathbb{R}^{d}$ and $\forall h\in H$ it satisfies the following conditions:

• •

  If $\|u-v\|_{2}\leq R$, then $\operatorname{Pr}[h(u)=h(v)]\geq P_{1}$, and

• •

  If $\|u-v\|_{2}>cR$, then $\operatorname{Pr}[h(u)=h(v)]\leq P_{2}$

Where for a proper LSH $c>1$ and $P_{1}>P_{2}$. The original LSH for Euclidean distance uses,

where $\lfloor\cdot\rfloor$ is the floor function, $a\in\mathbb{R}^{d}$ is a random vector generated from a $p$-stable distribution (e.g., Normal or Cauchy distributions) [17], and $b$ is a real number chosen uniformly from $[0,\omega)$ such that $\omega$ is the width of the hash bucket. Let $r=\|u-v\|_{p}$ and $f_{p}(t)$ denote the probability density function of the absolute value of the p-stable distribution, then this family of hash functions leads to the following probability of collision:

To amplify the gap between $P_{1}$ and $P_{2}$ one can concatenate several such hash functions. Let $\mathcal{G}:=\{g:\mathbb{R}^{d}\rightarrow U^{k}\}$ denote the concatenation of $k$ randomly selected hash functions, i.e., $g(u)=(h_{1}(u),...,h_{k}(u))$. In practice, multiple such hash functions $g_{1},\ldots,g_{T}$ are often used. For points, $u$ and $v$, within $R$-distance from one another, the probability that they collide at least in one of $g_{j}$s is: $1-(1-P_{1}^{k})^{T}$.

Another commonly used, and related, family of hash functions consist of random projections and thresholding, i.e., $h_{a,b}(u)=\operatorname{sgn}(a^{T}u+b)$ [14]. Using $k$ such random projections (i.e., binary codes of length $k$) provides the following collision probability:

Since the conception of its idea [21, 17], many variants of LSH have been proposed. However, these approaches are not designed to handle set queries. Here, we extend LSH to set-structured data via Sliced-Wasserstein Embeddings.

The idea of Sliced-Wasserstein Embeddings (SWE) is rooted in Linear Optimal Transport [65, 30, 41] and was first introduced in the context of pattern recognition from 2D probability density functions (e.g., images) [29] and more recently in [56]. Our work extends the framework to d-dimensional distributions with the specific application of set retrieval in mind. Consider a set of probability measures $\{\mu_{i}\}_{i=1}^{I}$ with densities $\{p_{i}\}_{i=1}^{I}$, and recall that we use $\mu_{i}$ to represent the i’th set $X_{i}=\{x_{n}^{i}\}_{n=0}^{N_{i}-1}$. At a high level, SWE can be thought as a set-2-vector operator, $\phi$, such that:

For convenience, we use $\mu_{i}^{\theta}:=g_{\theta\#}\mu_{i}$ to denote the slice of measure $\mu_{i}$ with respect to $g_{\theta}$. Also, let $\mu_{0}$ denote a reference measure, with $\mu_{0}^{\theta}$ being its corresponding slice. The optimal transport coupling (i.e., Monge coupling) between $\mu_{i}^{\theta}$ and $\mu_{0}^{\theta}$ can be written as

and recall that $F_{\mu_{i}^{\theta}}^{-1}$ is the quantile function of $\mu_{i}^{\theta}$. Now, letting $id$ denote the identity function, we define the so-called cumulative distribution transform (CDT) [46] of $\mu_{i}^{\theta}$ as

For a fixed $\theta$, we can show that $\phi^{\theta}(\mu_{i})$ satisfies (see supplementary material):

1. C1.

   The weighted $\ell_{2}$-norm of the embedded slice, $\phi_{\theta}(\mu_{i})$, satisfies:

   	$\displaystyle\|\phi_{\theta}(\mu_{i})\|_{\mu_{0}^{\theta},2}$	$\displaystyle=\left(\int_{\mathbb{R}}\|\phi_{\theta}(\mu_{i}(t))\|_{2}^{2}d\mu_{0}^{\theta}(t)\right)^{\frac{1}{2}}$
   		$\displaystyle=\mathcal{W}_{2}(\mu_{i}^{\theta},\mu_{0}^{\theta}),$

   As a corollary we have $\|\phi_{\theta}(\mu_{0})\|_{\mu_{0}^{\theta},2}=0$.

2. C2.

   The weighted $\ell_{2}$ distance between two embedded slices satisfies:

   $\displaystyle\|\phi_{\theta}(\mu_{i})-\phi_{\theta}(\mu_{j})\|_{\mu_{0}^{\theta},2}=\mathcal{W}_{2}(\mu_{i}^{\theta},\mu_{j}^{\theta}).$

   (8)

It follows from C1 and C2 that:

For probability measure $\mu_{i}$, we then define the mapping to the embedding space via,

Finally, we can re-weight the embedding (according to $d\mu^{\theta}_{0}$) such that the weighted $\ell_{2}$ in 8 becomes the $\ell_{2}$ distance as in 5. Next, we describe this reweighting and other implementation considerations in more details.

The GSW distance in Eq. (3) relies on integration on $\Omega_{\theta}$ (e.g., $\mathbb{S}^{d-1}$ for linear slices), which cannot be directly calculated. Following the common practice in the literature [52, 12, 32, 31, 39], here, we approximate the integration on $\theta$ via a Monte-Carlo (MC) sampling of $\Omega_{\theta}$. Let $\Theta_{L}=\{\theta_{l}\sim\sigma(\theta)\}_{l=1}^{L}$ denote a set of $L$ parameters sampled independently and uniformly from $\Omega_{\theta}$. We assume an empirical reference measure, $\mu_{0}=\frac{1}{M}\sum_{m=1}^{M}\delta(x-x_{m}^{0})$ with $M$ samples. The MC approximation is written as:

Finally, our SWE embedding is calculated via:

which satisfies:

As for the approximation error, we rely on Theorem 6 in [44], which uses Hölder’s inequality and the results on the moments of the Monte Carlo estimation error to obtain:

The upper bound indicates that the approximation error decreases with $\sqrt{L}$. The numerator, however, is implicitly dependent on the dimensionality of input space. Meaning that a larger number of slices, $L$, is needed for higher dimensions.

Here we review the algorithmic steps to obtain SWE. We consider $X_{i}=\{x^{i}_{n}\}_{n=0}^{N_{i}-1}$ as the input set with $N_{i}$ elements, and $X_{0}=\{x^{0}_{m}\}_{m=0}^{M-1}$ denote the reference set of $M$ samples where in general $M\neq N_{i}$. For a fixed slicer $g_{\theta}$ we calculate $\{g_{\theta}(x^{i}_{n})\}_{n=0}^{N_{i}-1}$ and $\{g_{\theta}(x^{0}_{m})\}_{m=0}^{M-1}$ and sort them increasingly. Let $\pi_{i}$ and $\pi_{0}$ denote the permutation indices (obtained from argsort). Also, let $\pi_{0}^{-1}$ denote the ordering that permutes the sorted set back to the original ordering. Then we numerically calculate the Monge coupling $T^{\theta}_{i}$ via:

where $F_{\mu_{0}^{\theta}}(x_{m}^{0})=\frac{\pi^{-1}_{0}[m]+1}{M}$, assuming that the indices start from 0. Here $F_{\mu_{i}^{\theta}}^{-1}$ is calculated via interpolation. In our experiments we used linear interpolation similar to [38]. Note that the dimensionality of the Monge coupling is only a function of the reference cardinality, i.e., $T_{i}^{\theta}\in\mathbb{R}^{M}$. Consequently, we write:

and repeat this process for $\theta\in\{\theta_{l}\sim\sigma(\theta)\}_{l=1}^{L}$, while we emphasize that this process can be parallelized. The final embedding is achieved via weighting and concatenating $\phi_{\theta_{l}}$s as in Eq. (11), where the coefficient $\frac{1}{\sqrt{LM}}$ allows us to simplify the weighted Euclidean distance, $\|\cdot\|_{\mu_{0},2}$, to Euclidean distance, $\|\cdot\|_{2}$. Algorithm 1 summarizes the embedding process, and Figure 1 provides a graphical depiction of the process. Lastly, the SWE’s computational complexity for a set with cardinality $|X|=N$ is $\mathcal{O}(LN(d+logN))$, where we assumed the cardinality of the reference set is of the same order as $N$. Note that $\mathcal{O}(LNd)$ is the cost of slicing and $\mathcal{O}(LNlog(N))$ is the sorting and interpolation cost to calculate Eq. (13).

Our proposed Set LOcality Sensitive Hashing (SLOSH) leverages the Sliced-Wasserstein Embedding (SWE) to embed training sets, $\mathcal{X}_{train}=\{X_{i}|X_{i}=\{x^{i}_{n}\in\mathbb{R}^{d}\}_{n=0}^{N_{i}-1}\}$, into a vector space where we can use Locality Sensitive Hashing (LSH) to perform ANN search. We treat each input set $X_{i}$ as a probability measure $\mu_{i}(x)=\frac{1}{N_{i}}\sum_{n=1}^{N_{i}}\delta(x-x_{n}^{i})$. For a reference set $X_{0}$ with cardinality $|X_{0}|=M$ and $L$ slices, we embed the input set $X_{i}$ into a $(\mathbb{R}^{LM})$-dimensional vector space using Algorithm 1. With abuse of notation we use $\phi(\mu_{i})$ and $\phi(X_{i})$ interchangeably.

Given SWE, a family of $(R,c,P_{1},P_{2})$-sensitive LSH functions, $H$, will induce the following conditions,

• •

  If $\widehat{\mathcal{GSW}}_{2,L}(X_{i},X_{j})\leq R$, then $Pr[h(\phi(X_{i}))=h(\phi(X_{j}))]\geq P_{1}$, and

• •

  If $\widehat{\mathcal{GSW}}_{2,L}(X_{i},X_{j})>cR$, then $Pr[h(\phi(X_{i}))=h(\phi(X_{j}))]\leq P_{2}$

As mentioned in Section 2, for amplifying the gap between $P_{1}$ and $P_{2}$, one uses $g(X_{i})=[h_{1}(\phi(X_{i})),...,h_{k}(\phi(X_{i}))]$, which results in a code length, $k$, for each input set, $X_{i}$. Finally, if $\widehat{\mathcal{GSW}}_{2,L}(X_{i},X_{j})\leq R$, by using $T$ such codes, $g_{t}$ for $t\in\{1,...,T\}$, of length $k$, we can ensure collision at least in one of $g_{t}$s with probability $1-(1-P_{1}^{k})^{T}$.

Let ${X}_{i}=\{x^{i}_{n}\in\mathbb{R}^{d}\}_{n=0}^{N_{i}-1}$ be the input set with $N_{i}$ elements. We denote $[X_{i}]_{k}=\{[x^{i}_{n}]_{k}\}_{n=0}^{N_{i}-1}$ as the set of all elements along the $k$’th dimension, $k\in\{1,2,...d\}$. Below we provide a quick overview of the baseline approaches, which provide different set-2-vector mechanisms.

Generalized-Mean Pooling (GeM) [53] was originally proposed as a generalization of global mean and global max pooling on Convolutional Neural Network (CNN) features to boost image retrieval performance. Given the input ${X}_{i}$, GeM calculates the (p-th)-moment of each feature, $f^{(p)}\in\mathbb{R}^{d}$, as:

When pooling parameter $p=1$, we end up with average pooling. While as $p\to\infty$, we get max pooling. In practice, we found that a concatenation of higher-order GeM features, i.e., $\phi_{GeM}(X_{i})=[f^{(1)};...;f^{(p)}]\in\mathbb{R}^{pd}$, leads to the best performance, where $p$ is GeM’s hyper-parameter.

Covariance Pooling [1, 64] presents another way to capture second-order statistics and provide more informative representations. It was shown that this mechanism can be applied to CNN features as an alternative to global mean/max pooling to generate state-of-the-art results on facial expression recognition tasks[1]. Given input set ${X}_{i}$, the unbiased covariance matrix is computed by:

where $\overline{\mu}_{i}=\frac{1}{N_{i}}\sum_{n=1}^{N_{i}}x^{i}_{n}$. The output matrix can be further regularized by adding a multiple of the trace to diagonal entries of the covariance matrix to ensure symmetric positive definiteness (SPD), $C_{\lambda}=C+\lambda trace(C)\textbf{I}$ where $\lambda$ is a regularization hyper-parameter and I is the identity matrix. Covariance pooling then uses $\phi_{Cov}(X_{i})=\text{flatten}(C_{\lambda})$.

Featurewise Sort Pooling (FSPool) [71] is a powerful technique for learning representations from set-structured data. In short, this approach is based on sorting features along all elements of a set, $[X_{i}]_{k}$:

The fixed-dimensional representation is then obtained via an interpolation along the $N_{i}$ dimension of $f$. More precisely, a continuous linear operator $W$ is learned and the inner product between this continuous operator and $f$ is evaluated at $M$ fixed points (i.e., leading to weighted summation over $d$), resulting in a $M$-dimensional embedding.

Given that we are interested in a data-independent set representation, we cannot rely on learning the continuous linear operator $W$. Instead, we perform interpolation along the $N_{i}$ axis on $M$ points and drop the inner product all together to obtain a $(M\times d)$-dimensional data-independent set representation. The mentioned variation of FSPool is similar to the Sliced-Wasserstein Embedding when $L=d$ and $\Theta_{L}=I_{d\times d}$, i.e., axis-aligned projections.

Point Cloud MNIST 2D [34, 18] consists of 60,000 training samples and 10,000 testing samples. Each sample is a 2-dimensional point cloud derived from an image in the original MNIST dataset. The sets have various cardinalities in the range of $|X_{i}|\in[34,351]$.

ModelNet40 [68] contains 3-dimensional point clouds converted from 12,311 CAD models in 40 common object categories. We used the official split with 9,843 samples for training and 2,468 samples for testing. We sample $N_{i}$ points uniformly and randomly from the mesh faces of each object, where $N_{i}=\lfloor n_{i}\rfloor,n_{i}\sim\mathcal{N}(512,64)$. To avoid any orientation bias, we randomly rotate the point clouds by $\rho\in[0,45]$ degrees around the x-axis. Finally, we normalize each sample to fit within the unit cube to avoid any scale bias and to enforce the methods to rely on shapes.

The Oxford Buildings Dataset [49] has 5,062 images containing eleven different Oxford landmarks. Each landmark has five corresponding queries leading to 55 queries over which an object retrieval system can be evaluated. In our experiments, we use a pretrained VGG16 [57] on ImageNet1k [54] as a feature extractor, and use the features in the last convolutional layer as a set representation for an input image. We resize the images without cropping, which leads to varied input size images, and therefore gives set representations with varied cardinalities. We further perform a dimensionality reduction on these features to obtain sets of features in an $(d=8)$-dimensional space.

For all datasets, we first calculate the set-2-vector embeddings for all baselines and SWE. Then, we apply Locality-Sensitive Hashing (LSH) to the embedded sets and report Precision@k and accuracy (based on majority voting) for all the approaches on the test sets. We use the FAISS library [24], developed by Facebook AI Research, for LSH. For all methods, we use a hash code length of $1024$, and we report our results for $k=4,8,\text{and}~{}16$. For SLOSH, we consider three different settings for the number of slices, namely, $L<d$, $L=d$, and $L>d$. We repeat the experiments five times per method and report the mean Precision@k and accuracy in Table 1. In all experiments, the best hyper-parameters are selected for each approach based on cross validation. We see that SLOSH provides a consistent lead on all datasets, especially, when $L>d$. Figure 2 provide set retrieval examples on the three data sets. Next, we provide a sensitivity analysis of our approach.

Sensitivity to code length. For all datasets, we study the sensitivity of the different embeddings to the hashing code-length. We vary the code length from $16$ to $1024$, and report the average of Precision@k over five runs. Figure 3 shows the outcome of this study. It can be seen that methods that encode higher statistical orders of the underlying distribution gain performance as a function of the hash code length. In addition, we see a strong performance gap between SWE and other set-2-vector approaches, which points at the more descriptive nature of our proposed method.

Sensitivity to the number of slices. Next, we study the sensitivity of SLOSH to the choice of number of slices, $L$, for the three datasets. We measure the average Precision@k over five runs for various number of slices and for different code-lengths and report our results in Figure 4. As expected we see a performance increase with respect to increasing $L$.

Sensitivity to the reference. Finally, we perform a study on the sensitivity to the choice of the reference set, $X_{0}$. We measure the performance of SLOSH on the three datasets and for various code-lengths, when the reference set is: 1) calculated using K-Means on the elements of all sets, 2) a random set from the dataset, 3) sampled from a uniform distribution, and 4) sampled from the normal distribution. We see that SLOSH’s performance depends on the reference choice. However, we point out that our formulation implies that as $L\rightarrow\infty$ the embedding becomes independent of the choice of the reference, and the observed sensitivity could be due to using finite $L$.

Here we include the proof for the C1 and C2 conditions covered in Section 4. Recall that $\mu_{i}$ represents a probability measure, $\mu_{i}^{\theta}$ represents $g_{\theta\#}\mu_{i}$, where $g_{\theta}:\mathbb{R}^{d}\rightarrow\mathbb{R}$ with some regularity constraints, and that we define the cumulative distribution transform (CDT) [46] of $\mu_{i}^{\theta}$ as

where $T_{i}^{\theta}$ is the Monge map/coupling, and $id$ denotes the identity function. For a fixed $\theta$, here we prove that $\phi^{\theta}(\mu_{i})$ satisfies the following conditions:

C1. The weighted $\ell_{2}$-norm of the embedded slice, $\phi_{\theta}(\mu_{i})$, satisfies:

C2. The weighted $\ell_{2}$ distance satisfies:

As a corollary of C1 and C2 we have:

For the sake of completeness, here we include the computational complexities of the baseline methods used in our paper. In Table 2, we provide the computational complexity of embedding a set $X=\{x_{n}\in\mathbb{R}^{d}\}_{n=1}^{N}$ into a vector space.