Efficient (restricted) matrix multiplication implementation using indicator matrix

username: agentdr
cooked: <p>Hi all,</p>
<p>I have a dot product which I aim to optimise by using some a priori information.<br>
I have two matrices, i.e. Q (n x d) and K (d x n), and since n can be very large this results in quadratic complexity. But I also have a symmetric indicator matrix A (n x n) where each row has knn ones and n-knn zeros (knn &lt; n, always). I use this matrix A to select the relevant columns in K for each row in Q. So now, my question is, how can I efficiently implement the outlined code part below? And, how can I store A as efficient as possible in terms of memory?</p>
<pre><code class="lang-auto">import torch
import time
import random
from sklearn.neighbors import kneighbors_graph

n = 1024
d = 512
knn = 64
q = torch.randn(n, d)
k = torch.randn(d, n)
coords = torch.tensor([(random.random(), random.random()) for _ in range(n)])
A = torch.tensor(kneighbors_graph(coords, knn, mode=‘connectivity’, include_self=False).todense()).bool()

t1_start = time.time()

#start optimise; goal: this should be more efficient than full dot product

outs = [q[i].matmul(k[:,A[i]]) for i in range(n)]
outs = torch.stack(outs)

#end optimise

t1_stop = time.time()
time_need = t1_stop - t1_start
print(time_need) # 0.1741


# this is the full dot product 
t1_start = time.time()
outs=q.matmul(k)
t1_stop = time.time()
time_need = t1_stop - t1_start
print(time_need) #0.00365</code></pre>

username: KFrank
cooked: <p>Hi Agent!</p>
<aside class="quote no-group" data-username="agentdr" data-post="1" data-topic="144525" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/agentdr/48/46828_2.png" class="avatar"> agentdr:</div>
<blockquote>
<p>how can I efficiently implement the outlined code part below?</p>
</blockquote>
</aside>
<p>You can package the set of dot products into a single <a href="https://pytorch.org/docs/1.10.0/generated/torch.einsum.html" rel="noopener nofollow ugc">einsum()</a><br>
operation (“Einstein summation,” a generalized matrix multiplication),<br>
thereby removing the loop (expressed as a python list comprehension),<br>
by “expanding” the matrices, at the cost of materializing a <code>d x n**2</code><br>
expanded version of <code>k</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.10.2'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; n = 7
&gt;&gt;&gt; d = 5
&gt;&gt;&gt; r = 3   # knn
&gt;&gt;&gt;
&gt;&gt;&gt; q = torch.randn (n, d)
&gt;&gt;&gt; k = torch.randn (d, n)
&gt;&gt;&gt;
&gt;&gt;&gt; # build (non-symmetric) boolean indicator matrix
...
&gt;&gt;&gt; A = torch.randn (n, n).argsort (dim = 1) &lt; r
&gt;&gt;&gt;
&gt;&gt;&gt; torch.all (A.sum (dim = 1) == r)   # verify correct number of "True"s in rows of A
tensor(True)
&gt;&gt;&gt;
&gt;&gt;&gt; # version with loop
...
&gt;&gt;&gt; outs = [q[i].matmul(k[:,A[i]]) for i in range(n)]
&gt;&gt;&gt; outs = torch.stack(outs)
&gt;&gt;&gt;
&gt;&gt;&gt; # loop-free version, but with d x n**2 "expanded" matrix
...
&gt;&gt;&gt; qex = q.unsqueeze (1).expand (n, r, d).reshape (r * n, d)
&gt;&gt;&gt; kex = k.unsqueeze (1).expand (d, n, n).reshape (d, n * n)[:, A.view (n * n)]
&gt;&gt;&gt; outsB = torch.einsum ('ij, ji -&gt; i', qex, kex).view (n, r)
&gt;&gt;&gt;
&gt;&gt;&gt; torch.allclose (outs, outsB)   # cross-check results
True
&gt;&gt;&gt;
&gt;&gt;&gt; # shapes of "expanded" matrices and result
...
&gt;&gt;&gt; qex.shape
torch.Size([21, 5])
&gt;&gt;&gt; kex.shape
torch.Size([5, 21])
&gt;&gt;&gt;
&gt;&gt;&gt; outsB
tensor([[-0.5084,  3.1422,  0.9380],
        [ 0.0939,  2.3767,  0.5234],
        [-3.2668, -0.2849, -1.3499],
        [-2.0864,  1.1269,  0.9776],
        [ 0.1827,  0.2465, -1.3500],
        [ 3.4298,  2.5439,  0.3097],
        [ 0.6107, -1.6878, -2.8169]])
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>And, how can I store A as efficient as possible in terms of memory?</p>
</blockquote>
</aside>
<p>Assuming that <code>knn</code> is significantly less than <code>n</code>, the information in <code>A</code> will<br>
be stored most efficiently as a list (tensor) of the indices of its <code>True</code><br>
elements.  If you flatten <code>A</code> first in to a vector (of length <code>n**2</code>), you save<br>
space by using only a single index per element (and this structure is<br>
also convenient for the above “expanded” matrix approach):</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; # express indicator matrix more compactly as list of (flattened) indices
...
&gt;&gt;&gt; Aind = torch.where (A.view (n * n))[0]
&gt;&gt;&gt; Aind
tensor([ 2,  3,  5,  9, 12, 13, 16, 19, 20, 21, 26, 27, 29, 31, 33, 35, 37, 41,
        43, 47, 48])
&gt;&gt;&gt; kexB = k.unsqueeze (1).expand (d, n, n).reshape (d, n * n)[:, Aind]   # expanded k
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (kex, kexB)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>