Efficient batch dot product

username: smonsays
cooked: <p>I have given a batch of row vectors stored in the matrix <code>U</code>, a batch of column vectors stored in the matrix <code>V</code> and a single matrix <code>M</code>. For each row vector <code>u in U</code> and each column vector <code>v in V</code> I want to compute the sum of the matrix product <code>u *M*v</code> for each batch.</p>
<p>How can I efficiently implement this (potentially using <a href="https://pytorch.org/docs/stable/torch.html#torch.bmm" rel="nofollow noopener">bmm()</a>, <a href="https://pytorch.org/docs/stable/torch.html#torch.matmul" rel="nofollow noopener">matmul()</a> or maybe even <a href="https://pytorch.org/docs/stable/torch.html#torch.einsum" rel="nofollow noopener">einsum</a>)?</p>
<p>Here is a small toy example doing what I want to do with a for loop:</p>
<pre><code class="lang-python">import torch

U = torch.arange(1,10).reshape(3,3)
V = torch.arange(1,10).reshape(3,3)
M = torch.tensor([1, 2, 3]).repeat(3,1)

result = 0
for u,v in zip(U.t(), V):
    result += torch.matmul(torch.matmul(u,V),v)
</code></pre>
<pre><code class="lang-python">result:
tensor(4545)
</code></pre>
<p>I know there is <code>torch.bmm()</code> to perform batch matrix matrix multiplication. If there was something similar for a batch vector dot product (e.g. <code>torch.bvv()</code>) I could do <code>bvv(matmul(U,M),V)</code> .</p>

username: InnovArul
cooked: <p>Isn’t it nothing but twice <code>bmm()</code>?</p>
<pre><code class="lang-auto">import torch
U = torch.arange(1,10).reshape(1, 3,3) 
V = torch.arange(1,10).reshape(1, 3,3) 
M = torch.tensor([1, 2, 3]).repeat(3,1).view(1,3,3) 
result = torch.bmm(U, M).bmm(V)
</code></pre>
<pre><code class="lang-auto">result: 
tensor([[[ 180,  216,  252],
         [ 450,  540,  630],
         [ 720,  864, 1008]]])
</code></pre>

username: smonsays
cooked: <aside class="quote no-group" data-username="InnovArul" data-post="2" data-topic="41382">
<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/innovarul/48/5282_2.png" class="avatar"> InnovArul:</div>
<blockquote>
<p>Isn’t it nothing but twice <code>bmm()</code> ?</p>
</blockquote>
</aside>
<p>I don’t think that is correct. Note that the <code>result</code> should be a scalar.</p>

username: colesbury
cooked: <pre><code class="lang-python">((U.t() @ V) * V).sum()
</code></pre>
<p>(EDIT: The <code>@</code> operator is equivalent to torch.matmul)</p>

username: InnovArul
cooked: <p>In that case, you can call <code>.sum()</code> on <code>result</code> to get a scalar.  Basically, every element of <code>result</code> matrix contains the multiplication result of a combination of <code>u, M, v</code>.<br>
Maybe, you can work out the math and the <code>result.sum()</code> to see if it is correct.</p>

username: tom
cooked: <p><code>torch.einsum('bi,ij,bj', U, M, V)</code> if you want the sum, <code>'bi,ij,bj-&gt;b'</code> if you prefer the batch items separately. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Best regards</p>
<p>Thomas</p>

username: smonsays
cooked: <aside class="quote no-group" data-username="tom" data-post="6" data-topic="41382">
<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/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p><code>torch.einsum('bi,ij,bj', U, M, V)</code> if you want the sum, <code>'bi,ij,bj-&gt;b'</code> if you prefer the batch items separately. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Best regards</p>
<p>Thomas</p>
</blockquote>
</aside>
<p>Perfect, thanks! Einsum is really neat, I took the time to get familiar with it and came up with the same result.</p>

username: zeligism
cooked: <p>I faced a similar problem and noticed that a faster way of doing <code>torch.einsum('bi,ij,bj-&gt;b', U, M, V)</code> is <code>torch.sum(U @ M * V, dim=1)</code>.</p>

username: tom
cooked: <p>For this particular case, the specialized <code>bilinear</code> function might be a good thing to use.</p>

username: zeligism
cooked: <p>Thanks for pointing out the <code>bilinear</code> function, it’s very useful and I wasn’t aware of it.</p>
<p>Though, in my case, I had to do a bit of squeeze/unsqueeze-ing to get it to work. The dimensions are exactly as in my einsum example, namely (B,M), (B,N), and (M,N) for U, V, and W, respectively. Calling <code>bilinear(U,V,W)</code> in this case requires W to have a dim of (Y,M,N), where Y is the number of out-features, so I had to call it like this <code>bilinear(U,V,W.unsqueeze(0)).squeeze()</code> to make it equivalent to <code>torch.sum(U @ W * V, dim=1)</code>.</p>
<p>In addition to that, <code>torch.sum(U @ W * V, dim=1)</code> seems to be about ~1.5 times faster than <code>bilinear(U,V,W.unsqueeze(0)).squeeze()</code> in my case, though I’m not sure why.</p>

username: Gao_Haibin
cooked: <p>(U * V).sum(axis=-1)</p>