Understanding the differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks

TL;DR

The paper introduces the Dynamical Systems Framework (DSF) to unify attention, State Space Models, and Recurrent Neural Networks as backbone sequence mixers, addressing the quadratic complexity of softmax attention in long contexts. By reformulating all architectures as a linear time-varying recurrence, DSF enables principled theoretical comparisons and reveals how normalization, state expansion, and parameter coupling drive performance differences. The authors provide both analytical results and empirical validations on MQAR, LRA, and WikiText-103, showing, for example, that linear attention can approach softmax performance with sufficient state expansion and proper normalization, while SSMs like S6 benefit from input-dependent normalization and recurrent structure. The DSF framework thus offers a unified lens to guide the principled development of future hybrid and more scalable foundation models, while also highlighting limitations related to efficiency guarantees and the scope of empirical validation.

Abstract

Softmax attention is the principle backbone of foundation models for various artificial intelligence applications, yet its quadratic complexity in sequence length can limit its inference throughput in long-context settings. To address this challenge, alternative architectures such as linear attention, State Space Models (SSMs), and Recurrent Neural Networks (RNNs) have been considered as more efficient alternatives. While connections between these approaches exist, such models are commonly developed in isolation and there is a lack of theoretical understanding of the shared principles underpinning these architectures and their subtle differences, greatly influencing performance and scalability. In this paper, we introduce the Dynamical Systems Framework (DSF), which allows a principled investigation of all these architectures in a common representation. Our framework facilitates rigorous comparisons, providing new insights on the distinctive characteristics of each model class. For instance, we compare linear attention and selective SSMs, detailing their differences and conditions under which both are equivalent. We also provide principled comparisons between softmax attention and other model classes, discussing the theoretical conditions under which softmax attention can be approximated. Additionally, we substantiate these new insights with empirical validations and mathematical arguments. This shows the DSF's potential to guide the systematic development of future more efficient and scalable foundation models.

Figures (7)

• Figure 1: Comparison of linear attention and softmax attention on two MQAR tasks $\{(L=256, \textrm{KV-pairs}=16), (L=512, \textrm{KV-pairs}=64)\}$, fixed model size $d=512$, and varying state expansion $n$. We report the best result from a learning rate sweep in $\texttt{np.logspace}(-4, -2, 4)$.
• Figure 2: Model accuracy with increasing model size $d$ for different models: softmax, linear, and normalized attention, S6, and SSD. The MQAR task is $(L=512, \textrm{KV-pairs}=64)$, we fix $n=128$, and report the best performance of a learning rate sweep in $\texttt{np.logspace}(-4, -2, 4)$.
• Figure 3: Comparison of qLSTM \ref{['eqn:qlstm']} and a qLSTM variant where the original state transition $\Lambda_i$ is replaced by \ref{['eq:rev-sigmoid']}.
• Figure 4: Visual representation of the dimensions considered in the DSF.
• Figure 5: Model accuracy with increasing model size $d$ for different models: softmax, linear, and normalized attention, S6, and SSD. The MQAR task is $(L=512, \textrm{KV-pairs}=64)$, we fix $n=128$, and report the best performance of a learning rate sweep in $\texttt{np.logspace}(-4, -2, 4)$. Solid lines are the average accuracy over 10 different seeds, while the shaded area show the standard deviation.

Theorems & Definitions (5)

• Remark 1
• Lemma 1
• proof
• Lemma 2
• proof