In this task, the transformer learns to copy a random sequence of symbols. This benchmark requires very sparse random attention and is a hard task for dense attentions like linear attention.

The transformer reads, one element at a time, a sequence like the following

0 t1 t2 t3 t4 0 t1 t2 t3 t4

and must predict the next element. The loss is taken in the second half of the sequence as t1-t4 are completely random tokens but they are repeated.

The installation requirements are the following:

• torch
• pytorch-fast-transformers

They can be installed in most systems via

pip install torch pytorch-fast-transformers

The main.py has thorough command line help that can be invoked via the --help command line argument. The attention implementation is chosen via the --attention_type argument.

python main.py  # Runs with the causal-linear attention_type
python main.py --attention_type full  # Uses softmax attention
python main.py --attention_type reformer  # Uses reformer or lsh attention
                                          # see also --rounds for changing the
                                          # hashing rounds

Other things to try are increasing the maximum length of the sequence

python main.py --sequence_length 512

or playing around with the transformer parameters like --rounds and --bits for LSH attention.