MultiHeadAttention Attention_mask [Keras、Tensorflow] 示例

我正在努力掩盖 MultiHeadAttention 层的输入。我正在使用 Keras 文档中的 Transformer Block 进行自我关注。到目前为止，我在网上找不到任何示例代码，如果有人能给我一个代码片段，我将不胜感激。

变压器块来自this https://keras.io/examples/nlp/text_classification_with_transformer/ page:

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)

屏蔽的文档可以在下面找到this https://www.tensorflow.org/api_docs/python/tf/keras/layers/MultiHeadAttention link:

Attention_mask：形状为 [B, T, S] 的布尔掩码，可防止 注意某些位置。布尔掩码指定哪个查询 元素可以关注哪些关键元素，1表示关注，0表示关注 表示不注意。可能会针对丢失的批次进行广播 尺寸和头部尺寸。

我唯一可以运行的是在图层类外部创建的掩码作为 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个观察，因此它没有任何意义。 除此之外，我认为这没有正确屏蔽输入 - 一般来说，考虑到注意力掩码的形状（观察、序列长度、序列长度），我对如何屏蔽感到非常困惑。我的输入的形状是（观察、序列长度、特征）。该输入被零填充，但是，当涉及到转换器块时，它已经通过了嵌入层和 CNN。 我尝试了各种方法来编写函数，该函数在使用不同的 Tensor 或 Keras 对象进行训练时创建掩码。然而我每次都会遇到错误。

我希望更熟悉 Tensorflow/Keras 的人能够提供一个例子。 或者有人告诉我，考虑到我的架构，屏蔽是没有用的。该模型表现良好。然而，我希望屏蔽可以帮助加快计算速度。 但令我烦恼的是我无法理解它。

也许有点晚了，但对于任何最终在这篇文章中寻找解决方案的人来说，这可能会有所帮助。

使用 Transformer 的典型场景是在 NLP 问题中，其中有成批的句子（为了简单起见，我们假设它们已经被标记化）。考虑以下示例：

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

如您所见，我们有两个长度不同的句子。为了在张量流模型中学习它们，我们可以用一个特殊的标记填充最短的一个，比方说'[PAD]'，然后按照您的建议将它们输入到 Transformer 模型中。因此：

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

还假设我们已经有了从某些语料库中提取的标记词汇表，例如词汇表1000令牌，我们可以定义一个StringLookup将我们的一批句子转换为给定词汇的数字投影的层。我们可以指定使用哪个令牌masking.

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]令牌映射到0词汇中的价值。

典型的下一步是将这个张量输入到Embedding层，像这样：

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

这里的关键是论证mask_zero。根据文档 https://www.tensorflow.org/api_docs/python/tf/keras/layers/Embedding，这个论证的意思是：

布尔值，输入值 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] where B意味着批量大小（示例中为 2），T意味着查询大小（在我们的示例中为 7），以及S意味着key size（如果我们使用自注意力，则再次为 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 layer.