How do I obtain multiple heads via `multi_head_attention_forward` when using `nn.TransformerEncoder`?

sgbaird: <p>Maintainership of the <a href="https://github.com/anthony-wang/CrabNet/tree/master/crabnet" rel="noopener nofollow ugc">CrabNet</a> repository is being transferred to me, and I’m very interested in removing the hacky workaround described in<br>
<a href="https://github.com/anthony-wang/CrabNet#important---if-you-want-to-reproduce-the-publication-figures-1-and-2" class="inline-onebox" rel="noopener nofollow ugc">GitHub - anthony-wang/CrabNet: Predict materials properties using only the composition information!</a> so that I can keep the user-friendliness and include plotting functions as class methods without asking people to edit PyTorch source code.</p>
<blockquote>
<p>To properly export the attention heads from the PyTorch <code>nn.MultiheadAttention</code> implementation within the transformer encoder layer, you will need to manually modify some of the source code of the PyTorch library.<br>
This applies to PyTorch v1.6.0, v1.7.0, and v1.7.1 (potentially to other untested versions as well).</p>
<blockquote>
<p>For this, open the file:<br>
<code>C:\Users\{USERNAME}\Anaconda3\envs\{ENVIRONMENT}\Lib\site-packages\torch\nn\functional.py</code><br>
(where <code>USERNAME</code> is your Windows user name and <code>ENVIRONMENT</code> is your conda environment name (if you followed the steps above, then it should be <code>crabnet</code>))</p>
<p>At the end of the function defition of <code>multi_head_attention_forward</code> (line numbers may differ slightly):</p>
<pre data-code-wrap="python"><code class="lang-python">L4011 def multi_head_attention_forward(
# ...
# ... [some lines omitted]
# ...
L4291    if need_weights:
L4292        # average attention weights over heads
L4293        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
L4294        return attn_output, attn_output_weights.sum(dim=1) / num_heads
L4295    else:
L4296        return attn_output, None
</code></pre>
<p>Change the specific line</p>
<pre data-code-wrap="python"><code class="lang-python">return attn_output, attn_output_weights.sum(dim=1) / num_heads
</code></pre>
<p>to:</p>
<pre data-code-wrap="python"><code class="lang-python">return attn_output, attn_output_weights
</code></pre>
<p>This prevents the returning of the attention values as an average value over all heads, and instead returns each head’s attention matrix individually.<br>
For more information see:</p>
<ul>
<li><a href="https://github.com/pytorch/pytorch/issues/34537" class="inline-onebox" rel="noopener nofollow ugc">nn.MultiHeadAttention should be able to return attention weights for each head. · Issue #34537 · pytorch/pytorch · GitHub</a></li>
<li><a href="https://github.com/pytorch/pytorch/issues/32590" class="inline-onebox" rel="noopener nofollow ugc">[FYI] MultiheadAttention / Transformer · Issue #32590 · pytorch/pytorch · GitHub</a></li>
<li><a href="https://discuss.pytorch.org/t/getting-nn-multiheadattention-attention-weights-for-each-head/72195/" class="inline-onebox">Getting nn.MultiHeadAttention attention weights for each head</a></li>
</ul>
</blockquote>
</blockquote>
<p>The relevant code in the CrabNet repository is as follows:</p>
<h3><a name="encoder-__init__py-1" class="anchor" href="#encoder-__init__py-1"></a>Encoder <code>__init__.py</code></h3>
<aside class="onebox githubblob" data-onebox-src="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/crabnet/kingcrab.py#L151-L157">
  <header class="source">

      <a href="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/crabnet/kingcrab.py#L151-L157" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/crabnet/kingcrab.py#L151-L157" target="_blank" rel="noopener nofollow ugc">anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/crabnet/kingcrab.py#L151-L157</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="151" style="counter-reset: li-counter 150 ;">
          <li>if self.attention:</li>
          <li>    encoder_layer = nn.TransformerEncoderLayer(self.d_model,</li>
          <li>                                               nhead=self.heads,</li>
          <li>                                               dim_feedforward=2048,</li>
          <li>                                               dropout=0.1)</li>
          <li>    self.transformer_encoder = nn.TransformerEncoder(encoder_layer,</li>
          <li>                                                     num_layers=self.N)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<h3><a name="encoder-forward-2" class="anchor" href="#encoder-forward-2"></a>Encoder <code>forward</code></h3>
<aside class="onebox githubblob" data-onebox-src="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/crabnet/kingcrab.py#L176-L177">
  <header class="source">

      <a href="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/crabnet/kingcrab.py#L176-L177" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/crabnet/kingcrab.py#L176-L177" target="_blank" rel="noopener nofollow ugc">anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/crabnet/kingcrab.py#L176-L177</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="176" style="counter-reset: li-counter 175 ;">
          <li>x = self.transformer_encoder(x_src,</li>
          <li>                             src_key_padding_mask=src_mask)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<h3><a name="plotting-using-hook-3" class="anchor" href="#plotting-using-hook-3"></a>Plotting using hook</h3>
<aside class="onebox githubblob" data-onebox-src="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/Paper_FIG_1.py#L31-L56">
  <header class="source">

      <a href="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/Paper_FIG_1.py#L31-L56" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/Paper_FIG_1.py#L31-L56" target="_blank" rel="noopener nofollow ugc">anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/Paper_FIG_1.py#L31-L56</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="31" style="counter-reset: li-counter 30 ;">
          <li>class SaveOutput:</li>
          <li>    def __init__(self):</li>
          <li>        self.outputs = []</li>
          <li></li>
          <li>    def __call__(self, module, module_in, module_out):</li>
          <li>        self.outputs.append(module_out)</li>
          <li></li>
          <li>    def clear(self):</li>
          <li>        self.outputs = []</li>
          <li></li>
          <li>save_output = SaveOutput()</li>
          <li>hook_handles = []</li>
          <li></li>
          <li>for layer in model.model.modules():</li>
          <li>    if isinstance(layer, torch.nn.modules.activation.MultiheadAttention):</li>
          <li>        # print('isinstance')</li>
          <li>        handle = layer.register_forward_hook(save_output)</li>
          <li>        hook_handles.append(handle)</li>
          <li></li>
          <li>len(save_output.outputs)</li>
      </ol>
    </code></pre>


  This file has been truncated. <a href="https://github.com/anthony-wang/CrabNet/blob/94184a9f6166cefeb64b9c4faadbe61b87b3fe56/Paper_FIG_1.py#L31-L56" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>The hacky workaround that I’m trying to remove is the editing of the PyTorch source code. As long as users can still install via <code>pip</code> or <code>conda</code>, and users don’t have to worry about the change mentioned above (i.e. editting the source code, or some other hacky workaround), I’m game (some of the intended audience will have limited programming experience). If no easy solutions exist as-is, then I can move this into a GitHub issue as a feature request. See <a href="https://github.com/pytorch/pytorch/issues/32590#issuecomment-1050988587" class="inline-onebox" rel="noopener nofollow ugc">[FYI] MultiheadAttention / Transformer · Issue #32590 · pytorch/pytorch · GitHub</a></p>
<p>Happy to hear any other general suggestions or criticisms (e.g. you’re going about this the wrong way). My experience with transformers is mostly limited to CrabNet, and I haven’t had much success with PyTorch hooks in the past. I’d be happy to use a different approach/module etc. as long as the basic structure and functionality (i.e. transformer architecture) can be retained.</p>

jasperhyp: <p>I just wanna say your post has helped me a lot!  Thanks!</p>

FALL-ML: <p>Hah, I was thinking through this exact problem all day and just saw your forum post! I think part of the problem lies in the way that the <a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/transformer.html#TransformerEncoder" rel="noopener nofollow ugc">TransformerEncoderLayer</a> is written. In the self-attention block, the needs_weight argument is hardcoded as False:</p>
<pre><code class="lang-auto">x = self.self_attn(x, x, x,
                   attn_mask=attn_mask,
                   key_padding_mask=key_padding_mask,
                   need_weights=False)[0]
</code></pre>
<p>For newer versions of Pytorch, the MultiheadAttention module has a flag in the forward pass that allows you to turn off weight averaging (average_attn_weights: bool = True). However, this doesn’t matter if you’re using the TransformerEncoderLayer because the hardcoded need_weights is just going to stop any weight return. If it wasn’t hardcoded, you could have used a default override statement. For example:</p>
<pre><code class="lang-auto"># the four defaults are: 
# key_padding_mask: Optional[Tensor] = None
# need_weights: bool = True
# attn_mask: Optional[Tensor] = None
# average_attn_weights: bool = True
nn.MultiheadAttention.forward.__defaults__ = (None, True, None, False)
</code></pre>
<p><strong>Two possible solutions.</strong> So I think this leaves you with two options. One is to write your own versions of the TransformerEncoderLayer and the MultiheadAttention classes. I think this is a painful option that has the added drawback of missing out on future improvements to these two classes. Option number two is to monkey patch the <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4905" rel="noopener nofollow ugc">multi_head_attention_forward</a> function in functional.py. This is not elegant, but it gets around users with limiting programming experience having to manually edit the PyTorch source code.</p>
<p><strong>The monkey patch solution.</strong> In the imports for train_crabnet.py, import the inspect package.</p>
<pre><code class="lang-auto">import inspect
</code></pre>
<p>Then implement these four lines anywhere after importing Model but somewhere before your first hook:</p>
<pre><code class="lang-auto">source = inspect.getsource(torch.nn.functional.multi_head_attention_forward)
new_source = source.replace('attn_output_weights.sum(dim=1) / num_heads', 'attn_output_weights')
new_source = new_source.replace('return attn_output, None', 'return attn_output, attn_output_weights')
exec(new_source, torch.nn.functional.__dict__)
</code></pre>
<ul>
<li>Line 1 grabs the source code for the multi_head_attention_forward function in functional.py.</li>
<li>Line 2 finds the line where attention head averaging occurs and replaces it with no averaging.</li>
<li>Line 3 finds the return line (which is separate in later versions of Pytorch) and replaces it with a return that includes the attention weights. <strong>Note:</strong> This line might not be necessary in earlier versions of Pytorch, but it is necessary in the version I’m currently using (1.12.0). It should be harmless if it isn’t necessary.</li>
<li>Line 4 re-executes the source code with the changes from lines 2-3.</li>
</ul>
<p>After making these changes your save_output.output will return a list with length equal to the number of batches. Each entry in the list (e.g. save_output.output[0]) is a tuple of length 2, where the second entry is your non-averaged attention weights. For example, I currently have 4x heads, batch sizes of 1024, and 5 features:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/8/288aea1213ba3ef4ff44f11a80d9cfc46d0924fb.png" alt="image" data-base62-sha1="5MEIJ7rIvpG99Dk4g81ZNtNoPnJ" width="248" height="210"></p>
<p>You can see that the first dimension is 4096 (batch_size * heads). Whereas the second and third dimension represent the square attention weight map (5 x 5).</p>
<p>With the monkey patch commented out, the TransformerEncoderLayer ensures that None is returned in the second entry of the tuple.</p>