矩阵之间的行向量两两距离

给定数据矩阵\(A\in R^{n\times d}\)和矩阵\(B\in R^{m\times d}\) ，A和B中每一行都是一个数据点，现在要去求A中所有元素和B中所有元素之间的欧氏距离。即计算矩阵\(D =(d_{ij} = ||a_{i,:} - b_{j, :}||) \in R^{n\times m}\)。
直接去做\(n\times m\)的两层循环当然可以得到结果，但是既然数据都已经放在矩阵里了，它自然可以“并行”地去计算，好利用一些现有的矩阵计算库（比如numpy、torch）等的优化。

具体来讲，把欧氏距离的运算展开

观察可以知道，式子中的三项分别等于A中第i行的平方和、B中第j行的平方和、A中第i行和B中第j行的内积，这样的话，可以分别对A和B做逐元素平方求和，而向量两两做内积，得到的结果其实就是A与B转置相乘。矩阵相乘作为一个经典操作那是被优化得好好的，比起自己去做两层循环方便得多。

整个过程就成了\(RowSum(A .* A)^T + RowSum(B .* B) - 2A\cdot B^T\) 。

这里有个问题，A平方和之后转置得到的是\(n\times 1\) 的向量，B平方和得到的是\(1\times m\)的向量，要把它们都扩展成\(n\times m\)的向量才好相加。幸运的是，现在常见的矩阵运算库基本都是有“广播”功能的。当向量\(v \in R^{n\times 1}\)和向量\(u\in R^{1\times m}\) 相加时，它们分别会被自动扩展成为矩阵\(V\in R^{n\times m}\)和\(U\in R^{n\times m}\) ，显然可以想到V就是把v复制了m列，U是把u复制了n行。通过自动广播，维度的处理就完全不是问题了。

具体代码如下

# numpy version
A2 = (A ** 2).sum(axis=1).reshape((-1, 1))	# 转为n*1
B2 = (B ** 2).sum(axis=1) # 这里不用reshape了，(n,)和(1,n)的向量在计算过程中效果是一样的
D = A2 + B2 - 2 * A.dot(B.transpose())
# D = np.clamp(D, a_min=0) # 有时候会出现精度问题导致结果小于0 ，我习惯于在这里进行一次校正
D = D.sqrt()	# 开方得到欧氏距离
# ...

# torch version
A2 = (A ** 2).sum(dim=1).reshape((-1, 1))
B2 = (A ** 2).sum(dim=1)
D = A2 + B2 - 2 * A.mm(B.t())
# ...

实际场景经常会遇到一个点矩阵\(A\in R^{a\times d}\) 需要求所有数据点之间的两两距离，此时就可以用上面的方式去计算A和A自身的两两距离。

批量计算点集与点集之间的两两距离

今天在做实验时遇到这样一个场景，\(X\in R^{B\times n\times d}\) 和\(Y\in R^{B\times m\times d}\)，其实这也是玩深度学习时常会出现到的情况，就是数据集被划分为一个一个的Batch，这里我想对每一个batch b分别计算\(X[b,:, :]\) 和\(Y[b,:, :]\) 这两个矩阵中的两两距离，放在常用的计算库（如torch、paddle等）中，一样是可以转换为矩阵运算的，只是需要做一点点的reshape而已。

# paddle version
X2 = (X ** 2).sum(axis=2).reshape((batch_size, 1, -1))
Y2 = (Y ** 2).sum(axis=2).reshape((batch_size, -1, 1)) # 由于多了第一维，这里的reshape不能偷懒不写了
D = X2 + Y2 - 2 * X.bmm(Y.transpose((0, 2, 1)))

之所以能这么做，得益于很多库实现了bmm运算，即对每个batch分别进行矩阵乘法（这个函数应用于张量的时代，在numpy里还没有）。

这一需求其实在pytorch中有一个特定的函数torch.cdist完成了。然而很不幸的百度paddle框架中并没有类似函数，所以只能自己实现了一版。当然torch内部似乎是用了contiguous函数去做了些矩阵结构的调整（从GitHub issue中看到），具体实现细节倒没有去看。