MultiHeadAttention 层

[源代码]

MultiHeadAttention 类

tf_keras.layers.MultiHeadAttention(
    num_heads,
    key_dim,
    value_dim=None,
    dropout=0.0,
    use_bias=True,
    output_shape=None,
    attention_axes=None,
    kernel_initializer="glorot_uniform",
    bias_initializer="zeros",
    kernel_regularizer=None,
    bias_regularizer=None,
    activity_regularizer=None,
    kernel_constraint=None,
    bias_constraint=None,
    **kwargs
)

MultiHeadAttention 层。

这是"Attention is all you Need"（Vaswani 等人，2017 年）论文中描述的多头注意力机制的实现。如果 query、key、value 相同，则这是自注意力。query 中的每个时间步都关注 key 中的相应序列，并返回一个固定宽度的向量。

该层首先对 query、key 和 value 进行投影。这些（有效地）是长度为 num_attention_heads 的张量列表，其中相应的形状为 (batch_size, <query dimensions>, key_dim)、(batch_size, <key/value dimensions>, key_dim)、(batch_size, <key/value dimensions>, value_dim)。

然后，对查询和键张量进行点积并缩放。然后对这些进行 softmax 处理以获得注意力概率。然后，值张量通过这些概率进行插值，然后连接回单个张量。

最后，最后一个维度为 value_dim 的结果张量可以进行线性投影并返回。

当在自定义层内使用 MultiHeadAttention 时，自定义层必须实现自己的 build() 方法，并在其中调用 MultiHeadAttention 的 _build_from_signature()。这使得在加载模型时能够正确恢复权重。

示例

对两个序列输入执行 1D 交叉注意力，并带有注意力掩码。返回头部的附加注意力权重。

>>> layer = MultiHeadAttention(num_heads=2, key_dim=2)
>>> target = tf.keras.Input(shape=[8, 16])
>>> source = tf.keras.Input(shape=[4, 16])
>>> output_tensor, weights = layer(target, source,
...                                return_attention_scores=True)
>>> print(output_tensor.shape)
(None, 8, 16)
>>> print(weights.shape)
(None, 2, 8, 4)

在 5D 输入张量的轴 2 和 3 上执行 2D 自注意力。

>>> layer = MultiHeadAttention(
...     num_heads=2, key_dim=2, attention_axes=(2, 3))
>>> input_tensor = tf.keras.Input(shape=[5, 3, 4, 16])
>>> output_tensor = layer(input_tensor, input_tensor)
>>> print(output_tensor.shape)
(None, 5, 3, 4, 16)

参数

• num_heads：注意力头的数量。
• key_dim：查询和键的每个注意力头的尺寸。
• value_dim：值每个注意力头的尺寸。
• dropout：Dropout 概率。
• use_bias：布尔值，表示密集层是否使用偏置向量/矩阵。
• output_shape：除了 batch 和 sequence 维度之外，输出张量的预期形状。如果未指定，则投影回查询特征维度（查询输入的最后一个维度）。
• attention_axes：应用注意力的轴。None 表示在除 batch、heads 和 features 之外的所有轴上应用注意力。
• kernel_initializer：密集层核的初始化器。
• bias_initializer：密集层偏置的初始化器。
• kernel_regularizer：密集层核的正则化器。
• bias_regularizer：密集层偏置的正则化器。
• activity_regularizer：密集层活动的正则化器。
• kernel_constraint：密集层核的约束。
• bias_constraint：密集层核的约束。

调用参数

• query：形状为 (B, T, dim) 的查询 Tensor。
• value：形状为 (B, S, dim) 的值 Tensor。
• key：形状为 (B, S, dim) 的可选键 Tensor。如果未提供，则将使用 value 作为 key 和 value，这是最常见的情况。
• attention_mask：形状为 (B, T, S) 的布尔掩码，可防止对某些位置进行注意力。布尔掩码指定了哪些查询元素可以关注哪些键元素，1 表示注意力，0 表示无注意力。可以对缺失的 batch 维度和 head 维度进行广播。
• return_attention_scores：一个布尔值，指示如果为 True，输出应为 (attention_output, attention_scores)，如果为 False，则为 attention_output。默认为 False。
• training：一个 Python 布尔值，指示层应以训练模式（添加 dropout）还是推理模式（无 dropout）运行。将使用父层/模型的训练模式，或者在没有父层时为 False（推理）。
• use_causal_mask：一个布尔值，指示是否应用因果掩码以防止 token 关注未来的 token（例如，在解码器 Transformer 中使用）。

返回

• attention_output：计算结果，形状为 (B, T, E)，其中 T 表示目标序列形状，E 是查询输入的最后一个维度（如果 output_shape 为 None）。否则，多头输出将被投影到 output_shape 指定的形状。
• attention_scores：[可选] 在注意力轴上的多头注意力系数。