如何在PyTorch中高效计算张量中元素的成对差值？

问题描述 #

给定1D输入张量：

tensor([1, 2, 3, 4, 5, 6.])

需要生成如下预期结果：

tensor([-1., -2., -3., -4., -5., -1., -2., -3., -4., -1., -2., -3., -1., -2., -1.])

已尝试方法的局限 #

• torch.pdist方法：计算1-范数距离得到的是绝对值，无法直接得到带符号的差值：

import torch
x = torch.tensor([[1], [2], [3], [4], [5], [6.]])
torch.pdist(x, p=1)
# 输出：tensor([1., 2., 3., 4., 5., 1., 2., 3., 4., 1., 2., 3., 1., 2., 1.])

• 广播生成矩阵：通过广播得到的差值矩阵包含无意义的下三角和对角线元素，不符合输出要求：

x = torch.tensor([1, 2, 3, 4, 5, 6.])
x[:, None] - x[None, :]
# 输出：
# tensor([[ 0., -1., -2., -3., -4., -5.],
#        [ 1.,  0., -1., -2., -3., -4.],
#        [ 2.,  1.,  0., -1., -2., -3.],
#        [ 3.,  2.,  1.,  0., -1., -2.],
#        [ 4.,  3.,  2.,  1.,  0., -1.],
#        [ 5.,  4.,  3.,  2.,  1.,  0.]])

• 其他备选方案：scipy的squareform函数可用性存疑，不确定是否会破坏损失函数可微性；Python循环逐个计算拼接的方法速度慢，还可能干扰梯度计算，不推荐。

解决方案 #

利用PyTorch的掩码操作提取上三角区域（排除对角线）的差值，再展平即可，全程支持梯度计算：

import torch

x = torch.tensor([1, 2, 3, 4, 5, 6.], requires_grad=True)
# 生成广播差值矩阵
diff_matrix = x[:, None] - x[None, :]
# 生成上三角掩码（仅保留i < j的位置）
mask = torch.triu(torch.ones_like(diff_matrix), diagonal=1)
# 提取目标元素并展平
result = diff_matrix[mask.bool()].flatten()

print(result)
# 输出：tensor([-1., -2., -3., -4., -5., -1., -2., -3., -4., -1., -2., -3., -1., -2., -1.],
#        grad_fn=<ReshapeAliasBackward0>)

方案说明 #

1. torch.triu(..., diagonal=1)生成上三角掩码，仅保留矩阵中索引i < j的位置为1，其余为0；
2. 通过布尔掩码过滤出需要的差值元素，展平后得到预期的一维张量；
3. 所有操作均基于PyTorch原生张量运算，自动支持梯度反向传播，不会影响损失函数的可微性，且运算效率远高于Python循环。

内容的提问来源于stack exchange，提问作者dawid