Matrix multiplication with values only at masked entries

hjain: <p>I want to multiply two dense matrices A(1000000,300) and B(1000000,300). The resulting matrix will be of size million x million, which offcourse won’t fit into memory and will take forever to compute. However, I want values only at few elements which are specified in a sparse matrix C(1000000,1000000).<br>
Is there an efficient(possibly that utilizes GPUs) way to do this on Pytorch?</p>

munkiti: <p>Did you try just iterating over the required elements in C matrix and take dot product of corresponding row of A and column of B? Of-course i believe there can more efficient solutions, but i think this is the most basic one.</p>
<p>While iterating you can also use multiprocessing to run the computations on various GPU’s for various required elements in C.</p>

hjain: <p>yes but that takes a lot of time. I wanted a vectorized solution that uses inbuilt functions</p>

richard: <p>There’s a way to do this if we restrict your mask M:<br>
Let’s say you have a mask <code>M</code> such that no two 1’s in the mask are in the same row or column.</p>
<p>Then the following should equivalent to <code>(z @ y) * M</code>, where the <code>@</code> sign is matrix multiplication:<br>
<code>(z.t() * (y @ M.t())).to_dense().sum(dim=0).diag() @ M</code></p>
<p>I think pytorch does support sparse x dense -&gt; sparse via <code>torch.mm</code>. If you need a dense x sparse -&gt; sparse (because <code>M</code> will probably be sparse), you can use the identity <code>AB = ( AB )^T ^T = (B^T A^T)^T</code>.</p>
<p>I haven’t actually tried out the expression above for your use case, though.</p>

hjain: <p>Thanks a lot Richard! Unfortunately, the mask I have contains multiple ones in each row. For instance for an M of size million x million there are about 300 ones in each row.<br>
Is there any solution for this?</p>

hjain: <p>Any more suggestions to solve this??</p>

tom: <p>If you have your mask 1 locations as a <code>rows</code> and <code>columns</code> <code>LongTensors</code>, so for example (your matrices are larger, but for demonstration purposes)</p>
<pre><code class="lang-auto">A = torch.randn(5,10)
B = torch.randn(8,10)
rows = torch.arange(0, 16, out=torch.LongTensor())%5
cols = (torch.arange(0, 16, out=torch.LongTensor())/2)%8
idxes = torch.stack([rows,cols], dim=0)
mask = torch.sparse.FloatTensor(idxes, torch.ones(16)).to_dense()
</code></pre>
<p>how about doing</p>
<pre><code class="lang-auto">vals = (A[rows]*B[cols]).sum(1)
res = torch.sparse.FloatTensor(idxes, vals)
</code></pre>
<p>Then you see that the difference to the masked product is 0 up to machine precision:</p>
<pre><code class="lang-auto">print((res.to_dense()-(A@B.t()) * mask).abs().max())
</code></pre>
<p>Now this will have quite a bit of memory overhead, but you could split it into chunks of nonzero entries and cat them together afterwards if needed.</p>
<p>Best regards</p>
<p>Thomas</p>

shenshanf: <p>I implemented what you said using numba cuda. see：<a href="https://github.com/shenshanf/matmul_with_mask/" class="inline-onebox" rel="noopener nofollow ugc">GitHub - shenshanf/matmul_with_mask: matrix multiplication with sparse mask that indicate the operation position</a></p>

shenshanf: <p>see this api function: torch.sparse.sampled_addmm<br>
<a href="https://pytorch.org/docs/stable/generated/torch.sparse.sampled_addmm.html#torch.sparse.sampled_addmm" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.sparse.sampled_addmm.html#torch.sparse.sampled_addmm</a></p>