Luna: Linear Unified Nested Attention

TL;DR

<3-5 sentence high-level summary> Luna tackles the quadratic bottleneck of Transformer attention by introducing Linear Unified Nested Attention, which uses two nested linear attention passes and an auxiliary fixed-length input $P$ to pack and unpack context. This design yields total linear time and space complexity, while enabling variable-length sequences and autoregressive (causal) attention through a carefully designed causal variant. Across long-context sequence modeling, neural machine translation, and large-scale masked language modeling and fine-tuning, Luna achieves competitive or superior accuracy with substantial efficiency gains, particularly when projecting to small lengths like $l=16$. The approach also demonstrates that encoding contextual information in $P$ enhances performance, offering a practical, scalable alternative to traditional full-attention transformers for long sequences and pretraining regimes.

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

Figures (2)

• Figure 1: Trade-off between accuracy (y-axis), speed (x-axis) and memory (cir-radius) on LRA.
• Figure 2: Illustration of the architecture of one Transformer encoder layer (left) versus one Luna encoder layer (right).