Masking queries in torch. MultiheadAttention

Adamits: <p>Hi, I am trying to use torch. MultiheadAttention for the following use case:</p>
<p>I have documents of |Q| queries, and sentences of length |K| (here, K==V). I would like for each Q to attend to all K, and ultimately, I will combine the |Q| context vectors. If I am batching these inputs, I understand that I can pass <code>key_padding_mask=|B|x|K|</code> where <code>|B|</code> is my batch size, to the model to mask the appropriate keys in each batch. I also understand that I can use <code>attn_mask=|Q|x|K|</code> to mask which k a given output (in my case I want |Q| outputs) attends to.</p>
<p>However, because I have a variable length |Q|, how can I mask this batchwise? Can I simply pass <code>attn_mask=|B|x|Q|x|K|</code> and get the desired effect? Essentially, certain entire distributions over |K| are invalid because that q is masked. I guess I can ust 0 out those vectors after the computation? Is my understanding even correct, or is the implementation of MultiheadAttention not valid for my use-case?</p>

Adamits: <p>It’s actually pretty surprising to me that you cannot mask the query with this function. Am I not using it as intended?</p>

Adamits: <p>Nevermind, this does exist, I was just getting an error because your first dimension for <code>attn_mask</code> needs to be batch_size * num_heads. I just had <code>batch_size</code>. This is explained in <code>F.multi_head_attention_forward</code>, but not in the <code>MultiheadAttention</code> <code>forward</code> method.</p>

Adamits: <p>Note: this feature was merged  on Jan 23 - <a href="https://github.com/pytorch/pytorch/pull/31996" rel="nofollow noopener">https://github.com/pytorch/pytorch/pull/31996</a> - and is not in the pip distribution of pytorch</p>
<p>EDIT: Note that this is available in the most recent pip release, 1.5: <a href="https://github.com/pytorch/pytorch/releases" rel="nofollow noopener">https://github.com/pytorch/pytorch/releases</a></p>
<aside class="onebox githubpullrequest">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/pull/31996" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Pull Request">
    <svg width="60" height="60" class="github-icon" viewbox="0 0 12 16" aria-hidden="true"><path d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/pull/31996" target="_blank" rel="nofollow noopener">Support 3D attention mask in MultiheadAttention.</a>
    </h4>

    <div class="branches">
      <code>pytorch:master</code> ← <code>pytorch:3d_attn_mask</code>
    </div>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-01-09" data-time="18:23:55" data-timezone="UTC">06:23PM - 09 Jan 20 UTC</span>
      </div>

      <div class="user">
        <a href="https://github.com/zhangguanheng66" target="_blank" rel="nofollow noopener">
          <img alt="zhangguanheng66" src="https://avatars0.githubusercontent.com/u/6156351?v=4" class="onebox-avatar-inline" width="20" height="20">
          zhangguanheng66
        </a>
      </div>

      <div class="lines" title="19 commits changed 3 files with 57 additions and 22 deletions">
        <a href="https://github.com/pytorch/pytorch/pull/31996/files" target="_blank" rel="nofollow noopener">
          <span class="added">+57</span>
          <span class="removed">-22</span>
        </a>
      </div>
    </div>

  </div>
</div>
  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>