Extracting differences from Attention Masks and Paddings, in order to get logit differences

Alex_Ge: <p>For better or worse, I have the following scenario for a custom Transformer based model:</p>
<ul>
<li>logits [batch, sequence, features]</li>
<li>ideals [batch, sequence, features]</li>
</ul>
<p>Then I have the attention masks (which I use to test predictability of the model)</p>
<ul>
<li>attn_masks [batch, sequence]</li>
</ul>
<p>And finally I have the padding masks (which are pretty standard, padding what is not part of the sequence)</p>
<ul>
<li>paddings [batch, sequence]</li>
</ul>
<p>I’ve managed to write a very ugly, slow, monstrosity which does two things:</p>
<ol>
<li>it calculates the R2 score by looking at the values from the start of the attention mask, all the way to the end of the sequence</li>
<li>it calculates an F1 score, by extracting the values from the start of the attention mask, up to the start of the padding (and binarizes it)</li>
</ol>
<p>I accept full responsibility the code may be buggy.<br>
The main logic is the following (I use pytorch-lighting so the data arrive in <code>outputs</code> I then gather the data and concatenate it into single tensors):</p>
<pre><code class="lang-auto">i = 0
for attn_mask, pad_mask in zip(torch.unbind(attns), torch.unbind(pads)):
    idx  = (attn_mask == True).nonzero(as_tuple=False)
    end = (pad_mask == True).nonzero(as_tuple=False)

    if len(end) &gt; 0:
        z = torch.min(end).item()
    else:
        z = self.max_seq

    if len(idx) &gt; 0:
        k = torch.min(idx).item()
        t_x = ideals[i, k:]
        t_y = logits[i, k:]

        masked_x = ideals[i, k - 1:z]
        masked_x_hat = ideals[i, k - 1:z]

    i += 1
</code></pre>
<p>I then use <code>t_y</code> and <code>t_x</code> to calculate an R2 score, and <code>masked_x</code> and <code>masked_x_hat</code> to binarise and calculate F1 score on a specific feature.</p>
<p>The problem is:</p>
<ol>
<li>
<code>i</code> is the batch (I know its horrible) and the loop above unrolls for each and every single entry in the test batch.</li>
<li>the way <code>idx</code> and <code>end</code> are computer are per-sample in the batch. Same for <code>z</code>.</li>
</ol>
<p>I suspect there is a way to vectorize this and speed it up considerably.<br>
Playing around with dummy data, I managed to get as far as below:</p>
<p>Assume that it’s a 2 item batch, and a 5-length sequence.</p>
<pre><code class="lang-auto">pads = torch.Tensor([[False, False, False, False, True], [False, False, False, False, False]])
attns =  torch.Tensor([[False, False, True, True, True], [False, False, False, True, False]])
vals  = torch.rand(2, 5)
diffs = torch.eq(attns, pads)
(diffs == False).nonzero(as_tuple=False)
</code></pre>
<p>Basically, in this example, for the R2 Score I want the values starting from the first <code>True</code> to the end of the sequence and for the F1 Score, I want the values starting from the first <code>True</code> in attention, up to the first <code>True</code> in the padding (or the end of the sequence).</p>
<p>I have no clue what built-in <code>torch</code> methods exist to achieve this without doing the loop. The data is in the GPU if it matters. Any help is much appreciated!</p>