numerically stable way to multiply log probability matrices in numpy

Question

I need to take the matrix product of two NumPy matrices (or other 2d arrays) containing log probabilities. The naive way `np.log(np.dot(np.exp(a), np.exp(b)))` is not preferred for obvious reasons. Using ``` from scipy.misc import logsumexp res = np.zeros((a.shape[0], b.shape[1])) for n in range(b.shape[1]): # broadcast b[:,n] over rows of a, sum columns res[:, n] = logsumexp(a + b[:, n].T, axis=1) ``` works but runs about 100 times slower than `np.log(np.dot(np.exp(a), np.exp(b)))` Using ``` logsumexp((tile(a, (b.shape[1],1)) + repeat(b.T, a.shape[0], axis=0)).reshape(b.shape[1],a.shape[0],a.shape[1]), 2).T ``` or other combinations of tile and reshape also work but run even slower than the loop above due to the prohibitively large amounts of memory required for realistically sized input matrices. I am currently considering writing a NumPy extension in C to compute this, but of course I'd rather avoid that. Is there an established way to do this, or does anybody know of a less memory…

Accepted Answer

`logsumexp` works by evaluating the right-hand side of the equation

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

I.e., it pulls out the max before starting to sum, to prevent overflow in `exp`. The same can be applied before doing vector dot products:

```
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) }
```

but by taking a different turn in the derivation, we obtain

```
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)])
```

The final form has a vector dot product in its innards. It also extends readily to matrix multiplication, so we get the algorithm

```
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
```

This creates two `A`-sized temporaries and two `B`-sized ones, but one of each can be eliminated by

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

and similarly for `B`. (If the input matrices may be modified by the function, all the temporaries can be eliminated.)

numerically stable way to multiply log probability matrices in numpy

Solution

Problem