[Paper Review] Luna: Linear Unified Nested Attention

Xuezhe Ma, Xiang Kong+5 more|arXiv (Cornell University)|Jun 3, 2021

Topic ModelingComputer Science48 references49 citations

TL;DR

Luna introduces a linear-time attention mechanism by using two nested linear attention functions with an extra fixed-length input sequence P, enabling efficient long-sequence modeling while retaining competitive performance across benchmarks.

ABSTRACT

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

Motivation & Objective

• Motivate the need to scale Transformers to long sequences due to quadratic attention complexity.
• Propose Luna to achieve linear time/space via pack and unpack attention with a fixed-length input P.
• Enable variable-length and autoregressive attention while maintaining efficiency and contextual capacity.
• Demonstrate Luna's effectiveness across long-context modeling, machine translation, and large-scale pretraining/fine-tuning.

Proposed method

• Introduce pack attention that compresses the context C into a fixed-length sequence P using the query P.
• Introduce unpack attention that expands the packed representation back to the original query length using X as the query in Attn(X, Y_P).
• Define LunaAttn(X, P, C) that outputs both Y_X and Y_P to carry contextual information through layers.
• Stack Luna layers with LayerNorm and FFN, applying normalization to Y_X and Y_P and FFN to Y_X only.
• Optionally share W_Q, W_K, W_V within a layer to reduce parameters (W_K tied to W_V).
• Present a Luna causal attention variant that preserves linear complexity with tailored activation functions (elu+1 or softplus) for packing, and softmax for unpacking.

Experimental results

Research questions

• RQ1 How can attention be reformulated to achieve linear time/space without sacrificing long-range contextual modeling?
• RQ2 Can an extra fixed-length input sequence P capture sufficient contextual information for variable-length inputs?
• RQ3 Does Luna support causal attention and autoregressive decoding while remaining efficient?
• RQ4 How does Luna perform on long-context benchmarks, MT, and large-scale pretraining relative to strong baselines?

Key findings

• Luna achieves competitive or better performance than full-attention baselines and other efficient attention methods across long-context sequence modeling, MT, and pretraining tasks.
• Luna demonstrates linear time and memory complexity due to pack/unpack attention, with efficiency gains shown against several baselines.
• Small projection lengths (e.g., l = 16) provide strong performance in LRA benchmarking and MT translation tasks.
• Contextual information encoded into P improves performance versus non-contextual P in MT experiments.
• In MT translation on WMT’14 EN→DE, Luna with softplus activation and shared keys/values shows competitive BLEU scores, with contextual P providing a notable boost over non-contextual P.
• Pretraining experiments indicate Luna can match or exceed baseline models on some QA and classification tasks, particularly with smaller data regimes; full-model capacity affects large-data results.