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

Thank you very much for the helpful feedback and pointers! Changes or clarifications (as deemed appropriate) for every single point raised will be incorporated in the final version of the paper. In what follows we provide a brief discussion on the points that were raised in the review:

Prior studies: Thank you for the feedback and pointing us at Oren et al. (2024). We agree that connections between the three model classes have been studied in previous work and are well-established. Our intent is to translate the three model classes specifically as dynamical systems - which is achieved by the DSF - in order to use methods from control theory to analyze the models and to obtain a single form of representation to easily compare certain parameterization choices between the model classes. In the revised version, we made sure to better highlight this intent and we appropriately cited Oren et al. (2024). We also acknowledge that the infinite representation of softmax-attention has been discussed previously in Katharopoulos et al. (2020) but also in Nauen et al. (2024) and Choromanski et al. (2021) among others. We believe the discussion on representation of softmax-attention is important in the context of the DSF, which is why we included it. In the revised manuscript, we made sure to include additional related works and improve the presentation to better reflect previous works.

Empirical novelty: We acknowledge that the insight on increased performance with larger state sizes is well-established in the literature, we believe it is beneficial to show that insights provided by the DSF are consistent with previous results. In the revised manuscript, we improved the presentation of this insight to clearly state that we are only recovering existing results and cite the relevant works more prominently. For the insights on RNNs, the main difference between the proposed qLSTM variant and the RG-LRU is the shared parameterization in state transition $A$ and input matrix $B$. While RG-LRU shares parameters in both matrices, the qLSTM does not (as briefly mentioned in lines 263-265). Interestingly, any attention-based model also shows this coupling in the DSF. In the revised manuscript, we added a discussion of this point for all three model classes (please also see the global response) and plan to include an empirical analysis for the camera-ready version.

Task complexity & modalities: Thank you for the feedback and the suggestion. We agree that more complex and also diverse tasks can strengthen the insights we provide. While we believe in-depth analysis of a specific insight (e.g. the proposed normalized attention) including showing performance on complex tasks are more suited for subsequent works (we only have limited computational resources), we agree that other task modalities offer more insights. Therefore, to improve the empirical validation of our findings, we include experiments on the LRA benchmark, which also contains image tasks, in the revised manuscript. As an example, we included the results of the LRA image task in the attached pdf. In doing this we generated an additional insight (performance gap between SSM models and transformers on LRA) as discussed in the general response.

Attention and state size: Thank you for the excellent suggestion. Preliminary experiments we ran show better performance of the other models with their state size increased to the size of the KV cache, but not outperforming softmax attention. Hence, it seems that there is more to softmax attention than just the size of the KV cache. Additionally, we believe there is more potential to such an analysis, since it relates to the question of how information is compressed in the state of the DSF. This is important since running these models with the size of the KV cache is not sustainable in practical applications with larger sequence lengths. We think such an analysis could lead to further follow-up works investigating different aspects of this question. Therefore, we will conduct full experiments for the revised version and perform an in-depth analysis of the results.

Minor: Thank you for pointing out the two inconsistencies. It is true that the matrix in Eq. (12) is not $L \times L$, but should be a block matrix consisting of $L \times L$ blocks. We will correct this in the revised manuscript. We will also improve the language for practical separability, since it is indeed misleading in the current form.

Thank you for the clarifying comments and additional results. I believe after the response to the reviewers this is a stronger paper and of interest to the community. I have increased my score.

Thank you very much for the helpful feedback and pointers! Changes or clarifications (as deemed appropriate) for every single point raised will be incorporated in the final version of the paper. In what follows we provide a brief discussion on the points that were raised in your review:

High-level discussion: Thank you for this suggestion! We indeed believe that this would help to make the paper accessible to a broader audience. To address this comment, we plan to add an introduction to the concept of dynamical systems and how it applies in the context of foundation models. We will also delve into the specific dynamical system parametrization chosen by the DSF, which is the canonical state space representation. This choice will be motivated by two reasons: (i) this exists a wide body of literature on state space model representations for dynamical systems, and (ii) it encompasses transformers, RNNs and SSMs in a suitable fashion that allows for further analysis. We will also better explain how expressing foundation models in the DSF allows for a direct comparison between models and enables the use of control theoretical tools for their study.

Prior theoretical comparative studies: Thank you for pointing this out, we acknowledge that the statement in the introduction is not correct. We intended to convey that to the best of our knowledge so far there has been no unified approach to theoretically analyze transformers, SSMs, and RNNs, but only empirical studies benchmarking specific models of these classes against each other. In the revised version, we changed this to reflect prior theoretical work comparing various models with attention, including properly referencing the papers that are mentioned in the review.

Empirical novelty: Thank you for pointing us at Arora et al. (2024). We agree that the insight “increased state size leads to better performance” is not novel and has been shown in previous work. We intended to include this insight, since we believe it is beneficial to show that insights generated from the DSF are consistent with the literature. In the revised manuscript, we appropriately cite Arora et al. (2024) and improve the presentation to reflect that this insight has been reported before. Additionally, we improved the presentation of the overall document to better highlight the other two insights generated from the DSF (i.e. normalized attention and the S6-inspired state transition for qLSTMs) as well as the additional insights we added (please see the global response for more details). This change puts more focus on the new insights and explicitly presents the state size finding as established knowledge.

Lemma 2: Thank you for the feedback, we acknowledge that the language in and around Lemma 2 is imprecise. As correctly pointed out, Lemma 2 should state that expressivity with a larger state size is at least as good as with lower state size, i.e., not strictly better. We reformulated Lemma 2 and the text surrounding it to reflect this. Additionally, we now explicitly state that this only holds if everything else is held constant as mentioned in the review.

Results discussion: Thank you for the feedback. We agree that the key take-aways from Fig. 2 are lacking and are not appropriately discussed in the text. In the revised manuscript, we added the following discussion points:

• The results in Fig. 2 and Tab. 1 support the conclusions that (i) performance of linear attention can be improved by using a suitable normalization function and (ii) this normalization function does not need to use the keys and queries but can be learnt independently.
• The difference in parameterization between the proposed normalized attention and SSD mainly lies in the recurrent nature of the normalization parameters in transition matrix $A$ and input matrix $B$ for normalized attention. With the addition of the LRA experiments in the revised manuscript and the performance gap between SSMs and transformers, this warrants the question on the role of this recurrent normalization. In the revised manuscript we added a discussion on this question as detailed in the global response.

Typos: Thank you! We will amend these in the revised version.

Thank you very much for the helpful feedback and pointers! Changes or clarifications for every single point raised will be incorporated in the final version of the paper. We provide a brief discussion on the points that were raised in the review:

Benchmark: Our main goal was to show-case some selected insights the DSF can bring to all three investigated model classes, i.e., transformers, SSMs, and RNNs. Since softmax-attention-based models are known to be the best performing models for language modalities and considerable research focuses on sub-quadratic models to close the performance gap to softmax-attention, we opted to benchmark on language tasks. However, we added the LRA benchmark to the experimental validation of our insights (see global response). The LRA results are consistent with existing empirical results, which show that attention-based models perform worse than SSM-based models on the LRA. Why this is the case is still an open research question and warrants more detailed investigation, but the DSF points to a potential hypothesis for this performance gap: the recurrent normalization discussed in Section 4.2 of the paper. Along with the LRA results we added a discussion on this point to the revised paper.

Insights: Thank you for your feedback. We agree that the paper only showed a few theoretical insights. Therefore, we have extended on these insights and added several more as discussed in the general response. Also thank you for the specific suggestions, which we discuss below.

• Finite-dimensional approximation: In the revised manuscript, we expanded on the discussion in Section 3.2.1 on using a Taylor approximation of softmax-attention. Specifically for the Taylor approximation, there exist error bounds that can be used to bound the error between the approximation and softmax-attention on a given interval, e.g., the Lagrange bound (see e.g. Section 3.3.2. in Chevillard et al. (2008)). The Lagrange error bound defines the worst-case error of the approximation on the given interval and thus provides a quantitative metric of how well softmax-attention can be approximated. Additionally, we extended the discussion of existing works that use a finite-dimensional approximation of softmax.
• Kernel method: Thank you for the excellent question. The attention function $\zeta$ in Eq. (13) (later rewritten in Lemma 1) can be interpreted as a kernel if $\psi = \phi$ , with softmax-attention using a specific type of exponential kernel function, i.e., $\exp(x^\top y)$. This was noticed and further analyzed in Tsai et al. (2019) and Katharopoulos et al. (2020). In the revised manuscript, we added a remark on this and refer to the mentioned papers for more details.

Training dynamics: While training dynamics (especially convergence) can be studied empirically using experiments, the DSF also allows theoretical analysis. As discussed in Example 2 of Dörfler et al. (2024), a gradient-based optimization algorithm (e.g. SGD) can be interpreted and written as a dynamical system. Using this perspective together with the DSF allows interpretation of the training dynamics as two interacting dynamical systems. Thus, the training dynamics can be analyzed theoretically with tools from control theory, e.g., via Lyapunov theory for convergence/stability of the training. However, we believe this question requires an in-depth investigation and additional empirical validation, which are out of scope for this paper.

Stability: The exploding/vanishing gradient problem is linked to the eigenvalues of state transition matrix $A$ for SSMs (Orvieto et al. (2023)). Therefore, SSM models actively restrict the eigenvalues to the range [0,1]. Using the DSF the same argumentation can be made for transformers and RNNs. In the revised paper, we included an additional discussion on this.

Long-range dependencies: A model's abiltiy to capture long-range dependencies are due to two factors: (i) the modulus of the eigenvalues in dynamic matrix $A$, and (ii) the dimension of the state $x$ (in DSF representation). Other factors include the compression of the input into the state by the encoding architecture, but this is out of the scope of our analysis that focuses on the recurrent block of learning architectures. Given two trained models with numerical values for $A$ and a sufficiently large state dimension, the DSF would be capable of predicting the performance in long-range contexts by looking at the eigenvalues of $A$. Moreover, theory on system theoretic model reduction (e.g. Obinata & Anderson (2012)) provides principled ways of comparing the ability of a given state dimension to compress past information. Thus, the DSF can be leveraged to provide a theoretical analysis on the long-range capabilities of the analyzed models. To complement our analysis, we will incorporate this discussion into the revised manuscript.

Task suitability: We believe that the question of architecture suitability for specific tasks is a broad open research question, and many factors need to be considered to provide a clear answer. Specifically, the recurrent models presented in this paper are only one aspect of deep learning architectures, but the encoders, decoders, skip connections, etc. also play a role in task suitability. For this reason, we do not aim to answer this question in full in this paper. However, motivated by the comment we have experimentally explored different data modalities. Our findings illustrate (see attached pdf and paper experiments) that transformer architectures (specially with softmax attention) tend to perform better on language tasks and SSMs on tasks that require long-range dependencies such as the LRA benchmark. However, the performance of a model is highly dependent on the specific formulation of the task.

Thank you very much for the helpful feedback and pointers! Changes or clarifications (as deemed appropriate) for every single point raised will be incorporated in the final version of the paper. In what follows we provide a brief discussion on the points that were raised in the review:

Empirical evaluation: We agree that the empirical evaluations of our insights were somewhat limited. To improve this we added additional experiments on the LRA benchmark, which also includes image tasks to broaden the empirical evaluation. The partial results (due to time constraints), i.e., only the sequential CIFAR-10 (image) task are shown in Tab. 1 of the attached pdf, and the results of the full benchmark will be added to the revised manuscript.

Statistical significance: Thank you for this comment! We are including error bars to Fig. 2 and 3 in the paper. These results are shown (as shaded regions) in Fig. 1 and 2 of the attached pdf.

Theoretical results: We acknowledge that Lemma 2 is somewhat straightforward and that this effect has already been shown empirically. However, this finding is a trivial consequence of the DSF and we wanted to highlight that the DSF provides insights that are consistent with the literature.

Theoretical sections: Thank you for this comment! We acknowledge that the theoretical sections might be dense for an audience unfamiliar with dynamical systems. To improve upon this point, we plan to incorporate a brief introduction to the conceptual basics of dynamical systems in the introduction of the revised manuscript. We will also extend Section 3.1 with a more detailed presentation to make it more amenable for a broader audience.

Practical impact: We added several more insights and analyses to the revised document as stated in the global response and we believe that the proposed work enables leveraging system theoretic insights in order to generate more specific practical impacts in future works.

Training and inference efficiency: Although the primary goal of the DSF is to carry out theoretical comparison and obtain insights on the mathematical representations of foundation models, we believe it can also help in providing pointers for efficient training and inference. For instance, as discussed below in the response to the question regarding efficient implementations, the DSF allows identification of efficient algorithms to implement a specific model. Additionally, the DSF naturally allows analysis of the eigenvalues of the state transition matrix $A$, which are linked to the exploding/vanishing gradient problem (Orvieto et al. (2023)). We included a discussion on this to the revised version. On a more theoretical level, it is possible to view the training process as an interaction of two dynamical systems by formulating the optimization algorithm as a separate dynamical system (see e.g. Example 2 in Dörfler et al. (2024)). In the revised version, we added a remark on this, but we think elaborating on this requires a more in-depth analysis that is out of scope of the paper.

Scalability: In order to explore how the insights (e.g. normalized attention) scale to billions of parameters, an empirical evaluation is vital. Due to the extensive computational demands of this task, we believe this analysis is out of the scope for this paper. In a theoretical analysis, normalized attention scales linearly with context length, similarly to linearized attention and SSD. As such, we anticipate that it will scale similarly to other SSM models that have been demonstrated at a larger scale.

Softmax approximation: Thank you for the excellent question! We believe this is a very interesting research question that warrants further investigation, but has no straightforward answer. However, in the special case of polynomial kernel functions (in Lemma 1), softmax-attention can be exactly recovered using infinitely many polynomials. Therefore, in the limit these polynomial kernels can recover the softmax performance with appropriate weights as there is an exact correspondence.

Alternative normalization schemes: Yes, we also considered other normalization schemes to the one shown in Eq. (21), i.e.,

$

\eta(u_i) &= e^{\textrm{ReLU}(W_\eta u_i)} \\\\
\eta(u_i) &= \sigma(W_\eta u_i)

$

We observed that they perform similarly to the one presented in the paper. We included the results on these additional normalization schemes to the appendix of the revised manuscript including a brief discussion.

Long-term memory: The long-term memory capabilities of a model from a control theoretic perspective depend on the eigenvalues of the state transition matrix $A$ and the size of the state $x$ and thus the size of $A$ (Orvieto et al. (2023)). Therefore, any model reformulated in the DSF can be analyzed in this way. Since the proposed state transition is the same state transition as in S6, the two models have the same theoretical long-term memory capabilities. These capabilities then solely depend on the hyper-parameter $n$ (state size) and the initialization of $A$. However, we agree that this question warrants a more detailed analysis and thorough empirical validation, which is out of scope of this paper.

Efficient implementations: Thank you for the excellent question. Yes, the DSF can be leveraged to identify and also develop new efficient implementations. In the revised version, we expand the first paragraph in Section 3.2 that only briefly mentions that the DSF can be used for this. In the appendix, we also added a few examples of how this can be done, e.g., the proposed normalized attention can be efficiently implemented using flash linear attention.