Luna: Linear Unified Nested Attention

By Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, Luke S. Zettlemoyer. 2021

140 citations recorded by the source corpus.

The quadratic computational and memory complexities of the Transformer's attention mechanism have limited its scalability for modeling long sequences. In this paper, we propose Luna, a linear unified nested attention mechanism that 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 sequence into a sequence of fixed length. Then, the packed sequence is unpacked using the second attention function. As compared to a more traditional attention mechanism, Luna introduces an additional sequence with a fixed length as input and an additional corresponding output, which allows Luna to perform attention operation linearly, while also storing adequate contextual information. We perform extensive evaluations on three benchmarks of sequence modeling tasks: long-context sequence modeling, neural machine translation and masked language modeling for large-scale pretraining. Competitive or even better experimental results demonstrate both the effectiveness and efficiency of Luna compared to a variety

Subfields: Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Model Architectures

1. Christopher R'e
2. Tri Dao
3. Yi Tay
4. Mostafa Dehghani
5. Xipeng Qiu
6. Donald Metzler
7. Stefano Ermon
8. Atri Rudra

Citations in the corpus

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, background): “Other approaches include compressing along the sequence dimension to attend to multiple tokens at once [50, 55, 76, 86].”
• Self-attention Does Not Need $O(n^2)$ Memory (2021, methodology): “…and architectures with more favorable complexity classes (Kitaev et al., 2020; Roy et al., 2021; Zaheer et al., 2020; Choromanski et al., 2020; Wang et al., 2020; Ren et al., 2021; Child et al., 2019; Tay et al., 2021; Wang et al., 2020; Ma et al., 2021; Shen et al., 2021; Qiu et al., 2020).”
• Perceiver IO: A General Architecture for Structured Inputs & Outputs (2021, background): “Cross-attention itself is in widespread use as a mechanism for allowing information from a source to influence processing of a target domain of a different size or structure [18, 20, 53, 88, 64, 68, 34, 51].”
• Efficient Transformers: A Survey (2020, background): “, 2020b) O(N logN) 7 LP Luna (Ma et al., 2021) O(kN) X M TokenLearner (Ryoo et al.”
• General-purpose, long-context autoregressive modeling with Perceiver AR (2022, methodology): “, 2019) and Luna (Ma et al., 2021), that use differently shaped query and key-value inputs to limit the cost of attention and produce “downsampled” intermediate representations that are smaller than the input’s size.”
• Harnessing the Power of LLMs in Practice: A Survey on ChatGPT and Beyond (2023, methodology): “There is a push towards devising efficient attention mechanisms, such as linear nested attention [80], Transient Global attention [44], CoLT5 [4], and Synthesizer [111].”
• LongT5: Efficient Text-To-Text Transformer for Long Sequences (2021, unknown): “Another type of approaches is leveraging the low-rank approximation of the attention matrix, such as Lin-former (Wang et al., 2020), Performer (Choroman-ski et al., 2021), Random Feature Attention (Peng et al., 2021), and LUNA (Ma et al., 2021).”
• A Survey of Transformers (2021, background): “Set Transformer [70] and Luna [90] use a number of external trainable global nodes to summarize information from inputs and then the summarized representations serve as a compressed memory that the inputs attend to.”
• Scatterbrain: Unifying Sparse and Low-rank Attention Approximation (2021, background): “• Luna [43] (concurrent work): they use a fixed-length extra sequence and two consecutive attention steps: the context sequence attends to the extra sequence, and then the query sequence attends to the extra sequence.”
• RWKV: Reinventing RNNs for the Transformer Era (2023, methodology): “…including sparse attention (Beltagy et al., 2020; Kitaev et al., 2020; Guo et al., 2022), approximating the full attention matrix (Wang et al., 2020; Ma et al., 2021; Choromanski et al., 2020), combining chunked attention with gating (Ma et al., 2023) and other efficient methods (Katharopoulos et…”
• Simplified State Space Layers for Sequence Modeling (2022, background): “[39] Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer.”
• FLatten Transformer: Vision Transformer using Focused Linear Attention (2023, background): “Although enjoying linear computational complexity, various previous works have also proved that simply replacing Softmax attention with linear attention usually results in severe performance drop [34, 2, 7, 27].”