Relative position/type embeddings implementation

username: mrpatekful
cooked: <p>Hi, I am trying to implement a relative type embedding for transformer based dialogue models, similarily to relative position embedding in <a href="https://arxiv.org/pdf/1803.02155.pdf" rel="noopener nofollow ugc">https://arxiv.org/pdf/1803.02155.pdf</a></p>
<p>According to the article the the usual way of computing self attention:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/3/030e826cff5f450d8d3320c2e5866da81b577734.png" alt="Screenshot from 2020-04-12 11-55-31" data-base62-sha1="r2waleLEvyZmG3InyXPaiNPNdi" width="173" height="55"></p>
<pre data-code-wrap="python"><code class="lang-python">e = torch.matmul(query, key.T)
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/4/04fef6309f88fb02ae2d4f9a41f0520468428833.png" alt="Screenshot from 2020-04-12 11-57-31" data-base62-sha1="Ica5EqzUU986h8SCXjOtzzjuy7" width="158" height="53"></p>
<pre data-code-wrap="python"><code class="lang-python">a = torch.nn.functional.softmax(e, dim=-1)
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/7/d746a95d2fc41e58f876bf493c40a359a5a635f2.png" alt="image" data-base62-sha1="uIq2Y32m4ph7FnRQUbF461rWbiG" width="139" height="53"></p>
<pre data-code-wrap="python"><code class="lang-python">z = torch.matmul(a, value)
</code></pre>
<p>is modified to incorporate (by addition) a <code>[batch_size, seq_len, seq_len, embed_dim]</code> sized tensor with the relative position distance embeddings for every position pair in the final <code>z</code> vector. As the position values are the same for the batches, this can be simplified to <code>[seq_len, seq_len, embed_dim]</code> tensor, therefore sparing computation costs.</p>
<p>The modified equation to incorporate the pos embed matrix in self attention is then:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/6/06b67c3d67a00d637e9239cf5b4f72a82225987b.png" alt="Screenshot from 2020-04-12 12-04-55" data-base62-sha1="XnPWviBdQRBEuu1PyFV34d9W5l" width="185" height="57"></p>
<p>where <code>e</code> can be rewritten as the following to achieve the said optimization of removing the unnecessary broadcasting of the batch dimension:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/7/6732e76bafeeb55727f8615461b8ab4aa3d01386.png" alt="Screenshot from 2020-04-12 12-07-37" data-base62-sha1="eIWeVqkOTZv3A9TuAoAM5VHkVLg" width="225" height="59"></p>
<p>This basically means there are two terms, the first is the regular <code>torch.matmul(query, key.T)</code> product and</p>
<pre data-code-wrap="python"><code class="lang-python">torch.matmul(q, pos_embed_mat.T)
</code></pre>
<p>The equation for the <code>e</code> tensor in pytorch then can be written as:</p>
<pre data-code-wrap="python"><code class="lang-python">e = torch.matmul(query, key.T) + torch.matmul(q, pos_embed_mat.T)
</code></pre>
<p>The final output is then:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/5/2536ffa13e6c1966ba55e16f190db71641213976.png" alt="Screenshot from 2020-04-12 12-04-10" data-base62-sha1="5jdv11fJQt1z58KfsDkXD2A3cxw" width="178" height="55"></p>
<pre data-code-wrap="python"><code class="lang-python">a = torch.nn.functional.softmax(e, dim=-1)

z = torch.matmul(a, value) + torch.matmul(a, pos_embed)
</code></pre>
<p>The above code snippets are simplified version of the real code, as these do not take into account the head dimensions and the required various reshape operations to ensure the correct tensor sizes for the matrix products. I use the following function, which produces correct output for correctly shaped input tensors:</p>
<pre data-code-wrap="python"><code class="lang-python">def relative_attn_inner(x, y, pos_embed):
    """
    x: [batch_size, heads, length, head_dim] or [batch_size, heads, length, length]
    y: [batch_size, heads, head_dim, length] or [batch_size, heads, length, head_dim]
    pos_embed: [length, length, head_dim] or [length, head_dim, length]
    """
    batch_size, heads, length, _ = x.size()

    xy = torch.matmul(x, y)

    x_t = x.permute(2, 0, 1, 3)
    x_t_r = x_t.reshape([length, heads * batch_size, -1])
    x_tz = torch.matmul(x_t_r, pos_embed)
    x_tz_r = x_tz.reshape([length, batch_size, heads, -1])
    x_tz_r_t = x_tz_r.permute([1, 2, 0, 3])

    return xy + x_tz_r_t
</code></pre>
<p>The whole attention mechanism is then:</p>
<pre data-code-wrap="python"><code class="lang-python">def relative_attn(q, k, v, pos_embed):
    """
    q: [batch_size, heads, length, head_dim]
    k: [batch_size, heads, length, head_dim]
    v: [batch_size, heads, length, head_dim]
    pos_embed: [length, length, head_dim]
    """
    d = value.size(-1) ** 0.5
    e = relative_attn_inner(q, k.transpose(2, 3), pos_embed.transpose(1, 2)) / d
    a = torch.nn.functional.softmax(e, dim=-1)
    z = relative_attn_inner(a, v, pos_embed)

    return z
</code></pre>
<p>This all works fine, however I am trying to replace the <code>pos_embed</code> tensor with a <code>role_embed</code> tensor, where the elements of the matrix are not the pairwise relative distances of the input tokens, but the 1 or 0 values, whether the given element at position <code>i, j</code> belongs to an utterance spoken by the same person in a context of several turns of dialogs between two agents.</p>
<p><code>Hi how are you [EOS] Hi I am fine [EOS]</code></p>
<p>is encoded as</p>
<p><code>1, 1, 1, 1, 1, 0, 0, 0, 0, 0</code><br>
<code>1, 1, 1, 1, 1, 0, 0, 0, 0, 0</code><br>
<code>1, 1, 1, 1, 1, 0, 0, 0, 0, 0</code><br>
<code>1, 1, 1, 1, 1, 0, 0, 0, 0, 0</code><br>
<code>1, 1, 1, 1, 1, 0, 0, 0, 0, 0</code><br>
<code>0, 0, 0, 0, 0, 1, 1, 1, 1, 1</code><br>
<code>0, 0, 0, 0, 0, 1, 1, 1, 1, 1</code><br>
<code>0, 0, 0, 0, 0, 1, 1, 1, 1, 1</code><br>
<code>0, 0, 0, 0, 0, 1, 1, 1, 1, 1</code><br>
<code>0, 0, 0, 0, 0, 1, 1, 1, 1, 1</code></p>
<p>This is then fed to a role embedding layer to produce the same <code>[batch_size, length, length, head_dim]</code> sized tensor as the initial <code>pos_embed</code>. However in the case of my role embeddings, the batch dimension can not be simplified, as the embeddings differ for every elment of a batch.</p>
<p><strong>My problem is the above code works for only <code>[length, length, head_dim]</code> sized tensors as this <code>x_tz = torch.matmul(x_t_r, pos_embed)</code> product will have an adittional batch size dimension.</strong> Could someone that is interested in this topic help me find a way to rewrite the above equations to fix this problem?</p>
<p>I appreciate any form of help, and also here is a colab link to play with the above code:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1cFLuRm3zvts3L82VQ4-R7Rzhv_nowlhS">
  <header class="source">
      <img src="https://ssl.gstatic.com/colaboratory-static/common/f6c736cbde16632a9fcd16d9ab1970b1/img/favicon.ico" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1cFLuRm3zvts3L82VQ4-R7Rzhv_nowlhS" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img width="260" height="260" src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar">

<h3><a href="https://colab.research.google.com/drive/1cFLuRm3zvts3L82VQ4-R7Rzhv_nowlhS" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

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

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