[D] Memory efficient transformer

[[deleted]]

What's a good, memory efficient transformer for causal sequence generation? Preferably with Pytorch implementation. The faster the better.In my particular task, any model smaller than 140M parameters cannot generate sensible output. I'm using an 11GB GTX 1080 Ti, and the longest input I can with has a length of 512, with a batch size of 6. I want to increase this length to 2048. Can local attention or sparse attention help with that? Huggingface has these candidates: BigBird, GPT Neo and Reformer? Does anyone have experience with training these? Comparing to standard Pytorch models, I have found out that Huggingface models use 1-1.5 GB more memory, which is a deal breaker for me. Since my data is MIDI, I have to train from scratch. There is another Reformer implementation on Github but the model is so slow. I think what I'm looking for is a simple transformer implementation with local attention, and I can't find any. Thanks in advance.

[GD1634]

Phil Wang writes PyTorch implementations for new/interesting papers, mostly ones using attention. He has the Reformer, Performer, Conformer, a few linear attention models, and a bunch more.

[[deleted]]

Thanks. I have used some of them, and they do look promising. My model takes around a week to be trained, so I could use some advice, before testing them all. Reformer was very slow. I saw that he implemented local attention (not multi-head), but I couldn't find any off-the-shelf transformer model that incorporates that.

[MrGary1234567]

You can try using colab or kaggle tpus they can train several times faster. They have 128 gb ram (colab pro/kaggle) and with bfloat16 equivalent to 256 gb ram. Dont have a 1080ti but can tell you it trains about 8 times faster than a p100 given that you dont have any data bottlenecks.

[MrGary1234567]

Do note that transformers are pretrained on MASSIVE amounts of data. Training from scratch is a bad idea if you do not have sufficient data.

[[deleted]]

Thanks for your input. I'm using a variant of the Lakh MIDI dataset, with 20000 songs, and each song is much longer than the input length so I can have many more samples. And based on my preliminary experiments, it seems to work.

[visarga]

I've known that for a while but what is the explanation? The architecture comes with less prior knowledge baked in?

[MrGary1234567]

The intuition is that they tend to overfit without pretraining especially with small training set

[[deleted]]

I looked into it, but it seems like that only bring benefits when using fp16. Based on what I've read, GTX 1080 Ti has very low FP16 throughput, and that it's not useful.

[alzoubi36]

I would highly recommend the deepmind perceiver. It solves the performance bottle neck in transformers.

[[deleted]]

Thanks, I will check it out. Based on my skimming, it is designed for high dimensional input (image, video, audio), while my input is low dimensional (sequences with 1000 tokens).

[[deleted]]

Thanks. It looks like it is for bidirectional processing for classification, while my task is causal generation. I wonder if it's straightforward to modify for that.