Module fast_transformers.attention

Implementations of different types of attention mechanisms.

Expand source code
#
# Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
# Written by Angelos Katharopoulos <angelos.katharopoulos@idiap.ch>,
# Apoorv Vyas <avyas@idiap.ch>
#

"""Implementations of different types of attention mechanisms."""


from .attention_layer import AttentionLayer
from .full_attention import FullAttention
from .linear_attention import LinearAttention
from .causal_linear_attention import CausalLinearAttention
from .clustered_attention import ClusteredAttention
from .improved_clustered_attention import ImprovedClusteredAttention
from .reformer_attention import ReformerAttention
from .conditional_full_attention import ConditionalFullAttention
from .exact_topk_attention import ExactTopKAttention
from .improved_clustered_causal_attention import ImprovedClusteredCausalAttention
from .local_attention import LocalAttention

Sub-modules

fast_transformers.attention.attention_layer

The base attention layer performs all the query key value projections and output projections leaving the implementation of the attention to the inner …

fast_transformers.attention.causal_linear_attention

Implement causally masked linear attention.

fast_transformers.attention.clustered_attention

Implement clustered self attention.

fast_transformers.attention.conditional_full_attention

Implement a self attention that delegates to full attention or another attention depending on the input sequence length.

fast_transformers.attention.exact_topk_attention

Implement the oracle top-k attention. The top-k keys are exact ones. MultiHeadAttention module. Note that this module is to be used in conjuction with …

fast_transformers.attention.full_attention

Implement the full attention similar to the one implemented by PyTorch's MultiHeadAttention module. Note that this module is to be used in conjuction …

fast_transformers.attention.improved_clustered_attention

Implement improved clustered self attention.

fast_transformers.attention.improved_clustered_causal_attention

Implement improved clustered causal self attention.

fast_transformers.attention.linear_attention

Implement unmasked linear attention.

fast_transformers.attention.local_attention

Implement local context attention.

fast_transformers.attention.reformer_attention

Implement the Reformer attention from the paper "Reformer the efficient transformer".