MEMO : Accelerating Transformers with Memoization on Big Memory Systems

Abstract

Transformers gain popularity because of their superior prediction accuracy and inference throughput. However, the transformer is computation-intensive, causing a long inference time. The existing work to accelerate transformer inferences has limitations because of the changes to transformer architectures or the need for specialized hardware. In this paper, we identify the opportunities of using memoization to accelerate the attention mechanism in transformers without the above limitation. Built upon a unique observation that there is a rich similarity in attention computation across inference sequences, we build an attention database upon the emerging big memory system. We introduce the embedding technique to find semantically similar inputs to identify computation similarity. We also introduce a series of techniques such as memory mapping and selective memoization to avoid memory copy and unnecessary overhead. We enable 21% performance improvement on average (up to 68%) with the TB-scale attention database and with ignorable loss in inference accuracy.

LUNet: An enhanced upsampling fusion network with efficient self-attention for semantic segmentation

“…Recent studies have suggested various strategies [13,16,[37][38][39][40] to improve the computational efficiency of self-attention. These strategies include: (1) Diminishing the dimension of input by utilizing low-rank or compressed embeddings [41][42][43]; (2) Implementing local or sparse self-attention mechanisms, such as convolutional layers [41,44] or shifted window [36,45], which can focus on nearby or relevant elements.…”