KerasHub:预训练模型 / API 文档 / 建模层 / RotaryEmbedding 层

RotaryEmbedding 层

[源]

RotaryEmbedding

keras_hub.layers.RotaryEmbedding(
    max_wavelength=10000, scaling_factor=1.0, sequence_axis=1, feature_axis=-1, **kwargs
)

旋转位置编码层。

此层使用旋转矩阵对绝对位置信息进行编码。它通过混合使用具有几何级增长波长的正弦和余弦函数来计算旋转编码。定义和提出于 RoFormer:增强型 Transformer 与旋转位置嵌入。输入必须是一个具有序列维度和特征维度的张量。通常,这将是一个形状为 (batch_size, sequence_length, feature_length)(batch_size, sequence_length, num_heads, feature_length) 的输入。此层将返回一个已应用旋转嵌入的新张量。

参数

  • max_wavelength:int。正弦/余弦曲线的最大角波长。
  • scaling_factor:float。用于缩放 token 位置的缩放因子。
  • sequence_axis:int。输入张量中的序列轴。
  • feature_axis:int。输入张量中的特征轴。
  • **kwargs:传递给 keras.layers.Layer 的其他关键字参数,包括 nametrainabledtype 等。

调用参数

  • inputs:要应用嵌入的张量输入。它可以具有任何形状,但必须同时包含序列轴和特征轴。旋转嵌入将应用于 inputs 并返回。
  • start_index:一个整数或整数张量。计算旋转嵌入的起始位置。这在缓存解码(cached decoding)期间很有用,其中每个位置在循环中单独预测。

示例

batch_size = 16
feature_length = 18
sequence_length = 256
num_heads = 8

# No multi-head dimension.
tensor = np.ones((batch_size, sequence_length, feature_length))
rot_emb_layer = RotaryEmbedding()
tensor_rot = rot_emb_layer(tensor)

# With multi-head dimension.
tensor = np.ones((batch_size, sequence_length, num_heads, feature_length))
tensor_rot = rot_emb_layer(tensor)

参考