MultiHeadAttention Attention_mask [Keras、Tensorflow] 示例

Question

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

本页的变压器块：

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)

屏蔽的文档可以在此链接下找到：

focus_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 的人能够提供一个例子。或者有人告诉我，考虑到我的架构，屏蔽是没有用的。该模型表现良好。然而，我希望屏蔽可以帮助加快计算速度。但令我烦恼的是我无法理解它。

Answer 1

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

使用 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层，将我们的一批句子转换为给定词汇表的数字投影。并且我们可以指定使用哪个 token 来进行屏蔽。

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。根据文档，这个参数意味着：

布尔值，输入值 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表示密钥大小（如果我们使用自注意力，则再次为 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图层生成的蒙版。