Batch matrix multiplication of 3D tensors

shon711: <p>Hello.</p>
<p>I have series of matrix multiplication in a for loop structure, I want to transform it to one “big” matrix to do all the multiplication together to better utilize the GPU.</p>
<p>Here is the current implementation:</p>
<p>The model input <code>x, y</code> in shape of <code>[batch_size, k, config.hidden_size]</code>.<br>
For each category id <code>[0, 1, 2, 3]</code> we compute:</p>
<ol>
<li>non-linear reps for <code>x</code> and <code>y</code> → <code>cat_x_reps</code>, <code>cat_y_reps</code>
</li>
<li>multiple <code>cat_x_reps</code>, <code>cat_y_reps</code> with a bilinear matrix<br>
<code>cat_x_reps</code> x <code>C_i</code> x <code>cat_y_reps</code>
</li>
</ol>
<pre><code class="lang-auto">import torch
from torch import nn
from transformers.activations import ACT2FN

class FullyConnectedLayer(Module):
    def __init__(self, config, input_dim, output_dim, dropout_prob):
        super(FullyConnectedLayer, self).__init__()

        self.input_dim = input_dim
        self.output_dim = output_dim
        self.dropout_prob = dropout_prob

        self.dense = Linear(self.input_dim, self.output_dim)
        self.layer_norm = LayerNorm(self.output_dim, eps=config.layer_norm_eps)
        self.activation_func = ACT2FN[config.hidden_act]
        self.dropout = Dropout(self.dropout_prob)

    def forward(self, inputs):
        temp = inputs
        temp = self.dense(temp)
        temp = self.activation_func(temp)
        temp = self.layer_norm(temp)
        temp = self.dropout(temp)
        return temp


class MyClass(nn.Module):
    def __init__(self, config, args):
        super().__init__(config)
     
        self.ffnn_size = args.ffnn_size
        self.num_heads = args.num_heads
        
        self.x_cat_mlp = nn.ModuleList([FullyConnectedLayer(config, config.hidden_size, self.ffnn_size, args.dropout_prob)
                                        for _ in range(self.num_heads)])
        self.y_cat_mlp = nn.ModuleList([FullyConnectedLayer(config, config.hidden_size, self.ffnn_size, args.dropout_prob)
                                        for _ in range(self.num_heads)])
        
        self.classifiers = nn.ModuleList([Linear(self.ffnn_size, self.ffnn_size) 
                                          for _ in range(self.num_heads)])

def forward(self, x_reps, y_reps):
         # x,  y -&gt;  [batch_size, k, config.hidden_size]
        
        cat_logit_list = []
        for cat_id in range(self.num_heads):
            cat_x_reps = self.x_cat_mlp[cat_id](x)
            cat_y_reps = self.y_cat_mlp[cat_id](y)

            temp = self.classifiers[cat_id](cat_x_reps)                                    # [batch_size, k, ffnn_size]
            cat_logits = torch.matmul(temp, cat_y_reps.permute([0, 2, 1])).  # [batch_size, k, k]
            cat_logit_list.append(cat_logits)

        return torch.stack(cat_logit_list, dim=1)
</code></pre>
<p>I was able to do it for the attribute <code>self.x_cat_mlp</code> because all the operations on the same input <code>(x, y)</code>.</p>
<p>I defined:</p>
<pre><code class="lang-auto">self.all_cats_size = self.ffnn_size * self.num_cats
self.x_cat_mlp = FullyConnectedLayer(config, config.hidden_size, self.all_cats_size, args.dropout_prob)
x_reps = self.x_cat_mlp(x)                                                          # [batch, k, ffnn * num_cats]
cat_x_reps = x_reps.view((batch_size, self.num_cats, k, self.ffnn_size))            # [batch, num_cats, k, ffnn]
</code></pre>
<p>I am struggling to do it for <code>self.classifiers</code>. the input is different each time, i.e <code>cat_x_reps</code> is different for each category</p>

KFrank: <p>Hi Shon!</p>
<aside class="quote no-group quote-modified" data-username="shon711" data-post="1" data-topic="153644" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/eada6e/48.png" class="avatar"> shon711:</div>
<blockquote>
<p>I want to transform it to one “big” matrix to do all the multiplication together<br>
…<br>
2. multiple <code>cat_x_reps</code>, <code>cat_y_reps</code> with a bilinear matrix<br>
<code>cat_x_reps</code> x <code>C_i</code> x <code>cat_y_reps</code></p>
<pre><code class="lang-auto">...
            temp = self.classifiers[cat_id](cat_x_reps)                                    # [batch_size, k, ffnn_size]
            cat_logits = torch.matmul(temp, cat_y_reps.permute([0, 2, 1])).  # [batch_size, k, k]
            cat_logit_list.append(cat_logits)
</code></pre>
</blockquote>
</aside>
<p>My interpretation of the above is that:</p>
<pre data-code-wrap="python"><code class="lang-python">            temp = self.classifiers[cat_id](cat_x_reps)
</code></pre>
<p>is the application of a <code>Linear</code> to <code>cat_x_reps</code>, and corresponds to<br>
<code>C_i</code> x <code>cat_y_reps</code> in your summary.</p>
<p>Note that applying a <code>Linear</code> is not just a matrix multiplication, but<br>
also entails adding the <code>bias</code> term:</p>
<pre data-code-wrap="python"><code class="lang-python">linear (t) == t @ linear.weight.T + linear.bias
</code></pre>
<p>So describing <code>C_i</code> x <code>cat_y_reps</code> as just matrix multiplication is an<br>
over-simplification.</p>
<p>Leaving the issue of the <code>bias</code> term aside for the moment, you can<br>
<code>stack()</code> the <code>weight</code> tensors of your <code>classifiers</code> (along <code>dim = 1</code>, if<br>
your choose), <code>stack()</code> the <code>cat_x_reps</code> and <code>stack()</code> the <code>cat_y_reps</code>,<br>
and then use <code>torch.einsum()</code> to compute all of the three-term products,<br>
<code>cat_x_reps</code> x <code>C_i</code> x <code>cat_y_reps</code>, for all values of <code>i</code> (which I take to be<br>
<code>cat_id</code>) all at once in one “batch.”</p>
<p>Now for the <code>bias</code> term:</p>
<p><code>cat_x_reps</code> x (<code>C_i</code> x <code>cat_y_reps</code> + <code>bias_i</code>) is equal to<br>
<code>cat_x_reps</code> x <code>C_i</code> x <code>cat_y_reps</code> + <code>cat_x_reps</code> x <code>bias_i</code> so you can<br>
compute the three-term product as above, and then <code>stack()</code> the <code>bias_i</code>,<br>
compute <code>cat_x_reps</code> x <code>bias_i</code> for all <code>i</code> all at once (using, for example,<br>
<code>torch.einsum()</code>), and then add it to the three-term product to get the<br>
final result.</p>
<p>Best.</p>
<p>K. Frank</p>

shon711: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>,<br>
Thank you for your response!</p>
<p>Just to verify:<br>
After stacking, the shapes are:</p>
<pre><code class="lang-auto">all_X  [btach_size, num_cats, k,     ffnn]
all_C  [btach_size, num_cats, ffnn,  ffnn]
all_Y  [btach_size, num_cats, k,     ffnn]
</code></pre>
<p>Can you elaborate how to use <code>torch.einsum()</code> for these shapes? we need to have:</p>
<pre><code class="lang-auto">for i in range(num_cat):
    x = all_X[:, i, :, :]                               # [btach_size, k,     ffnn]
    c = all_C[:, i, :, :]                               # [btach_size, ffnn,  ffnn]
    y = all_Y[:, i, :, :]                               # [btach_size, k,     ffnn]

    temp = torch.matmul(x, c)                           # [batch_size, k, ffnn_size]
    logits = torch.matmul(temp, y.permute([0, 2, 1])).  # [batch_size, k, k]
</code></pre>
<p>If we will do<br>
<code>all_X</code> x <code>all_C</code> x <code>all_Y</code><br>
it include all the different combinations, for example: <code>X_i</code> x <code>C_i</code> x <code>Y_j</code> which I don’t need to perform.</p>

KFrank: <p>Hi Shon!</p>
<aside class="quote no-group" data-username="shon711" data-post="3" data-topic="153644" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/eada6e/48.png" class="avatar"> shon711:</div>
<blockquote>
<pre><code class="lang-auto">all_X  [btach_size, num_cats, k,     ffnn]
all_C  [btach_size, num_cats, ffnn,  ffnn]
all_Y  [btach_size, num_cats, k,     ffnn]
</code></pre>
</blockquote>
</aside>
<p>It’s not entirely clear to me what the values <code>num_cats</code> and <code>k</code> are.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">for i in range(num_cat):
    x = all_X[:, i, :, :]                               # [btach_size, k,     ffnn]
    c = all_C[:, i, :, :]                               # [btach_size, ffnn,  ffnn]
    y = all_Y[:, i, :, :]                               # [btach_size, k,     ffnn]

    temp = torch.matmul(x, c)                           # [batch_size, k, ffnn_size]
    logits = torch.matmul(temp, y.permute([0, 2, 1])).  # [batch_size, k, k]
</code></pre>
</blockquote>
</aside>
<p>Let me assume that you want the shape of your final “all the multiplication<br>
together” result to be <code>[batch_size, num_cats, k, k]</code>.</p>
<p>You can compute your three-term product with <code>einsum()</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">result = torch.einsum ('bnkf, bnfg, bnlg -&gt; bnkl', all_X, all_C, all_Y)
</code></pre>
<p>Make sure you understand how <code>einsum()</code> works so that you are<br>
<em>contracting</em> (“multiplying” together) the desired pairs of indices.</p>
<p>If <code>all_C</code> is constructed from the <code>weight</code>s of <code>Linear</code>s, you may need<br>
to transpose the last two dimensions (the <code>ffnn</code> dimensions) of <code>all_C</code>,<br>
either explicitly or by swapping the symbolic indices in the <code>einsum()</code><br>
expression, in order to replicate the result you would get by applying<br>
the <code>Linear</code> to <code>all_Y</code>.</p>
<p>(This again ignores any <code>bias</code> terms that may have appeared in your<br>
<code>Linear</code>s.)</p>
<p>Best.</p>
<p>K. Frank</p>

shon711: <p><a class="mention" href="/u/kfrank">@KFrank</a><br>
Thanks ! this is working, WOW <code>einsum</code> such a powerful method !</p>
<p><code>k</code> is the sequence length.<br>
<code>num_cats</code> is the number of “learning” matrices we have.</p>
<p>You right, I want <code>[batch_size, num_cats, k, k]</code></p>
<p>I took your note about the weights’s dim swap.<br>
In addition, <code>all_C</code> is the learnable matrices and its shape is <code>[num_cats, ffnn, ffnn]</code></p>
<p>I am a bit struggling to add the bias term as you suggested:<br>
after calc the three-term product I got a tensor of shape <code>logits = [batch_size, num_cats, k, k]</code></p>
<p>now I need to do</p>
<blockquote>
<p><code>cat_x_reps</code> x (<code>C_i</code> x <code>cat_y_reps</code> + <code>bias_i</code> ) is equal to<br>
<code>cat_x_reps</code> x <code>C_i</code> x <code>cat_y_reps</code> + <code>cat_x_reps</code> x <code>bias_i</code> so you can<br>
compute the three-term product as above, and then <code>stack()</code> the <code>bias_i</code> ,<br>
compute <code>cat_x_reps</code> x <code>bias_i</code> for all <code>i</code> all at once (using, for example,<br>
<code>torch.einsum()</code> ), and then add it to the three-term product to get the<br>
final result.</p>
</blockquote>
<p>So I stacked all the bias term → <code>all_bias = [num_cats, ffnn]</code><br>
And sum:<br>
<code>logits + torch.einsum('bnkf, nf -&gt; bnk', all_x, all_bias).unsqueeze(-1)</code></p>
<p>but this isn’t reproduce the same performance (not so far from the for loop method but still behind).</p>
<p>BTW, I ran the for loop method with <code>bias=False</code> when initizle the Linear attribute, then calc the performance(loss).<br>
Then I did the same only with the 3-term product and I got exact the same loss.<br>
So I only struggling with the bias. any idea why?</p>
<p><a class="mention" href="/u/kfrank">@KFrank</a> Thanks in advance you helped me a lot!!</p>

KFrank: <p>Hi Shon!</p>
<aside class="quote no-group" data-username="shon711" data-post="5" data-topic="153644" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/eada6e/48.png" class="avatar"> shon711:</div>
<blockquote>
<p>but this isn’t reproduce the same performance (not so far from the for loop method but still behind).</p>
<p>BTW, I ran the for loop method with <code>bias=False</code> when initizle the Linear attribute, then calc the performance(loss).<br>
Then I did the same only with the 3-term product and I got exact the same loss.</p>
</blockquote>
</aside>
<p>I’m not sure what you’re asking.</p>
<p>Is your issue that <code>einsum()</code> runs more slowly than the loop (but gives the<br>
same answer)?</p>
<p>Or is <code>einsum()</code> giving you different values than the loop does?</p>
<p>If the latter, check that any differences aren’t due to expected floating-point<br>
round-off error.  You might try comparing the results with <a href="https://pytorch.org/docs/1.11/generated/torch.allclose.html" rel="noopener nofollow ugc">torch.allclose()</a>, or<br>
performing the computations with double-precision to see if that reduces the<br>
discrepancies.</p>
<p>Best.</p>
<p>K. Frank</p>

shon711: <p>I found the issue,<br>
The linear transformation with the for loop performed on <code>X</code> and the result multiplied by <code>Y</code>.</p>
<p>So instead of<br>
<code>cat_x_reps</code> x <code>C_i</code> x <code>cat_y_reps</code> + <code>cat_x_reps</code> x <code>bias</code><br>
I need to do:<br>
<code>cat_x_reps</code> x <code>C_i</code> x <code>cat_y_reps</code> + <code>bias</code> x <code>cat_y_reps</code></p>
<p>Now I can reproduce the same loss score.</p>
<p>Again, thank you very 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>