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

Question

Woodbury matrix identity指出某些矩阵的 k 阶校正的逆可以通过对原始矩阵的逆进行 k 阶校正来计算。

如果A是 p × p通过 UCV 进行秩校正的全秩矩阵哪里U是 p × k , C是 k × k ，和V是 k × p ，那么 Woodbury 恒等式为:

(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 进行比较时，我的 Woodbury 函数并不像我预期的那么快。我预计反转时间会随着维度 p 呈立方增长。 。但我的结果是:

我的问题是:为什么 Woodbury 方法这么慢？这仅仅是因为我将 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

正如您所提到的，在 A 是对角线的情况下，使用 Woodbury 技术可以更快地反转 A + UCV。也就是说，在 Woodbury 公式中，乘以 A^{-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矩阵!是的，对角线是唯一包含非零的部分，但所有零元素的算术仍将在这个稠密矩阵中进行!相反，您应该使用 Numpy 的广播功能来避免这种不必要的工作。 (或者，使用 Scipy 的 sparse module 作为稀疏对角矩阵。)

例如，

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)。对于其他替换产品也是如此。

我在我的(稍微快一点)机器上重新运行了你的计时并生成了以下图:

更清晰的性能区别!