MultiHeadAttention attention_mask [Keras, Tensorflow] 示例

Question

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

来自this页面的变压器块：

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链接下找到：

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

我唯一可以运行的是在图层类之外创建为 numpy 数组的掩码：

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

然后在调用层时输入，transformer 块中唯一的变化是：

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

但是，当在拟合时给定 batch_size 时，这当然不起作用，并且仅适用于我的记忆中的 5 次观察，所以它没有任何意义。 除此之外，我认为这并没有正确屏蔽输入 - 一般来说，考虑到 attention_mask 的形状（观察值、sequence_length、sequence_length），我对如何屏蔽感到很困惑。我输入的形状是（观察、序列长度、特征）。该输入被零填充，但是，当涉及到转换器块时，它已经通过了嵌入层和 CNN。 我尝试了各种方法来编写一个函数，该函数在使用不同的 Tensor 或 Keras 对象进行训练时创建掩码。但是我每次都遇到错误。

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

Answer 1

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

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

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

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

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

另外假设我们已经有一个从一些语料库中提取的标记词汇表，例如1000 标记的词汇表，我们可以定义一个StringLookup 层，将我们的一批句子转换成给定词汇表的数字预测。我们可以指定哪个令牌用于ma​​sking。

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

布尔值，输入值 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 表示 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 层生成的掩码。