From EMD and WMD to WRD: Similarity Calculation of Text Vector Sequences

Optimal Transport

Suppose at positions i=1,2,\dots,n, we have distributions of earth with quantities p_1,p_2,\dots,p_n. For simplicity, let the total amount of earth be 1, i.e., \sum p_i = 1. Now, we want to move the earth to positions j=1,2,\dots,n', where the required quantities are q_1, q_2, \dots, q_{n'}. The cost of moving earth from i to j is d_{i,j}. We seek the plan with the minimum cost and the corresponding minimum cost value.

This is a classic optimal transport problem. We denote the optimal plan as \gamma_{i,j}, representing the amount of earth moved from i to j in this plan. Obviously, we have the constraints: \sum_j \gamma_{i,j}=p_i,\quad \sum_i \gamma_{i,j}=q_j,\quad \gamma_{i,j} \geq 0 \label{eq:cond} Thus, our optimization problem is: \min_{\gamma_{i,j} \geq 0} \sum_{i,j} \gamma_{i,j} d_{i,j}\quad \text{s.t.} \quad \sum_j \gamma_{i,j}=p_i, \sum_i \gamma_{i,j}=q_j

Reference Implementation

Although it looks complex, a careful observation reveals that the above formula is actually a linear programming problem—finding the extremum of a linear function under linear constraints. Since scipy comes with a linear programming solver linprog, we can use it to implement a function for calculating the Wasserstein distance:

import numpy as np
from scipy.optimize import linprog

def wasserstein_distance(p, q, D):
    """Solve Wasserstein distance via linear programming
    p.shape=[m], q.shape=[n], D.shape=[m, n]
    p.sum()=1, q.sum()=1, p in [0,1], q in [0,1]
    """
    A_eq = []
    for i in range(len(p)):
        A = np.zeros_like(D)
        A[i, :] = 1
        A_eq.append(A.reshape(-1))
    for i in range(len(q)):
        A = np.zeros_like(D)
        A[:, i] = 1
        A_eq.append(A.reshape(-1))
    A_eq = np.array(A_eq)
    b_eq = np.concatenate([p, q])
    D = D.reshape(-1)
    # Remove the last redundant constraint to avoid contradictions from float errors
    result = linprog(D, A_eq=A_eq[:-1], b_eq=b_eq[:-1])
    return result.fun

Readers may notice that when passing the constraints, A_eq=A_eq[:-1], b_eq=b_eq[:-1] is used, which removes the last constraint. This is because 1=\sum_{i=1}^n p_i = \sum_{j=1}^{n'} q_j, so the equality constraints in [eq:cond] are inherently redundant. In actual calculations, floating-point errors might cause redundant constraints to contradict each other, leading to failure in solving the linear program. Thus, removing the last redundant constraint reduces the possibility of errors.

Basic Form

Suppose there are two sentences s = (t_1,t_2,\dots,t_n) and s' = (t'_1, t'_2, \dots, t'_{n'}). After some mapping (such as Word2Vec or BERT), they become corresponding vector sequences (\boldsymbol{w}_1,\boldsymbol{w}_2,\dots,\boldsymbol{w}_n) and (\boldsymbol{w}'_1, \boldsymbol{w}'_2, \dots, \boldsymbol{w}'_{n'}). Now we find a way to use the Wasserstein distance to compare the similarity of these two sequences.

According to the previous section, the Wasserstein distance requires three quantities: p_i, q_j, d_{i,j}. We define them one by one. Since there is no prior knowledge, we can simply set p_i \equiv \frac{1}{n} and q_j \equiv \frac{1}{n'}. Now only d_{i,j} remains. Obviously, d_{i,j} represents some difference between the vector \boldsymbol{w}_i from the first sequence and the vector \boldsymbol{w}'_j from the second sequence. For simplicity, we can use the Euclidean distance \left\Vert \boldsymbol{w}_i - \boldsymbol{w}'_j\right\Vert. Thus, the degree of difference between the two sentences can be expressed as: \min_{\gamma_{i,j} \geq 0} \sum_{i,j} \gamma_{i,j} \left\Vert \boldsymbol{w}_i - \boldsymbol{w}'_j\right\Vert\quad \text{s.t.} \quad \sum_j \gamma_{i,j}=\frac{1}{n},\sum_i \gamma_{i,j}=\frac{1}{n'} This is the Word Mover’s Distance (WMD). It can be understood as the shortest path to transform one sentence into another, and in a sense, it can also be seen as a smooth version of Edit Distance. In practice, WMD is usually calculated after removing stop words.

Reference Implementation

The reference implementation is as follows:

def word_mover_distance(x, y):
    """Reference implementation of WMD (Word Mover's Distance)
    x.shape=[m,d], y.shape=[n,d]
    """
    p = np.ones(x.shape[0]) / x.shape[0]
    q = np.ones(y.shape[0]) / y.shape[0]
    D = np.sqrt(np.square(x[:, None] - y[None, :]).sum(axis=2))
    return wasserstein_distance(p, q, D)

Lower Bound Formula

In a retrieval scenario, calculating WMD for an input sentence against every sentence in a database and then sorting them is computationally expensive. Therefore, we should try to reduce the number of WMD calculations, for example, by filtering out samples using simpler and more efficient metrics before calculating WMD for the remaining ones.

Fortunately, we can derive a lower bound formula for WMD, which the original paper calls Word Centroid Distance (WCD): \begin{aligned} \sum_{i,j} \gamma_{i,j} \left\Vert \boldsymbol{w}_i - \boldsymbol{w}'_j\right\Vert &= \sum_{i,j} \left\Vert \gamma_{i,j}(\boldsymbol{w}_i - \boldsymbol{w}'_j)\right\Vert\\ &\geq \left\Vert \sum_{i,j}\gamma_{i,j}(\boldsymbol{w}_i - \boldsymbol{w}'_j)\right\Vert\\ &= \left\Vert \sum_i\left(\sum_j\gamma_{i,j}\right)\boldsymbol{w}_i - \sum_j\left(\sum_i\gamma_{i,j}\right)\boldsymbol{w}'_j\right\Vert\\ &= \left\Vert \frac{1}{n}\sum_i\boldsymbol{w}_i - \frac{1}{n'}\sum_j\boldsymbol{w}'_j\right\Vert\\ \end{aligned} In other words, WMD is greater than the Euclidean distance between the average vectors of the two sentences. So, when retrieving sentences with small WMD, we can first use WCD to filter out sentences with large distances and then use WMD for the remaining ones.

Basic Form

First, WRD proposes the view that "the norm of a word vector is positively correlated with the importance of the word" and validates this view through experimental results. In fact, this view is consistent with the author’s previously proposed simpler GloVe model; see "A More Unique Word Vector Model (V): Interesting Results". In WMD, p_i, q_j also represent the importance of a word in a sentence in some sense, so we can set: \begin{aligned} &p_i = \frac{\left\Vert \boldsymbol{w}_i\right\Vert}{Z}, &Z=\sum_{i=1}^n \left\Vert\boldsymbol{w}_i\right\Vert\\ &q_j = \frac{\left\Vert \boldsymbol{w}'_j\right\Vert}{Z'}, &Z'=\sum_{j=1}^{n'}\left\Vert\boldsymbol{w}'_j\right\Vert \end{aligned} Then d_{i,j} uses the cosine distance: d_{i,j}=1 - \frac{\boldsymbol{w}_i\cdot \boldsymbol{w}'_j}{\left\Vert\boldsymbol{w}_i\right\Vert\times \left\Vert\boldsymbol{w}'_j\right\Vert} Resulting in: \min_{\gamma_{i,j} \geq 0} \sum_{i,j} \gamma_{i,j} \left(1 - \frac{\boldsymbol{w}_i\cdot \boldsymbol{w}'_j}{\left\Vert\boldsymbol{w}_i\right\Vert\times \left\Vert\boldsymbol{w}'_j\right\Vert}\right)\quad \text{s.t.} \quad \sum_j \gamma_{i,j}=\frac{\left\Vert \boldsymbol{w}_i\right\Vert}{Z},\sum_i \gamma_{i,j}=\frac{\left\Vert \boldsymbol{w}'_j\right\Vert}{Z'} This is the Word Rotator’s Distance (WRD). Since the metric used is cosine distance, the transformation between two vectors is more like a rotation rather than a move, hence the name. Also, because it uses cosine distance, its result is within [0, 2], making it relatively easier to perceive the degree of similarity.

Reference Implementation

The reference implementation is as follows:

def word_rotator_distance(x, y):
    """Reference implementation of WRD (Word Rotator's Distance)
    x.shape=[m,d], y.shape=[n,d]
    """
    x_norm = (x**2).sum(axis=1, keepdims=True)**0.5
    y_norm = (y**2).sum(axis=1, keepdims=True)**0.5
    p = x_norm[:, 0] / x_norm.sum()
    q = y_norm[:, 0] / y_norm.sum()
    D = 1 - np.dot(x / x_norm, (y / y_norm).T)
    return wasserstein_distance(p, q, D)

def word_rotator_similarity(x, y):
    """1 - WRD
    x.shape=[m,d], y.shape=[n,d]
    """
    return 1 - word_rotator_distance(x, y)

Lower Bound Formula

Similar to WMD, we can also derive a lower bound formula for WRD: \begin{aligned} 2\sum_{i,j} \gamma_{i,j} \left(1 - \frac{\boldsymbol{w}_i\cdot \boldsymbol{w}'_j}{\left\Vert\boldsymbol{w}_i\right\Vert\times \left\Vert\boldsymbol{w}'_j\right\Vert}\right)&=\sum_{i,j} \gamma_{i,j} \left\Vert \frac{\boldsymbol{w}_i}{\left\Vert \boldsymbol{w}_i\right\Vert} - \frac{\boldsymbol{w}'_j}{\left\Vert \boldsymbol{w}'_j\right\Vert}\right\Vert^2 \\ &\geq \left\Vert \sum_{i,j}\gamma_{i,j}\left(\frac{\boldsymbol{w}_i}{\left\Vert \boldsymbol{w}_i\right\Vert} - \frac{\boldsymbol{w}'_j}{\left\Vert \boldsymbol{w}'_j\right\Vert}\right)\right\Vert^2\\ &= \left\Vert \sum_i\left(\sum_j\gamma_{i,j}\right)\frac{\boldsymbol{w}_i}{\left\Vert \boldsymbol{w}_i\right\Vert} - \sum_j\left(\sum_i\gamma_{i,j}\right)\frac{\boldsymbol{w}'_j}{\left\Vert \boldsymbol{w}'_j\right\Vert}\right\Vert^2\\ &= \left\Vert \frac{1}{Z}\sum_i\boldsymbol{w}_i - \frac{1}{Z'}\sum_j\boldsymbol{w}'_j\right\Vert^2\\ \end{aligned} The inequality is based on Jensen’s inequality (or a generalized version of the basic inequality). However, this part of the content did not appear in the WRD paper; it was supplemented by the author.