#

相対的多面的注意

これは、論文「Transformer-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], [9, 0, 7]] 下の三角形をマスクするのが理想的ですが、この目的には問題ありません。

23def shift_right(x: torch.Tensor):

#

0 の列を連結する

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}$ 。

第2用語と第3用語では、相対位置エンコーディングが導入されています。 $B Q_{i}^{⊤} U_{j}^{K}$ So は、 $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} Z$

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()