Linearizing Transformer with Key-Value Memory
This work addresses efficiency and accuracy issues in transformer models for sequence generation tasks, offering a practical improvement over existing variants.
The authors tackled the performance drop and inefficiency of existing linear-time transformers on short sequences by proposing MemSizer, which combines low-rank projection with recurrent-style incremental computation, achieving linear time and constant memory at inference while improving accuracy in tasks like machine translation and text summarization.
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.