Linearizing Transformer with Key-Value Memory

TL;DR

MemSizer is a novel transformer variant that combines low-rank projection and kernel-based methods to achieve linear time complexity, constant memory, and improved performance on sequence generation tasks.

Contribution

It introduces MemSizer, a new approach that enhances efficiency and accuracy of transformers by integrating low-dimensional projections with recurrent-style incremental computation.

Findings

MemSizer achieves linear inference time and constant memory.

It outperforms vanilla and other efficient transformers in translation, summarization, and language modeling.

MemSizer maintains high accuracy even with short sequence generation.

Abstract

Efficient transformer variants with linear time complexity have been developed to mitigate the quadratic computational overhead of the vanilla transformer. Among them are low-rank projection methods such as Linformer and kernel-based Transformers. Despite their unique merits, they usually suffer from a performance drop comparing with the vanilla transformer on many sequence generation tasks, and often fail to obtain computation gain when the generation is short. We propose MemSizer, an approach towards closing the performance gap while improving the efficiency even with short generation. It projects the source sequences into lower dimension representations like Linformer, while enjoying efficient recurrent-style incremental computation similar to kernel-based transformers. This yields linear computation time and constant memory complexity at inference time. MemSizer also employs a lightweight multi-head mechanism which renders the computation as light as a single-head model. We demonstrate that MemSizer provides an improved balance between efficiency and accuracy over the vanilla transformer and other efficient transformer variants in three typical sequence generation tasks, including machine translation, abstractive text summarization, and language modeling.