State-space models can learn in-context by gradient descent

1. The paper addresses a theoretical statement about the transformer model, specifically that it can learn in-context through gradient descent. However, this concept is already well-known and documented. In theoretical research, it is widely recognized that transformers, convolutional models, and recurrent models, such as state-space models, are universal approximators and are capable of learning continuous target relationships. Therefore, demonstrating the same for state-space models does not appear to offer significant theoretical advancement. If the authors wish to underscore the importance of this work, I would recommend showing that previous work in approximation theory does not extend to the in-context learning case. Without this distinction, the contribution seems to fall within a subset of known results that hold limited value for theoretical study.

The reviewer mentions that the ability of transformers and other universal models to perform ICL is well-documented, and thus demonstrating the same for SSMs might not constitute a significant theoretical advancement. We respectfully clarify that our work provides a mechanistic explanation specifically tied to gradient descent dynamics in SSMs, which distinguishes it from prior general approximation results. While it is true that many architectures can approximate target functions, we focus on elucidating how SSMs can explicitly emulate gradient-based learning mechanisms. Our goal is not to provide a general approximation result but a concrete and direct construction to show that ICL can be implement using GD in SSMs.

2. The notion of in-context learning, as presented, lacks practical interest. Simply stating that a model "can learn" through in-context learning is insufficient, as the same argument could be made for various methods, including Newton's or quasi-Newton's methods. There is no compelling reason for practitioners to assume that, when state-space models engage in in-context learning, the behavior in terms of loss convergence or asymptotic rates would align with that of gradient descent. Clarifying this distinction would strengthen the paper’s contribution. Could you provide empirical comparisons of convergence rates or asymptotic behavior between your method and alternatives like Newton's or quasi-Newton's methods.

We are not aware of any specific direct construction showing that SSMs in the general form commonly used can learn using Newton’s or quasi-Newton’s methods. We do in fact demonstrate in Figs 3 & 4 that SSMs trained from a random initialisation do exhibit convergence to the same values as would GD. We plan to look into the similar empirical comparisons for Newton’s and quasi-Newton’s methods in the future work.

3. The paper introduces a state-space model with local self-attention, which is not a commonly adopted approach in practice. It would be more beneficial to align the framework with models that are widely used in real-world applications. A method akin to the linear transformer might be more appropriate and could provide a better point of reference for practical utility.

Local self attention is more commonly referred to as sliding window attention in literature (we have updated our paper to reflect this). It is in fact used in major high-performance models such as Griffin [1] and Samba [2].

[1] S. De et al., ‘Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models’, Feb. 29, 2024, arXiv:2402.19427. doi: 10.48550/arXiv.2402.19427.

[2] L. Ren, Y. Liu, Y. Lu, Y. Shen, C. Liang, and W. Chen, ‘Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling’, Jun. 11, 2024, arXiv:2406.07522. doi: 10.48550/arXiv.2406.07522.

We thank the reviewer for recognizing the soundness and clarity of the paper and the empirical evidence supporting the papers claims in small-scale settings.

We also thank the reviewer for bringing up a very closely related paper in Zucchet et al. 2023. However, there are critical differences between their paper and ours as we explain below. To summarize, ours is a more explicit, direct and parsimonious construction showing that SSMs can do ICL by GD while theirs is very indirect.

In detail: In Zucchet et al., they show the equivalence of gated linear recurrent networks and linear self-attention. It is true that linear self-attention has been shown to be able to do ICL by GD [1], and so, indirectly, they show that gated linear recurrent networks can perform ICL by GD. In our case, we consider the general formulation of SSMs and show a correspondence directly between them and ICL via GD, leading to a simpler and more understandable construction.

Our construction is more parsimonious and efficient due to the following:

1. a model using linear attention model requires $3f^2$ input parameters for the separate embeddings for query, key and value, whereas in our case we only need $f^2$ input parameters since we have only one embedding.
2. To train the parameters of a linear-regression model with $f \times f$ parameters the GD-SSM requires $f^2$ recurrent neurons with $f^2$ recurrent parameters since the recurrent matrix is diagonal. Whereas a linear self-attention as constructed in [1] requires $12f^2$ parameters in the linear self-attention.
3. The RNN construction in Zucchet et al. uses $O(f^4)$ neurons to represent the $3f^2$ parameters in the linear self-attention layer, further increasing the size and redundancy of the model.

Our construction distils out the exact inductive bias required for ICL. Moreover, the that the simple form of linear transformers used in [1] still lag behind state-space models in task performance.

While the end result is the same, the method used to achieve this is completely different — ours is more direct.

We have added a discussion of this to our paper.

[1] J. von Oswald et al., ‘Transformers Learn In-Context by Gradient Descent’, in Proceedings of the 40th International Conference on Machine Learning, PMLR, Jul. 2023, pp. 35151–35174. Accessed: Apr. 16, 2024. [Online]. Available: https://proceedings.mlr.press/v202/von-oswald23a.html

We thank the reviewer for their constructive points, and appreciate the reviewer recognizing the novelty of our perspective on SSMs as gradient accumulators, the soundness of our mathematical theory, and the clarity of our step-by-step theoretical exposition.

Responses to weaknesses:

While the paper does a good job of explaining the theory and formulation of GD-SSM and empirically validating GD-SSM on regression tasks, I didn't feel like it shed that much insight into the currently used SSM variants used in practice.

We agree that the constructed approach doesn’t exactly match existing methods, although the use of sliding window attention is becoming more common on highly-performant models [1,2] and 2-D state and input-dependent output computation is pretty standard through gating [3]. To explore why Griffin and Mamba do so badly, we explored their performance a bit more (see Table 1 in appendix). We found that adding additional layers (our earlier experiments were with 1-layer), increasing the model dimension, and increasing the number of heads improves performance . We hypothesize that this is due to the fact that the architecture doesn’t exactly match GD, and additional layers are required to play the role of the sliding window attention to perform time-mixing (access adjacent elements in the sequence).

Our construction distils out the exact inductive bias required for ICL.

Related to the above, the experimental section is light on architectural details, making it hard to determine what exactly is being compared empirically and what conclusions can be drawn. Please include more experimental details including the architectures and hyperparameters.

We’ve added more details in the Appendix about the details of the experiments and architectural details. Please let us know if any specific aspects still remain unclear.

The paper is often unclear on the terminology of local self-attention and local linear self-attention. The formulation in Section 3.2 appears to only require local linear self-attention, yet other times in the paper local self-attention is used. In the related works section the two are contrasted. I would recommend being very explicit and consistent on this point as the two are very different.

We agree with the reviewer, and have changed all references to “Local Self-attention” into "Sliding Window Attention” which is more explicit.

The paper is limited to regression style in-context learning. This is interesting and amenable to theoretical analysis, but also limits the impact of the investigation. See https://arxiv.org/abs/2401.12973 and https://arxiv.org/abs/2402.04248 for other papers that have empirically investigated other types of in-context learning in different architectures.

We agree with the reviewer, and plan to look at other ICL tasks in future work.

[1] S. De et al., ‘Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models’, Feb. 29, 2024, arXiv:2402.19427. doi: 10.48550/arXiv.2402.19427.

[2] L. Ren, Y. Liu, Y. Lu, Y. Shen, C. Liang, and W. Chen, ‘Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling’, Jun. 11, 2024, arXiv:2406.07522. doi: 10.48550/arXiv.2406.07522.

[3] A. Gu and T. Dao, ‘Mamba: Linear-Time Sequence Modeling with Selective State Spaces’, Dec. 01, 2023, arXiv: arXiv:2312.00752.

[4] T. Dao and A. Gu, ‘Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality’, in Proceedings of the 41st International Conference on Machine Learning, PMLR, Jul. 2024, pp. 10041–10071.

[5] J. Park et al., ‘Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks’, Feb. 06, 2024, arXiv: arXiv:2402.04248.

We thank the reviewer for recognizing the strengths of the paper including its clarity, ability to bridge key domains and helping enhance the understanding of inductive biases for ICL tasks. We respond to listed weaknesses below:

(1) The authors have overlooked important related work in this domain. For example, [1] presented at NeurIPS 2016, demonstrates that RNNs can perform gradient descent during the forward pass.

We appreciate the reviewer bringing up [1] (Andrychowicz et al., NeurIPS 2016). However, we believe this work is completely unrelated to ours. Reference [1] employs meta-learning to train two levels of LSTMs: one LSTM generating weight updates which are used to update the “lower level” LSTM (see Fig. 3 in [1]). This meta-learning framework fundamentally differs from our work, where we construct a state-space model (SSM) that implicitly performs gradient descent as part of its forward pass dynamics to perform in-context learning. We do not use meta-learning, and we do not use LSTMs producing updates for another LSTM.

Additionally, Section 3.1 in the current paper shares several similarities with Section 2 of [1]. To be clear, this is not an accusation of plagiarism, but rather an indication of missing key references. While there is a difference in scope (RNNs versus SSMs), this oversight reduces the originality of contribution #1. I kindly ask the authors to specify the principal differences between the approach taken by [1] and the approach used by SSMs in their work to better highlight the novel contributions.

We would appreciate a more detailed explanation of the reviewer’s concern on the similarity between Section 2 of [1] and our Section 3.1. To our understanding [1] does not consider linear regression at all, and in Section 2, formulates a meta-learning loss. Whereas, in our Section 3.1, we start from linear regression and construct an SSM that can perform in-context learning. Our approach provides a mechanistic explanation for how in-context learning can emerge in SSMs without requiring meta-learned updates or external optimization signals.

(2) Overlooks simple alternative approaches: While the theoretical analysis is accurate, the authors overlook significant alternative approaches. References [2-4] demonstrate the connection between S6 layers and attention, showing that S6 can be more expressive than attention without the softmax. Additionally, various ICL studies for transformers omit the softmax (see [5-6] as an example), allowing direct conclusions that could extend to SSMs. Given the extensive exploration of ICL capabilities in transformers, discussions on the ICL potential of SSMs should consider this reduction approach. I recommend that the authors evaluate whether this approach could yield additional theoretical developments to strengthen their analysis. Moreover, I suggest that the authors explicitly state the advantages of their approach compared to the proposed alternatives. For instance, Section 3 introduces specific circuits that cannot be achieved through simple reductions.

We recognize the relevance of studies [2-6] demonstrating connections between SSMs, attention mechanisms, and gradient descent. While those studies propose indirect reduction-based approaches, our GD-SSM construction explicitly and directly encodes gradient descent into the recurrent dynamics of SSMs, without relying on architectural simplifications. Therefore, we contend that our approach is simpler.

While it is true that linear self-attention can be written as a recurrent network, our approach demonstrates a more direct connection between state-space models and ICL by GD. Note that the simple form of linear transformers used in [5] still lag behind state-space models in task performance. Moreover our construction is more parameter efficient than linear self-attention — a model using linear attention model requires $3mf$ input parameters for the separate embeddings for query, key and value, whereas in our case we only need $mf$ input parameters since we have only one embedding.

References [2-4] deals with selective state-space models with non-linear input-dependence of the SSM parameters. Our construction suggests that such non-linear Input dependence, while powerful, may not be necessary for in-context learning. The remaining studies are empirical, and don’t show a direct theoretical construction, but are nonetheless interesting to understand if the ICL mechanism in the wild corresponds to the one we propose.

As suggested, we will expand our discussion to explicitly outline the advantages of our approach over reduction-based methods, such as greater interpretability and architectural alignment with currently used SSM and clarify that the unique circuits introduced in Section 3 cannot be achieved through these reductions alone. This addition will provide a clearer context for how GD-SSM complements and extends prior work.