# Efficient Distance Matrix Computation

**URL:** <https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065>\
**Category:** Uncategorized\
**Created:** [October 25, 2017, 7:21pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065 "2017-10-25T19:21:03Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![AtheMathmo](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/f475e1/32.png) [@AtheMathmo](https://discuss.pytorch.org/u/AtheMathmo)\
**Post date:** [October 25, 2017, 7:21pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/1 "2017-10-25T19:21:03Z")

</div>

Hi All,

For the project I’m working on right now I need to compute distance matrices over large batches of data. I have two matrices X and Y, where X is nxd and Y is mxd. Then the distance matrix D is nxm and contains the squared euclidean distance between each row of X and each row of Y.

So far I’ve implemented this in a few different ways but each has their issues and I’m hoping someone more experienced with pytorch might be able to help me get an implementation that matches my needs. I’ll go through each approach and the related issues below.

### Approach A - Direct Expanded

```
def expanded_pairwise_distances(x, y=None):
    '''
    Input: x is a Nxd matrix
           y is an optional Mxd matirx
    Output: dist is a NxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:]
            if y is not given then use 'y=x'.
    i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2
    '''
    if y is not None:
         differences = x.unsqueeze(1) - y.unsqueeze(0)
    else:
        differences = x.unsqueeze(1) - x.unsqueeze(0)
    distances = torch.sum(differences * differences, -1)
    return distances

```

This approach adds extra dimensions to compute the difference between all combinations of rows and columns at once. This requires a lot of memory and is slow.

### Approach B - Quadratic Expansion

```
def pairwise_distances(x, y=None):
    '''
    Input: x is a Nxd matrix
           y is an optional Mxd matirx
    Output: dist is a NxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:]
            if y is not given then use 'y=x'.
    i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2
    '''
    x_norm = (x**2).sum(1).view(-1, 1)
    if y is not None:
        y_norm = (y**2).sum(1).view(1, -1)
    else:
        y = x
        y_norm = x_norm.view(1, -1)

    dist = x_norm + y_norm - 2.0 * torch.mm(x, torch.transpose(y, 0, 1))
    return dist

```

This approach requires less memory and is faster than the above. However this suffers from rounding errors and in my setup leads to numerical instability. This seems to be a well documented issue with this approach and existing libraries will use thresholding to default to the direct computation when needed.

### Approach C - Direct row-wise

```
def row_pairwise_distances(x, y=None, dist_mat=None):
    if y is None:
        y = x
    if dist_mat is None:
        dtype = x.data.type()
        dist_mat = Variable(torch.Tensor(x.size()[0], y.size()[0]).type(dtype))

    for i, row in enumerate(x.split(1)):
        r_v = row.expand_as(y)
        sq_dist = torch.sum((r_v - y) ** 2, 1)
        dist_mat[i] = sq_dist.view(1, -1)
    return dist_mat

```

The above is numerically stable and has a lower memory footprint than A but it is far slower than approach B. I think that it should be possible to modify C to be faster but I don’t know enough about the autograd engine and pytorch data storage to make the tweaks myself. Does anybody have any suggestions?

---

<div class="post-metadata">

**Author:** ![chenchr](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/chenchr/32/10407_2.png) [@chenchr](https://discuss.pytorch.org/u/chenchr)\
**Post date:** [December 4, 2017, 4:52pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/2 "2017-12-04T16:52:32Z")

</div>

Hello. I am doing a project that need to calculate this distance too. Firstly I want to express my gratitude to you for sharing such elaborate approach.  
I firstly just deal with the calculation by approch A and `out of memory` occurs, in the pytorch document it seems that there are not good solution to date.  
I had no idea calculating (a-b)^2 by a^2+b^2-2ab, I think this is the best way for speed and memory consideration by using ready-made pytorch api.  
I think the best way to deal with numerical stability may be writing a cffi extension and using cuda to calculating the distance directly.

---

<div class="post-metadata">

**Author:** ![AtheMathmo](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/f475e1/32.png) [@AtheMathmo](https://discuss.pytorch.org/u/AtheMathmo)\
**Post date:** [December 4, 2017, 6:38pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/3 "2017-12-04T18:38:00Z")

</div>

Glad there is something helpful here for you! In the end I managed to get a suitable solution by adapting Approach B to clamp the distance matrix to be positive and ensure the diagonal is zero when needed:

```
def pairwise_distances(x, y=None):
    '''
    Input: x is a Nxd matrix
           y is an optional Mxd matirx
    Output: dist is a NxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:]
            if y is not given then use 'y=x'.
    i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2
    '''
    x_norm = (x**2).sum(1).view(-1, 1)
    if y is not None:
        y_t = torch.transpose(y, 0, 1)
        y_norm = (y**2).sum(1).view(1, -1)
    else:
        y_t = torch.transpose(x, 0, 1)
        y_norm = x_norm.view(1, -1)
    
    dist = x_norm + y_norm - 2.0 * torch.mm(x, y_t)
    # Ensure diagonal is zero if x=y
    # if y is None:
    # dist = dist - torch.diag(dist.diag)
    return torch.clamp(dist, 0.0, np.inf)

```

For my code this was enough to keep things stable during optimization. I think that in general you might be correct that we’d need a cffi extension.

---

<div class="post-metadata">

**Author:** ![ruthcfong](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/r/ed655f/32.png) [@ruthcfong](https://discuss.pytorch.org/u/ruthcfong)\
**Post date:** [December 7, 2017, 5:48pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/4 "2017-12-07T17:48:36Z")

</div>

@AtheMathmo: For pytorch version 0.2.0\_4 (I’m not sure about the latest pytorch version), I found that `torch.clamp(dist, 0.0, np.inf)` didn’t deal with `nan` values that sometimes arises along the diagonal when running `pairwise_distances(x, y=None)`. To mitigate this issue, I replaced the line with the following:

`dist[dist != dist] = 0 # replace nan values with 0`

---

<div class="post-metadata">

**Author:** ![fusheng\_hao](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/f/b19c9b/32.png) [@fusheng\_hao](https://discuss.pytorch.org/u/fusheng_hao)\
**Post date:** [December 15, 2017, 12:52am UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/5 "2017-12-15T00:52:02Z")

</div>

I expand both x ( n x d ) and y ( m x d) as follows:

```
n = x.size(0)
m = y.size(0)
d = x.size(1)

x = x.unsqueeze(1).expand(n, m, d)
y = y.unsqueeze(0).expand(n, m, d)

```

Then calculate dist[i,j] = ||x[i,:]-y[j,:]||^2 in the following way:

```
dist = torch.pow(x - y, 2).sum(2) 

```

Both x and y are torch.Tensor. Is this exactly what you want? Or is there a more efficient way?

---

<div class="post-metadata">

**Author:** ![ans](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ans/32/3331_2.png) [@ans](https://discuss.pytorch.org/u/ans)\
**Post date:** [February 9, 2018, 1:48pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/6 "2018-02-09T13:48:50Z")

</div>

Thank you so much !  
With your function I was able, moving the matrix on cuda, to  
compute pairwise distances 210/220 times faster with respect  
to scipy.spatial.distance.pdist !  
I checked this on matrices from 1000 to 10000 samples of flattened  
cifar10 images.  
There is a small error with respect to the scipy function, I estimated  
the relative error (max / median, taken over each pairwise distance)  
to be around 1e-5, which is fine for me.

---

<div class="post-metadata">

**Author:** ![derek\_kim](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/derek_kim/32/8299_2.png) [@derek\_kim](https://discuss.pytorch.org/u/derek_kim)\
**Post date:** [August 1, 2018, 2:58am UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/7 "2018-08-01T02:58:15Z")

</div>

I think your approach is almost identical to the author’s first approach. Redundant.

---

<div class="post-metadata">

**Author:** ![marx\_was\_right](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/marx_was_right/32/4954_2.png) [@marx\_was\_right](https://discuss.pytorch.org/u/marx_was_right)\
**Post date:** [August 28, 2018, 4:02pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/8 "2018-08-28T16:02:15Z")

</div>

> [@fusheng\_hao](#):
>
> dist = torch.pow(x - y, 2).sum(2)

This answer is MEGA AESTHETIC.  
thanks!!

also shout out to @AtheMathmo for posing the similar problem I had in terms that helped me understand my own situation

---

<div class="post-metadata">

**Author:** ![iamlockelightning](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/iamlockelightning/32/6832_2.png) [@iamlockelightning](https://discuss.pytorch.org/u/iamlockelightning)\
**Post date:** [October 11, 2018, 7:49am UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/9 "2018-10-11T07:49:44Z")

</div>

Hi, thank you for your **Approach B - Quadratic Expansion**. Pretty impressive!  
But may I ask is there a similar way for calculating the **L1 distance** (||x - y||) efficiently?

---

<div class="post-metadata">

**Author:** ![ehsanmok](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ehsanmok/32/15849_2.png) [@ehsanmok](https://discuss.pytorch.org/u/ehsanmok)\
**Post date:** [November 16, 2018, 10:47pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/10 "2018-11-16T22:47:19Z")

</div>

Note that methods above are not as numerically stable as `torch.norm(input[:, None] - input, dim=2, p=2)` or [pdist](https://github.com/pytorch/pytorch/issues/9406) which has recently been implemented for faster computations of the distance matrix containing upper triangular values.

Example:

```auto
a = torch.tensor([[1.4335, -1.0990, -0.8586],
        [2.1553, 2.7028, -0.8020],
        [1.0524, 0.1599, -0.0374]])

def dist(x):
    assert len(x.size()) == 2
    norm = (x ** 2).sum(1).view(-1, 1)
    dn = (norm + norm.view(1, -1)) - 2.0 * (x @ x.t())
    return dn.sqrt()

def dist2(x):
    return th.norm(x[:, None] - x, dim=2, p=2)

```

then

```auto
dist(a)
tensor([[9.7656e-04, 3.8701e+00, 1.5507e+00], # --> entry (0,0) MUST be zero
        [3.8701e+00, 0.0000e+00, 2.8752e+00],
        [1.5507e+00, 2.8752e+00, 0.0000e+00]])

dist2(a) 
tensor([[0.0000, 3.8701, 1.5507],
        [3.8701, 0.0000, 2.8752],
        [1.5507, 2.8752, 0.0000]])

torch.pdist(a) --> tensor([3.8701, 1.5507, 2.8752])

```

---

<div class="post-metadata">

**Author:** ![dashesy](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/dashesy/32/4869_2.png) [@dashesy](https://discuss.pytorch.org/u/dashesy)\
**Post date:** [December 21, 2018, 12:24am UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/11 "2018-12-21T00:24:50Z")

</div>

Can we use the currently available `torch.pairwise_distance`, or does it effectively use your first approach?

---

<div class="post-metadata">

**Author:** ![Haydnspass](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/haydnspass/32/3780_2.png) [@Haydnspass](https://discuss.pytorch.org/u/Haydnspass)\
**Post date:** [February 21, 2019, 10:18am UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/12 "2019-02-21T10:18:49Z")

</div>

No, I don’t think so. What we have here is numpy / scipy cdist (so the two sets of points do not have to be the same ones, i.e. do not need to be of an equal number) whereas torch pairwise\_distance does require that.  
It fails for something like (First dimension is batch size, N, D)

```auto
x = torch.zeros((3, 5, 2))
y = torch.zeros((3, 2, 2))

x[0, 0, 0] = 1
torch.pairwise_distance(x, y, p=2)

```

However, you may vote: [https://github.com/pytorch/pytorch/issues/15253](https://github.com/pytorch/pytorch/issues/15253)

---

<div class="post-metadata">

**Author:** ![Shihab\_Shahriar](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/shihab_shahriar/32/10025_2.png) [@Shihab\_Shahriar](https://discuss.pytorch.org/u/Shihab_Shahriar)\
**Post date:** [September 5, 2019, 6:56pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/13 "2019-09-05T18:56:58Z")

</div>

GOD I love this answer, thanks man.

---

<div class="post-metadata">

**Author:** ![Shivam\_Chandhok](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/shivam_chandhok/32/10384_2.png) [@Shivam\_Chandhok](https://discuss.pytorch.org/u/Shivam_Chandhok)\
**Post date:** [October 7, 2019, 2:44pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/14 "2019-10-07T14:44:53Z")

</div>

Can this be done for a custom distance function like wasserstein distance

---

<div class="post-metadata">

**Author:** ![usman095](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/usman095/32/20361_2.png) [@usman095](https://discuss.pytorch.org/u/usman095)\
**Post date:** [December 30, 2019, 8:32am UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/15 "2019-12-30T08:32:12Z")

</div>

Correction:

```auto
dist = dist - torch.diag(dist.diag)

```

should be

```auto
dist = dist - torch.diag(dist.diag())

```

---

<div class="post-metadata">

**Author:** ![nihilistsumo](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nihilistsumo/32/32904_2.png) [@nihilistsumo](https://discuss.pytorch.org/u/nihilistsumo)\
**Post date:** [December 31, 2020, 2:36pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/16 "2020-12-31T14:36:06Z")

</div>

Very helpful. Now can this be done batch wise? I mean, if a is a 3D tensor of shape batch x m x hidden size, how should we apply dist2 so that we get a tensor back of shape batch x m x m?

---

<div class="post-metadata">

**Author:** ![francois-rozet](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/francois-rozet/32/29950_2.png) [@francois-rozet](https://discuss.pytorch.org/u/francois-rozet)\
**Post date:** [March 7, 2021, 2:24pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/17 "2021-03-07T14:24:51Z")

</div>

For those looking at this post nowadays (2021), [`torch.cdist`](https://pytorch.org/docs/stable/generated/torch.cdist.html#torch.cdist) is available.

---

<div class="post-metadata">

**Author:** ![sgbaird](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sgbaird/32/40284_2.png) [@sgbaird](https://discuss.pytorch.org/u/sgbaird)\
**Post date:** [August 18, 2021, 6:31pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/18 "2021-08-18T18:31:22Z")

</div>

Only for 1-, 2-, 3- etc. norms from what it looks like. Anyone have a typical speed gain for `torch.cdist`?

---

<div class="post-metadata">

**Author:** ![sgbaird](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sgbaird/32/40284_2.png) [@sgbaird](https://discuss.pytorch.org/u/sgbaird)\
**Post date:** [September 4, 2021, 6:50pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/19 "2021-09-04T18:50:48Z")

</div>

I think that [scipy.stats.wasserstein\_distance](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.wasserstein_distance.html) would be a good starting point for this. The [source code](https://github.com/scipy/scipy/blob/v1.7.1/scipy/stats/stats.py#L8245-L8319) mostly uses standard NumPy functionality for which I think there are compatible PyTorch functions. Not exactly sure how that would translate to the .view() approach of B, though. If generating the pairwise distance matrix is the main desired output, I have a working Numba implementation that is ~130x faster than using `cdist(x, y, metric=scipy.stats.wasserstein_distance)`

---

<div class="post-metadata">

**Author:** ![11126](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/11126/32/30126_2.png) [@11126](https://discuss.pytorch.org/u/11126)\
**Post date:** [October 1, 2021, 4:11pm UTC](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065/20 "2021-10-01T16:11:25Z")

</div>

May I ask for the numba solution you mentioned? I’m not familiar with numba and so far only managed to get a 2x boost

[Next page](https://discuss.pytorch.org/t/efficient-distance-matrix-computation/9065.md?page=2)