如何用torch.cdist正确计算相近向量的批量L2距离？

1. 使用compute_mode='donot_use_mm_for_euclid_dist'是否正确？ #

是的，这是处理相近向量场景的正确做法。

矩阵乘法方式计算L2距离依赖公式||a-b||² = ||a||² - 2a·b + ||b||²，当a和b向量非常接近时，||a||²与2a·b - ||b||²的数值会极其接近，浮点数相减时会丢失大量有效精度，也就是所谓的灾难性抵消，甚至可能出现计算结果为负的情况（后续会被钳位为0，但精度已经损失）。

而compute_mode='donot_use_mm_for_euclid_dist'会强制torch.cdist采用直接计算向量差的平方和的逻辑：先计算a - b，再对差值的每个元素平方、求和，最后开根号。这种方式从根源上避免了相近向量的抵消问题，能保证计算精度。

2. 计算相近向量批量成对L2距离的正确方法 #

目前针对这类场景，推荐以下两种方案：

方案一：使用torch.cdist（官方首选） #

继续使用torch.cdist并指定compute_mode='donot_use_mm_for_euclid_dist'，这是PyTorch官方为精度敏感场景提供的最优解，既保证了数值稳定性，也经过了性能优化，不会比矩阵乘法方式有明显的性能差距。

修正你代码中小问题（b是普通张量，没有weight属性）后的示例：

import torch
import torch.nn as nn
import numpy as np

nr_units = 200
z_dim = 128

a = nn.Embedding(num_embeddings=nr_units, embedding_dim=z_dim)
a.weight.data.uniform_(-np.sqrt(1 / z_dim), np.sqrt(1 / z_dim))

b = torch.rand(1024, z_dim)

# 正确调用torch.cdist
dist_cdist = torch.cdist(
    x1=a.weight,  # a是Embedding层，取其weight张量
    x2=b,
    p=2,
    compute_mode='donot_use_mm_for_euclid_dist'
)

方案二：手动实现数值稳定的计算逻辑 #

如果需要自定义扩展，也可以手动实现直接计算向量差的逻辑，和cdist指定参数后的内部逻辑一致：

# 扩展维度实现广播计算
a_expanded = a.weight.unsqueeze(1)  # shape: (nr_units, 1, z_dim)
b_expanded = b.unsqueeze(0)          # shape: (1, 1024, z_dim)

# 计算向量差的平方和再开根号
diff = a_expanded - b_expanded
dist = torch.sqrt(torch.sum(torch.square(diff), dim=-1))  # shape: (nr_units, 1024)

这种手动实现完全避免了灾难性抵消问题，精度和cdist的指定模式一致，适合需要自定义中间步骤的场景。

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