Feedback Transformer - Pytorch
Simple implementation of Feedback Transformer in Pytorch. They improve on Transformer-XL by having each token have access to the representations of all previous layers through time. This is achieved by aggregating the outputs of all layers into a shared memory, which each token across layers can attend to at each time step.
The main drawback is longer training time, due to its non-parallel nature. But I thought I'd build it to further exploration and research into this line of work.
I also took the liberty to add some various enhancements, including pre-normalization, GLU gated feedforwards, as well as simplified T5 relative positional embeddings.
Install
$ pip install feedback-transformer-pytorch
Usage
import torch
from feedback_transformer_pytorch import FeedbackTransformer
model = FeedbackTransformer(
num_tokens = 20000, # number of tokens
dim = 512, # dimension
depth = 6, # depth
seq_len = 2, # the sequence length of each segment or window
mem_len = 256, # length of the memory buffer
dim_head = 64, # dimension of each head
heads = 8, # number of heads
attn_dropout = 0.1, # attention dropout
ff_dropout = 0.1 # feedforward dropout
).cuda()
x = torch.randint(0, 20000, (2, 64)).cuda()
model(x) # (2, 64, 20000)
If you would like to have fine control over the memory (when to detach, etc), you can do it with some extra keyword arguments on .forward
import torch
from feedback_transformer_pytorch import FeedbackTransformer
model = FeedbackTransformer(
num_tokens = 20000,
dim = 512,
depth = 6,
seq_len = 32,
mem_len = 256
).cuda()
x1 = torch.randint(0, 20000, (2, 32)).cuda()
x2 = torch.randint(0, 20000, (2, 32)).cuda()
x3 = torch.randint(0, 20000, (2, 32)).cuda()
out1, mem1 = model(x1, return_memory = True)
out2, mem2 = model(x2, memory = mem1, return_memory = True)
out3, mem3 = model(x3, memory = mem2, return_memory = True) # (2, 32, 20000)
Citations
@misc{fan2021addressing,
title = {Addressing Some Limitations of Transformers with Feedback Memory},
author = {Angela Fan and Thibaut Lavril and Edouard Grave and Armand Joulin and Sainbayar Sukhbaatar},
year = {2021},
eprint = {2002.09402},
archivePrefix = {arXiv},
primaryClass = {cs.LG}
}
46 Sep 20, 2022
2 Mar 15, 2022
0 Jul 21, 2022
28 Dec 13, 2022
0 Sep 24, 2021
3 Aug 17, 2022
34 Oct 02, 2022
74 Dec 05, 2022
4 Feb 09, 2022
145 Sep 28, 2022
8 Sep 05, 2022
51 Nov 29, 2022
70 Nov 16, 2022
607 Jan 07, 2023
485 Dec 26, 2022
64 Nov 16, 2022
15 Oct 17, 2022
119 Jan 08, 2023
14 Dec 27, 2022
177 Dec 01, 2022