How to calculate pair-wise differences between two tensors in a vectorized way?

I have two tensors of shape (4096, 3) and (4096,3). What I’d like to do is calculate the pairwise differences between all of the individual vectors in those matrices, such that I end up with a (4096, 4096, 3) tensor. This can be done in for-loops, but I’d like to do a vectorized approach. NumPy lets you do some broadcasting approaches, but I’m not sure how to do the same for PyTorch. Any help would be greatly appreciated!

1 Like

I’m not sure if I understand it correctly, but I think this will do it:

a = torch.randn(5, 3)
b = torch.randn(5, 3)

res = a.unsqueeze(1) - b
# res[i] corresponds to (a[i] - b)
8 Likes

you can use the code from here for pairwise L2 distance.

1 Like

Thanks, this is exactly what I was looking for! Out of curiosity, why is it that the extra dimension is added to the middle, so it becomes (5, 1, 3)?

Lastly, a follow-up question, is it possible to apply this batchwise, so if you had for example:

a = torch.randn(1, 5, 3)
b = torch.randn(1, 5, 3)

What seems to be working is just doing res = a.unsqueeze(2) - b, but I’m not sure if that’s 100% correct (although the values match with both methods when I use two fixed arrays).

Thank you for linking the code!

This shouldn’t work for batch sizes > 1.
However, this code should work:

a = torch.randn(2, 5, 3)
b = torch.randn(2, 5, 3)

res = a.unsqueeze(2) - b.unsqueeze(1)
# res[batch_idx, i] corresponds to (a[batch_idx, i] - b[batch_idx])

The extra dimension is added so that the tensors will automatically be broadcasted.
Basically this would be equal to (a.unsqueeze(1) - b.unsqueeze(0)).

4 Likes

Great, thanks so much!

1 Like

Great answer!
How would I achieve pairwise concatenation?
The ‘res.shape’ should be torch.Size([2, 5, 5, 6]) or torch.Size([2, 10, 3, 3]).

Could you give an example of this operation or post some pseudo code for it, please?

Let’s assume that we have a tensor with shape [batch_size, channel_dim, spatial_dim]:

a = torch.randn(2, 3, 5)
tensor([[[-0.8526,  1.0994, -0.2659, -0.3834,  1.3758],
         [ 0.4018, -0.9505,  0.0167, -0.2450, -1.1768],
         [ 0.9758, -0.7121,  0.7073, -2.0957, -0.8515]],

        [[ 0.4197, -0.0777, -0.0946,  0.1594, -0.0315],
         [-0.1422, -0.8626,  2.9035, -0.5575, -0.4993],
         [-0.7888, -0.5184, -0.7460, -0.4184, -0.2338]]])

The ‘res’ should have the shape of torch.Size([2, 6, 5, 5]) in this case, and given any pair of spatial locations, the pairwise concatenation result should be:

res[0, :, 2, 3] = torch.cat([a[0, :, 2], a[0, :, 3]], dim=-1) 
                = tensor([-0.2659,  0.0167,  0.7073, -0.3834, -0.2450, -2.0957])