The top_k_attention from ag1988

ag1988 / top_k_attention Goto Github PK

View Code? Open in Web Editor NEW

The accompanying code for "Memory-efficient Transformers via Top-k Attention" (Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, Jonathan Berant. SustaiNLP 2021).

Python 97.22% Shell 2.78%

top_k_attention's Introduction

Memory-efficient Transformers via Top-k Attention

This repository contains the accompanying code for the paper:

"Memory-efficient Transformers via Top-k Attention." Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, Jonathan Berant. In SustaiNLP, 2021. [PDF]

Structure

The repository contains:

our implementation/benchmarking of top-k attention (in nocache_attention dir)
unifiedqa/T5 finetuning/inference using our top-k attention at feed-forward layers (in unifiedqa dir)

Coming Soon:

BERT QA model with top-k attention
T5 multi-head layers with top-k attention (current code is only for FF layers)

Citation

@inproceedings{gupta2021memoryefficient,
  title={Memory-efficient Transformers via Top-k Attention}, 
  author={Ankit Gupta and Guy Dar and Shaya Goodman and David Ciprut and Jonathan Berant},
  booktitle = {Proceedings of SustaiNLP: Workshop on Simple and Efficient Natural Language Processing},
  year={2021},
  publisher = {Association for Computational Linguistics}
}

top_k_attention's People