Speed difference in torch.einsum and torch.bmm when adding an axis

username: hugoych
cooked: <p>Hi,<br>
When using self-attention, I found it’s common usage to use <code>torch.einsum</code> such as follows:</p>
<pre><code class="lang-auto">queries  =  torch.normal(0, 1, (b, h, q, d)).to('cuda') 
keys  =  torch.normal(0, 1, (b, h, q, d)).to('cuda')  # Because self-attention k == q
pre_softmax = torch.einsum("bhqd,bhkd-&gt;bhqk", queries, keys)
</code></pre>
<p>If my understanding is correct in that full self-attention example, we perform <code>b*h*q*q</code> operations.</p>
<p>If one wants to reduce the numbers of keys <strong>for each query</strong> to a fixed value <code>k &lt;&lt; q</code>, then theoretically the number of operations reduces to <code>b*h*q*k</code> which should therefore reduce both time and memory complexity. Because one would have <code>k</code> separate keys for each query , the implementation using <code>torch.einsum</code> would look like the following:</p>
<pre><code class="lang-auto">queries  =  torch.normal(0, 1, (b, h, q, d)).to('cuda') 
keys  =  torch.normal(0, 1, (b, h, q, k, d)).to('cuda') 
pre_softmax = torch.einsum("bhqd,bhqkd-&gt;bhqk", queries, keys)
</code></pre>
<p>In practice, memory consumption does decrease by a factor <code>q/k</code> but not time complexity which becomes much larger as k grows.</p>
<p>On the other, if we simulate a scenario where the <code>k</code> keys are <strong>shared</strong> across all the <code>q</code> queries such that the code would look like the following :</p>
<pre><code class="lang-auto">queries  =  torch.normal(0, 1, (b, h, q, d)).to('cuda') 
keys  =  torch.normal(0, 1, (b, h, k, d)).to('cuda') 
pre_softmax = torch.einsum("bhqd,bhkd-&gt;bhqk", queries, keys)
</code></pre>
<p>We then obtain both linear improvement in time and memory as expected.</p>
<p>If we have a 1000 queries I plotted the following performance depending on <code>k</code>:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/6/360d084f6c4a715108c1018d237c32fbed62d34a.png" data-download-href="https://discuss.pytorch.org/uploads/default/360d084f6c4a715108c1018d237c32fbed62d34a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/6/360d084f6c4a715108c1018d237c32fbed62d34a_2_517x344.png" alt="image" data-base62-sha1="7I9HJgIyOWvw0tXqrNHmnsPX2Yy" width="517" height="344" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/6/360d084f6c4a715108c1018d237c32fbed62d34a_2_517x344.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/6/360d084f6c4a715108c1018d237c32fbed62d34a_2_775x516.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/3/6/360d084f6c4a715108c1018d237c32fbed62d34a.png 2x" data-dominant-color="F6F7F6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">819×545 32.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I was wondering what explained such difference in the implementation of <code>einsum</code> and I assume <code>torch.bmm</code> when the number of operations is the same?  Is there a known way to get around this issue?</p>
<p>Thanks in advance !</p>
<p>Here is <a href="https://colab.research.google.com/gist/hugoych/1e69ab7388803896b5bd23dec413088e/torch_bmm_with-replicated_axis.ipynb" rel="noopener nofollow ugc">gist</a>  from a colab to replicate plot and with more example using also bmm directly.</p>

username: eqy
cooked: <p>It appears that some variant of this is a known issue <a href="https://github.com/pytorch/pytorch/issues/60295" rel="noopener nofollow ugc">Optimize torch.einsum · Issue #60295 · pytorch/pytorch (github.com)</a> but I am not sure if anyone is actively working on it.</p>

username: hugoych
cooked: <p>Hi, thanks for the answer <a class="mention" href="/u/eqy">@eqy</a>!</p>
<p>I’ve tried to see if the issue was related by running additional experiments on the CPU to compare with the two suggested improvements.</p>
<p>From the following code, I believe we can conclude the issue isn’t related.</p>
<pre><code class="lang-auto">b, h, q, k, d = 8, 1, 1000, 100, 32
test_mat = torch.normal(0, 1, (b, h, q, d))
test_mat_u = test_mat.unsqueeze(-2) # (b, h, q, 1 ,d)
test_mat_q_k = test_mat_u.permute(0,1,3,2,4)[...,:k,:].tile((1, 1, q, 1, 1)) # (b, h, q, k, d)
test_mat_k = test_mat[...,:k,:] # (b, h, k, d)


print('### SEPERATE KEYS ###')

print("\n Timing contract('bhqd,bhqkd-&gt;bhqk', test_mat, test_mat_q_k, backend='torch') : 25.6M FLOPs")
%timeit contract("bhqd,bhqkd-&gt;bhqk", test_mat, test_mat_q_k, backend='torch')

print("\n Timing torch.einsum('bhqd,bhqkd-&gt;bhqk', test_mat, test_mat_q_k)  : 25.6M FLOPs")
%timeit torch.einsum("bhqd,bhqkd-&gt;bhqk", test_mat, test_mat_q_k)

print("\n Timing np.einsum('bhqd,bhqkd-&gt;bhqk',  test_mat.numpy(), test_mat_q_k.numpy())  : 25.6M FLOPs")
%timeit np.einsum("bhqd,bhqkd-&gt;bhqk", test_mat.numpy(), test_mat_q_k.numpy(), optimize='optimal')

print("\n Timing torch.matmul(test_mat_u, test_mat_q_k.transpose(-2,-1)).view(b,h,q,k) : 25.6M FLOPs")
%timeit torch.matmul(test_mat_u, test_mat_q_k.transpose(-2,-1)).view(b,h,q,k)


print('\n ### SHARED KEYS ###')

print("\n Timing contract('bhqd,bhkd-&gt;bhqk', test_mat, test_mat_k, backend='torch') : 25.6M FLOPs")
%timeit contract("bhqd,bhkd-&gt;bhqk", test_mat, test_mat_k, backend='torch')

print("\n Timing torch.einsum('bhqd,bhkd-&gt;bhqk', test_mat, test_mat_k)  : 25.6M FLOPs")
%timeit torch.einsum("bhqd,bhkd-&gt;bhqk", test_mat, test_mat_k)

print("\n Timing np.einsum('bhqd,bhkd-&gt;bhqk',  test_mat.numpy(), test_mat_k.numpy())  : 25.6M FLOPs")
%timeit np.einsum("bhqd,bhkd-&gt;bhqk", test_mat.numpy(), test_mat_k.numpy(), optimize='optimal')

print("\n  Timing torch.matmul(test_mat, test_mat_k.transpose(-2,-1)).view(b,h,q,k) : 25.6M FLOPs")
%timeit torch.matmul(test_mat, test_mat_k.transpose(-2,-1)).view(b,h,q,k)


print('\n ### ALL KEYS ###')

print("\n Timing torch.einsum('bhqd,bhkd-&gt;bhqk', test_mat, test_mat)  : 256M FLOPs")
%timeit torch.einsum("bhqd,bhkd-&gt;bhqk", test_mat, test_mat)
</code></pre>
<p>It returns:</p>
<pre><code class="lang-auto">### SEPERATE KEYS ###

 Timing contract('bhqd,bhqkd-&gt;bhqk', test_mat, test_mat_q_k, backend='torch') : 25.6M FLOPs
100 loops, best of 5: 16.4 ms per loop

 Timing torch.einsum('bhqd,bhqkd-&gt;bhqk', test_mat, test_mat_q_k)  : 25.6M FLOPs
100 loops, best of 5: 12.5 ms per loop

 Timing np.einsum('bhqd,bhqkd-&gt;bhqk',  test_mat.numpy(), test_mat_q_k.numpy())  : 25.6M FLOPs
100 loops, best of 5: 11.8 ms per loop

 Timing torch.matmul(test_mat_u, test_mat_q_k.transpose(-2,-1)).view(b,h,q,k) : 25.6M FLOPs
100 loops, best of 5: 12.4 ms per loop

 ### SHARED KEYS ###

 Timing contract('bhqd,bhkd-&gt;bhqk', test_mat, test_mat_k, backend='torch') : 25.6M FLOPs
1000 loops, best of 5: 1.34 ms per loop

 Timing torch.einsum('bhqd,bhkd-&gt;bhqk', test_mat, test_mat_k)  : 25.6M FLOPs
1000 loops, best of 5: 1.27 ms per loop

 Timing np.einsum('bhqd,bhkd-&gt;bhqk',  test_mat.numpy(), test_mat_k.numpy())  : 25.6M FLOPs
100 loops, best of 5: 13 ms per loop

  Timing torch.matmul(test_mat, test_mat_k.transpose(-2,-1)).view(b,h,q,k) : 25.6M FLOPs
1000 loops, best of 5: 1.27 ms per loop

 ### ALL KEYS ###

 Timing torch.einsum('bhqd,bhkd-&gt;bhqk', test_mat, test_mat)  : 256M FLOPs
100 loops, best of 5: 10.6 ms per loop

</code></pre>
<p>For a better contraction path, <code>opt_einsum</code> now supports torch. However, as shown above, using <code>opt_einsum.contract</code> doesn’t improve over <code>torch.einsum</code>.</p>
<p>For the second suggestion on discontiguous cases, I directly compared to <code>np.einsum</code> as it’s supposed t be handled in <code>numpy</code>. Unfortunately, similarly <code>np.einsum</code> doesn’t handle the additional dimension better. Moreover, <code>np.einsum</code> seems to implement the “shared” cases as the separate resulting in a similar run-time ( see above ).</p>
<p>If those two suggestions from the issue you gave seem to be a possible improvement for <code>torch.einsum</code>, this example seems to fall under neither specific case. Moreover, the issue seems to come from <code>torch.matmul</code> itself as it does show the same pattern in runtime, even though the FLOPs are similar between “shared” and “separate” keys cases.</p>