tfm.nlp.layers.ReuseTransformer

Transformer layer.

tfm.nlp.layers.ReuseTransformer(
    num_attention_heads,
    inner_dim,
    inner_activation,
    head_size=None,
    output_range=None,
    kernel_initializer='glorot_uniform',
    bias_initializer='zeros',
    kernel_regularizer=None,
    bias_regularizer=None,
    activity_regularizer=None,
    kernel_constraint=None,
    bias_constraint=None,
    use_bias=True,
    norm_first=False,
    norm_epsilon=1e-12,
    output_dropout=0.0,
    attention_dropout=0.0,
    inner_dropout=0.0,
    attention_initializer=None,
    attention_axes=None,
    reuse_attention=0,
    use_relative_pe=False,
    pe_max_seq_length=512,
    layer_idx=None,
    max_reuse_layer_idx=None,
    **kwargs
)

This layer implements the ReuseTransformer Encoder from "Leveraging redundancy in attention with Reuse Transformers". (https://arxiv.org/abs/2110.06821)

Args
num_attention_heads	Number of attention heads.
inner_dim	The output dimension of the first Dense layer in a two-layer feedforward network.
inner_activation	The activation for the first Dense layer in a two-layer feedforward network.
head_size	Projection size of heads.
output_range	the sequence output range, [0, output_range) for slicing the target sequence. None means the target sequence is not sliced.
kernel_initializer	Initializer for dense layer kernels.
bias_initializer	Initializer for dense layer biases.
kernel_regularizer	Regularizer for dense layer kernels.
bias_regularizer	Regularizer for dense layer biases.
activity_regularizer	Regularizer for dense layer activity.
kernel_constraint	Constraint for dense layer kernels.
bias_constraint	Constraint for dense layer kernels.
use_bias	Whether to enable use_bias in attention layer. If set False, use_bias in attention layer is disabled.
norm_first	Whether to normalize inputs to attention and intermediate dense layers. If set False, output of attention and intermediate dense layers is normalized.
norm_epsilon	Epsilon value to initialize normalization layers.
output_dropout	Dropout probability for the post-attention and output dropout.
attention_dropout	Dropout probability for within the attention layer.
inner_dropout	Dropout probability for the first Dense layer in a two-layer feedforward network.
attention_initializer	Initializer for kernels of attention layers. If set None, attention layers use kernel_initializer as initializer for kernel.
attention_axes	axes over which the attention is applied. None means attention over all axes, but batch, heads, and features.
reuse_attention	reuse_attention: An integer specifying number of heads to reuse. -1 for all heads.
use_relative_pe	whether to use relative position bias.
pe_max_seq_length	used to set the size of the relative positin encodings.
layer_idx	the idx of this layer.
max_reuse_layer_idx	layer idx (if passed) greater than this value will not reuse attention scores from previous layers.
**kwargs	keyword arguments.

Methods

call

View source

call(
    inputs
)

Transformer self-attention encoder block call.

Args
inputs	a single tensor or a list of tensors. input tensor as the single sequence of embeddings. [input tensor, attention mask] to have the additional attention mask. [query tensor, attention mask, attention scores] to have additional attention scores for reuse computation. If attention scores is None, the reuse_attention flag will be ignored.

Returns
An output tensor with the same dimensions as input/query tensor. Attention scores if return_attention_scores is true.