Spaces:
Running
Running
File size: 5,800 Bytes
9d61c9b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
import math
from typing import Tuple
import torch
from torch import nn
from torch.nn import Module
import torch.nn.functional as F
class RelativeMultiHeadAttention(Module):
r"""Multi-head attention with relative positional encoding.
This concept was proposed in the
[Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://arxiv.org/abs/1901.02860)
Args:
d_model (int): The dimension of model
num_heads (int): The number of attention heads.
Inputs: query, key, value, pos_embedding, mask
- **query** (batch, time, dim): Tensor containing query vector
- **key** (batch, time, dim): Tensor containing key vector
- **value** (batch, time, dim): Tensor containing value vector
- **pos_embedding** (batch, time, dim): Positional embedding tensor
- **mask** (batch, 1, time2) or (batch, time1, time2): Tensor containing indices to be masked
Returns:
- **outputs**: Tensor produces by relative multi head attention module.
Note: `d_model` should be divisible by `num_heads` in other words `d_model % num_heads` should be zero.
"""
def __init__(
self,
d_model: int = 512,
num_heads: int = 16,
):
super().__init__()
assert d_model % num_heads == 0, "d_model % num_heads should be zero."
self.d_model = d_model
self.d_head = int(d_model / num_heads)
self.num_heads = num_heads
self.sqrt_dim = math.sqrt(d_model)
self.query_proj = nn.Linear(d_model, d_model)
self.key_proj = nn.Linear(d_model, d_model, bias=False)
self.value_proj = nn.Linear(d_model, d_model, bias=False)
self.pos_proj = nn.Linear(d_model, d_model, bias=False)
self.u_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
self.v_bias = nn.Parameter(torch.Tensor(self.num_heads, self.d_head))
torch.nn.init.xavier_uniform_(self.u_bias)
torch.nn.init.xavier_uniform_(self.v_bias)
self.out_proj = nn.Linear(d_model, d_model)
def forward(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
pos_embedding: torch.Tensor,
mask: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
r"""Function applies multi-head attention along with relative positional encoding to the inputs. It restructures the input queries, keys, and values according to individual attention heads, applies biases, calculates content and position scores, and combines these to get the final score. A softmax activation is applied over the final score, followed by the calculation of context (contextual representation of input).
Performs the forward pass on the queries, keys, values, and positional embeddings with a mask.
Args:
query (torch.Tensor): The input tensor containing query vectors.
key (torch.Tensor): The input tensor containing key vectors.
value (torch.Tensor): The input tensor containing value vectors.
pos_embedding (torch.Tensor): The positional embedding tensor.
mask (torch.Tensor): The mask tensor containing indices to be masked.
Returns:
Tuple[torch.Tensor, torch.Tensor]: The context and attention tensors.
Tensor produces by relative multi head attention module.
"""
batch_size = query.shape[0]
query = self.query_proj(query).view(batch_size, -1, self.num_heads, self.d_head)
key = (
self.key_proj(key)
.view(batch_size, -1, self.num_heads, self.d_head)
.permute(0, 2, 1, 3)
)
value = (
self.value_proj(value)
.view(batch_size, -1, self.num_heads, self.d_head)
.permute(0, 2, 1, 3)
)
pos_embedding = self.pos_proj(pos_embedding).view(
batch_size, -1, self.num_heads, self.d_head,
)
u_bias = self.u_bias.expand_as(query)
v_bias = self.v_bias.expand_as(query)
a = (query + u_bias).transpose(1, 2)
content_score = a @ key.transpose(2, 3)
b = (query + v_bias).transpose(1, 2)
pos_score = b @ pos_embedding.permute(0, 2, 3, 1)
pos_score = self._relative_shift(pos_score)
score = content_score + pos_score
score = score * (1.0 / self.sqrt_dim)
score.masked_fill_(mask, -1e9)
attn = F.softmax(score, -1)
context = (attn @ value).transpose(1, 2)
context = context.contiguous().view(batch_size, -1, self.d_model)
return self.out_proj(context), attn
def _relative_shift(self, pos_score: torch.Tensor) -> torch.Tensor:
r"""The main idea of relative positional encoding is that the attention score doesn't only depend on the query and the key, but also on the relative position of the key with respect to the query. This becomes particularly useful when working with sequences of tokens, like in NLP tasks, as it helps the model to be aware of the position of the words (or tokens) in the sentence.
Performs the relative shift operation on the positional scores.
Args:
pos_score (torch.Tensor): The positional scores tensor.
Returns:
torch.Tensor: The shifted positional scores tensor.
"""
batch_size, num_heads, seq_length1, seq_length2 = pos_score.size()
zeros = torch.zeros(
(batch_size, num_heads, seq_length1, 1), device=pos_score.device,
)
padded_pos_score = torch.cat([zeros, pos_score], dim=-1)
padded_pos_score = padded_pos_score.view(
batch_size, num_heads, seq_length2 + 1, seq_length1,
)
return padded_pos_score[:, :, 1:].view_as(pos_score)
|