Masked Matrix multiplication

user_9926472 0
Viewed 743

I have a A = N x M matrix and another array B = N x P x M where P is typically 9 or 15. For each vector a from A, it has to be multiplied with each pi from B of the same row, to get an output of dimensions N x P.

I am using numpy and Python and would be performing this operation on a GPU.

For a tiny example, let N=4, M=5, P=3.

Let A be:

array([[0.18503431, 0.2628188 , 0.26343728, 0.8356702 , 0.47581551],
       [0.70827725, 0.04006919, 0.58975722, 0.90874113, 0.43946412],
       [0.40669507, 0.63328008, 0.95832881, 0.59041436, 0.63578578],
       [0.12129919, 0.74470057, 0.62271405, 0.97760796, 0.6499647 ]])

Let B be:

array([[[4.29031165e-01, 6.17324572e-01, 6.54726975e-02, 1.72218768e-02, 3.53970827e-01],
        [3.38821841e-01, 3.80128792e-01, 7.70995505e-01, 7.38437494e-03, 5.87395036e-02],
        [4.75661932e-01, 3.75617802e-01, 1.28564731e-01, 3.66302247e-01, 6.70953890e-01]],

       [[8.96228996e-02, 1.67135584e-02, 4.56921778e-01, 8.25731354e-01, 7.66242539e-01],
        [5.16651815e-01, 4.27179773e-01, 9.34673912e-01, 2.04687170e-01, 7.68417953e-01],
        [5.90980849e-01, 5.03013376e-01, 8.41765736e-02, 8.08221224e-01, 7.76765422e-01]],

       [[3.25802668e-01, 8.58148960e-01, 9.47505735e-01, 1.01405305e-01, 8.34114717e-01],
        [1.65308159e-01, 9.74572631e-01, 2.69886016e-01, 7.44036253e-02, 4.73350521e-01],
        [8.59030672e-01, 3.96972621e-01, 7.34687493e-01, 2.84647032e-02, 7.19723378e-01]],

       [[1.35751242e-01, 1.74882898e-01, 5.48875709e-01, 7.33443675e-01, 4.05282650e-01],
        [8.41298770e-01, 6.24323279e-01, 5.83482185e-01, 4.28514313e-01, 1.96797205e-01],
        [7.93345700e-04, 3.01441721e-01, 7.59451146e-01, 9.09102382e-01, 7.11518948e-01]]])

This is how I want my output to be:

[[np.dot(a[0], b[0][0]), np.dot(a[0], b[0][1]), np.dot(a[0], b[0][2])],
[np.dot(a[1], b[1][0]), np.dot(a[1], b[1][1]), np.dot(a[1], b[1][2])],
[np.dot(a[2], b[2][0]), np.dot(a[2], b[2][1]), np.dot(a[2], b[2][2])],
[np.dot(a[3], b[3][0]), np.dot(a[3], b[3][1]), np.dot(a[3], b[3][2])]]

Doing this manually gives:

[[0.44169455751462816, 0.3998276862221848, 0.845960080871557],
[1.4207326179275017, 1.4579799277670968, 1.564201768913105],
[2.174162453912622, 1.287925491552765, 1.779226448174152],
[1.4689343122491012, 1.4771555510001255, 2.0487088726424365]]

Since I want to do this on a GPU, this obviously calls for converting my problem into matrix multiplication (this is true if I dont use a GPU as well for that matter). But I do not exactly know how to convert it to that.

One idea that I had was to reshape B to be Q x M where Q=NxP. And then perform some sort of sparse multiplication, where for every row i of a boolean sparse matrix, I turn on (0:P) + P*ith elements. (Drawing it out makes sense), however I certainly feel that there is a much more elegant way to do this as creating sparse matrices and performing the operations might take time, and that the sparsity of my matrix is not random at all.

How can I solve this elegantly.

Note that I cannot do some operations such as broadcasting/repeating my A matrix P times and performing the huge matrix multiplication and picking out relevant values, since typically N and M will be quite huge (2000ish and 256 respectively), but P will be quite small, hence doing a global matrix multiplication for all vectors means I will be doing >95% unnecessary computations!.

pytorch matrix numpy python
2 Answers

You can use einsum here to compute this efficiently.

np.einsum('ij,ikj->ik', A, B) # or torch.einsum

>>> np.allclose(np.einsum('ij,ikj->ik', A, B), manual)
True
user_34832030 
(np.expand_dims(A, 1) * B).sum(axis=2)

sort of a "manual" matrix multiplication. First, A is reshaped to N x 1 x M. Then, it is element-wise multiplied by B. Numpy will broadcast A shape to match B, so at the end we'll get N x P x M. each row will contain A[i] * B[i], all that remains to do is to sum along the last axis to get a dot product.

user_2716410
Related