MultiHeadAttention attention_mask [Keras，Tensorflow]示例

Question

我正在努力为MultiHeadAttention层隐藏我的输入。我正在使用来自Keras文档的转换器块进行自我关注。到目前为止，我无法在网上找到任何示例代码，如果有人能给我一个代码片段，我会很感激的。

来自这页面的转换器块：

class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super(TransformerBlock, self).__init__()
        self.att = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.ffn = keras.Sequential(
            [layers.Dense(ff_dim, activation="relu"), layers.Dense(embed_dim),]
        )
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

    def call(self, inputs, training):
        attn_output = self.att(inputs, inputs)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(inputs + attn_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        return self.layernorm2(out1 + ffn_output)

在这链接下可以找到掩蔽文档：

attention_mask: B，T，S形状的布尔掩码，防止注意某些位置。布尔掩码指定哪些查询元素可以处理哪些关键元素，1表示注意，0表示不注意。对于缺少的批处理尺寸和头尺寸，可以进行广播。

唯一可以运行的是在layer类之外创建一个掩码，作为numpy数组：

mask = np.ones((observations, sequence_length, sequence_length))
mask[X[:observations,:,0]==0]=0

然后在调用层时输入，变压器块中唯一的更改是：

def call(self, inputs, mask, training):
    attn_output = self.att(inputs, inputs, attention_mask=mask)

然而，这当然不起作用，当给一个batch_size时，拟合，并且只工作对我的记忆的5个观察，所以它没有任何意义。除此之外，我不认为这是正确地掩盖输入--一般来说，考虑到attention_mask的形状(观察、sequence_length、sequence_length)，我对如何掩蔽感到非常困惑。我输入的形状是(观察，sequence_length，特性)。这个输入是由零填充，然而，当谈到变压器块，它已经通过一个嵌入层和CNN。我尝试过各种编写函数的方法，它在使用不同的张量或Keras对象进行训练时创建掩码。然而，我每次都会遇到错误。

我希望更流利的Tensorflow/Keras语言能够提供一个例子。或者有人告诉我，考虑到我的建筑，掩蔽是无用的。这个模型表现很好。然而，我希望掩蔽能帮助加快计算速度。这让我觉得我不能把我的头绕过去。

Answer 1

也许现在有点晚了，但是对于那些在这篇文章上寻找解决方案的人来说，这可能会有所帮助。

使用转换器的典型场景是NLP问题，其中有大量句子(让我们假设它们已经被标记为简单性)。请考虑以下示例：

sentences = [['Lorem', 'ipsum', 'dolor', 'sit', 'amet'], ['Integer', 'tincidunt', 'in', 'arcu', 'nec', 'fringilla', 'suscipit']]

如你所见，我们有两句不同长度的句子。为了在tensorflow模型中向它们学习，我们可以使用一个特殊的令牌来填充最短的令牌，比如'[PAD]'，并按照您的建议将它们输入变压器模型。因此：

sentences = tf.constant([['Lorem', 'ipsum', 'dolor', 'sit', 'amet', '[PAD]', '[PAD]'], ['Integer', 'tincidunt', 'in', 'arcu', 'nec', 'fringilla', 'suscipit']])

另外，假设我们已经有了从某些语料库中提取的标记词汇表，例如1000标记的词汇表，我们可以定义一个StringLookup层，该层可以将我们的一批句子转换为给定词汇表的它们的数值选择。我们可以指定用于掩蔽的令牌。

lookup = tf.keras.layers.StringLookup(vocabulary=vocabulary, mask_token='[PAD]')
x = lookup(sentences)
# x is a tf.Tensor([[2, 150, 19, 997, 9, 0, 0], [72, 14, 1, 1, 960, 58, 87]], shape=(2, 7), dtype=int64)

我们可以看到，[PAD]令牌映射到词汇表中的值。

典型的下一步是将该张量输入到Embedding层，如下所示：

embedding = tf.keras.layers.Embedding(input_dim=lookup.vocabulary_size(), output_dim=64, mask_zero=True)

这里的关键是参数mask_zero。根据文档，这个论点意味着：

布尔值，无论输入值0是否是一个特殊的“填充”值，都应该隐藏.

这允许embedding层为后续层生成一个掩码，以指示哪些位置应该参加，哪些位置不应该参加。可以通过以下方式访问此掩码：

mask = embedding.compute_mask(sentences)
# mask is a tf.Tensor([[True, True, True, True, True, False, False], [True, True, True, True, True, True, True]], shape=(2, 7), dtype=bool)

嵌入的张量具有如下形式：

y = embedding(sentences)
# y is a tf.Tensor of shape=(2, 7, 64), dtype=float32)

为了将mask使用到MultiHeadAttention层中，必须对掩码进行整形以满足形状要求，根据文档，[B, T, S]是B表示批处理大小(在示例中是2)，T是指查询大小(在我们的示例中是7)，S是指键大小E 229(如果我们正在使用自我关注的话，也是7)。此外，在一个多头注意层，我们必须照顾头部的数量H。使用此输入创建兼容掩码的最简单方法是通过广播：

mask = mask[:, tf.newaxis, tf.newaxis, :]
# mask is a tf.Tensor of shape=(2, 1, 1, 7), dtype=bool) -> [B, H, T, S]

然后，我们最终可以按以下方式提供MultiHeadAttention层：

mha = tf.keras.layers.MultiHeadAttention(num_heads=4, key_dim=64)
z = mha(y, y, attention_mask=mask)

因此，为了使用带有掩码的TransformerBlock层，您应该向call方法添加一个mask参数，如下所示：

def call(self, inputs, training, mask=None):
    attn_output = self.att(inputs, inputs, attention_mask=mask)
    ...

在调用MultiHeadAttention层的层/模型中，必须传递/传播使用Embedding层生成的掩码。