#

相对多头注意力

这是 paper Transfor mer-XL：PyTorch 中固定长度上下文之外的细心语言模型中相对多头关注的实现。

16import torch
17from torch import nn
18
19from labml.logger import inspect
20from labml_nn.transformers.mha import MultiHeadAttention

#

此方法将矩阵的 $i^{t h}$ 行按 $i$ 列移动。

如果输入为[[1, 2 ,3], [4, 5 ,6], [7, 8, 9]] ，则移位的结果将为[[1, 2 ,3], [0, 4, 5], [6, 0, 7]] 。理想情况下，我们应该掩盖下三角形，但这对我们的目的来说是可以的。

23def shift_right(x: torch.Tensor):

#

连接一列零

33    zero_pad = x.new_zeros(x.shape[0], 1, *x.shape[2:])
34    x_padded = torch.cat([x, zero_pad], dim=1)

#

重塑并从末端移除多余的元素

37    x_padded = x_padded.view(x.shape[1] + 1, x.shape[0], *x.shape[2:])
38    x = x_padded[:-1].view_as(x)

#

41    return x

#

相对多头注意模块

我们重写了多头注意模块，因此我们只需要编写该get_scores 方法即可。

44class RelativeMultiHeadAttention(MultiHeadAttention):

#

52    def __init__(self, heads: int, d_model: int, dropout_prob: float = 0.1):

#

线性变换不需要偏差，因为我们在计算分数时会明确包含偏差。但是，有偏见value 可能是有道理的。

56        super().__init__(heads, d_model, dropout_prob, bias=False)

#

相对位置的数量

59        self.P = 2 ** 12

#

键相对于查询的相对位置嵌入。我们需要 $2 P$ 嵌入，因为键可以在查询之前或之后。

63        self.key_pos_embeddings = nn.Parameter(torch.zeros((self.P * 2, heads, self.d_k)), requires_grad=True)

#

键相对于查询的相对位置嵌入偏差。

65        self.key_pos_bias = nn.Parameter(torch.zeros((self.P * 2, heads)), requires_grad=True)

#

查询的位置嵌入与查询的位置无关

67        self.query_pos_bias = nn.Parameter(torch.zeros((heads, self.d_k)), requires_grad=True)

#

获取相对注意力分数

绝对关注

A_{j}^{ab s} = l i n_{q} (X_{i}^{q} + P_{i})^{⊤} l i n_{k} (X_{j}^{k} + P_{j}) = A Q_{i}^{⊤} K_{j} + B Q_{i}^{⊤} U_{j}^{K} + C U_{i}^{Q}^{⊤} K_{j} + D U_{i}^{Q}^{⊤} U_{j}^{K}

其中 $Q_{i}, K_{j}$ ，是原始嵌入的线性变换 $X_{i}^{q}, X_{j}^{k}$ ， $U_{i}^{Q}, U_{j}^{K}$ 是绝对位置编码的线性变换 $P_{i}, P_{j}$ 。

他们认为，无论查询的位置如何，对给定键的关注都应该相同。因此， $C U_{i}^{Q}^{⊤} K_{j}$ 用常量替换 $C v^{⊤} K_{j}$ 。

对于第二项和第三项，引入了相对位置编码。因此 $B Q_{i}^{⊤} U_{j}^{K}$ ，替换 $B Q_{i}^{⊤} R_{i - j}$ 为 $D U_{i}^{Q}^{⊤} U_{j}^{K}$ 和 $D S_{i - j}$ 。

A_{i, j}^{re l} = A Q_{i}^{⊤} K_{j} + B Q_{i}^{⊤} R_{i - j} + C v^{⊤} K_{j} + D S_{i - j}

69    def get_scores(self, query: torch.Tensor, key: torch.Tensor):

#

$R_{k}$

108        key_pos_emb = self.key_pos_embeddings[self.P - key.shape[0]:self.P + query.shape[0]]

#

$S_{k}$

110        key_pos_bias = self.key_pos_bias[self.P - key.shape[0]:self.P + query.shape[0]]

#

$v^{⊤}$

112        query_pos_bias = self.query_pos_bias[None, None, :, :]

#

$(A + C)_{i, j} = Q_{i}^{⊤} K_{j} + v^{⊤} K_{j}$

117        ac = torch.einsum('ibhd,jbhd->ijbh', query + query_pos_bias, key)

#

$B^{'}_{i, k} = Q_{i}^{⊤} R_{k}$

119        b = torch.einsum('ibhd,jhd->ijbh', query, key_pos_emb)

#

$D^{'}_{i, k} = S_{k}$

121        d = key_pos_bias[None, :, None, :]

#

移动行 $(B^{'} + D^{'})_{i, k}$ 以获取 $(B + D)_{i, j} = (B^{'} + D^{'})_{i, i - j}$

124        bd = shift_right(b + d)

#

移除多余的头寸

126        bd = bd[:, -key.shape[0]:]

#

返回总和 $A Q_{i}^{⊤} K_{j} + B Q_{i}^{⊤} R_{i - j} + C v^{⊤} K_{j} + D S_{i - j}$

134        return ac + bd

#

137def _test_shift_right():
138    x = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
139    inspect(x)
140    inspect(shift_right(x))
141
142    x = torch.arange(1, 6)[None, :, None, None].repeat(5, 1, 1, 1)
143    inspect(x[:, :, 0, 0])
144    inspect(shift_right(x)[:, :, 0, 0])
145
146    x = torch.arange(1, 6)[None, :, None, None].repeat(3, 1, 1, 1)
147    inspect(x[:, :, 0, 0])
148    inspect(shift_right(x)[:, :, 0, 0])
149
150
151if __name__ == '__main__':
152    _test_shift_right()