[Paper Review] Random Feature Attention

Hao Peng, Nikolaos Pappas+4 more|arXiv (Cornell University)|Mar 3, 2021

Topic Modeling70 references122 citations

TL;DR

Rfa replaces softmax attention with a linear-time, linear-space random feature approximation, with an optional gate for recency bias, achieving comparable or better performance to strong transformers and faster decoding in MT."

ABSTRACT

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.

Motivation & Objective

• Motivate scalable attention for long sequences in transformers.
• Propose Random Feature Attention (Rfa) as a linear-time, linear-space substitute for softmax attention.
• Incorporate an optional gating mechanism to introduce recency bias during learning.
• Demonstrate Rfa's effectiveness across language modeling, machine translation, and long text classification.

Proposed method

• Derive an unbiased kernel-based approximation of exp(q·k/σ^2) using random feature maps φ to approximate softmax attention.
• Rewrite attention as φ(q)ᵀ S / (φ(q)· z) with S and z accumulated via φ(k)⊗v and φ(k), enabling linear-time computation.
• Introduce Rfa-Gate, a recurrent-like gating mechanism g_t to smoothly decay history and encode recency bias.
• Allow Rfa to be a drop-in replacement for softmax attention with minimal parameter increase (≈0.1%).
• Explore Gaussian and arc-cosine random feature maps as φ, with normalization considerations for q and k norms.

Experimental results

Research questions

• RQ1 How can attention be approximated to scale linearly with sequence length without sacrificing performance?
• RQ2 Does random feature-based attention (Rfa) match or exceed the performance of standard softmax attention across language modeling, translation, and long-sequence classification?
• RQ3 Can a gating mechanism in Rfa capture recency bias and improve performance on tasks requiring locality?
• RQ4 What are the speedup and memory benefits of Rfa during decoding and for long inputs compared to vanilla transformers?

Key findings

• Rfa achieves comparable or better perplexity than the base transformer in WikiText-103, with gating providing notable improvements.
• In machine translation benchmarks, all Rfa variants decode at least ~1.8× faster than the Base transformer, with BLEU scores comparable to Base.
• On long text classification tasks, Rfa is competitive in accuracy and offers speed/memory advantages over several efficient transformer variants.
• Rfa demonstrates substantial decoding speedups (up to 12× for 2048-length outputs) and lower memory usage for longer sequences.
• The Gaussian feature map generally yields more stable training and better performance than arc-cosine in language modeling experiments.
• The gating variant (Rfa-Gate) provides benefits in language modeling, especially on WikiText-103.