矩阵之间的行向量两两距离
给定数据矩阵\(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
)等的优化。
具体来讲,把欧氏距离的运算展开
\[\begin{aligned} ||a_{i,:}-b_{j,:}||^2 & =(a_{i,:}-b_{j,:})^2 \\ & = a_{i,:}^2 + b_{j,:}^2 - 2a_{i,:}\cdot b_{j,:} \end{aligned} \]观察可以知道,式子中的三项分别等于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中看到),具体实现细节倒没有去看。