[Paper Review] Attention as an RNN

Leo Feng, Frederick Tung+4 more|arXiv (Cornell University)|May 22, 2024

Educational and Psychological AssessmentsPsychology3 citations

TL;DR

This paper proposes Aaren, a novel attention mechanism that reformulates self-attention as a differentiable RNN, enabling both parallel training like Transformers and efficient, constant-memory inference with incremental token updates. Aaren achieves Transformer-level performance across 38 sequential modeling tasks while reducing computational and memory complexity from quadratic to linear scaling.

ABSTRACT

The advent of Transformers marked a significant breakthrough in sequence modelling, providing a highly performant architecture capable of leveraging GPU parallelism. However, Transformers are computationally expensive at inference time, limiting their applications, particularly in low-resource settings (e.g., mobile and embedded devices). Addressing this, we (1) begin by showing that attention can be viewed as a special Recurrent Neural Network (RNN) with the ability to compute its extit{many-to-one} RNN output efficiently. We then (2) show that popular attention-based models such as Transformers can be viewed as RNN variants. However, unlike traditional RNNs (e.g., LSTMs), these models cannot be updated efficiently with new tokens, an important property in sequence modelling. Tackling this, we (3) introduce a new efficient method of computing attention's extit{many-to-many} RNN output based on the parallel prefix scan algorithm. Building on the new attention formulation, we (4) introduce extbf{Aaren}, an attention-based module that can not only (i) be trained in parallel (like Transformers) but also (ii) be updated efficiently with new tokens, requiring only constant memory for inferences (like traditional RNNs). Empirically, we show Aarens achieve comparable performance to Transformers on $38$ datasets spread across four popular sequential problem settings: reinforcement learning, event forecasting, time series classification, and time series forecasting tasks while being more time and memory-efficient.

Motivation & Objective

• To address the high inference cost of Transformers in low-resource settings such as mobile and embedded devices.
• To bridge the gap between the parallelizability of Transformers and the efficient, incremental inference of RNNs.
• To develop a new attention formulation that supports both parallel training and constant-memory, streaming inference.
• To introduce a method that computes attention's many-to-many RNN output efficiently using parallel prefix scans.
• To demonstrate that the resulting Aaren module matches Transformer performance while significantly reducing time and memory complexity.

Proposed method

• Reformulates standard self-attention as a special case of a many-to-one RNN, enabling a recurrent interpretation of attention computation.
• Introduces a novel RNN-based formulation of attention using the parallel prefix scan algorithm to compute the many-to-many RNN output efficiently.
• Derives a differentiable, exact re-implementation of softmax-based attention as an RNN, avoiding the exponential timestamp bias seen in prior approximations.
• Designs Aaren as a learnable RNN module that maintains a constant hidden state size and supports incremental updates with new tokens.
• Employs a learnable initial hidden state vector $ q $ to enable the RNN to capture long-range dependencies without storing past activations.
• Leverages efficient parallel prefix sum algorithms (e.g., Hillis-Steele) to compute attention weights in $ O(N \log N) $ time, enabling scalable inference.

Experimental results

Research questions

• RQ1 Can standard self-attention be formally reinterpreted as a recurrent neural network?
• RQ2 Why do attention-based models like Transformers fail to support efficient incremental inference despite their parallel training capability?
• RQ3 Is it possible to design an attention mechanism that combines the parallelizability of Transformers with the constant-memory efficiency of RNNs?
• RQ4 Can a new attention formulation based on parallel prefix scans achieve competitive performance with Transformers while reducing computational complexity?
• RQ5 How does the proposed Aaren module compare to Transformers in terms of inference efficiency and model performance across diverse sequential tasks?

Key findings

• Aaren achieves performance comparable to Transformers on 38 diverse sequential modeling datasets across reinforcement learning, event forecasting, time series classification, and time series forecasting.
• The cumulative inference time and memory usage of Aaren scale linearly with sequence length, while Transformers with KV-caching scale quadratically.
• Aaren requires only a marginal 0.016% increase in parameters (3,152,896 vs. 3,152,384) compared to equivalent Transformers, making the efficiency gains cost-effective.
• The model maintains constant memory usage during inference, enabling deployment on low-resource devices such as mobile and embedded systems.
• The parallel prefix scan-based attention computation enables efficient, differentiable RNN-style inference while preserving the expressivity of full-attention mechanisms.
• Empirical results confirm that Aaren outperforms or matches Transformers in both speed and memory efficiency across all evaluated tasks.