用于快速矩阵求逆的伍德伯里恒等式——比预期慢

Question

的Woodbury的矩阵身份指出一些矩阵的秩-K校正的逆可以通过执行秩k值校正到原始矩阵的逆矩阵来计算。

如果A是一个p × p满秩矩阵，通过UCV其中Uis p × k、Cisk × k和Vis 进行秩校正k × p，则伍德伯里恒等式为：

(A + UCV)^{-1} = A^{-1} - A^{-1} U (C^{-1} + V A^{-1} U)^{-1} V A^{-1}

关键是不是反转p × p矩阵，而是反转k × k矩阵。在许多应用中，我们可以假设k < p. A在某些情况下，反转可能很快，例如如果A是对角矩阵。

我在这里实现了这个，假设A是对角线，那C就是身份：

def woodbury(A, U, V, k):
    A_inv = np.diag(1./np.diag(A))  # Fast matrix inversion of a diagonal.
    B_inv = np.linalg.inv(np.eye(k) + V @ A_inv @ U)
    return A_inv - (A_inv @ U @ B_inv @ V @ A_inv)

理智检查了我的实施

n = 100000
p = 1000
k = 100
A = np.diag(np.random.randn(p))
U = np.random.randn(p, k)
V = U.T
M = U @ V + A
M_inv = woodbury(A, U, V, k)
assert np.allclose(M @ M_inv, np.eye(p))

但是当我实际将其与 进行比较时numpy.linalg.inv，我的伍德伯里函数并没有我预期的那么快。我希望反转的时间随着 th 维度的三次方增长p。但我的结果是：

我的问题是：为什么伍德伯里方法这么慢？仅仅是因为我将 Python 代码与 LAPACK 进行比较还是有其他事情发生？

编辑：我对 einsum() 与广播的实验

我实现了三个版本：（1）使用einsum和einsum_path，（2）根据接受的答案使用广播，以及（3）使用两者。这是我使用的实现，使用einsum优化einsum_path：

def woodbury_einsum(A, U, V, k):
    A_inv = np.diag(1./np.diag(A))
    tmp   = np.einsum('ab,bc,cd->ad',
                      V, A_inv, U,
                      optimize=['einsum_path', (1, 2), (0, 1)])
    B_inv = np.linalg.inv(np.eye(k) + tmp)
    tmp   = np.einsum('ab,bc,cd,de,ef->af',
                      A_inv, U, B_inv, V, A_inv,
                      optimize=['einsum_path', (0, 1), (0, 1), (0, 2), (0, 1)])
    return A_inv - tmp

结果在这里：

因此，避免使用对角矩阵进行矩阵乘法的计算成本比使用 优化矩阵乘法的顺序和内存占用更快einsum()。

Answer 1

正如您所提到的，在diagonal的情况下，A + UCV通过使用 Woodbury 技术可以更快地进行反转。也就是说，在伍德伯里公式中，您的乘法应该及时发生，而不是因为您所做的只是缩放右/左项的行/列。AA^{-1}O(p x m)O(p x m x p)

但是，这不是您在以下代码中所做的！

def woodbury(A, U, V, k):
    A_inv = np.diag(1./np.diag(A))
    B_inv = np.linalg.inv(np.eye(k) + V @ A_inv @ U)
    return A_inv - (A_inv @ U @ B_inv @ V @ A_inv)

你A_inv是一个完整的p x p矩阵！是的，对角线是唯一包含非零元素的部分，但所有零元素的算术运算仍将在这个密集矩阵中进行！相反，您应该使用 Numpys 广播功能来避免这种不必要的工作。（或者，作为使用 Scipysparse模块的稀疏对角矩阵。）

例如，

def efficient_woodbury(A, U, V, k):
    A_inv_diag = 1./np.diag(A)  # note! A_inv_diag is a vector!
    B_inv = np.linalg.inv(np.eye(k) + (V * A_inv_diag) @ U)
    return np.diag(A_inv_diag) - (A_inv_diag.reshape(-1,1) * U @ B_inv @ V * A_inv_diag)

该产品V * A_inv_diag与您的相同，V @ A_inv但仅在O(p x k)一段时间内运行，而不是O(p x k x p). 其他被替换的产品也是如此。

我在我的（稍微快一点）机器上重新运行了你的计时并产生了以下图：

更清晰的性能区别！