AlibiBias 层

[源代码]

AlibiBias

keras_hub.layers.AlibiBias(alibi_bias_max=8, **kwargs)

一个将 alibi 偏差添加到注意力分数的层。

此层将 alibi 偏差添加到注意力分数。Alibi 偏差是一种线性、非学习的偏差。在 Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation 中定义和形式化。

此层将注意力分数作为输入,并返回添加 alibi 偏差后的注意力分数。输出的形状将与输入相同。

参数

  • alibi_bias_max: int。此值将用于计算每个头的斜率。头的斜率是一个几何序列,从 2**(-alibi_bias_max/num_heads) 开始,并使用相同的值作为其比率。默认为 8。
  • **kwargs: 传递给 keras.layers.Layer 的其他关键字参数,包括 nametrainabledtype 等。

调用参数

  • attention_scores: 乘以 Transformer 多头注意力层的查询和键以添加 alibi 偏差的结果。形状为 (batch_size, num_heads, query_length, key_length)

示例

query_length = 10
key_length = 10
num_heads = 4
batch_size = 2
hidden_dim = 8

# Create new alibi layer.
alibi_layer = keras_hub.layers.AlibiBias()

query = np.zeros((batch_size, num_heads, query_length, hidden_dim))
key = np.zeros((batch_size, num_heads, hidden_dim, key_length))

attention_scores = keras.ops.matmul(query, key)

# Add alibi bias to attention scores.
attention_scores = alibi_layer(attention_scores)

参考文献