Keras 2 API 文档 / Layers API / 注意力层 / Attention layer

注意力层

[源代码]

Attention

tf_keras.layers.Attention(use_scale=False, score_mode="dot", **kwargs)

点积注意力层,又称 Luong 风格注意力。

输入是形状为 [batch_size, Tq, dim]query 张量,形状为 [batch_size, Tv, dim]value 张量,以及形状为 [batch_size, Tv, dim]key 张量。计算遵循以下步骤:

  1. 计算形状为 [batch_size, Tq, Tv] 的得分,作为 query-key 点积:scores = tf.matmul(query, key, transpose_b=True)
  2. 使用得分计算形状为 [batch_size, Tq, Tv] 的分布:distribution = tf.nn.softmax(scores)
  3. 使用 distributionvalue 进行线性组合,得到形状为 [batch_size, Tq, dim] 的结果:return tf.matmul(distribution, value)

参数

  • use_scale:如果为 True,将创建一个标量变量来缩放注意力分数。
  • dropout:0 到 1 之间的浮点数。用于遗漏(drop)注意力得分的单元的比例。默认为 0.0。
  • score_mode:用于计算注意力得分的函数,可以是 {"dot", "concat"} 之一。"dot" 指的是 query 和 key 向量之间的点积。"concat" 指的是 query 和 key 向量拼接后的双曲正切值。

调用参数

  • inputs:以下张量的列表
    • query:形状为 [batch_size, Tq, dim] 的 Query Tensor
    • value:形状为 [batch_size, Tv, dim] 的 Value Tensor
    • key:形状为 [batch_size, Tv, dim] 的可选 Key Tensor。如果未提供,则将 value 同时用于 keyvalue,这是最常见的情况。
  • mask:以下张量的列表
    • query_mask:形状为 [batch_size, Tq] 的布尔掩码 Tensor。如果提供,则在 mask==False 的位置输出将为零。
    • value_mask:形状为 [batch_size, Tv] 的布尔掩码 Tensor。如果提供,将应用掩码,使得在 mask==False 的位置处的值不会对结果做出贡献。
  • return_attention_scores:布尔值,如果为 True,则将注意力分数(经过掩码和 softmax 后)作为额外的输出参数返回。
  • training:Python 布尔值,指示该层是应以训练模式(添加 dropout)还是推理模式(无 dropout)运行。
  • use_causal_mask:布尔值。设置为 True 用于 decoder self-attention。添加一个掩码,使得位置 i 不能 attending 到位置 j > i。这可以防止信息从未来流向过去。默认为 False

输出

Attention outputs of shape `[batch_size, Tq, dim]`.
[Optional] Attention scores after masking and softmax with shape
    `[batch_size, Tq, Tv]`.

queryvaluekey 的含义取决于应用。例如,在文本相似度的情况下,query 是第一段文本的序列嵌入,value 是第二段文本的序列嵌入。key 通常与 value 是同一个张量。

以下是一个在 CNN+Attention 网络中使用 Attention 的代码示例:

# Variable-length int sequences.
query_input = tf.keras.Input(shape=(None,), dtype='int32')
value_input = tf.keras.Input(shape=(None,), dtype='int32')

# Embedding lookup.
token_embedding = tf.keras.layers.Embedding(input_dim=1000, output_dim=64)
# Query embeddings of shape [batch_size, Tq, dimension].
query_embeddings = token_embedding(query_input)
# Value embeddings of shape [batch_size, Tv, dimension].
value_embeddings = token_embedding(value_input)

# CNN layer.
cnn_layer = tf.keras.layers.Conv1D(
    filters=100,
    kernel_size=4,
    # Use 'same' padding so outputs have the same shape as inputs.
    padding='same')
# Query encoding of shape [batch_size, Tq, filters].
query_seq_encoding = cnn_layer(query_embeddings)
# Value encoding of shape [batch_size, Tv, filters].
value_seq_encoding = cnn_layer(value_embeddings)

# Query-value attention of shape [batch_size, Tq, filters].
query_value_attention_seq = tf.keras.layers.Attention()(
    [query_seq_encoding, value_seq_encoding])

# Reduce over the sequence axis to produce encodings of shape
# [batch_size, filters].
query_encoding = tf.keras.layers.GlobalAveragePooling1D()(
    query_seq_encoding)
query_value_attention = tf.keras.layers.GlobalAveragePooling1D()(
    query_value_attention_seq)

# Concatenate query and document encodings to produce a DNN input layer.
input_layer = tf.keras.layers.Concatenate()(
    [query_encoding, query_value_attention])

# Add DNN layers, and create Model.
# ...