Fast 1-Wasserstein Distance Approximation Using Greedy Strategies
Guillaume Houry

Han Bao

Han Zhao

Makoto Yamada

ENS Paris-Saclay

Kyoto University

UIUC

OIST

Word Mover Distance (WMD) computes similarity measures between text documents based on
optimal transport. Its computation can be fastened using tree embeddings of the vocabulary,e.g. Tree Wasserstein distance or Flowtree algorithms.
The latter relies on a greedy matching of the embedded words in the documents to compare. It
provides accurate approximations of WMD, but its complexity increases linearly with the treeheight. Our work aims to further fasten the Flowtree algorithm by proposing a new approachbased on 1D representation of the tree. Our main contributions are:•A method that embeds trees into 1D, reducing tree matching to 1D matching;•Two algorithms, 1D k-Greedy and 1D-ICT, using greedy strategies in 1D.

The Flowtree pipeline, and our proposed approach based on 1D embedding of the tree and greedy matching.

Input

Data modelingDocuments

BoW vectors

Tree embs.

𝐷𝑜𝑐1𝜇𝜇෤

1D embs.

1D k-Greedy𝜈෤𝜇′෤𝜈′෤

1D-ICT

𝐷𝑜𝑐2Background: Word Mover Distance and Tree
Wasserstein DistanceApproximation algorithm

𝜈OutputFlowtree෩
𝑊(𝜇,𝜈)

Greedy heuristics for 1D matching
We proposed two heuristics to compute greedy matchings from a 1D embedding:Word Mover’s Distance (WMD) [1] measures similarity between text documents,modeled as bag of words vectorized using neural embeddings. It is defined as
the Wasserstein distance between uniform distributions over these point sets.
Definition. Let 𝑋, 𝑑 be a finite metric space and μ, ν two distributions over 𝑋,with support Ex = (𝑥1 , … , 𝑥𝑛 ) and Ey = (𝑦1 , … , 𝑦𝑚 ). The Wasserstein distance
between μ and ν is the cost of the optimal transport between the distributions:

𝐸𝑞. 1• In 1D 𝒌-greedy, we restrict the number of NN updates to 𝑘; remaining points
are matched in arbitrary order. The complexity is 𝑶 𝒌 + 𝟏 𝒏 + 𝒎 𝒍𝒐𝒈 𝒏 ;• In 1D-ICT, we relax the optimal transport constraints using the idea of theIterative Constrained Transfer algorithm [4]. The complexity is 𝑶 𝒏 + 𝒎 .
Algorithm: 1D 𝒌-Greedy.

Algorithm: 1D-ICT.

Compute the 1D-NN map between 𝐸𝑥 and 𝐸𝑦
For 𝑘෨ = 0 𝑡𝑜 𝑘:

Compute the 1D-NN map between 𝐸𝑥 and 𝐸𝑦

Sort 𝐸𝑥 by incr. value of 1D-NN dist. in 𝐸𝑦
For 𝑥 ∈ 𝐸𝑥 :(1)
𝑦 = NN of 𝑥 in 𝐸𝑦

With:

Several heuristics aim to reduce the cost of Wasserstein computation, like Tree
Wasserstein Distance (TWD) [2] which approximates the true distance 𝒅 by a෩ in Eq. 1, with an tree embedding of the vocabulary 𝑋.
tree distance 𝒅For 𝑖 = 1 𝑡𝑜 𝑛:
𝐸𝑦 = {𝑦1 , … , 𝑦𝑚 }𝑦 = 1D-NN of 𝑥 in 𝐸𝑦While 𝜇 𝑥 ≠ 0:
u = min(𝜇 𝑥 , 𝜈 𝑦 )෩ += 𝑢 ⋅ 𝑑(𝑥, 𝑦)
𝑊𝜇 𝑥 −= 𝑢Remove 𝑦 from 𝐸𝑦

u = min(𝜇 𝑥 , 𝜈 𝑦 )෩ += 𝑢 ⋅ 𝑑(𝑥, 𝑦)
𝑊𝜇 𝑥 −= 𝑢 𝑎𝑛𝑑 𝜈 𝑦 −= 𝑢If 𝜈 𝑦 = 0:
Remove 𝑦 from 𝐸𝑦
Update the 1D-NN of 𝑥 in 𝐸𝑦
If 𝜇 𝑥 = 0:
Remove 𝑥 from 𝐸𝑥
If 𝑘෨ = 𝑘 and 𝜇 𝑥 ≠ 0:

Update the 1D-NN of 𝑥 in 𝐸𝑦
෩Return 𝑊

In both algorithms, the nearest neighbors arecomputed using a 1D embedding of 𝑋 ; 𝑑෩ approximates
denotes the true distance in 𝑋. 𝑊the true Wasserstein distance in 𝑋 based on
the NN relation in the 1D embedding space.

Repeat from (1)෩Return 𝑊

Experiments
Greedy algorithms with 1D Tree embeddings
Experiments show that the accuracy of our algorithms is similar to the Flowtreebaseline, while being faster for deeper tree embeddings.

Although fast to compute, TWD provides rough estimates of WMD. Flowtree [3]෩ on the tree
improves the accuracy of TWD by computing the optimal matching Πሚ then by evaluating the Wasserstein distance on the true distance 𝑑.
distance 𝑑,

From trees to 1D greedy matching
෩ on a tree is computed using a greedy algorithm, whose
The optimal matching Πcomplexity is 𝑶( 𝒏 + 𝒎 𝒉) where 𝑛 and 𝑚 are the number of words in thecompared documents and ℎ is the height of the embedded tree. Our motivation
was to find faster algorithm that scales better with the tree height.

From trees to 1D embeddings
First, we observed that it is possible to embed a tree in a real line, so that greedy
matching in the tree is equivalent to a 1D greedy matching. This embedding isobtained by adjusting the spaces Δ0 , Δ1 , Δ2 … between each node.

Greedy algorithms with sliced embeddings
These algorithms can also be usedwith other 1D embeddings, e.g.
sliced projections [5]. The accuracy
is then lower, but the method is more
polyvalent – as it does not requires
to fix the vocabulary 𝑿.
Accuracy can be improved by taking
the maximum value over 𝑁 random
projections.

References
Second, we noticed that a nearest neighbor (NN) mapping between two sets of
1D points is always non-decreasing, hence fast to compute and update.

[1] Kusner, Matt, et al. "From word embeddings to document distances." International conferenceon machine learning. PMLR, 2015.
[2] Le, Tam, et al. "Tree-sliced variants of Wasserstein distances." Advances in neural information
processing systems 32, 2019.
[3] Backurs, Arturs, et al. "Scalable nearest neighbor search for optimal transport." International
Conference on machine learning. PMLR, 2020.
[4] Atasu, Kubilay, and Thomas Mittelholzer. "Linear-complexity data-parallel earth mover’s
distance approximations." International Conference on Machine Learning. PMLR, 2019.

Greedy methods match pairs of points by increasing order of NN distance.
Therefore, we can use this property combined with other heuristics, to design
faster greedy algorithms.

[5] Rabin, Julien, et al. "Wasserstein barycenter and its application to texture mixing." Scale Spaceand Variational Methods in Computer Vision: Third International Conference, SSVM 2011, EinGedi, Israel, May 29–June 2, 2011, Revised Selected Papers 3. Springer Berlin Heidelberg, 2012.