Compute matrix multiplication only for given coordinates
In PyTorch, I would like to compute

E * A.mm(B)


where E can be a very sparse matrix consisting of 0's and 1's. In other words, I want to compute A.mm(B) and then leave only the certain coordinates. Is there a way to compute such a sparse matrix efficiently? I have full control over matrix representations.

Also, in most cases, E consists only of 1's, so I would like this case also to be handled efficiently.

You don't need an element-wise multiplication for that, as E is essentially a boolean matrix, that is used as a mask to select the values, where E is 1, and discard the values where E is 0.

C = A.mm(B)

# Ensure that E is a boolean matrix to only keep values where E is True,
# otherwise the 0s and 1s would be treated as indices to select the values.
C = C[E.to(torch.bool)]


If you want to avoid the entire matrix multiplication and only compute the values you would be masking afterwards, you need to manually select the values for A and B that produce the desired values in C.

The matrix multiplication C = AB, where A is an m x n matrix and B  an n x p matrix, produces an m x p matrix C, whose values are obtained by multiplying the i-th row of A with the j-th column of B element-wise and taking the sum of these n products. Formally:



Given E, an m x p matrix, that determines which elements of C are required, the index pairs of the required elements are given as follows:



# Indices of required elements (i.e. indices of non-zero elements of E)
# Separate the tensor of (i, j) pairs, into a pair of tensors,
# containing the indices i and j respectively.
indices_i, indices_j = E.nonzero().unbind(dim=1)

# Select all needed rows of A and the needed columns of B
A = A[indices_i]
B = B[:, indices_j]

# Calculate the values
# B is transposed to change the column vectors to row vectors
# such that the two can be multiplied element-wise.
C = torch.sum(A * B.transpose(0, 1), dim=1)




Is it more efficient to selectively calculate the values you want compared to performing the entire matrix multiplication and then only keep the values you want?

The answer is a resounding No. The matrix multiplication is highly optimised, much more optimised than doing the steps manually with operations that themselves are well optimised. Especially, when E contains mostly 1s, then you're basically re-implementing the matrix multiplication, which is guaranteed to be less efficient. Even for the case where E contains mostly 0s, the matrix multiplication is just faster.

To support my claims, I've timed them. For convenience I did it in IPython, which has the built-in %timeit command.

In [1]: import torch
   ...:
   ...:
   ...: def masked(A, B, E):
   ...:     C = A.mm(B)
   ...:     return C[E]
   ...:
   ...:
   ...: def selective(A, B, E):
   ...:     indices_i, indices_j = E.nonzero().unbind(dim=1)
   ...:     return torch.sum(A[indices_i] * B[:, indices_j].transpose(0, 1), dim=1)
   ...:
   ...:
   ...: A = torch.rand(1200, 1000)
   ...: B = torch.rand(1000, 1100)
   ...: # Only 10% of the elements are 1
   ...: E_mostly_zeros = torch.rand(1200, 1100) &lt; 0.1
   ...: # 90% of the elements are 1
   ...: E_mostly_ones = torch.rand(1200, 1100) &lt; 0.9

In [2]: # All close instead of equal to account for floating point errors
   ...: torch.allclose(masked(A, B, E_mostly_ones), selective(A, B, E_mostly_ones))
Out[2]: True

In [3]: # All close instead of equal to account for floating point errors
   ...: torch.allclose(masked(A, B, E_mostly_zeros), selective(A, B, E_mostly_zeros))
Out[3]: True

In [4]: %timeit masked(A, B, E_mostly_ones)
8.16 ms ± 20.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [5]: %timeit selective(A, B, E_mostly_ones)
2.09 s ± 11.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [6]: %timeit masked(A, B, E_mostly_zeros)
5.73 ms ± 24.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

] In [7]: %timeit selective(A, B, E_mostly_zeros)
266 ms ± 3.36 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)


The matrix multiplication is staggeringly fast, being over 256x faster when E contains 90% ones (8.16ms vs 2090ms), and over 46x faster when E contains only 10% ones (5.73ms vs 266ms).
