3 ms·
Check out "Attention as an RNN" by Feng et al (2024), with Bengio as a co-author. https://arxiv.org/pdf/2405.13956 https://arxiv.org/pdf/2405.13956 Abstract: T
by bravura 2y ago
Check out "Attention as an RNN" by Feng et al (2024), with Bengio as a co-author. https://arxiv.org/pdf/2405.13956 https://arxiv.org/pdf/2405.13956
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 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 many-tomany RNN output based on the parallel prefix scan algorithm. Building on the
new attention formulation, we (4) introduce 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.