add shaws relative positional embedding

Author: lucidrains

conformer/conformer.py

@@ -77,7 +77,8 @@ def __init__(
         dim,
         heads = 8,
         dim_head = 64,
-        dropout = 0.
+        dropout = 0.,
+        max_pos_emb = 512
     ):
         super().__init__()
         inner_dim = dim_head * heads
@@ -87,17 +88,28 @@ def __init__(
         self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
         self.to_out = nn.Linear(inner_dim, dim)
 
+        self.max_pos_emb = max_pos_emb
+        self.rel_pos_emb = nn.Embedding(2 * max_pos_emb + 1, dim_head)
+
         self.dropout = nn.Dropout(dropout)
 
     def forward(self, x, context = None, mask = None, context_mask = None):
-        device, h, has_context = x.device, self.heads, exists(context)
+        n, device, h, max_pos_emb, has_context = x.shape[-2], x.device, self.heads, self.max_pos_emb, exists(context)
         context = default(context, x)
 
         q, k, v = (self.to_q(x), *self.to_kv(context).chunk(2, dim = -1))
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
 
         dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
 
+        # shaw's relative positional embedding
+        seq = torch.arange(n, device = device)
+        dist = seq[:, None] - seq[None, :]
+        dist = dist.clip(-max_pos_emb, max_pos_emb) + max_pos_emb
+        rel_pos_emb = self.rel_pos_emb(dist).to(q)
+        pos_attn = einsum('b h n d, n r d -> b h n r', q, rel_pos_emb) * self.scale
+        dots = dots + pos_attn
+
         if exists(mask) or exists(context_mask):
             mask = default(mask, lambda: torch.ones(*x.shape[:2], device = device))
             context_mask = default(context_mask, mask) if not has_context else default(context_mask, lambda: torch.ones(*context.shape[:2], device = device))