Python numpy中对数概率矩阵乘法的数值稳定方法

我需要取两个NumPy矩阵（或其他2d数组）的矩阵积，其中包含对数概率。出于明显的原因，不推荐使用简单的方法

np.log（np.dot（np.exp（a），np.exp（b））

使用

工作正常，但运行速度比

np.log（np.dot（np.exp（a），np.exp（b））慢100倍左右。

使用

或者平铺和重塑的其他组合也可以工作，但由于实际大小的输入矩阵需要大量内存，因此运行速度甚至比上面的循环还要慢

我目前正在考虑用C编写一个NumPy扩展来计算这个，但当然我宁愿避免这样。有没有一种既定的方法可以做到这一点，或者有没有人知道一种执行这种计算的内存密集度较低的方法

编辑： 感谢larsmans提供此解决方案（有关推导，请参见下文）：

使用iPython的magic

%timeit

函数将此方法与上面发布的方法（

logdot_old

）进行快速比较，得出以下结果：

In  [1] a = np.log(np.random.rand(1000,2000))

In  [2] b = np.log(np.random.rand(2000,1500))

In  [3] x = logdot(a, b)

In  [4] y = logdot_old(a, b) # this takes a while

In  [5] np.any(np.abs(x-y) > 1e-14)
Out [5] False

In  [6] %timeit logdot_old(a, b)
1 loops, best of 3: 1min 18s per loop

In  [6] %timeit logdot(a, b)
1 loops, best of 3: 264 ms per loop

---

显然拉斯曼的方法抹杀了我的

您正在访问

res

和

b

的列，这些列的性能较差。要尝试的一件事是将它们存储在中。

logsumexp

通过计算等式的右侧来工作

log(∑ exp[a]) = max(a) + log(∑ exp[a - max(a)])

也就是说，它在开始求和之前提取最大值，以防止

exp

中溢出。在进行矢量点积之前，也可以应用同样的方法：

log(exp[a] ⋅ exp[b])
 = log(∑ exp[a] × exp[b])
 = log(∑ exp[a + b])
 = max(a + b) + log(∑ exp[a + b - max(a + b)])     { this is logsumexp(a + b) }

但是通过在推导中进行不同的转向，我们得到

log(∑ exp[a] × exp[b])
 = max(a) + max(b) + log(∑ exp[a - max(a)] × exp[b - max(b)])
 = max(a) + max(b) + log(exp[a - max(a)] ⋅ exp[b - max(b)])

最终形式的内部有一个向量点积。它也很容易扩展到矩阵乘法，所以我们得到了算法

def logdotexp(A, B):
    max_A = np.max(A)
    max_B = np.max(B)
    C = np.dot(np.exp(A - max_A), np.exp(B - max_B))
    np.log(C, out=C)
    C += max_A + max_B
    return C

这将创建两个

A

大小的临时变量和两个

B

大小的临时变量，但每个临时变量中的一个可以通过

exp_A = A - max_A
np.exp(exp_A, out=exp_A)

同样，对于

B

。（如果该函数可以修改输入矩阵，则可以消除所有临时变量。）

假设A.shape==（n，r）和

B.shape==（r，m）

。在计算矩阵乘积

C=A*B

时，实际上有

n*m

和。在日志空间中工作时，为了获得稳定的结果，需要在每个求和中使用logsumexp技巧。幸运的是，使用numpy广播可以很容易地分别控制A和B的行和列的稳定性

代码如下：

def logdotexp(A, B):
    max_A = np.max(A,1,keepdims=True)
    max_B = np.max(B,0,keepdims=True)
    C = np.dot(np.exp(A - max_A), np.exp(B - max_B))
    np.log(C, out=C)
    C += max_A + max_B
    return C

注意：

def logdotexp_less_stable(A, B):
    max_A = np.max(A)
    max_B = np.max(B)
    C = np.dot(np.exp(A - max_A), np.exp(B - max_B))
    np.log(C, out=C)
    C += max_A + max_B
    return C

print('old method:')
print(logdotexp_less_stable([[0,0],[0,0]], [[-1000,0], [-1000,0]]))
print('new method:')
print(logdotexp([[0,0],[0,0]], [[-1000,0], [-1000,0]]))

这背后的推理与FredFoo的答案相似，但他对每个矩阵使用了一个最大值。由于他没有考虑每一个<代码> n*m < /代码>求和，所以最后一个矩阵的一些元素可能仍然不稳定，如一个注释中所提到的。 使用@identity-m反例与当前接受的答案进行比较：

def logdotexp_less_stable(A, B):
    max_A = np.max(A)
    max_B = np.max(B)
    C = np.dot(np.exp(A - max_A), np.exp(B - max_B))
    np.log(C, out=C)
    C += max_A + max_B
    return C

print('old method:')
print(logdotexp_less_stable([[0,0],[0,0]], [[-1000,0], [-1000,0]]))
print('new method:')
print(logdotexp([[0,0],[0,0]], [[-1000,0], [-1000,0]]))

哪张照片

old method:
[[      -inf 0.69314718]
 [      -inf 0.69314718]]
new method:
[[-9.99306853e+02  6.93147181e-01]
 [-9.99306853e+02  6.93147181e-01]]

---

Fred Foo目前接受的答案以及Hassan的答案在数值上不稳定（Hassan的答案更好）。稍后将提供一个哈桑回答失败的输入示例。我的执行情况如下：

将numpy导入为np
从scipy.special导入logsumexp
def logmatmulexp（log_A:np.ndarray，log_B:np.ndarray）->np.ndarray:
“”“给定形状为ϴ×R的矩阵log_A和形状为R×I的矩阵log_B，计算
（log_A.exp（）@log_B.exp（））.log（）以数值稳定的方式。
具有O（ϴRI）时间复杂度和空间复杂度
ϴ，R=对数ϴA.形状
I=原木形状[1]
断言log_B.shape==（R，I）
log_A_expanded=np.broadcast_to（np.expanded_dims（log_A，2），（ϴ，R，I））
log_B_expanded=np.broadcast_to（np.expanded_dims（log_B，0），（ϴ，R，I））
log_pairwise_products=log_A_expanded+log_B_expanded#shape:（ϴ，R，I）
返回logsumexp（对数成对乘积，轴=1）

就像Hassan和Fred Foo的答案一样，我的答案具有时间复杂性O（ϴRI）。他们的答案具有空间复杂度O（ϴR+RI）（我实际上对此不确定），而不幸的是，我的答案具有空间复杂度O（ϴRI）-这是因为numpy可以将一个ϴ×R矩阵乘以一个R×I矩阵，而无需分配额外的大小为ϴ×R×I的数组空间复杂度不是我的方法的内在属性——我认为如果你用循环写出它，你可以避免这种空间复杂度，但不幸的是，我不认为你可以用stock numpy函数做到这一点

我已经检查了我的代码实际运行的时间，它比常规矩阵乘法慢20倍

以下是您如何知道我的答案在数值上是稳定的：

显然，除回流线以外的所有线在数值上都是稳定的

已知

logsumexp

函数在数值上是稳定的

因此，我的

logmatmulexp

函数在数值上是稳定的

我的实现还有一个很好的特性。如果不使用numpy，而是在pytorch中编写相同的代码，或者使用另一个具有自动微分功能的库，您将自动获得数值稳定的反向传递。下面是我们如何知道反向过程在数值上是稳定的：

我的代码中的所有函数在任何地方都是可微的（不像

np.max

）

显然，除了返回线之外，所有线的反向传播在数值上都是稳定的，因为那里绝对没有发生任何奇怪的事情

通常pytorch的开发人员知道他们在做什么。所以，他们实现了反向传递就足够信任他们了

def logdotexp_less_stable(A, B):
    max_A = np.max(A)
    max_B = np.max(B)
    C = np.dot(np.exp(A - max_A), np.exp(B - max_B))
    np.log(C, out=C)
    C += max_A + max_B
    return C

print('old method:')
print(logdotexp_less_stable([[0,0],[0,0]], [[-1000,0], [-1000,0]]))
print('new method:')
print(logdotexp([[0,0],[0,0]], [[-1000,0], [-1000,0]]))

old method:
[[      -inf 0.69314718]
 [      -inf 0.69314718]]
new method:
[[-9.99306853e+02  6.93147181e-01]
 [-9.99306853e+02  6.93147181e-01]]