Handling matrix multiplication in log space in Python

Question

I am implementing a Hidden Markov Model and thus am dealing with very small probabilities. I am handling the underflow by representing variables in log space (so x → log(x)) which has the side effect that multiplication is now replaced by addition and addition is handled via numpy.logaddexp or similar.

Is there an easy way to handle matrix multiplication in log space?

Answer 1

This is the best way I could come up with to do it.

from scipy.special import logsumexp
def log_space_product(A,B):
    Astack = np.stack([A]*A.shape[0]).transpose(2,1,0)
    Bstack = np.stack([B]*B.shape[1]).transpose(1,0,2)
    return logsumexp(Astack+Bstack, axis=0)

The inputs A and B are the logs of the matrices A₀ and B₀ you want to multiply, and the functions returns the log of A₀B₀. The idea is that the i,j spot in log(A₀B₀) is the log of the dot product of the ith row of A₀ and the jth column of B₀. So it is the logsumexp of the ith row of A plus the jth column of B.

In the code, Astack is built so the i,j spot is a vector containing the ith row of A, and Bstack is built so the i,j spot is a vector containing the jth column of B. Thus Astack + Bstack is a 3D tensor whose i,j spot is the ith row of A plus the jth column of B. Taking logsumexp with axis = 0 then gives the desired result.