Fast and Numerically Stable Pairwise Distance Algorithms

Question

I'm looking for resources on fast, numerically stable pairwise euclidean distance algorithms. In particular, suppose $A \in \mathbb{R}^{M \times D}$ and $B \in \mathbb{R}^{N \times D}$ are two sets of row vectors. I would like to compute the matrix,

$$X \in \mathbb{R}^{M \times N}, \quad X_{i,j} = \|A_i - B_j\|_2^2.$$

So far I have found two methods:

Method 1 - The simple approach is to loop over each vector $A_i$ and each vector $B_j$ and populate $X$ with the corresponding sum of squares.

Method 2 - Take advantage of the fact that,

$$\|A_i - B_j\|_2^2 = \langle A_i - B_j, A_i - B_j \rangle = \|A_i\|_2^2 + \|B_j\|_2^2 - 2 \langle A_i, B_j \rangle.$$

Taking advantage of efficient code for computing each of these three terms, Method 2 is almost an order of magnitude faster. However, it is less numerically stable than Method 1. For example, Method 2 can output negative distances.

Are there alternate approaches to computing pairwise distances which are faster than Method 1 but guarantee (at the very least) all non-negative distances? have better numerical stability?

Code Demonstration - Below is an example comparison written in Python. Scipy's cdist() function is, effectively, an implementation of Method 1 whereas cdist_fast() below is an implementation of Method 2:

# experiment.py
import numpy as np
import time
from scipy.spatial.distance import cdist


def cdist_fast(XA, XB):
    XA_norm = np.sum(XA**2, axis=1)
    XB_norm = np.sum(XB**2, axis=1)
    XA_XB_T = np.dot(XA, XB.T)
    distances = XA_norm.reshape(-1,1) + XB_norm - 2*XA_XB_T
    return distances


def main():
    M,N = 5000, 128
    XA = np.random.randn(M,N)

    t = time.time()
    distances_cdist = cdist(XA, XA, metric='sqeuclidean')
    time_cdist = time.time() - t

    t = time.time()
    distances_cdist_fast = cdist_fast(XA, XA)
    time_cdist_fast = time.time() - t

    print(f'time_cdist = {time_cdist:.3f} s')
    print(f'time_cdist_fast = {time_cdist_fast:.3f} s')

    # check validity of results
    assert np.allclose(distances_cdist, distances_cdist_fast)

    # check that the results are non-negative
    try:
        assert (distances_cdist >= 0.0).all()
    except AssertionError:
        print('Numerical instability in cdist()')

    try:
        assert (distances_cdist_fast >= 0.0).all()
    except AssertionError:
        print('Numerical instability in cdist_fast()')


if __name__ == '__main__':
    main()

Script output on a 3.1 GHz Intel Core i7:

$ python experiment.py
time_cdist = 3.457 s
time_cdist_fast = 0.625 s
Numerical instability in cdist_fast()

Answer 1

I would strongly advise against using the method 2: $$\|A_i - B_j\|_2^2 = \langle A_i - B_j, A_i - B_j \rangle = \|A_i\|_2^2 + \|B_j\|_2^2 - 2 \langle A_i, B_j \rangle.$$ whether you are uing the absolute value on top of it or not. Because you are going to get your positive (but wrong) distances; however, you are still lacking any numerical stability.

Now, for 5000 vectors in $\mathbb R^{128}$, your (scipy) stable version is only ~6 times slower than the fast one. I don't think you are able to find a numerically stable version that is faster than provided by scipy.cdist. I would make sure that it is actually a bottleneck in your code. Especially since method 1 naturally sums only positive entries, an ideal and cheap scenario to not deal with the stability!

You definitely can try using Kahan or Neumaier summation, but I would be surprised that if they are employed (note, Neumaier has to be applied for each entry individually), the aforementioned method 2 would be any faster than scipy.cdist, most likely, at least tens time slower.

However, if the positiveness of the distance is the only thing you care, you certainly can use an absolute value on top of your cdist_fast(.,.) function.

Answer 2

The computed values are negative only when $A_i = B_j + O(\varepsilon)$, where $\varepsilon$ is the machine precision. The negative values are round-off errors.

If you return the absolute value of your computation, $\left| ||A_i||_2^2 + ||B_j||_2^2 - 2 <A_i,B_j> \right|$, you get the speed-up and positive numbers.