The currently accepted answer by Fred Foo, as well as Hassan's answer, are numerically unstable (Hassan's answer is better). An example of an input on which Hassan's answer fails will be provided later. My implementation is as follows:
import numpy as np
from scipy.special import logsumexp
def logmatmulexp(log_A: np.ndarray, log_B: np.ndarray) -> np.ndarray:
"""Given matrix log_A of shape ϴ×R and matrix log_B of shape R×I, calculates
(log_A.exp() @ log_B.exp()).log() in a numerically stable way.
Has O(ϴRI) time complexity and space complexity."""
ϴ, R = log_A.shape
I = log_B.shape[1]
assert log_B.shape == (R, I)
log_A_expanded = np.broadcast_to(np.expand_dims(log_A, 2), (ϴ, R, I))
log_B_expanded = np.broadcast_to(np.expand_dims(log_B, 0), (ϴ, R, I))
log_pairwise_products = log_A_expanded + log_B_expanded # shape: (ϴ, R, I)
return logsumexp(log_pairwise_products, axis=1)
Just like Hassan's answer and Fred Foo's answer, my answer has time complexity O(ϴRI). Their answers have space complexity O(ϴR+RI) (I am not actually sure about this), while mine unfortunately has space complexity O(ϴRI) - this is because numpy can multiply a ϴ×R matrix by a R×I matrix without allocating an additional array of size ϴ×R×I. Having O(ϴRI) space complexity is not an immanent property of my method - I think if you write it out using cycles, you can avoid this space complexity, but unfortunately I don't think you can do that using stock numpy functions.
I have checked how much actual time my code runs, it's 20 times slower than regular matrix multiplication.
Here's how you can know that my answer is numerically stable:
- Clearly, all lines other than the return line are numerically stable.
- The
logsumexp function is known to be numerically stable.
- Therefor, my
logmatmulexp function is numerically stable.
My implementation has another nice property. If instead of using numpy you write the same code in pytorch or using another library with automatic differentiation, you will get a numerically stable backward pass automatically. Here's how we can know the backward pass will be numerically stable:
- All functions in my code are differentiable everywhere (unlike
np.max)
- Clearly, back propagating through all lines except the return line is numerically stable, because absolutely nothing weird is happening there.
- Usually the developers of pytorch know what they're doing. So it's enough to trust them that they implemented backward pass of logsumexp in a numerically stable way.
- Actually the gradient of logsumexp is the softmax function (for reference google "softmax is gradient of logsumexp" or see https://arxiv.org/abs/1704.00805 proposition 1). It's known that softmax can be calculated in a numerically stable way. So the pytorch devs probably just use softmax there (I haven't actually checked).
Below is the same code in pytorch (in case you need backpropagation). Due to how pytorch backpropagation works, during forward pass it will save the log_pairwise_products tensor for the backward pass. This tensor is large, and you probably don't want it to be saved - you can just recalculate it once again during backward pass. In such case I suggest you use checkpointing - it's really easy - see the second function below.
import torch
from torch.utils.checkpoint import checkpoint
def logmatmulexp(log_A: torch.Tensor, log_B: torch.Tensor) -> torch.Tensor:
"""Given matrix log_A of shape ϴ×R and matrix log_B of shape R×I, calculates
(log_A.exp() @ log_B.exp()).log() and its backward in a numerically stable way."""
ϴ, R = log_A.shape
I = log_B.shape[1]
assert log_B.shape == (R, I)
log_A_expanded = log_A.unsqueeze(2).expand((ϴ, R, I))
log_B_expanded = log_B.unsqueeze(0).expand((ϴ, R, I))
log_pairwise_products = log_A_expanded + log_B_expanded # shape: (ϴ, R, I)
return torch.logsumexp(log_pairwise_products, dim=1)
def logmatmulexp_lowmem(log_A: torch.Tensor, log_B: torch.Tensor) -> torch.Tensor:
"""Same as logmatmulexp, but doesn't save a (ϴ, R, I)-shaped tensor for backward pass.
Given matrix log_A of shape ϴ×R and matrix log_B of shape R×I, calculates
(log_A.exp() @ log_B.exp()).log() and its backward in a numerically stable way."""
return checkpoint(logmatmulexp, log_A, log_B)
Here's an input on which Hassan's implementation fails but my implementation gives the correct output:
def logmatmulexp_hassan(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
log_A = np.array([[-500., 900.]], dtype=np.float64)
log_B = np.array([[900.], [-500.]], dtype=np.float64)
print(logmatmulexp_hassan(log_A, log_B)) # prints -inf, while the correct answer is approximately 400.69.
N * Kfeature matrix, withN >> K, and aK * Mweight matrix, withKandMof roughly the same size. – Splasherscipy.misc.logsumexpis doing what you think it is - according to the docs theb=parameter is actually a scaling factor forexp(a), i.e.np.log(np.sum(b*np.exp(a))). – Elevatorbin my code is a numpyndarray, not a parameter tologsumexp. The 1 corresponds to theaxisparameter, I've clarified this in the example above. – Splasherlogsumexpis usually used when, e.g., calculating the log-normalizer given a log-odds parameter vector. Are you sure that what you're doing makes sense? – Mimosaceouslogdotfunction. I am NOT trying to clarify @mart's idea. My usage is based on the fact that if you have a matrix A_{z,y} = p(Z=z|Y=y) and a matrix B_{y,x} = p(Y=y|X=x), and (X⊥Z|Y), then p(Z=z|X=x) = (AB)_{z,x}. Now, if I store logarithms of probabilities as A'_{z,y} = log p(Z=z|Y=y) and B'_{y,x} = log p(Y=y|X=x), then we have log p(Z=z|X=x) = logdot(A', B')_{z,x}. Personally I am planning to use this to build a categorical kinda directed graphical model. If I am wrong, please say where. – Luthern