tfm.nlp.layers.RelativePositionBias

Relative position embedding via per-head bias in T5 style.

tfm.nlp.layers.RelativePositionBias(
    num_heads: int,
    relative_attention_num_buckets: int = 32,
    relative_attention_max_distance: int = 128,
    bidirectional: bool = True,
    embeddings_initializer: Optional[Initializer] = None,
    **kwargs
)

Reference implementation in MeshTF: https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L1000

This layer implements the relative position bias used in "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer" (https://arxiv.org/abs/1910.10683)

Methods

call

View source

call(
    query: tf.Tensor, key: tf.Tensor
)

Implements the forward pass.

Args
query query input tensor shape [batch, query length, hidden size].
key key input tensor shape [batch, key length, hidden size].

Returns
A tensor in shape of [batch, heads, query length, key length].