PaperlayerView exact revision on arXiv

Source-linked AI summary

Random Feature Attention

Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, Lingpeng Kong

arXiv:2103.02143v2cs.CL

TL;DR

Long-sequence transformers are limited by attention's quadratic time and space costs. The paper proposes RFA, which uses random features to approximate softmax attention with linear complexity and optional recency-biased gating. Across language modeling, machine translation, and long text classification, RFA matches or exceeds strong baselines while decoding twice as fast in machine translation.

  • Problem

    Transformer attention models pairwise interactions but has quadratic time and memory costs, limiting efficiency for long sequences.

  • Method

    RFA approximates softmax attention with random feature methods, supports linear time and space complexity, and optionally learns recency bias through gating.

  • Results

    RFA achieves comparable or better performance than strong baselines across three task types, including decoding around 2× faster than a transformer in machine translation.

  • Takeaways & Limitations

    RFA is a drop-in softmax-attention substitute whose efficiency improvements become more pronounced on longer sequences.

  • Takeaways & Limitations

    RFA trains around 15% slower than the baseline because of feature-map overhead.

Abstract

from arXiv · show

Transformers are state-of-the-art models for a variety of sequence modeling tasks. At their core is an attention function which models pairwise interactions between the inputs at every timestep. While attention is powerful, it does not scale efficiently to long sequences due to its quadratic time and space complexity in the sequence length. We propose RFA, a linear time and space attention that uses random feature methods to approximate the softmax function, and explore its application in transformers. RFA can be used as a drop-in replacement for conventional softmax attention and offers a straightforward way of learning with recency bias through an optional gating mechanism. Experiments on language modeling and machine translation demonstrate that RFA achieves similar or better performance compared to strong transformer baselines. In the machine translation experiment, RFA decodes twice as fast as a vanilla transformer. Compared to existing efficient transformer variants, RFA is competitive in terms of both accuracy and efficiency on three long text classification datasets. Our analysis shows that RFA's efficiency gains are especially notable on long sequences, suggesting that RFA will be particularly useful in tasks that require working with large inputs, fast decoding speed, or low memory footprints.

1 INTRODUCTION

Transformers rely on attention to model pairwise interactions, but its quadratic time and memory costs make long-sequence processing expensive. RFA addresses this by providing a linear-time, linear-space attention variant with optional recency-biased gating and competitive task performance.

  • Attention models pairwise interactions regardless of distance, but gives transformers quadratic time and memory costs, especially on long sequences.
  • RFA uses random feature methods to approximate softmax attention while scaling linearly in sequence length for both time and space.It approximates exp(x · y) with a feature-map inner product.
  • RFA and its gated variant are drop-in substitutes for softmax attention and add less than 0.1% more parameters.
  • RFA achieves comparable performance to vanilla transformer baselines across language modeling, machine translation, and long text classification.It outperforms a recent related efficient-attention approach in the reported tasks.
  • RFA decodes around 2× faster than a transformer baseline on machine translation without accuracy loss.The analysis reports a 12× decoding speedup using less than 10% of the memory for 2,048-length outputs.

2 BACKGROUND

Standard attention combines values using softmax weights over queries and keys, but full-sequence computation requires quadratic space and can require quadratic time during autoregressive decoding. Random feature methods approximate kernel evaluations through feature-map inner products, enabling an efficient softmax-attention approximation whose variance decreases as feature dimension increases.

  • For each query, attention linearly combines value vectors using weights produced by a softmax over keys.
  • A single-query attention computation takes O(M) time and space, while full-sequence space is O(MN) and autoregressive decoding can be quadratic in sequence length.
  • Random Fourier features transform vectors so their inner product approximates a desired shift-invariant kernel evaluation.The feature map uses sine and cosine components of randomly sampled projections.
  • The estimation variance decreases inversely with the random feature dimension D.
  • Random feature methods approximate exp(⟨x,y⟩) and thereby provide an efficient approximation to softmax attention.

3 MODEL

RFA replaces softmax attention with a random-feature approximation that preserves a drop-in interface while reducing attention computation to linear time and space. Its causal recurrence supports learned recency bias through gating, with experiments indicating efficiency gains for long sequences and competitive performance.

  • 3.1 RANDOM FEATURE ATTENTION: RFA approximates exp(x · y) with random feature maps, eliminating nonlinear interactions between transformed queries and keys.The approximation uses φ(x) · φ(y) as a kernel estimate of exp(x · y).
  • 3.1 RANDOM FEATURE ATTENTION: Causal RFA maintains a recurrent state of weighted values and feature sums, allowing attention over prefixes without increasing time or memory overhead.The states S_t and z_t summarize history and are updated incrementally at each timestep.
  • 3.2 RFA-GATE: LEARNING WITH RECENCY BIAS: RFA-GATE exponentially decays historical state using learned scalar gates, favoring more recent context.The gates are computed from the timestep input and applied to both the value state and normalization state.
  • 3.3 DISCUSSION: The method assumes normalized queries and keys for its basic formulation, although this norm-1 constraint is not mandatory and has little preliminary performance impact when σ is set or learned properly.
  • 3.3 DISCUSSION: RFA supports Gaussian and arc-cosine random-feature variants, which achieve similar performance in the reported experiments.
  • 3.4 COMPLEXITY ANALYSIS: RFA yields its strongest practical efficiency benefits on long or sequentially decoded inputs: 1.9× faster machine-translation decoding, 5.3× speedup at length 4,000, and lower memory overhead when 2D ≪ M.At 512-token context, the analysis reports no speedup; the benefit of feature maps larger than 2d is marginal.

4 EXPERIMENTS

Experiments evaluate RFA variants on language modeling, machine translation, and long text classification, using controlled transformer-based comparisons. RFA matches or improves accuracy while offering substantial decoding and efficiency benefits, with results depending on task and variant.

  • Experimental setup: RFA replaces softmax attention in transformer-based models and is evaluated on language modeling, machine translation, and long text classification.The experiments use comparable implementations, tuning, and training procedures; language-model RFA variants use 64-dimensional random feature maps.
  • Language modeling: RFA-GATE-Gaussian outperforms BASE by at least 1.2 test perplexity under both WikiText-103 model-size settings.Without gating, RFA variants outperform φelu but underperform BASE; gating improves RFA by more than 1.8 perplexity.
  • Language modeling: RFA-GATE-Gaussian gains more than 1.5 test perplexity when the last hidden state is passed between mini-batches.This stateful variant uses RFA’s RNN-style computation during training and evaluation.
  • Efficiency trade-offs: RFA trains around 15% slower than BASE because of additional feature-map overhead.This training-time cost contrasts with the decoding and long-sequence efficiency gains reported elsewhere.
  • Machine translation: All RFA machine-translation variants achieve similar BLEU performance to BASE while decoding more than 1.8× faster.The comparison covers three datasets, and the reported decoding speed is relative to BASE.
  • Long text classification: Across three long-text classification datasets, RFA outperforms the transformer by 0.3% average accuracy and speeds up by 1.1–5.3×.RFA is reported as the only model competitive in both accuracy and efficiency against comparable baselines.

5 ANALYSIS

RFA maintains nearly constant decoding speed and memory overhead as sequence length grows, unlike the baseline. Its analysis also examines evaluation settings and possible transfer from pretrained softmax transformers.

  • Experimental setting: The decoding comparison uses equal-length input and output sequences, six-layer encoders and decoders, greedy decoding, batch size 16, and a TPU v2.The compared models have the same size, with additional hyperparameters summarized in Appendix B.2.
  • Decoding efficiency: At sequence length 2,048, RFA decodes around 12× faster than the baseline while using less than 10% of its memory.Both RFA variants show nearly constant memory overhead across lengths, whereas the baseline slows and uses more memory for longer sequences.
  • Decoding efficiency: RFA-arccos slightly outperforms RFA-Gaussian in speed and memory efficiency when both use the same D.The φarccos feature map is half the size of φGaussian under this setting.
  • Additional analysis: Unconditional decoding shows overall speed and memory trends similar to those observed in sequence-to-sequence decoding.These results are reported in Figure 3 in Appendix C.1.
  • Experimental setting: The controlled batch size is less favorable to RFA, because its lower memory overhead could otherwise support larger batches and potentially greater speed gains.The authors retain equal batch sizes because this setting better reflects applications that generate one sequence at a time.
  • Additional analysis: An RFA model initialized from a pretrained softmax transformer achieves decent training loss after moderate fine-tuning, despite weak cross-evaluation between attention types.The analysis suggests transferring knowledge from a pretrained transformer to an RFA model as a potential application.

6 RELATED WORK

Efficient-transformer research commonly targets scaling to long sequences through sparse attention, compressed context, or memory modules. RFA instead attends beyond a fixed context window through stateful computation.

  • Efficiency goals: A shared motivation across efficient-transformer studies is scaling transformers to long sequences.The paper also identifies weight sharing, quantization, knowledge distillation, and adapters as orthogonal efficiency choices.
  • Sparse attention patterns: Sparse-attention methods improve efficiency by limiting attention’s reception field, using patterns that may be fixed in advance or learned from data.The passage notes that empirical verification remains incomplete for many such approaches.
  • Compressed context: Compressed-context methods reduce effective sequence length or store past context in a limited-size memory module.These approaches moderately increase overhead when accessing longer history.
  • Compressed context: RFA attends beyond a fixed context window through a stateful computation, reminiscent of recurrent neural-network language models.

7 CONCLUSION

The paper presents RFA as a linear-time, linear-space random-feature approximation to softmax attention with optional recency-biased gating. Across evaluated tasks, it matches or exceeds strong baselines and improves long-sequence decoding efficiency.

  • Contribution: RFA views softmax attention through kernel methods and approximates it with random feature methods.An optional gating mechanism provides a straightforward way to learn with recency bias.
  • Contribution: RFA has linear time and space complexity in sequence length and serves as a drop-in substitute for softmax attention in transformer models.
  • Results: Across language modeling, machine translation, and long text classification benchmarks, RFA achieves comparable or better performance than strong baselines.
  • Results: In machine translation, RFA decodes twice as fast, with further time and memory efficiency improvements possible for longer sequences.

Appendices

The appendices describe causal and cross RFA computation procedures using accumulated random-feature statistics. They also state the underlying attention setup and the random-feature approximation that enables efficient softmax attention.

  • Algorithms: Algorithms 1 and 2 specify causal and cross random feature attention computation procedures.
  • Causal RFA: Causal RFA maps each query and key to random features, accumulates key-value statistics and feature sums, then normalizes each output.The procedure iterates through the sequence and returns the collection of outputs.
  • Cross RFA: Cross RFA first accumulates key-value statistics and feature sums over encoder positions, then maps decoder queries and computes their outputs.The procedure returns the decoder output collection after processing all queries.
  • Kernel approximation: Random-feature methods derive an unbiased estimate of exp(⟨·,·⟩) and an efficient approximation to softmax attention.The appendix attributes the stated result to Yu et al. (2016).

A.3 DERIVATION OF CAUSAL RFA

Causal RFA computes attention over the prefix using recurrently maintained feature-weighted value and normalization states. The gated variant decays prefix keys and values, introducing a sequential dependency that can make training quadratic despite linear softmax normalization.

  • Causal RFA: Causal RFA restricts attention to prefix keys and values, enabling recurrent computation of its sufficient statistics.The derivation changes the attended keys and values to the prefix and computes both quantities recurrently.
  • Causal RFA: The recurrent states update by adding the current feature-weighted value and feature vector at each timestep.With S0 = 0 and z0 = 0, the updates are St = St−1 + φ(kt) ⊗vt and zt = zt−1 + φ(kt).
  • Without the norm constraint: Removing the unit-vector assumption adds a scalar norm-dependent term C(x) = exp(∥x∥2 /2σ2) to the attention computation.The resulting computation otherwise remains similar to the unit-norm derivation.
  • Gated RFA: Gated RFA decays all prefix keys and values before attention, so the next query cannot be processed until the current query finishes.This sequential dependency yields quadratic time in sequence length during language-model training, even though softmax normalization itself is linear.
  • Complexity comparison: In sequence-to-sequence models, RFA has lower space complexity and remains linear-time during teacher-forcing training and autoregressive decoding.The comparison assumes enough threads to parallelize softmax attention across timesteps during teacher forcing; decoding is separately characterized as linear for RFA and quadratic for softmax.

B EXPERIMENTAL DETAILS

The experiments use JAX implementations, controlled training procedures, and separate language-modeling and machine-translation configurations. Experimental comparisons vary random-matrix sampling while reporting dataset and model-training settings.

  • Implementation and datasets: The implementation is based on JAX, and dataset statistics are summarized separately for the experimental corpora.The dataset-statistics table distinguishes WikiText-103 split sizes in tokens from other datasets’ sizes in instances.
  • Random projections: Training samples a different random projection matrix for each attention head and draws matrices from an offline pool.The pool typically contains 200 matrices, avoiding Gaussian-sampling overhead during training; preliminary experiments favor this over one fixed projection.
  • Language modeling: Language-modeling models are trained for up to 150K Adam steps with early stopping on development perplexity.They use fixed sinusoidal position embeddings, no ℓ2 regularization, 16 TPU v3 accelerators for training, and a single TPU v2 accelerator for testing.
  • Machine translation: Machine-translation models are trained for up to 350K Adam steps with early stopping on development BLEU.The EN-DE and EN-FR experiments use different batch sizes, and no ℓ2 regularization or gradient clipping is used.

C.1 MORE RESULTS ON DECODING SPEED AND MEMORY OVERHEAD

Additional experiments examine RFA’s decoding and memory behavior, random-feature size, and transfer between softmax and RFA attention. They show that feature sizes must be sufficiently large for convergence and that pretrained attention parameters do not directly transfer, although RFA can recover pretraining loss through finetuning.

  • Decoding speed and memory: Figure 3 compares unconditional RFA decoding speed and memory overhead with softmax attention as output length varies.The experiment uses encoder-free models to simulate applications such as language-model sampling.
  • Random-feature size: RFA-Gaussian training fails to converge with feature sizes of 32 or 64 in cross attention and 32 in causal attention.Accuracy improves with 256 cross-attention features and 128 causal-attention features, while larger sizes provide marginal benefit.
  • Cross-function transfer: Without finetuning, softmax and RFA-Gaussian models initialized from one another achieve only 2.3 and 1.1 BLEU, respectively, versus more than 35.2 for pretrained models.This result suggests that RFA’s good performance is not due to learning to imitate softmax attention outputs.
  • Finetuning: Finetuned RFA reaches similar training loss to the pretrained model in roughly 1,500 steps, whereas resetting attention parameters reaches pretraining loss in less than 200 steps.The reset condition changes query, key, value, and output projections to random initialization while retaining the attention function.
  • Implications: RFA can recover pretraining loss, and finetuning costs much less computation than training from scratch, but retaining most pretrained knowledge remains a proposed future application.The paper identifies faster and more memory-friendly sampling as the intended benefit of such an RFA model.
Loading 2103.02143v2…