PaperlayerView exact revision on arXiv

Source-linked AI summary

Luna: Linear Unified Nested Attention

Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, Luke Zettlemoyer

arXiv:2106.01540v2cs.LGcs.CL

TL;DR

Transformer attention’s quadratic time and memory costs limit its use on long sequences. Luna replaces regular attention with nested linear operations that pack context into a fixed-length sequence and unpack it; across sequence-modeling tasks, it reports comparable or better performance with speed and memory gains.

  • Problem

    Transformer attention has quadratic time and memory complexity, limiting its scalability for modeling long sequences.

    Source: §1 · p. 1, §2.1 · p. 2

  • Method

    Luna approximates regular softmax attention with nested pack and unpack attention functions, using a fixed-length sequence to represent context.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    Across long-context modeling, machine translation, and large-scale pretraining and finetuning, Luna achieves comparable or better performance than strong baselines while gaining speed and memory efficiency.

    Source: §6 · p. 9

  • Takeaways & Limitations

    Luna is presented as a drop-in substitute for softmax attention that captures contextual information with linear attention operations.

    Source: §6 · p. 9

    Luna causal attention requires sequential computation at each timestep, despite retaining linear complexity.

    Source: §3.3 · p. 4

Abstract

from arXiv · show

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

1 Introduction

Transformer attention captures sequence context through pairwise interactions, but its quadratic time and memory costs restrict long-sequence applications. Luna uses nested linear attention and reports competitive or better performance with speed and memory gains.

1 Introduction

  • Transformer attention models pairwise interactions across sequence inputs, but its quadratic time and memory complexity restricts applications requiring longer inputs.Source: §1 · p. 1
  • Luna reports competitive or better performance and prominent speed and memory gains across three sequence-modeling tasks, including superior performance with projection length 16.Figure 1 compares accuracy, speed, and memory on LRA.Source: Figure 1 · p. 1, §1 · p. 2
  • Prior efficiency approaches include restricting attention with sparsity or approximating full-rank attention through low-rank projections and kernelization.The cited discussion says these methods' efficiency gains are less prominent for moderate sequences and their performance remains behind regular-attention Transformers.Source: §1 · pp. 1–2
  • Luna uses two nested attention functions to pack context into a fixed-length sequence and then unpack it, adding a fixed-length input and corresponding output.The extra input supports linear attention while storing contextual information.Source: §1 · p. 2

2 Background

The background defines attention as mapping queries and context to an output, with pairwise alignment scores imposing quadratic cost. It also lays out the Transformer layer’s residual, feed-forward, and normalization operations.

2.1 Attention

  • The attention function takes query and context sequences and returns an output sequence matching the query length; its projections produce queries, keys, and values.Source: §2.1 · p. 2
  • The attention matrix records alignment scores between query and context tokens, and calculating it takes O(nm) time and space.This becomes quadratic with sequence length when the two lengths scale together.Source: §2.1 · p. 2

2.2 Transformer Layers

  • A Transformer layer applies attention and a position-wise feed-forward network with residual connections and layer normalization after each residual addition.The described layer uses post-layer normalization rather than pre-layer normalization.Source: §2.2 · p. 2

3 Linear Unified Nested Attention (Luna)

Luna replaces full attention with two nested linear-efficiency attention operations: one packs context into a fixed-length sequence and the other unpacks it for the original query. The section develops this design for standard and causal attention and relates it to Linformer and Set Transformer.

3.1 Pack and Unpack Attention

  • Pack attention uses fixed-length P to compress context C into YP, and unpack attention uses the packed sequence to return an output aligned with the original query length.Their complexities are O(lm) and O(ln), respectively, where l is fixed.Source: §3.1 · p. 3
  • Luna carries contextual information forward by making YP an additional layer output that serves as the next layer's P input.The first layer initializes P with learnable positional embeddings.Source: §3.1 · pp. 3–4

3.2 Luna Layers

  • Luna layers apply layer normalization to both attention outputs but the feed-forward network only to YX, while sharing WQ and WK within each layer to reduce parameters.The paper evaluates the performance decline against Luna with full parameter sets.Source: §3.2 · p. 4, §3.1 · p. 4

3.3 Luna Causal Attention

  • Causal Luna requires P not to contain information about X that could leak future tokens, and its causal function restricts each timestep to past and present rows.For pack attention, the paper uses elu(·)+1 or softplus(·) rather than softmax; unpack attention can use softmax because it normalizes over the fixed-length dimension.Source: §3.3 · p. 4
  • Causal Luna retains O(ln) complexity but computes sequentially at each timestep; P may come from the Luna encoder in encoder-decoder models or be learned per layer in decoder-only models.Sinusoidal positional embeddings produced similar results to the learnable embeddings.Source: §3.3 · p. 4

3.4 Discussion

  • Compared with Linformer, Luna's pack attention supports variable-length inputs and contextualized projections, and experiments report improvements over non-contextual projections.Linformer's learned projection matrix depends on the input length.Source: §3.4 · p. 5
  • Luna's P resembles Set Transformer's inducing-point side memory, but Luna is a drop-in replacement for softmax attention rather than a model designed for permutation-invariant set inputs.The inducing points can be viewed as non-contextual versions of P.Source: §3.4 · p. 5

4 Experiments

Across long-context modeling, translation, and language-model pretraining, Luna shows competitive or stronger performance, with efficiency and design outcomes varying by setting. The experiments also identify contextual P, activation choice, and model capacity as relevant to results.

4.1 Long-Context Sequence Modeling

  • Luna significantly outperforms all baselines on average LRA accuracy, wins three of five tasks, and matches the best-performing model on the other two.The five tasks span 1K–8K-token inputs; Retrieval was trained for 20K steps because models had not converged after 5K.Source: §4.1 · pp. 5–6
  • At 4K input length, Luna-16 uses 10% of Transformer memory and runs 1.2–5.5 times faster; Luna retains superior LRA performance at projected length 16.Luna-16 is comparable in speed to Performer and Linformer; the reported speedup varies with sequence length.Source: §4.1 · p. 6

4.2 Machine translation

  • On WMT’14 EN→DE, Luna-16 matches RFA while Luna-32 remains below Transformer-base; softplus consistently beats elu(·) + 1.Sharing causes a marginal performance decline; shared Luna has 7% more parameters than Transformer and 5% more than RFA.Source: §4.2 · p. 7, Table 4 · p. 7
  • On WMT’14, Luna with contextual P significantly outperforms the baseline with non-contextual P.Both models use projected length 16, parameter sharing, and softplus activation.Source: §4.2 · p. 8

4.3 Masked Language Modeling for Large-Scale Pretraining

  • With 16GB pretraining data, Luna has similar or slightly better downstream results than other pretrained models and leads on QNLI and SST-2; with 160GB, it is slightly worse than RoBERTa.The authors suggest capacity may explain the larger-corpus gap; full-parameter Luna performs better than its parameter-shared version.Source: §4.3 · p. 8

5 Related Work

Related work surveys approaches to Transformer efficiency, including sparsity, kernel-based attention, recurrence, and other model or context-compression techniques. It also notes that Luna can use its P sequence as memory across segments.

5 Related Work

  • The surveyed efficiency approaches include weight sharing, quantization, sparse attention, side memory, and low-rank or compressed context.The overview directs readers to Tay et al. (2020b) for a detailed review.Source: §5 · pp. 8–9
  • Sparse attention restricts each token to a fixed range, while some methods use hashing, online clustering, or learned sorting to make the range more flexible.These methods include local, blockwise, strided, and compressed attention patterns.Source: §5 · p. 9
  • Kernel methods avoid explicitly computing the attention matrix by rewriting attention with kernels; they can also be viewed as low-rank context compression.The passage relates this view to Linformer and Luna.Source: §5 · p. 9
  • Transformer-XL (Dai et al., 2019) connects input blocks through recurrence; Luna can likewise use P as memory across segments.Compressive Transformer extends Transformer-XL by retaining fine-grained memory of past chunk activations.Source: §5 · p. 9

6 Conclusion

Luna is a linear drop-in substitute for softmax attention that retains contextual information with a fixed-length extra input. Across three sequence-modeling tasks, it achieves comparable or better performance than strong baselines while improving speed and memory efficiency.

6 Conclusion

  • Across long-context sequence modeling, neural machine translation, and large-scale pretraining and finetuning, Luna matches or exceeds strong baselines while gaining speed and memory efficiency.The conclusion characterizes Luna as a simple, efficient, effective drop-in substitute for regular softmax attention.Source: §6 · p. 9
  • Luna introduces an extra fixed-length input to capture adequate contextual information while performing attention operations linearly.The authors identify combining Luna with recurrence, using P as running memory across input segments, as future work.Source: §6 · p. 9
  • The authors identify applying Luna to long-input tasks such as document-level summarization and translation as a future direction.They also propose combining Luna with recurrence methods, where P could serve as running memory across input segments.Source: §6 · p. 9

A.1 Long-Context Sequence Modelling

For most long-context sequence-modeling tasks, the study follows Tay et al. (2021) configurations; Retrieval uses a longer training schedule to guarantee convergence.

A.1 Long-Context Sequence Modelling

  • Retrieval models are trained for 20k steps rather than the 5k steps prescribed by Tay et al. (2021) to guarantee convergence.The passage says other tasks closely follow Tay et al. (2021) for preprocessing, data splits, architecture, and batch size.Source: §A.1 · p. 14
  • The study mainly tunes learning rate, dropout, and attention dropout for these tasks.Batch size, layer count, and warmup steps follow Tay et al. (2021).Source: §A.1 · p. 14
  • Hyperparameters for the long-context sequence-modeling tasks are listed in Table 7.The passage identifies the table as containing task-model hyperparameters.Source: §A.1 · p. 14

A.2 Neural Machine Translation

The machine-translation experiments use a Transformer-base setup on WMT 2014 English–German, with standard validation and test splits and specified training and optimization configurations.

A.2 Neural Machine Translation

  • Table 7 is captioned as listing hyperparameters for models in the long-range arena tasks.The caption describes LR and Attn-Dropout as denoting learning, batch size, and attention dropout.Source: Table 7 · p. 14
  • WMT 2014 English–German experiments use 4.5M training sentence pairs, Newstest2013 for validation, and Newstest2014 for testing.The models use 512-dimensional embeddings, six-layer encoder and decoder stacks, eight attention heads, and 2048 feed-forward dimensions.Source: §A.2 · p. 14
  • Each model is trained for 500,000 updates with label smoothing of 0.1 and dropout of 0.1.Training uses distributed runs on eight NVIDIA Tesla V100 GPUs, with a maximum batch of 8192 tokens per GPU.Source: §A.2 · p. 14

A.3 Masked Language Modeling for Large-Scale Pretraining and Finetuing

The pretraining experiments use the masked-language-modeling objective on two corpora, followed by task-specific finetuning configurations for GLUE, RACE, and CSQA.

A.3 Masked Language Modeling for Large-Scale Pretraining and Finetuing

  • Pretraining uses the masked-language-modeling objective on a 16GB BookCorpus-and-English-Wikipedia corpus and a 160GB RoBERTa corpus.The larger corpus also includes CC-News, OpenWebText, and Stories; models are trained on 64 Tesla V100 GPUs.Source: §A.3 · p. 14
  • Pretraining uses Adam with a linear-decay learning-rate scheduler.Table 8 is identified as listing hyperparameters for pretraining Luna-128 on the two public corpora.Source: §A.3 · p. 14, Table 8 · p. 15
  • Finetuning follows the released RoBERTa finetuning script for different tasks, with main hyperparameters listed in Table 9.Table 9 covers finetuning Luna on GLUE, RACE, and CSQA.Source: §A.3 · p. 14, Table 9 · p. 15
Loading 2106.01540v2…