RotaryEmbedding layer

[源代码]

RotaryEmbedding 类

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

旋转位置编码层。

该层使用旋转矩阵对绝对位置信息进行编码。它使用具有几何递增波长的正弦和余弦函数的组合来计算旋转编码。在 RoFormer: Enhanced Transformer with Rotary Position Embedding 中定义和制定。输入必须是一个张量，其中包含一个序列维度和一个特征维度。通常，输入形状将是 (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 的其他关键字参数，包括 name、trainable、dtype 等。

调用参数

• inputs: 要应用嵌入的张量输入。它可以是任何形状，但必须包含序列轴和特征轴。旋转嵌入将应用于 inputs 并返回。
• start_index: 一个整数或整数张量。用于计算旋转嵌入的起始位置。这在缓存解码期间很有用，因为每个位置都是在循环中单独预测的。
• positions: 形状为 (sequence_length,) 或 (batch_size, sequence_length) 的张量。输入序列的自定义位置。如果指定，将使用此张量来计算旋转嵌入，并且 start_index 参数将被忽略。这对于非标准位置的情况很有用。

示例

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)

参考文献

• RoFormer: Enhanced Transformer with Rotary Position Embedding