Luna: Linear Unified Nested Attention

  • 2021-06-03 01:47:26
  • Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, Luke Zettlemoyer
  • 46

Abstract

The quadratic computational and memory complexities of the Transformer'sattention mechanism have limited its scalability for modeling long sequences.In this paper, we propose Luna, a linear unified nested attention mechanismthat approximates softmax attention with two nested linear attention functions,yielding only linear (as opposed to quadratic) time and space complexity.Specifically, with the first attention function, Luna packs the input sequenceinto a sequence of fixed length. Then, the packed sequence is unpacked usingthe second attention function. As compared to a more traditional attentionmechanism, Luna introduces an additional sequence with a fixed length as inputand an additional corresponding output, which allows Luna to perform attentionoperation linearly, while also storing adequate contextual information. Weperform extensive evaluations on three benchmarks of sequence modeling tasks:long-context sequence modeling, neural machine translation and masked languagemodeling for large-scale pretraining. Competitive or even better experimentalresults demonstrate both the effectiveness and efficiency of Luna compared to avariety

 

Quick Read (beta)

loading the full paper ...