计算（平方）马氏距离的矢量化代码

Question

编辑 2：这篇文章似乎已从 CrossValidated 移至 StackOverflow，因为它主要是关于编程的，但这意味着花哨的 MathJax 不再起作用了。希望这仍然是可读的。

假设我想用协方差矩阵 S 计算两个向量 x 和 y 之间的平方马氏距离。这是一个由

定义的相当简单的函数

M2(x, y; S) = (x - y)^T * S^-1 * (x - y)

使用 python 的 numpy 包我可以这样做

# x, y = numpy.ndarray of shape (n,)
# s_inv = numpy.ndarray of shape (n, n)
diff = x - y
d2 = diff.T.dot(s_inv).dot(diff)

或在 R 中作为

diff <- x - y
d2 <- t(diff) %*% s_inv %*% diff

在我的情况下，我被给予了

• m by n 矩阵 X
• n-维向量mu
• n by n 协方差矩阵 S

并且想找到m-维向量d这样

d_i = M2(x_i, mu; S)  ( i = 1 .. m )

其中x_i 是X 的第i 行。

在python中使用一个简单的循环并不难完成：

d = numpy.zeros((m,))
for i in range(m):
    diff = x[i,:] - mu
    d[i] = diff.T.dot(s_inv).dot(diff)

当然，鉴于外部循环发生在 python 中，而不是 numpy 库中的本机代码中，这意味着它的速度不如预期的快。 $n$ 和 $m$ 分别约为 3-4 和几十万，我在交互式程序中经常这样做，因此加速会非常有用。

在数学上，我能够使用基本矩阵运算来表述这一点的唯一方法是

d = diag( X' * S^-1 * X'^T )

在哪里

 x'_i = x_i - mu

这很容易编写一个矢量化版本，但不幸的是，计算一个超过 100 亿个元素的矩阵并且只取对角线的效率低下……我相信这个操作应该很容易用爱因斯坦表示法表达，因此希望可以使用numpy 的einsum 函数快速评估，但我什至还没有开始弄清楚这个黑魔法是如何工作的。

所以，我想知道：有没有更好的方法来数学地制定这个运算（就简单的矩阵运算而言），或者有人可以建议一些很好的矢量化（python 或 R）代码来有效地做到这一点？

奖励问题，勇敢者

我真的不想做一次，我想做k ~ 100 次。给定：

• m by n 矩阵X

• k by n 矩阵U

• n 和 n 协方差矩阵的集合，每个矩阵都表示为 S_j (j = 1..k)

通过k 矩阵D 找到m，这样

D_i,j = M(x_i, u_j; S_j)

其中i = 1..m、j = 1..k、x_i 是X 的第i 行，u_j 是j 的第U 行。

即，将以下代码矢量化：

# s_inv is (k x n x n) array containing "stacked" inverses
# of covariance matrices
d = numpy.zeros( (m, k) )
for j in range(k):
    for i in range(m):
        diff = x[i, :] - u[j, :]
        d[i, j] = diff.T.dot(s_inv[j, :, :]).dot(diff)

Answer 1

首先，看起来你可能会得到 S 然后反转它。你不应该那样做；它很慢而且数字不准确。相反，你应该得到 S 的 Cholesky 因子 L 使得 S = L L^T;那么

M^2(x, y; L L^T)
  = (x - y)^T (L L^T)^-1 (x - y)
  = (x - y)^T L^-T L^-1 (x - y)
  = || L^-1 (x - y) ||^2,

并且由于 L 是三角形 L^-1 (x - y) 可以有效地计算。

事实证明，scipy.linalg.solve_triangular 如果你正确地改造它，会很高兴地一次做一堆这样的事情：

L = np.linalg.cholesky(S)
y = scipy.linalg.solve_triangular(L, (X - mu[np.newaxis]).T, lower=True)
d = np.einsum('ij,ij->j', y, y)

稍微分解一下，y[i, j] 是 L^-1 (X_j - \mu) 的第 i 个分量。然后einsum 调用就可以了

d_j = \sum_i y_{ij} y_{ij}
    = \sum_i y_{ij}^2
    = || y_j ||^2,

我们需要的。

不幸的是，solve_triangular 不会向量化它的第一个参数，所以你可能应该在那里循环。如果 k 仅为 100 左右，那将不是什么大问题。

---

如果您实际上得到的是 S^-1 而不是 S，那么您确实可以更直接地使用 einsum 来执行此操作。由于 S 在您的情况下非常小，因此实际上反转矩阵然后执行此操作也可能会更快。但是，一旦 n 是一个重要的大小，您就会通过这样做而丢掉很多数值精度。

要弄清楚如何处理 einsum，请根据组件编写所有内容。我将直接进入奖金案例，为了符号方便，写为 S_j^-1 = T_j：

D_{ij} = M^2(x_i, u_j; S_j)
  = (x_i - u_j)^T T_j (x_i - u_j)
  = \sum_k (x_i - u_j)_k ( T_j (x_i - u_j) )_k
  = \sum_k (x_i - u_j)_k \sum_l (T_j)_{k l} (x_i - u_j)_l
  = \sum_{k l} (X_{i k} - U_{j k}) (T_j)_{k l} (X_{i l} - U_{j l})

所以，如果我们创建形状为(m, n) 的数组X、形状为(k, n) 的U 和形状为(k, n, n) 的T，那么我们可以这样写

diff = X[np.newaxis, :, :] - U[:, np.newaxis, :]
D = np.einsum('jik,jkl,jil->ij', diff, T, diff)

diff[j, i, k] = X_[i, k] - U[j, k].

Answer 2

Dougal 给出了一个出色而详细的答案，但我想我会分享一个小的修改，我发现它可以提高效率，以防其他人尝试实现这一点。开门见山：

Dougal 的方法如下：

def mahalanobis2(X, mu, sigma):
    L = np.linalg.cholesky(sigma)
    y = scipy.linalg.solve_triangular(L, (X - mu[np.newaxis,:]).T, lower=True)
    return np.einsum('ij,ij->j', y, y)

我尝试过的数学等价变体是

def mahalanobis2_2(X, mu, sigma):

    # Cholesky decomposition of inverse of covariance matrix
    # (Doing this in either order should be equivalent)
    linv = np.linalg.cholesky(np.linalg.inv(sigma))

    # Just do regular matrix multiplication with this matrix
    y = (X - mu[np.newaxis,:]).dot(linv)

    # Same as above, but note different index at end because the matrix
    # y is transposed here compared to above
    return np.einsum('ij,ij->i', y, y)

使用相同的随机输入对两个版本进行 20 倍的测试并记录时间（以毫秒为单位）。对于 X 作为 1,000,000 x 3 矩阵（mu 和 sigma 3 和 3x3），我得到：

Method 1 (min/max/avg): 30/62/49
Method 2 (min/max/avg): 30/47/37

第二个版本的速度大约提高了 30%。 我主要是在 3 或 4 个维度上运行它，但为了看看它是如何缩放的，我尝试将 X 设置为 1,000,000 x 100 并得到：

Method 1 (min/max/avg): 970/1134/1043
Method 2 (min/max/avg): 776/907/837

这几乎是相同的改进。

---

我在对 Dougal 的回答的评论中提到了这一点，但在此处添加以增加可见性：

上面的第一对方法采用单个中心点 mu 和协方差矩阵 sigma 并计算到 X 的每一行的平方马氏距离。我的额外问题是使用多组 @987654327 多次执行此操作@ 和sigma 并输出一个二维矩阵。上面的一组方法可以通过一个简单的 for 循环来实现这一点，但 Dougal 还发布了一个更聪明的例子，使用 einsum。

我决定通过使用它们来解决以下问题，将这些方法相互比较：给定kd 维正态分布（中心存储在k 的行中@d 矩阵U和k by d by d 数组S 的最后两个维度中的协方差矩阵，找到存储在n by d 矩阵@ 行中的n 点的密度987654342@.

多元正态分布的密度是点到均值的马氏距离平方的函数。 Scipy 有一个作为scipy.stats.multivariate_normal.pdf 的实现作为参考。我每次使用相同的随机参数对所有三种方法运行 10 次，使用 d=3, k=96, n=5e5。以下是结果，以点/秒为单位：

[Method]: (min/max/avg)
Scipy:                      1.18e5/1.29e5/1.22e5
Fancy 1:                    1.41e5/1.53e5/1.48e5
Fancy 2:                    8.69e4/9.73e4/9.03e4
Fancy 2 (cheating version): 8.61e4/9.88e4/9.04e4

Fancy 1 是上述两种方法中较好的一种，Fancy2 是 Dougal 的第二个解决方案。由于Fancy 2 需要计算所有协方差矩阵的逆矩阵，我还尝试了一个“作弊版本”，将它们作为参数传递，但看起来这并没有什么区别。我曾计划包含非矢量化实现，但速度太慢，需要一整天。

我们可以从中得出的结论是，使用 Dougal 的第一种方法比使用 Scipy 的方法快 20%。 不幸的是，尽管它很聪明，但第二种方法仅比第一的。可能还有一些其他优化可以完成，但这对我来说已经足够快了。

我还测试了它如何在更高维度上进行扩展。与d=100, k=96, n=1e4：

Scipy:                      7.81e3/7.91e3/7.86e3
Fancy 1:                    1.03e4/1.15e4/1.08e4
Fancy 2:                    3.75e3/4.10e3/3.95e3
Fancy 2 (cheating version): 3.58e3/4.09e3/3.85e3

Fancy 1这次似乎有了更大的优势。另外值得注意的是，Scipy 抛出了LinAlgError 8/10 次，可能是因为我随机生成的一些 100x100 协方差矩阵接近奇异（这可能意味着其他两种方法在数值上不太稳定，我实际上并没有检查结果）。