Fast execution of function which resembles matrix multiplication

Question

I am trying to make the following function as fast as possible on large matrices. The matrices can have dimensions in the range of 10K-100K. The matrix values are always between 0 and 1. The function resembles matrix multiplication, but with log operations in the inner loop. These log operations appear to be a bottleneck.

I tried various ways of using Numba and Cython. I also tried writing as much as I could with Numpy. I also experimented with doing fewer memory lookups, but this did not seem to give much advantage. The fastest version is below. High precision is greatly preferred, but if there is a way to increase speed at its expense, that would also be appreciated. Thank you for your feedback.

import numpy as np
import numba
from numba import njit, prange

@numba.jit(nopython=True, fastmath=True, parallel=True)
def f(A, B):

    len_A = A.shape[0]
    len_B = B.shape[0]
    num_factors = B.shape[1]        
    C = np.zeros((len_A, len_B))

    for i in prange(len_A):
        for j in prange(len_B):         
            for a in prange(num_factors):

                A_elem = A[i,a]
                B_elem = B[j,a]
                AB_elem = (A_elem + B_elem)/2

                C[i,j] += A_elem * np.log(A_elem/AB_elem) + \
                          B_elem * np.log(B_elem/AB_elem) + \
                          (1-A_elem) * np.log((1-A_elem)/(1-AB_elem)) + \
                          (1-B_elem) * np.log((1-B_elem)/(1-AB_elem))

    C = (np.maximum(C, 0)/2*num_factors)**0.5
    return C

#A = np.random.rand(10000, 10000)
#B = np.random.rand(10000, 10000)
#f(A, B)

Answer 1

I think you can rearrange your equation from

\$ AB_{ija} = \frac{1}{2}(A_{ia} + B_{ja}) \$

\$ \begin{align} C_{ij} = \sum_a & A_{ia}\cdot \log\left(\frac{A_{ia}}{AB_{ija}}\right) + B_{ja}\cdot \log\left(\frac{B_{ja}}{AB_{ija}}\right) + \\ &+ (1-A_{ia})\cdot \log\left(\frac{1-A_{ia}}{1-AB_{ija}}\right) + (1-B_{ja})\cdot \log\left(\frac{1-B_{ja}}{1-AB_{ija}}\right) \end{align} \$

to

\$ \begin{align} C_{ij} = \sum_a & A_{ia} \cdot \log A_{ia} + B_{ja}\cdot \log B_{ja} - 2AB_{ija}\cdot \log AB_{ija}\\ &+ (1-A_{ia}) \cdot \log (1-A_{ia}) + (1-B_{ja})\cdot \log (1-B_{ja}) - 2(1-AB_{ija})\cdot \log (1-AB_{ija})\\ = \sum_a & A_{ia}\cdot \log A_{ia} + \sum_a B_{ja}\cdot \log B_{ja} - 2\sum_a AB_{ija}\cdot \log AB_{ija} + \dots \end{align} \$

This will at least save you some divisions, which can be a bit slower. But the crucial part is to see that you can compute the sums separately. The sums with only \$A\$ or \$B\$ are a lot less computation, which you would otherwise repeat, and only the combined terms are (computationally) hard.

It also means that there is probably a numpy way to eliminate at least some of your for loops, which might give you a bit of a speedup, although numba will negate some of that (by already having sped up the for loops). The hard part is then to use the numpy broadcasting correctly to achieve this, which is left as an exercise for now :).

Answer 2

You can simplify this by rearranging the function as follows (writing z=(a+b)/2):

a log(a/z) + b log(b/z) + (1-a) log( (1-a)/(1-z) ) + (1-b) log( (1-b)/(1-z) )

= a log(a) + b log(b) + (1-a) log(1-a) + (1-b) log(1-b) - 2z log(z) - 2(1-z)log(1-z)

The first 4 terms can be evaluated outside the main loop, so that you only need to evaluate log(z) and log(1-z) within the inner loop. This should give close to a factor 2 speed up.

Since the remaining function to evaluate in the inner loop, z log(z) + (1-z)log(1-z) is a function of a single variable, you could possible replace this with a look-up table .... but I don't know how much you would gain from that.

Alternatively, use

-log(0.5) - 2*(z-0.5)**2 > -z log(z) - (1-z)log(1-z) > -2z * log(0.5)

to first get upper and lower bounds for each C[i,j] and then focus attention on those values for which the upper bound is greater than the maximum of the lower bounds.