Computing bilinear with block diagonal matrix efficiently

Kore_ana: <p>I want to compute $x^TAx$ where $x$ is a vector and A is a matrix (i.e. bilinear form). Although nn.Bilinear would work, in my case $A$ is block-diagonal. So, I wanted to create a <code>nn.Module</code> that would compute this using much less memory and computation by only remembering the block diagonal matrices.</p>
<p>For example, Below is an example of computing, using <code>nn.Bilinear(x,x,A,bias)</code>  (here x’s batch dim is 3 and output dim is 2</p>
<pre><code class="lang-auto">
x = tensor([[ 0.,  1.,  2.,  3.],
        [ 4.,  5.,  6.,  7.],
        [ 8.,  9., 10., 11.]])
A = tensor([[[ 0.,  1.,  0.,  0.],
         [ 2.,  3.,  0.,  0.],
         [ 0.,  0., -0., -1.],
         [ 0.,  0., -2., -3.]],

        [[ 0.,  2.,  0.,  0.],
         [ 4.,  6.,  0.,  0.],
         [ 0.,  0., -0., -2.],
         [ 0.,  0., -4., -6.]]])
b = tensor([0, 1])
</code></pre>
<p>where,<br>
<code>x.shape, weight.shape, bias.shape torch.Size([3, 4]) torch.Size([2, 4, 4]) torch.Size([2])</code><br>
and the result of <code>nn.Bilinear</code> is</p>
<pre><code class="lang-auto">tensor([[ -42.,  -83.],
        [-138., -275.],
        [-234., -467.]])
</code></pre>
<p>Initially, I tried the following approach : since the quadratic form of block diagonals can be made equal to the sum of smaller chunks of bilinear (i.e. if x = [x1,x2] and A = [[A1,0][0,A2]], then $x^TAx = x_1^TA_1x_1 + x_2^TA_2x_2$), i tried to divide the x into chunks equal to the size of the block matrices, then send reshape it s.t. the batch dimension now is multiplied by the number of block matrices, then computed the bilinear using the increase batch size, but it didn’t work as expected.</p>
<p>Any suggestions or ideas would be greatly appreciated!<br>
(I thought of  running a for loop of <code>torch.matmul</code> over each block diagonal could work, but I am afraid that would not be parallel and hence decrease speed)</p>

KFrank: <p>Hi Kore (aka Danny)!</p>
<aside class="quote no-group quote-modified" data-username="Kore_ana" data-post="1" data-topic="194429" 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/kore_ana/48/45740_2.png" class="avatar"> Kore_ana:</div>
<blockquote>
<p>Although nn.Bilinear would work, in my case $A$ is block-diagonal.<br>
…<br>
the result of <code>nn.Bilinear</code> is</p>
<pre><code class="lang-auto">tensor([[ -42.,  -83.],
        [-138., -275.],
        [-234., -467.]])
</code></pre>
</blockquote>
</aside>
<p>I don’t believe that there is any way (without a loop) to perform your<br>
computation with <code>Bilinear</code> that doesn’t have inefficiencies such as<br>
multiplying things with the zero blocks.</p>
<p>Probably the simplest way to perform the necessary tensor multiplications<br>
is to use <a href="https://pytorch.org/docs/2.1/generated/torch.einsum.html" rel="noopener nofollow ugc">einsum().</a></p>
<p>Let’s assume that you store just the non-zero blocks of <code>A</code> (or pre-process<br>
<code>A</code> to extract its non-zero blocks).  You can then perform the bilinear<br>
“contraction” with a <code>.reshape()</code> of <code>x</code> and <code>einsum()</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__

x = torch.tensor([[ 0., 1.,  2.,  3.],
                  [ 4., 5.,  6.,  7.],
                  [ 8., 9., 10., 11.]])
A = torch.tensor([[[ 0.,  1.,  0.,  0.],
                   [ 2.,  3.,  0.,  0.],
                   [ 0.,  0., -0., -1.],
                   [ 0.,  0., -2., -3.]],
                  [[ 0.,  2.,  0.,  0.],
                   [ 4.,  6.,  0.,  0.],
                   [ 0.,  0., -0., -2.],
                   [ 0.,  0., -4., -6.]]])
b = torch.tensor([0., 1.])

resultA = torch.nn.functional.bilinear (x, x, A, b)

print ('resultA = ...')
print (resultA)

A_block = torch.tensor([[[[ 0.,  1.],      # store just the non-zero blocks of A
                          [ 2.,  3.]],
                         [[-0., -1.],
                          [-2., -3.]]],
                        [[[ 0.,  2.],
                          [ 4.,  6.]],
                         [[-0., -2.],
                          [-4., -6.]]]])

resultB = torch.einsum ('mij, nijk, mik -&gt; mn', x.reshape (-1, 2, 2), A_block, x.reshape (-1, 2, 2)) + b

print ('resultB = ...')
print (resultB)

print ('torch.equal (resultA, resultB) =', torch.equal (resultA, resultB))
</code></pre>
<p>Here is the above script’s output:</p>
<pre><code class="lang-plaintext">resultA = ...
tensor([[ -42.,  -83.],
        [-138., -275.],
        [-234., -467.]])
resultB = ...
tensor([[ -42.,  -83.],
        [-138., -275.],
        [-234., -467.]])
torch.equal (resultA, resultB) = True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Kore_ana: <p>Thank you so much for your answer!</p>
<p>I actually solved it myself, using a similar method, where I added an additional dimension (dimension with # of blocks), then used torch.matmul twice. However, your solution seems to be more general (I think it might be helpful when I need to implement a convolution version of this operation)</p>
<p>If you don’t mind, I have just a few questions</p>
<ol>
<li>Is defining custom operations (i.e. instead of using things like <code>F.Linear</code>) detrimental to speed? Does it affect speed to the point where I need to try out different implementations of the same operation then measure their time and choose the best one?</li>
<li>Are there any other essential operations that I should know? (things like <code>einsum</code>?) It’s my first time making a custom <code>nn.Module</code>, and it seems that knowing which operations to use when is really important to writing code that does not use loops?</li>
<li>If I wanted to expand this operation to a matrix made of different block sizes, would it be efficient, or even possible?</li>
</ol>
<p>Again, thank you so much <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, Danny Han</p>

KFrank: <p>Hi Danny!</p>
<aside class="quote no-group" data-username="Kore_ana" data-post="3" data-topic="194429" 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/kore_ana/48/45740_2.png" class="avatar"> Kore_ana:</div>
<blockquote>
<p>I added an additional dimension (dimension with # of blocks), then used torch.matmul twice.</p>
</blockquote>
</aside>
<p>In general, you can replace <code>einsum()</code> with appropriately chosen<br>
<code>permute()</code>s, <code>matmul()</code>s, and <code>sum()</code>'s.  (This is, roughly speaking,<br>
what <code>einsum()</code> does under the hood.)</p>
<p>But <code>einsum()</code> is really useful and I find it easier to read than various<br>
<code>matmul()</code> implementations for complicated contractions (such as yours).</p>
<p>(Note, it goes the other way, too.  You can replace a <code>matmul()</code> that is<br>
performing a single matrix multiplication with <code>einsum()</code>.  But in that case<br>
I would find <code>matmul()</code> easier to read.)</p>
<aside class="quote no-group">
<blockquote>
<ol>
<li>Is defining custom operations (i.e. instead of using things like <code>F.Linear</code>) detrimental to speed?</li>
</ol>
</blockquote>
</aside>
<p>In general no, if you implement your custom operation with pytorch<br>
tensor operations, avoiding explicit python loops.  Of course, you can<br>
do things in a suboptimal way, but it’s up to you not to do that.</p>
<p>For most realistic applications, the big work is in performing tensor<br>
multiplications of (large) tensors or applying functions element-wise<br>
to (large) tensors.  Pytorch does these things with pytorch tensor<br>
operations that are optimized for gpu (and cpu) floating-point pipelines.</p>
<p>The python glue that you write to chain such operations together is,<br>
in a sense, inefficient, but for most practical problems, the time spent<br>
executing the “inefficient” python is small compared to the real work.</p>
<p>For cases where you use python to chain together a lot of small tensor<br>
operations, you might find that using things like <a href="https://pytorch.org/docs/2.1/generated/torch.compile.html" rel="noopener nofollow ugc">torch.compile</a> speeds<br>
things up dramatically.</p>
<aside class="quote no-group">
<blockquote>
<p>Does it affect speed to the point where I need to try out different implementations of the same operation then measure their time and choose the best one?</p>
</blockquote>
</aside>
<p>Sometimes yes.  If speed is very important to your use case, it can<br>
make sense to time alternative implementations.  Note that the fastest<br>
approach can depend (in unexpected ways) on the size of the tensors,<br>
gpu vs. cpu, single vs. double precision, and one hardware platform<br>
vs. another.</p>
<p>As an example, there are cases discussed in this forum where <code>einsum()</code><br>
significantly underperformed or outperformed an equivalent <code>matmul()</code><br>
implementation.  I’m not sure that this is still the case because where<br>
the discrepancy was significant, it was fair to call it a performance bug<br>
in <code>einsum()</code> or <code>matmul()</code>, so those bugs have likely been fixed.</p>
<p>But just do a good job writing your code the first time around, and don’t<br>
waste your time performance-tuning various pieces of code unless you’ve<br>
identified them as bottlenecks in your overall computation.  If you make<br>
something that accounts for only 1% of your runtime ten times faster,<br>
who cares?</p>
<aside class="quote no-group">
<blockquote>
<ol start="2">
<li>Are there any other essential operations that I should know? (things like <code>einsum</code>?)</li>
</ol>
</blockquote>
</aside>
<p><code>einsum()</code> is probably one of the best examples of a pytorch function<br>
that can be broadly useful, but is underused.</p>
<p><code>expand()</code> is useful when you need to repeat slices of a tensor (say for<br>
operations that don’t support broadcasting).  It gives you an expanded<br>
<em>view</em> of the tensor, rather than materializing the expanded tensor, which<br>
can offer a large performance benefit.</p>
<p>To the Forum:  Help out here!  What are some other underloved pytorch<br>
operations that people should put in their general-purpose toolkit?</p>
<aside class="quote no-group">
<blockquote>
<ol start="3">
<li>If I wanted to expand this operation to a matrix made of different block sizes, would it be efficient, or even possible?</li>
</ol>
</blockquote>
</aside>
<p>This would be harder.  Pytorch does not support “ragged” tensors, that<br>
is, tensors whose slices have differing shapes.  As a simple example,<br>
a “matrix” whose first row has five elements and whose second row has<br>
three elements would be ragged.</p>
<p>One approach is to pad the ragged slices up to some common uniform<br>
size.  (This is conceptually similar to your including zero-blocks in your<br>
full block-diagonal matrix.)</p>
<p>You <em>might</em> be able to package your matrix with blocks of different sizes<br>
as a <a href="https://pytorch.org/docs/2.1/sparse.html" rel="noopener nofollow ugc">sparse tensor,</a> but sparse tensors in pytorch are only partially<br>
supported, so you might not be able to perform all of the steps you want.</p>
<p>So, if worst comes to worst, you may have to write a python loop that<br>
processes the differently-sized blocks one by one (and maybe use<br>
<code>torch.compile</code> to undo some of the python-loop damage).</p>
<p>Best.</p>
<p>K. Frank</p>