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

## Abstract

Abstract Deep state-space models (Deep SSMs) have shown capabilities for in-context learning on autoregressive tasks, similar to transformers.
However, the architectural requirements and mechanisms enabling this in recurrent networks remain unclear.
This study demonstrates that state-space model architectures can perform gradient-based learning and use it for in-context learning.
We prove that a single structured state-space model layer, augmented with local self-attention, can reproduce the outputs of an implicit linear model with least squares loss after one step of gradient descent.
Our key insight is that the diagonal linear recurrent layer can act as a gradient accumulator, which can be ‘applied’ to the parameters of the implicit regression model.
We validate our construction by training randomly initialized augmented SSMs on simple linear regression tasks. The empirically optimized parameters match the theoretical ones, obtained analytically from the implicit model construction.
Extensions to multi-step linear and non-linear regression yield consistent results.
The constructed SSM encompasses features of modern deep state-space models, with the potential for scalable training and effectiveness even in general tasks.
The theoretical construction elucidates the role of local self-attention and multiplicative interactions in recurrent architectures as the key ingredients for enabling the expressive power typical of foundation models.

## 1 Introduction

The current generation of Large Language Models (LLMs) and foundation models are extremely capable and have started proliferating in several real-world applications.
These models are based on the transformer architecture (Vaswani et al., ), and a big part of their capability has been attributed to in-context learning (Wei et al., ; Lu et al., ).
In-context learning in transformers is relatively well studied (Wies et al., ; Pan et al., ; Guo et al., ; Garg et al., ).
A prominent explanation for the mechanism used by transformers to do in-context learning is that the model performs in-context learning by gradient descent (von Oswald et al., ; Akyürek et al., ).
But, the quadratic dependence of transformers on the input length makes them computationally expensive.
To mitigate this, there has been much work on alternatives based on recurrent networks, such as state-space models (Gu et al., ; Gu & Dao, ) and linear recurrent networks (Orvieto et al., ).
These recurrent models can perform inference efficiently since the computational complexity of recurrent networks is linear in the sequence length.
At the same time, linear recurrent networks allow parallelization across the sequence during training using associative scan.
The latest versions of these models are competitive with transformers at scale (Dao & Gu, ; De et al., ), and also capable of in-context learning (Grazzi et al., ).
However, the mechanism they use for in-context learning remains unclear.

This work focuses on the in-context learning mechanism of State-Space Models (SSMs).
Multiple variations of state-space models have recently shown competitive performance at scale (Gu & Dao, ; De et al., ), while earlier generations struggled with scaling (Smith et al., ; Poli et al., ).
All SSMs, linear attention models, and other linear recurrent networks share a common formalism of being linear recurrent networks interleaved with non-linear layers.
On the other hand, the ability to do in-context learning seems to be a hallmark of most recent scalable variants  (Grazzi et al., ).
Which features of these successful models contribute to in-context learning, as opposed to earlier variants?
Using a constructive approach, we pinpoint input-dependent input and output processing, as the key features required for in-context learning.

We show that SSMs with local self-attention, a form of input-dependent input processing, can perform in-context learning analogously to transformers, i.e. through gradient descent steps on an implicit linear regression problem.
The key insight we use is that the state of the recurrent network can be used to aggregate gradients of the parameters of the implicit linear model, which can later be ‘applied’ to the initial parameters of the implicit model.
Using the general SSM formalism, we show how to design the recurrent and output equations that enable them to do in-context learning.
Our construction, which we call GD-SSM, is not restricted to in-context learning tasks and performs well on general-purpose prediction problems.

In summary, our contributions are to show that:

- •
one-layer SSMs with diagonal recurrence, two-dimensional state, and input-dependent input and output processing can perform one step of minibatch gradient descent on a an implicit least-squares loss function;
- •
multi-step (minibatch) gradient descent can be achieved by stacking the 1-layer model;
- •
gradient descent for an implicit non-linear regression problem can be achieved by augmenting the SSMs with non-linearities;
- •
a randomly initialized model trained on regression tasks learns parameters that match our construction for in-context learning tasks based on regression.

## 2 Background

A sequence model operates on an input sequence ${\bm{S}}=\{{\bm{s}}_{t}\}_{t=1}^{T}\in\mathbb{R}^{T\times f}$, where $T$ is also referred to as the context-length and $f$ is the feature dimension.

Contemporary sequence models based on transformers interleave self-attention with MLP layers to perform sequence processing.
The self-attention performs sequence-mixing while the MLPs perform channel-mixing.
The most common form of self-attention has been the scaled dot-product self-attention, which embeds the sequence into a query ${\bm{Q}}={\bm{S}}{\bm{W}}_{Q}$, key ${\bm{K}}={\bm{S}}{\bm{W}}_{K}$ and value ${\bm{V}}={\bm{S}}{\bm{W}}_{V}$, where ${\bm{W}}_{Q},{\bm{W}}_{K}\in\mathbb{R}^{f\times m},{\bm{W}}_{V}\in\mathbb{R}^{
f\times d}$, and then calculates the output of attention as

$$ $\displaystyle\text{SA}({\bm{S}})=\text{softmax}\left(\frac{{\bm{Q}}{\bm{K}}^{T }}{\sqrt{m}}\right){\bm{V}}\,.$ (1) $$

Local self-attention (Beltagy et al., ) uses the same form as Eq. 1, but on an input sequence that is a subset of the full sequence, with a sliding window.

By discarding the softmax (and scaling for simplicity)
, self-attention can be written in a linear form

$$ $\displaystyle\text{LSA}({\bm{S}})$ $\displaystyle={\bm{Q}}{\bm{K}}^{T}{\bm{V}}\,.$ (2) $$

Deep SSMs are sequence models that replace the self-attention with a linear recurrent network, to perform the sequence-mixing.
In the most general form, the recurrent SSM block consists of a recurrent state ${\bm{Z}}_{t}\in\mathbb{R}^{d\times m}$ updated iteratively as

$$ $\displaystyle{\bm{Z}}_{t}={\bm{A}}({\bm{s}}_{t})*{\bm{Z}}_{t-1}+{\bm{B}}({\bm{ s}}_{t})\,,$ (3) $$

where ${\bm{A}},{\bm{B}}:\mathbb{R}^{f}\rightarrow\mathbb{R}^{d\times m}$, and $*$ is some multiplication operator (e.g. matrix or element-wise multiplication).

The output of the SSM is often calculated by an input-dependent linear transformation of the current state

$$ $\displaystyle{\bm{o}}_{t}={\bm{Z}}_{t}\,{\bm{U}}({\bm{s}}_{t})\,,$ (4) $$

combined with a non-linearity as

$$ $\displaystyle{\bm{o}}^{\prime}_{t}={\bm{U}}_{2}\left({\bm{o}}_{t}\odot\sigma({ \bm{U}}_{1}{\bm{o}}_{t})\right)\,,$ $$

where $\odot$ is elementwise multiplication.

Transformers with linear attention can also be written in this recurrent form Katharopoulos et al. ()

$$ $\displaystyle{\bm{Z}}_{t}={\bm{Z}}_{t-1}+{\bm{v}}({\bm{s}}_{t})\otimes{\bm{k}} ({\bm{s}}_{t})\,,$ (5) $$

where ${\bm{v}}({\bm{s}}_{t})={\bm{W}}_{V}^{T}{\bm{s}}_{t}\in\mathbb{R}^{d}$, ${\bm{k}}({\bm{s}}_{t})={\bm{W}}_{K}^{T}{\bm{s}}_{t}\in\mathbb{R}^{m}$ and $\otimes$ is the outer product.
This is a recurrent reformulation of the part of Eq. 2 containing ${\bm{V}}$ and ${\bm{K}}$ in recurrent form.

S4 (Gu et al., ) and multi-headed S5 (Smith et al., ) can also be written in this form

$$ $\displaystyle{\bm{Z}}_{t}={\bm{A}}{\bm{Z}}_{t-1}+{\bm{B}}{\bm{s}}_{t}\,,$ (6) $$

where ${\bm{A}}$ and ${\bm{B}}$ are learnable parameters of appropriate dimension, and ${\bm{Z}}_{t}$ consists of $m$ heads of dimension $d$ each (although multiple heads are not used in the original paper).

If local self-attention is used to process the input before being fed into the SSM, all instances of ${\bm{s}}_{t}$ would be replaced by a context vector

$$ $\displaystyle{\bm{c}}_{t}=\text{SA}({\bm{S}}^{\prime})\quad\text{or}\quad{\bm{ c}}_{t}=\text{LSA}({\bm{S}}^{\prime})\,,$ $$

where ${\bm{S}}^{\prime}\subset{\bm{S}}$ is a subsequence of the whole sequence, and (L)SA denotes the (linear) self-attention operation.

## 3 SSMs can emulate gradient descent on linear regression tasks

We will now show that an SSM as described in Section 2 can perform gradient descent on an implicit linear model to minimize a least squares loss (for particular choices of parameters).
Extensions to non-linear regression models are considered in Section 3.3.

Consider a linear regression problem.
The goal is to minimize the corresponding least squares loss using gradient descent.
The linear model will be the implicit model on which the SSM performs gradient descent.
Performing mini-batch (batch size $>$ 1) gradient descent on the parameters of this implicit model involves two steps: (i) to accumulate gradients of the loss with respect to the parameters, and (ii) apply the accumulated gradient to the initial value of the parameters of the linear model to calculate the updated parameters.
Predictions can be made with the updated parameters by combining them linearly with the input.

Assume the training samples for the linear regression problem are provided as a sequence of inputs and targets.
A large enough SSM can then accumulate the gradients of the loss function in its state *if a local self-attention-like layer processes the sequence inputs before the recurrence*(^1^11A SSM layer with input-dependent recurrence would be able to simulate a local self-attention layer, but with significantly increased computational and conceptual complexity.).

Given the accumulated gradients, the next-step emission of the SSM is equivalent to i) updating the parameters of the implicit model gradient and ii) computing the model output with the updated parameters.
Multiple steps of gradient descent can be achieved by stacking multiple layers, while nonlinearity in the implicit model can be handled by adding nonlinear input-output embedding layers.
We argue that the architecture that allows a single layer to perform gradient descent provides the inductive bias for the model to do in-context learning.

### 3.1 Single step 1-dimensional linear regression

Consider a linear regression model with 1-d output for simplicity

$$ $y={\bm{w}}^{T}{\bm{x}}\,,$ (7) $$

for parameter ${\bm{w}}\in\mathbb{R}^{f}$.
This is the implicit linear model we aim to reproduce for in the 1-dimensional target case.
Given dataset of $N$ samples $\mathcal{D}=\{\langle{\bm{x}}_{i},y_{i}\rangle\}_{i=0}^{N},{\bm{x}}\in\mathbb{
R}^{f},y\in\mathbb{R}$, the associated least squares loss is

$$ $\displaystyle\mathcal{L}(\mathcal{D};{\bm{w}})=\frac{1}{2N}\sum_{i=1}^{N}|| \hat{y}_{i}-y_{i}||_{2}^{2}=\frac{1}{2N}\sum_{i}\left({\bm{w}}^{T}{\bm{x}}_{i} -y_{i}\right)^{2}\,.$ (8) $$

The best fit for ${\bm{w}}$ is the minimum of $\mathcal{L}$ over ${\bm{w}}\in\mathbb{R}^{f}$.

The gradient of the loss calculated on the first $t$ samples of the dataset is

$$ $\displaystyle\nabla_{{\bm{w}}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{w}}_{0})$ $\displaystyle=\frac{1}{t}\sum_{i=1}^{t}\left({\bm{w}}_{0}^{T}\,{\bm{x}}_{i}-y_ {i}\right){\bm{x}}_{i}\,,$ $$

where $\mathcal{D}_{1:t}$ denotes the first $t$ samples in $\mathcal{D}$.
The unscaled gradient,

$$ $\displaystyle{\bm{g}}_{{\bm{w}}_{0}}(\mathcal{D}_{1:t})$ $\displaystyle=\sum_{i=1}^{t}\left({\bm{w}}_{0}^{T}\,{\bm{x}}_{i}-y_{i}\right){ \bm{x}}_{i}\,,$ $$

can be recursively calculated as

$$ $\displaystyle{\bm{g}}_{{\bm{w}}_{0}}(\mathcal{D}_{1:t})$ $\displaystyle={\bm{g}}_{{\bm{w}}_{0}}(\mathcal{D}_{1:t-1})+\left({\bm{w}}_{0}^ {T}\,{\bm{x}}_{t}-y_{t}\right){\bm{x}}_{t}\,.$ (9) $$

Scaling the ${\bm{g}}_{{\bm{w}}_{0}}(\mathcal{D}_{1:t})$ gives the mini-batch gradient $\nabla_{{\bm{w}}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{w}}_{0})=\frac{1}{t}\,g_{{
\bm{w}}_{0}}(\mathcal{D}_{1:t})$ for minibatch size $t$.

To make a prediction, $\hat{y}$, we apply the gradient to the parameters of the linear model in Eq. 7 and compute the corresponding output, i.e.

$$ $\displaystyle\hat{y}_{t+1}$ $\displaystyle=({\bm{w}}_{0}-\eta\nabla_{{\bm{w}}_{0}}\mathcal{L}(\mathcal{D}_{ 1:t};{\bm{w}}_{0}))^{T}\,{\bm{x}}_{t+1}\,,$ $\displaystyle=({\bm{w}}_{0}-\frac{\eta}{t}g_{{\bm{w}}_{0}}(\mathcal{D}_{1:t})) ^{T}\,{\bm{x}}_{t+1}\,.$ $$

When ${\bm{w}}_{0}={\bm{0}}$, this reduces to

$$ $\displaystyle\hat{y}_{t+1}$ $\displaystyle=-\frac{\eta}{t}{\bm{g}}_{{\bm{w}}_{0}}(\mathcal{D}_{1:t})^{T}\,{ \bm{x}}_{t+1}\,.$ (10) $$

#### Implementation as an SSM:

Equation 9, which is a linear recurrence equation, can be implemented by an appropriately constructed SSM.

###### Proposition 1

Given a diagonal linear recurrent layer, and tokens ${\bm{s}}_{j}={\bm{c}}_{j}=[{\bm{x}}_{j}\,y_{j},{\bm{x}}_{j+1}]$, for $j=1,\ldots,N$, and $[\ldots]$ concatenation, ${\bm{x}}_{j},y_{j}$ drawn from a linear model, one can construct recurrent matrix ${\bm{A}}({\bm{s}}_{j})$, input ${\bm{B}}({\bm{s}}_{j})$ and output matrix ${\bm{U}}({\bm{s}}_{j})$ such that each recurrent step for every token ${\bm{s}}_{j}$ produces $\hat{y}_{j+1}=-(\Delta{\bm{w}})^{T}{\bm{x}}_{j+1}$ as output, where $\Delta{\bm{w}}$ is one step of gradient descent, i.e. $\Delta{\bm{w}}=\eta\nabla_{\bm{w}}\mathcal{L}$.
The test input ${\bm{x}}_{N+1}$ is contained in token ${\bm{c}}_{N}$, and produces the test prediction $\hat{y}_{N+1}$.

Specifically, if we use ${\bm{z}}_{t}\in\mathbb{R}^{f}$ to denote the state of the recurrent network and let it directly correspond to the vector ${\bm{g}}_{{\bm{w}}_{0}}(\mathcal{D}_{1:t})$, the equivalent SSM layer is a linear recurrence equation,

$$ ${\bm{z}}_{t}={\bm{I}}\,{\bm{z}}_{t-1}+({\bm{w}}_{0}^{T}{\bm{x}}_{t}-y_{t})\,{ \bm{x}}_{t}\,.$ (11) $$

The state of the SSM, ${\bm{z}}_{t}$, represents the implicit linear regression problem through the *unscaled accumulated gradients* of the least squares loss with respect to the parameters ${\bm{w}}_{0}$ of the implicit linear model.

As linear regression is performed on the training dataset $\mathcal{D}=\{\langle{\bm{x}}_{i},y_{i}\rangle\}_{i=0}^{N}$, the SSM receives the training data in the form of a sequence as input.
In the most general case, this is a sequence ${\bm{s}}_{1}={\bm{x}}_{1},{\bm{s}}_{2}=[0,...,y_{1}],\ldots$ where $[\ldots]$ denoting concatenation and $y_{i}$ is padded with $f-1$ zeros for its dimensions to match that of ${\bm{x}}_{i}$.
This more general case is discussed in the next section.
But here, we will consider a case which simplifies our construction.

Let ${\bm{s}}_{1},{\bm{s}}_{2},\ldots$ be a sequence of constructed context vectors ${\bm{s}}_{t}={\bm{c}}_{t}$, where each ${\bm{c}}_{t}=[{\bm{x}}_{t}y_{t},{\bm{x}}_{t+1}]\in\mathbb{R}^{2f}$, and let us assume ${\bm{w}}_{0}={\bm{0}}$ for simplicity (^2^22The more general case is treated for the case of multi-step GD in Appendix A.2.).
If the sequence input weights ${\bm{\Psi}}\in\mathbb{R}^{2f\times f}$ are such that ${\bm{\Psi}}^{T}{\bm{c}}_{t}={\bm{x}}_{t}y_{t}$, Eq. 11 can be written as an SSM (Eq. 3), i.e.

$$ $\displaystyle{\bm{z}}_{t}$ $\displaystyle={\bm{I}}\,{\bm{z}}_{t-1}+{\bm{\Psi}}\,{\bm{c}}_{t}\,.$ (12) $$

A parameter matrix, ${\bm{\Psi}}$, satisfies equation 12 if all but the first $f$ diagonal entries are zero.
The state of this network is the unscaled gradient $t\,\nabla_{{\bm{w}}_{0}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{w}}_{0})$ and the state recursion accumulates the gradients as in step (i) above.

The accumulated gradient is then ‘applied’ to the implicit model’s initial parameters, ${\bm{w}}_{0}$, before computing the $(N+1)$-th output.
With ${\bm{w}}_{0}={\bm{0}}$, the output is

$$ $\displaystyle o_{t}$ $\displaystyle=\beta{\bm{z}}_{t}^{T}{\bm{\Theta}}{\bm{c}}_{t}\,,$ (13) $$

where $\beta=-\frac{\eta}{N}$, $\eta$ is the learning rate, and $N$ is the number of training points or, equivalently, the total length of the context.
The SSM final output, $o_{t}$ above, corresponds to a prediction of the trained linear model, $\hat{y}_{t+1}$ in Eq. 10, if ${\bm{\Theta}}$ obeys ${\bm{\Theta}}{\bm{c}}_{t}={\bm{x}}_{t+1}$.
It is easy to check that ${\bm{\Theta}}$ satisfies this condition if all but its last $f$ diagonal entries are zero (see Figure 3 for a concrete example).
Note that the output in Eq. 13 matches the general form of the SSM output in Eq. 4 (without the non-linearity).
hows that the following SSM

$$ $\displaystyle{\bm{c}}_{t}$ $\displaystyle=[{\bm{x}}_{t}y_{t},{\bm{x}}_{t+1}]\,,$ (14) $\displaystyle{\bm{z}}_{t}$ $\displaystyle={\bm{I}}\,{\bm{z}}_{t-1}+{\bm{\Psi}}\,{\bm{c}}_{t}\,,$ (15) $\displaystyle o_{t}$ $\displaystyle=\beta{\bm{z}}_{t}^{T}{\bm{\Theta}}{\bm{c}}_{t}=\hat{y}_{t+1}\,.$ (16) $$

can perform gradient descent on the parameters ${\bm{w}}_{0}$ of the implicit linear model and use this mechanism for in-context learning.
The specific structure of the SSM in equation 14 demonstrates the importance of multiplicative processing, for both the inputs and outputs.

###### Proposition 1

### 3.2 Single Step N-dimensional linear regression

In this section, we generalize the construction above to the N-dimensional case.
Without loss of generality, we assume the input and the target, ${\bm{x}}$ and ${\bm{y}}$, have both dimensions $f$ i.e. ${\bm{x}},{\bm{y}}\in\mathbb{R}^{f}$.
If this is not the case, the input and output dimensionality can be matched by defining appropriate embeddings(^3^33E.g. with dimensions $k,l$ into the same higher dimension $k+l=f$ by concatenation with appropriately sized zero vectors, or through a linear transformation to dimension $f$.).
We can then treat the N-dimensional system as $f$ 1-D linear regression problems, one for each element of ${\bm{y}}$.

###### Proposition 2

Given a diagonal linear recurrent layer augmented with local self-attention with sliding window of size 3, and tokens ${\bm{s}}_{2j}={\bm{x}}_{j}$ and ${\bm{s}}_{2j+1}={\bm{y}}_{j}$, for $j=1,\ldots,N$, ${\bm{x}}_{j},{\bm{y}}_{j}$ drawn from a linear model, one can construct recurrent matrix ${\bm{A}}({\bm{s}}_{j})$, input ${\bm{B}}({\bm{s}}_{j})$ and output matrix ${\bm{U}}({\bm{s}}_{j})$ such that each recurrent step for every token ${\bm{s}}_{j}$ produces $\hat{{\bm{y}}}_{j+1}=-(\Delta{\bm{W}})^{T}{\bm{x}}_{j+1}$ as output, where $\Delta{\bm{W}}$ is one step of gradient descent, i.e. $\Delta{\bm{W}}=\eta\nabla_{\bm{W}}\mathcal{L}$.
The test input ${\bm{x}}_{N+1}$ is contained in token ${\bm{s}}_{2N+2}$, and produces the test prediction $\hat{{\bm{y}}}_{N+1}$.

Similar to Eq. 11, we show the above by writing the SSM as

$$ $\displaystyle{\bm{Z}}_{t}={\bm{Z}}_{t-1}+{\bm{y}}_{t}\,{\bm{x}}_{t}^{T}\,,$ (17) $$

where ${\bm{Z}}_{t}$ corresponds to the parameters of the implicit linear model, ${\bm{W}}\in\mathbb{R}^{f\times f}$, and we assume, for simplicity, that ${\bm{W}}={\bm{0}}$(^4^44The more general case is discussed in Appendix A.2.).
The output is

$$ $\displaystyle{\bm{o}}_{t}=\beta{\bm{Z}}_{t}\,{\bm{x}}_{t+1}\,.$ (18) $$

See Appendix A.1 for the full derivation.

To see how this can be written in the form of Eq. 3, let the input sequence consist of the training dataset of the implicit linear regression problem (as before).
This time, we cast the training dataset into a standard sequence ${\bm{s}}_{1},{\bm{s}}_{2},\ldots$, where

$$ $\displaystyle{\bm{s}}_{2j}$ $\displaystyle={\bm{x}}_{j}\,,$ (19) $\displaystyle{\bm{s}}_{2j+1}$ $\displaystyle={\bm{y}}_{j}\,,$ (20) $$

and ${\bm{s}}_{i}$ denotes the $i$-th sequence element of the input, such that ${\bm{s}}_{2j+2}={\bm{x}}_{j+1}$.

At each step, the state update, in Eq. 17, and the output, in Eq. 18, include ${\bm{x}}_{t},{\bm{y}}_{t},{\bm{x}}_{t+1}$.
This necessitates the introduction of a quadratic (in ${\bf s}_{i}$) local self-attention mechanism.
Let ${\bm{C}}_{t}$ be the context matrix for the local self-attention, i.e. a sliding window of length three running through the sequence,

$$ $\displaystyle{\bm{C}}_{t}$ $\displaystyle=\begin{bmatrix}\vdots&\vdots&\vdots\\ {\bm{x}}_{t}&{\bm{y}}_{t}&{\bm{x}}_{t+1}\\ \vdots&\vdots&\vdots\end{bmatrix}\,.$ (21) $$

The local self-attention operation, ${\bm{C}}{\bm{Q}}{\bm{C}}^{T}$, is a truncated form of Eq. 2.
When ${\bm{Q}}=\left(\begin{smallmatrix}1\\
0\\
0\end{smallmatrix}\right)\left(\begin{smallmatrix}0&1&0\end{smallmatrix}\right
)=\left[\begin{smallmatrix}0&1&0\\
0&0&0\\
0&0&0\end{smallmatrix}\right]$, plugging ${\bm{C}}{\bm{Q}}{\bm{C}}^{T}$ into the second term in Eq. 17 produces an SSM with input ${\bm{y}}_{t}{\bm{x}}_{t}^{T}$.

The SSM can now be written as in Eq. 3, i.e.

$$ $\displaystyle{\bm{Z}}_{t}={\bm{Z}}_{t-1}+{\bm{C}}_{t}{\bm{Q}}{\bm{C}}_{t}^{T}\,,$ $$

with output

$$ $\displaystyle{\bm{o}}_{t}=\beta{\bm{Z}}_{t}\,{\bm{C}}_{t}\,{\bm{q}}\,,$ $$

where ${\bm{q}}=\left(\begin{smallmatrix}0\\
0\\
1\end{smallmatrix}\right)$ makes the output $\hat{{\bm{y}}}_{t+1}$ as in Eq. 10.
The output also matches the form of the general SSM output in Eq. 4.
The construction allows us to perform gradient descent on the parameters of an arbitrary dimensional implicit model.
We call this type of SSM a GD-SSM.

Eq. 3 shows that the recurrent state of the SSM is $2$-dimensional, as in most recently proposed SSMs (Dao & Gu, ).
The local self-attention reproduces the local linear self-attention of Eqs. 2 and 5.
As for the 1-dimensional case, the self-attention higher-order dependencies, ${\bm{C}}{\bm{Q}}{\bm{C}}^{T}$, is key for making the SSM learn (in-context) an implicit (linear) regression model.

###### Proposition 2

### 3.3 Generalising to any regression problem

#### Multi-step gradient descent:

The proposed construction can be extended to multi-step gradient descent.
Since each layer of a GD-SSM produces the parameters of the implicit linear model updated by one step of GD, this is equivalent to stacking together multiple layers.
In our derivations above, we assumed the initial parameter of the implicit linear model is ${\bm{0}}$.
In the multi-step GD, all layers other than the first will correspond to a *non-zero initialised implicit linear model*.
Technically, extra gradient steps in the implicit model introduce one additional term in the gradient accumulation equation.
Each of the two terms requires a separate (parallel) recurrence and is performed by a dedicated layer.
At the end, the states are combined to obtain the multiple-step GD update, with minor extra computational burdens.
See Appendix A.2 for a detailed construction of multi-step GD.

#### Non-linear regression:

Non-linear regression can be handled by adding MLP layers to the GD-SSM.
In the previous sections, we let GD-SSM accumulate the gradients of a linear regressor.
Additional MLP layers can learn to transform the state of these linear layers into quantities corresponding to the gradient of the implicit non-linear model.
See Appendix A.3 for a detailed explanation.

#### Regularisation terms in the loss:

As the recurrent layers only accumulate the gradients, we can separate the gradient calculation and accumulation from its application.
This has a practical advantage.
Any input-independent regularisation term, e.g. L2 norms, can be added to the model without changing the recurrence structure.
See Appendix A.4 for details.

## 4 Trained linear recurrent networks do emulate gradient descent on linear regression tasks

We investigated if the GD-SSM variant of the general SSM architecture does do gradient descent in-context learning.
To do this, we trained a randomly initialised model on various in-context learning tasks for linear and non-linear regression.
In each case, the sequence token input to the model consisted of the inputs ${\bm{x}}$ and target values ${\bm{y}}$ from the training dataset $\mathcal{D}$, and the model was expected to output the prediction for the query (test) ${\bm{x}}$ given in the last timestep.
The models were trained to minimize the mean squared error between the test prediction and target.

### 4.1 Single step 1-dimensional linear regression

Figure: Figure 1: Comparing one step of GD with a trained single layer GD-SSM for 1-dimensional regression: A: Trained single layer GD-SSM loss and GD-SSM loss with the parameters from our construction are identical. B: Cosine similarity and the L2 distance between models as well as their predictions. C: Comparison of loss between Gradient Descent (GD) and the SSM layer model for different input sizes $N=N_{x}$. D: The trained single 1-D SSM layer, and gradient descent show identically loss (in log-scale) when provided input data different than during training i.e. with scale of 1. We display the mean/std. or the single runs of 5 seeds.
Refer to caption: x1.png

We first tested the simplest case of one-step gradient descent on a linear regression problem with scalar predictions/targets.
This corresponds directly to the construction in Section 3.1.
To do this, similar to Garg et al. (); von Oswald et al. (), we randomly generated linear regression tasks consisting of training and test points, and trained the model to make a prediction for the test input using the training input as context.

We generated randomly sampled linear regression tasks $\tau$ in the following way:
Each task (context) $\tau$ consisted of a sequence of in-context training data $\mathcal{D}_{\tau}=\left\{\langle{\bm{x}}_{\tau,i},y_{\tau,i}\rangle\right\}_{
i=1}^{N}$ and test point $\langle{\bm{x}}_{\tau,N+1},y_{\tau,N+1}\rangle$.
To generate this, the ${\bm{x}}_{\tau,i}$s are sampled from a uniform distribution ${\bm{x}}_{\tau,i}\sim U(-1,1)^{f}$.
Then, for each task $\tau$, the parameters of its implicit linear model ${\bm{w}}_{\tau}$ is sampled from a normal distribution, so that each element $[{\bm{w}}_{\tau}]_{i}\sim\mathcal{N}(0,1)$.
This is used to calculate the $y_{\tau,i}$s for each corresponding ${\bm{x}}_{\tau,i}$ using $y_{\tau,i}={\bm{w}}_{\tau}^{T}{\bm{x}}_{\tau,i}$.

The sequence ${\bm{S}}=\{{\bm{s}}_{\tau,1},...,{\bm{s}}_{\tau,N}\}$ is constructed so that ${\bm{s}}_{\tau,t}={\bm{c}}_{\tau,t}=[{\bm{x}}_{\tau,t}y_{\tau,t},{\bm{x}}_{
\tau,t+1}]$, with $[\ldots]$ denoting the vector concatenation operation, and ${\bm{c}}_{t}$ is the constructed context vector.
Note that this includes the query ${\bm{x}}_{\tau,N+1}$ in ${\bm{c}}_{N}$.
We will use a more general construction in the next section.
The outputs of the GD-SSM at time $T=N$ is the prediction for ${\bm{x}}_{N+1}$ i.e. $SSM({\bm{S}}_{\tau})_{N}=\hat{y}_{\tau,N+1}$, with target $y_{\tau,N+1}={\bm{w}}_{\tau}{\bm{x}}_{\tau,N+1}$.
The model was trained to minimize the expected squared prediction error, averaged over linear regression tasks $\tau$:

$$ $\displaystyle\min_{\theta}\mathbb{E}_{\tau}\left[\left\|\hat{y}_{\theta}\left( {\bm{c}}_{\tau,1},\ldots,{\bm{c}}_{\tau,N}\right)-y_{\tau,\text{test}}\right\| ^{2}\right]\,,$ $$

where $\theta$ are the randomly initialised parameters of the GD-SSM.
We evaluate our model on multiple metrics:

- 1.
L2 norm between the difference in predictions $\left\|\hat{y}_{\theta}(x_{\tau,\text{test}})-\hat{y}_{\theta_{\text{GD}}}(x_{
\tau,\text{test}})\right\|_{2}$ where $\hat{y}_{\theta_{\text{GD}}}$ is the prediction from the GD based construction.
- 2.
Cosine similarity between the sensitivities $\frac{\partial\hat{y}_{\theta_{\text{GD}}}(x_{\tau,\text{test}})}{\partial x_{
\text{test}}}\quad\text{and}\quad\frac{\partial\hat{y}_{\theta}(x_{\tau,\text{
test}})}{\partial x_{\text{test}}}$.
- 3.
L2 norm between the sensitivites.
- 4.
Loss of the two models where validation data is sampled from $U(-\alpha,\alpha)^{n_{i}}$, with a different $\alpha$ than the one used in training.
- 5.
The loss of two models when trained for different number of feature dimensions $f$.

Figure 1 shows the results of this comparisons.
We find that these metrics show excellent agreement between the trained model and GD, over different hyperparameters.
To test if the trained network has learned a general purpose learning rule as opposed to fitting to the dataset, that is to test if the trained network generalises to linear regression tasks outside the training distribution, we draw values of inputs from $U(-\alpha,\alpha)$.
We see that our model generalises to both cases.

### 4.2 Single step N-dimensional linear regression

Figure: Figure 2: Comparing one step of GD with a trained single layer GD-SSM for N-dimensional regression: A: Trained single layer GD-SSM loss and GD-SSM loss with the parameters from our construction are identical. B: Cosine similarity and the L2 distance between models as well as their predictions. C: Comparison of loss between Gradient Descent (GD) and the SSM layer model for different input sizes $f$. D: The trained GD-SSM layer, and gradient descent show identically loss (in log-scale) when provided input data different than during training i.e. with scale of 1. We display the mean/std. or the single runs of 5 seeds.
Refer to caption: x2.png

Figure: Figure 3: Comparison of learned weights a GD-SSM with the GD-SSM parameters from our construction. A: Comparison of local self attention weights of trained GD-SSM with the GD-SSM parameters from our construction. B: Comparison of recurrent parameters of trained GD-SSM with the GD-SSM parameters from our construction. Since the recurrence parameters are tensors, for the ease of visualisation, each diagonal entry is the mean of the corresponding diagonal recurrence matrix. C: Comparison of skip connection weights of trained GD-SSM with the GD-SSM parameters from our construction.
Refer to caption: x3.png

In the more general case where we allow the targets ${\bm{y}}_{\tau,i}$ to be of arbitrary dimension, we also relax other assumptions in the form of the input.
Notably, instead of constructing the inputs in a particular way, we let the sequence be similar to a more natural sequence of ${\bm{x}}_{1},{\bm{y}}_{1},{\bm{x}}_{2},{\bm{y}}_{2},\ldots$, i.e. ${\bm{s}}_{2j}={\bm{x}}_{j};{\bm{s}}{2j+1}={\bm{y}}_{j}$.
This requires the introduction of a local self-attention mechanism with attention window of $3$ as discussed in Section 3.2.
The local self-attention mechanism allows the model to access inputs from multiple adjacent timesteps, and not just the current one.
We perform the same comparisons as in Section 4.1, and find that the model trained from a random initialisation has an excellent match with the construction, both in the output/loss metrics (Figure 2) and in the parameters it ends up with (Figure 3).

We compared loss metrics for GD-SSM with other standard models including one- and two-layer transformers, and S5 (Smith et al., ).
We found that, on exactly the same sequence token inputs, only the GD-SSM and a 2-layer transformer reached the same loss as actual GD, while other SSMs such as S5 struggle to perform ICL with 1 layer (Figure 4C).

Figure: Figure 4: A: Comparison of GD-SSM performance with single layer and two layers on regression task. B: Comparison of GD-SSM performance with and without MLP layers in single layer and multi-layer setup on non-linear regression task. C: Comparison with other models on 1-D regression. The GD-SSM model was evaluated with both 1-layer and 2-layer configurations, and the S5, Mamba, and Griffin models were included for comparison. TF refers to linear Transformer models, with both 1-layer and 2-layer variants tested to evaluate their performance. A similar plot for N-D regression is shown in Figure 6.
Refer to caption: x4.png

### 4.3 Multi-step and non-linear regression

To perform multi-step GD, we tested a GD-SSM with multiple-layers, where each layer of the GD-SSM does 1-step of GD as shown in Appendix A.2.
The token construction is identical to that used for N-dimensional linear regression, and each layer includes the local self-attention mechanism.
In Figure 4A, we show that that trained network is able to reach the same loss as the multi-step GD.

Finally, to be able to handle any regression task, we tested GD-SSM with an MLP layer on non-linear regression tasks.
Again, we see that the trained model is able to reach the same level of performance as performing GD directly on the dataset (Figure 4B).

## 5 Related work

In-context learning in transformers has been studied extensively, and various mechanisms have been proposed to explain it (Hendel et al., ; von Oswald et al., ; Akyürek et al., ).
Of those, the most prominent is that the self-attention mechanism performs gradient descent on a linear loss.
A construction with linear self-attention was demonstrated by von Oswald et al. ().
While linear-self attention can be written as an RNN (Katharopoulos et al., ), the results of von Oswald et al. () depends on the tokens being constructed in a specific way, which limits the generality of their construction.
Moreover, the most common type of self-attention used is softmax scaled dot-product self-attention rather than linear self-attention, and the basic linear self-attention mechanism used in von Oswald et al. () is not competitive with softmax self-attention.
Whereas, we show a construction that uses local-self attention with state-space models, which has been shown to be competitive with softmax self-attention transformers (De et al., ).

Liu et al. () formulate the SSM layer as an online optimization objective with exact solution.
Similarly, Sun et al. () use the state of the SSM to perform GD.
But both these papers consider online updates, that is minibatch size 1 updates.
One layer of their model is not capable of performing larger minibatch updates.
Moreover, without local self-attention, their online optimization objective is limited to using single sequence tokens with a single-layer.
And finally their goal is not to explain how in-context learning happens with existing architectures, but rather to develop entirely new SSM variants using the idea that an optimization process can be used for compressing information as well, which is the key property required of a recurrent network.

The combination of local-self attention with Mamba has been used in for improving performance in multiple instances, including most recently (De et al., ; Ren et al., ).
But these papers do not specifically construct their model for ICL, nor do they attempt to distill inductive biases required for ICL in state-space models.

## 6 Discussion

This work establishes a clear connection between state-space models (SSMs) and gradient-based in-context learning.
We have demonstrated, both theoretically and empirically, that SSMs can emulate gradient descent on implicit regression models, providing a mechanistic explanation for their in-context learning capabilities.
Our construction, GD-SSM, reveals the crucial role of architectural features such as local self-attention and multiplicative output interactions in enabling this behavior.
These findings not only explain the success of recent SSM variants in in-context learning tasks but also provide valuable insights for the design of future sequence models.

The alignment between our theoretical construction and the behavior of trained models suggests that the gradient descent mechanism may be a natural inductive bias in these architectures.
This understanding opens new avenues for analyzing and improving sequence models, potentially leading to more efficient and effective architectures for a wide range of tasks.

Future research could explore the implications of this mechanism in larger-scale models, more complex tasks, and real-world applications.
Additionally, investigating how this understanding can be leveraged to enhance the design and training of state-space models could yield significant advancements in the field of sequence modeling.

## Acknowledgments

We thank all the authors for their invaluable contributions and unwavering support throughout this research, and we are particularly grateful to Mark Schöne for his insightful feedback and constructive suggestions that greatly improved this work. Neeraj Mohan Sushma and Yudou Tian are funded by the German Federal Ministry of Education and Research (BMBF) project EVENTS (16ME0733). David Kappel is funded by the German Federal Ministry for Economic Affairs and Climate Action (BMWK) project ESCADE (01MN23004D). We sincerely thank the Gauss Centre for Supercomputing for supporting this project through the provision of computing time on the GCS Supercomputer JUWELS at the Jülich Supercomputing Centre (JSC).

## Appendix A Appendix

### A.1 N-D linear regression

The implicit linear model of N-dimensional ($N=f$) linear regression is

$$ $\hat{{\bm{y}}}={\bm{W}}^{T}{\bm{x}}\,,$ $$

for ${\bm{W}}\in\mathbb{R}^{f\times f}$ and $\hat{{\bm{y}}},{\bm{x}}\in\mathbb{R}^{f}$.
The loss is

$$ $\displaystyle\mathcal{L}(\mathcal{D};{\bm{W}}_{0})=\frac{1}{2N}\sum_{i}\left({ \bm{W}}_{0}^{T}{\bm{x}}_{i}-{\bm{y}}_{i}\right)^{T}\,\left({\bm{W}}_{0}^{T}{ \bm{x}}_{i}-{\bm{y}}_{i}\right)\,.$ (22) $$

The gradient of this loss calculated using the first $t$ samples $\mathcal{D}_{1:t}$ is

$$ $\displaystyle\nabla_{{\bm{W}}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{W}}_{0})$ $\displaystyle=\frac{1}{t}\sum_{i=1}^{t}\left.\frac{\partial}{\partial W}\left( {\bm{W}}^{T}\,{\bm{x}}_{i}-{\bm{y}}_{i}\right)^{T}\left({\bm{W}}^{T}\,{\bm{x}} _{i}-{\bm{y}}_{i}\right)\right|_{{\bm{W}}={\bm{W}}_{0}}\,.$ $$

For notational simplicity, we will write this as $f$ 1-D regression problems, one for each element of ${\bm{y}}$.
Using $[{\bm{W}}]_{:,i}$ to denote the $i$-th column of matrix ${\bm{W}}$, using $[{\bm{W}}_{0}]_{:,i}$ to denote the $i$-th column of matrix ${\bm{W}}_{0}$ and using $[{\bm{y}}_{t}]_{i}=[{\bm{W}}]_{:,i}^{T}\,{\bm{x}}_{t}$ to denote the $i$-th element of vector ${\bm{y}}_{t}$,

$$ $\displaystyle{\bm{g}}(\mathcal{D}_{1:t};[{\bm{W}}_{0}]_{:,i})$ $\displaystyle={\bm{g}}(\mathcal{D}_{1:t-1};[{\bm{W}}_{0}]_{:,i})+\left([{\bm{W }}_{0}]_{:,i}^{T}\,{\bm{x}}_{t}-[{\bm{y}}_{t}]_{i}\right){\bm{x}}_{t}\,,$ (23) $$

where $\frac{1}{t}{\bm{g}}(\mathcal{D}_{1:t};[{\bm{W}}_{0}]_{:,i})=\left[\nabla_{{\bm
{W}}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{W}}_{0})\right]_{:,i}$.

#### Implementation as a linear recurrent network:

The input is given as in Eqs. 19 &  20, with context matrix constructed as in Eq. 21.

The local self attention calculates ${\bm{C}}{\bm{Q}}{\bm{C}}^{T}$, which is then provided as input to the SSM layer.
If ${\bm{Q}}=\left(\begin{smallmatrix}1\\
0\\
0\end{smallmatrix}\right)\left(\begin{smallmatrix}0\\
1\\
0\end{smallmatrix}\right)^{T}=\left[\begin{smallmatrix}0&1&0\\
0&0&0\\
0&0&0\end{smallmatrix}\right]$, this corresponds to the input to the SSM being $[{\bm{y}}_{t}]_{i}{\bm{x}}_{t}$ (assuming ${\bm{W}}={\bm{0}}$).

Defining ${\bm{Z}}_{t}\in\mathbb{R}^{f\times f}$ as the state of the SSM, where $[{\bm{Z}}_{t}]_{:,i}={\bm{g}}(\mathcal{D}_{1:t};[{\bm{W}}_{0}]_{:,i})^{T}$ (note the transpose), and assuming ${\bm{W}}_{0}={\bm{0}}$, the equations Eq. A.2 for all $i$s (all columns) can be written as a single equation

$$ $\displaystyle{\bm{Z}}_{t}={\bm{Z}}_{t-1}+{\bm{C}}_{t}{\bm{Q}}{\bm{C}}_{t}^{T}\,,$ $$

with output

$$ $\displaystyle{\bm{o}}_{t}=-\frac{\eta}{N}{\bm{Z}}_{t}\,{\bm{C}}_{t}\,{\bm{q}}\,.$ $$

If ${\bm{q}}=\left(\begin{smallmatrix}0\\
0\\
1\end{smallmatrix}\right)$, this corresponds exactly to the output $\hat{{\bm{y}}}_{t+1}=-\eta\nabla_{{\bm{W}}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{
W}}_{0})^{T}\,{\bm{x}}_{t+1}$.
This provides a construction that allows us to perform gradient descent on the parameters of an arbitrary dimensional input model.

### A.2 Multi-step GD

###### Proposition 3

Given $1,\ldots,L$ diagonal linear recurrent layers, each augmented with local self-attention with sliding window of size 3, and sequence tokens ${\bm{s}}_{2j}={\bm{x}}_{j}$ and ${\bm{s}}_{2j+1}={\bm{y}}_{j}$, for $j=1,\ldots,N$, drawn from a linear model, for each layer $l$, one can construct pairs of recurrent matrices ${\bm{A}}^{(l,1)}({\bm{s}}_{j}),{\bm{A}}^{(l,2)}({\bm{s}}_{j})$, inputs ${\bm{B}}^{(l,1)}({\bm{s}}_{j}),{\bm{B}}^{(l,2)}({\bm{s}}_{j})$ and output matrices ${\bm{U}}^{(l,1)}({\bm{s}}_{j}),{\bm{U}}^{(l,2)}({\bm{s}}_{j})$ such that each recurrent step for every token ${\bm{s}}_{j}$ produces $\hat{{\bm{y}}}_{j+1}=(\Delta_{l}{\bm{W}})^{T}{\bm{x}}_{j+1}$ as output, where $\Delta_{l}{\bm{W}}$ is the update for $l$ steps of gradient descent, i.e. $\Delta_{l}{\bm{W}}=(({\bm{W}}_{0}-\eta\nabla_{\bm{W}}\mathcal{L}(\mathcal{D};{
\bm{W}}_{0}))-\eta\nabla_{\bm{W}}\mathcal{L}(\mathcal{D};({\bm{W}}_{0}-\eta
\nabla_{\bm{W}}\mathcal{L}(\mathcal{D};{\bm{W}}_{0})))...l\,\text{times})$.
The test input ${\bm{x}}_{N+1}$ is contained in token ${\bm{s}}_{2N+2}$, and produces the test prediction $\hat{{\bm{y}}}_{N+1}$.

We have so far assumed that ${\bm{W}}_{0}={\bm{0}}$ thus far, which for the first layer corresponds to initialising the parameters of the equivalent linear model to all zeros.
For multi-step GD, the second layer onwards have a non-zero initial value of parameters, so let’s derive a general form for a layer that performs GD with non-zero initialisation.

Starting from Eq. A.2, repeated below for convenience,

$$ $\displaystyle{\bm{g}}(\mathcal{D}_{1:t};[{\bm{W}}_{0}]_{:,i})$ $\displaystyle={\bm{g}}(\mathcal{D}_{1:t-1};[{\bm{W}}_{0}]_{:,i})+\left([{\bm{W }}_{0}]_{:,i}^{T}\,{\bm{x}}_{t}-[{\bm{y}}_{t}]_{i}\right){\bm{x}}_{t}\,,$ $$

we note that this accumulates two different components – $\left([{\bm{W}}_{0}]_{:,i}^{T}\,{\bm{x}}_{t}\,\right){\bm{x}}_{t}$ and $[{\bm{y}}_{t}]_{i}\,{\bm{x}}_{t}$.

We propose having two different heads (per layer) to accumulate these quantities separately (i) $[{\bm{y}}_{t}]_{i}\,{\bm{x}}_{t}$ and (ii) ${\bm{x}}_{t}\,{\bm{x}}_{t}$.

Defining ${\bm{Z}}_{t}\in\mathbb{R}^{m\times m}$ as the state of the SSM, where $[{\bm{Z}}_{t}]_{:,i}={\bm{g}}(\mathcal{D}_{1:t};[{\bm{W}}_{0}]_{:,i})^{T}$ (note the transpose), the recurrent network corresponding to the accumulation of (i) is the same as before:

$$ $\displaystyle{\bm{Z}}_{t}={\bm{Z}}_{t-1}+{\bm{C}}_{t}{\bm{Q}}{\bm{C}}_{t}^{T}\,.$ $$

For (ii), we can write an equivalent recurrent network

$$ $\displaystyle\tilde{{\bm{Z}}}_{t}=\tilde{{\bm{Z}}}_{t-1}+{\bm{C}}_{t}\tilde{{ \bm{Q}}}{\bm{C}}_{t}^{T}\,,$ $$

If $\tilde{{\bm{Q}}}=\left(\begin{smallmatrix}1\\
0\\
0\end{smallmatrix}\right)\left(\begin{smallmatrix}1\\
o\\
0\end{smallmatrix}\right)^{T}=\left[\begin{smallmatrix}1&0&0\\
0&0&0\\
0&0&0\end{smallmatrix}\right]$, this corresponds to the input to the SSM being ${\bm{x}}_{t}{\bm{x}}_{t}^{T}$.

The output at layer $l$ is

$$ $\displaystyle{\bm{o}}_{t}^{(l)}=\left({\bm{W}}_{l-1}+\Delta_{l}{\bm{W}}\right) \,{\bm{C}}_{t}\,{\bm{q}}\,,$ $$

where

$$ $\displaystyle\Delta_{l}{\bm{W}}=-\frac{\eta}{N}\left({\bm{W}}_{l-1}^{T}\, \tilde{{\bm{Z}}}_{t}^{(l)}-{\bm{Z}}_{t}^{(l)}\right)\,.$ $$

In summary, at each layer, there are two linear recurrent layers, and the output includes a multiplicative combination across layers and with the external input.

Since each recurrent layer is performing the same operation, one could also loop the output of each layer back to do multiple steps of GD.

###### Proposition 3

### A.3 Non-linear GD

Let us consider a non-linear regression problem with a least squares loss, where the output labels are again 1-dimensional for simplicity.
For a given dataset of $N$ samples $\mathcal{D}=\{<{\bm{x}}_{i},y_{i}>\}_{i=0}^{N},{\bm{x}}\in\mathbb{R}^{f},y\in
\mathbb{R}$,
predictions from a non-linear model are generated using

$$ $\hat{y}=g({\bm{w}}^{T}{\bm{x}})\,,$ $$

for ${\bm{w}}\in\mathbb{R}^{f}$ and where $g$ is some non-linear function such as sigmoid or an MLP.
This will be our implicit non-linear model for this 1-dimensional target case.

The best fit for ${\bm{w}}$ is found by minimizing the loss

$$ $\displaystyle\mathcal{L}(\mathcal{D};{\bm{w}}_{0})=\frac{1}{2N}\sum_{i=1}^{N}| |\hat{y}_{i}-y_{i}||_{2}^{2}=\frac{1}{2N}\sum_{i}\left(g({\bm{w}}_{0}^{T}{\bm{ x}}_{i})-y_{i}\right)^{2}\,.$ $$

The gradient of the loss calculated on the first $t$ samples of the dataset is

$$ $\displaystyle\nabla_{{\bm{w}}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{w}}_{0})$ $\displaystyle=\frac{1}{t}\sum_{i=1}^{t}\left(g({\bm{w}}_{0}^{T}\,{\bm{x}}_{i}) -y_{i}\right)g^{\prime}({\bm{w}}_{0}^{T}\,{\bm{x}}_{i}){\bm{x}}_{i}\,,$ $$

where $\mathcal{D}_{1:t}$ denotes the first $t$ samples in $\mathcal{D}$, and $g^{\prime}$ denotes the first derivative of $g$.

We define the term ${\bm{g}}$ but without containing $g$, as

$$ $\displaystyle{\bm{g}}_{{\bm{w}}_{0}}(\mathcal{D}_{1:t})$ $\displaystyle=\sum_{i=1}^{t}\left({\bm{w}}_{0}^{T}\,{\bm{x}}_{i}-y_{i}\right){ \bm{x}}_{i}\,,$ $$

which can be recursively calculated as before.
Note that this quantity is not the unscaled gradient anymore.
In this case we have the following Proposition:

###### Proposition 4

Given a diagonal linear recurrent layer augmented with local self-attention with sliding window of size 3, followed by a MLP layer, and tokens ${\bm{s}}_{2j}={\bm{x}}_{j}$ and ${\bm{s}}_{2j+1}={\bm{y}}_{j}$, for $j=1,\ldots,N$, drawn from a non-linear model, one can construct recurrent matrix ${\bm{A}}({\bm{s}}_{j})$, input ${\bm{B}}({\bm{s}}_{j})$ and output matrix ${\bm{U}}({\bm{s}}_{j})$ such that each recurrent step for every token ${\bm{s}}_{j}$ produces $\hat{{\bm{y}}}_{j+1}=-(\Delta{\bm{W}})^{T}{\bm{x}}_{j+1}$ as output, where $\Delta{\bm{W}}$ is one step of gradient descent, i.e. $\Delta{\bm{W}}=\eta\nabla_{\bm{W}}\mathcal{L}$.
The test input ${\bm{x}}_{N+1}$ is contained in token ${\bm{s}}_{2N+2}$, and produces the test prediction $\hat{{\bm{y}}}_{N+1}$.

Writing this as an SSM exactly as in 12 (with the same ${\bm{\Psi}}$)

$$ $\displaystyle{\bm{z}}_{t}$ $\displaystyle={\bm{I}}\,{\bm{z}}_{t-1}+{\bm{\Psi}}\,{\bm{c}}_{t}\,,$ $$

we now need to calculate the output as

$$ $\displaystyle\hat{y}_{t+1}$ $\displaystyle=-\eta\nabla_{{\bm{w}}_{0}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{w}} _{0})^{T}\,{\bm{x}}_{t+1}\,.$ $$

The interleaved MLP layer $f$ would need to learn a mapping such that

$$ $\displaystyle f\left(\sum_{i=1}^{t}\left({\bm{w}}_{0}^{T}\,{\bm{x}}_{i}-y_{i} \right){\bm{x}}_{i}\right)=\sum_{i=1}^{t}\left(g({\bm{w}}_{0}^{T}\,{\bm{x}}_{i })-y_{i}\right)g^{\prime}({\bm{w}}_{0}^{T}\,{\bm{x}}_{i}){\bm{x}}_{i}\,,$ $$

which it will be able to since it is an universal approximator.
This can then be used to calculate the output as

$$ $\displaystyle o_{t}$ $\displaystyle=f({\bm{z}}_{t})^{T}{\bm{\Theta}}{\bm{c}}_{t}\,.$ $$

###### Proposition 4

### A.4 Regularisation terms in the loss

For example, if we wished to change the loss function in Eq. 22 to include L2 regression,

$$ $\displaystyle\mathcal{L}(\mathcal{D};{\bm{W}})=\frac{1}{2N}\sum_{i=1}^{N}|| \hat{{\bm{y}}}_{i}-{\bm{y}}_{i}||_{2}^{2}+||{\bm{W}}||_{2}^{2}\,,$ $$

this will change the gradient to be

$$ $\displaystyle\nabla_{{\bm{w}}}\mathcal{L}(\mathcal{D}_{1:t};{\bm{W}}_{0})$ $\displaystyle=\frac{1}{t}\sum_{i=1}^{t}\frac{\partial}{\partial{\bm{W}}}||\hat {{\bm{y}}}_{i}-{\bm{y}}_{i}||_{2}^{2}+2{\bm{W}}_{0}\,.$ $$

To construct an SSM that does GD on this loss instead of the one in Eq. 8, all we’d have to do is change the $\Delta W$ in the output to be

$$ $\displaystyle\Delta_{l}{\bm{W}}=-\frac{\eta}{N}\left({\bm{W}}_{l-1}^{T}\, \tilde{{\bm{Z}}}_{t}^{(l)}-{\bm{Z}}_{t}^{(l)}+{\bm{W}}_{l-1}\right)\,.$ $$

### A.5 Experimental details of 1-d regression

The dataset construction and the model evaluation methods are same as that used in von Oswald et al. ().
Each task (context) $\tau$ consists of in-context training data $D_{\tau}=\left\{(x_{\tau,i},y_{\tau,i})\right\}_{i=1}^{N}$ and test point $(x_{\tau,N+1},y_{\tau,N+1})$.
At every optimization step of the model, we sample the regression parameters $W_{\tau}\sim\mathcal{N}(0,I)$.
We then sample $x_{\tau,i}\sim U(-1,1)^{f}$ and construct a scalar target $y_{\tau,i}=W_{\tau}x_{\tau,i}$, where ${f}$ is the feature size of the input.
To evaluate the model, a set of $10^{4}$ tasks are sampled and mean squared error is calculated.

To evaluate the proposed model, we conducted experiments using a constructed token dataset with an input feature size of 10. The model architecture is a single-layer, state space model (SSM) with a hidden dimension of 20. Initialization was performed with 2 blocks and an SSM latent size of 10. No activation function was used during training.

The training process spanned a maximum of 300,000 epochs with a batch size of 64. A cosine annealing schedule and linear warmup were utilized for the learning rate, beginning with an initial SSM learning rate of $1\times 10^{-4}$ for optimizing the SSM parameters using the AdamW optimizer. The global learning rate for the remaining parameters was set to $2\times 10^{-4}$, and these parameters were also optimized using AdamW. A weight decay of 0.05 was applied to regularize the model and prevent overfitting.

Figure: Figure 5: Visualisation of the trained parameters for two layer GD-SSM on linear regression task.
Refer to caption: x5.png

Figure: Figure 6: Comparison with other models on N-D regression. The GD-SSM model was evaluated with both 1-layer and 2-layer configurations, and the S5, Mamba, and Griffin models were included for comparison. TF refers to linear Transformer models, with both 1-layer and 2-layer variants tested to evaluate their performance.
Refer to caption: x6.png

### A.6 Experimental details of N-d regression

The setup is similar to 1-d regression, but we sample $n_{o}$ different regression parameters $W_{\tau}^{k}\sim\mathcal{N}(0,I)$, where $n_{o}$ is the dimension of vector $y_{\tau,i}$ and $1<=k<=n_{o}$. We then construct target $y_{\tau,i}^{k}=W_{\tau}^{k}x_{\tau,i}$ for all $k$.
For all our experiments, we choose $n_{o}=f$. We have evaluated the model based on all the experiments that are used for the 1-D regression evaluation.

To have the model prediction equivalent to one step gradient descent, we train the single layer GD-SSM using Adam optimizer, with an initial learning rate of 0.0001 for recurrent parameters with cosine annealing. For all the other parameters we double the learning rate that is used for recurrent parameters. In all our experiments, each optimization step contains 64 tasks.

### A.7 Experimental details of multi-step and non-linear regression

To emulate the multi-step regression task, we train the multi-layer GD-SSM architecture, without applying any non-linearity between the layers. In our experiments, we use two layer GD-SSM architecture. We use the same hyper-parameters that are used for the single layer GD-SSM. We compare the validation loss of the regression task between the learned model and the model based on construction. We observe the performance of two layer GD-SSM better than the GD-SSM initialised with the parameters from our construction, so we train the GD-SSM model initialised with the parameters from our construction 1000 steps to match the performance of a trained GD-SSM.

For non-linear regression tasks, we use the experiments from Finn et al. () and follow the data construction of von Oswald et al. (). To the output of the linear GD-SSM layer, we add a non-linear function, weighted sigmoid gated unit (Tanaka, ), as used in (Smith et al., ). For all the non-linear regression tasks, an input embedding layer is used in addition to the model architecture used for linear regression tasks. The hyper-parameters are the same that is used for all the previous tasks, and the trained model loss is compared with a model with GD-SSM layer/layers based on gradient descent construction and trained non-linearity layer/layers. Similar to multi-step linear regression, we observe the GD-SSM performance better than the GD-SSM initialised with the parameters from our construction.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: Please see weaknesses 1-3.
soundness: 3
strengths: (1) Enhances understanding of inductive biases in SSMs for ICL tasks.

(2) Bridges key domains such as SSMs, ICL, and mechanistic interpretability.

(3) Demonstrates generality by extending results to multi-step cases, multiple layers, and multi-dimensional data.

(4) Empirical analysis shows practical alignment with the theory-based construction in simple cases.

(5) Clarity: While some aspects could be improved (such as adding a figure to illustrate the main ideas in Section 3.1), the paper is clear, well-motivated, and easy to follow. It starts with simple cases and provides clear definitions, making it an enjoyable read!
summary: The paper investigates the ability of SSMs to perform ICL through gradient descent during the forward pass over recurrent steps. It provides a theoretical analysis of various ICL scenarios, demonstrating that simple SSMs, equipped with input and output-dependent processing, can accurately mimic gradient descent when data is presented as a sequence. This theory offers an explanation for the capacity of modern SSMs to execute ICL and outlines a network circuit capable of handling such tasks. Empirical experiments confirm that the proposed circuit can effectively learn and perform small-scale synthetic tasks in practice.
weaknesses: (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**. 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. 

(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. 

(3) **Understanding ICL in SSM variants can be enhanced** by examining their sub-components. Previous research indicates that S6 layers exhibit significantly better ICL capabilities than earlier SSMs [7]. However, while the authors highlight input and output-dependent processing as crucial features, they do not empirically ablate these features across various ICL tasks, nor do they provide a detailed theoretical analysis to substantiate this claim explicitly. I recommend adding a subsection that explores these aspects in depth. It is also important to note that input- and output-dependent processing can be implemented through various gating mechanisms. Hence, this claim could be considered somewhat ambiguous, as gated state-space models have been previously studied without demonstrating the same level of ICL capabilities as models like Mamba / S6.


(4) **The claims regarding GD-SSM** (“Our construction, which we call GD-SSM, is not restricted to in-context learning tasks and performs well on general-purpose prediction problems”) do not hold, and much **more empirical analysis is required to justify** them.

___

[1] Learning to learn by gradient descent by gradient descent. Andrychowicz et al.

[2] The Hidden Attention of Mamba Models. Ali et al.

[3] Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. Dao et al.

[4] Understanding the differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks. Sieber et al.

[5] Transformers Learn In-Context by Gradient Descent. Oswald et al. (see section 2)

[6] Why Can GPT Learn In-Context? Language Models Implicitly Perform Gradient Descent as Meta-Optimizers. Dai et al. (see section 3.1)

[7] Can mamba learn how to learn? a comparative study on in-context learning tasks. Park et al.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Can you explicitly define GD-SSM? Perhaps even provide pseudo-code? What are its learnable parameters? It is defined in 283 after the fact, but I think a more explicit definition would be helpful.

Questions related to experiments and first two bullets in Weaknesses section above:

2. How do the results in this paper explain the success of recent SSM variants (as claimed in lines 488)? Line 052-053 say : "Which features of these successful models contribute to in-context learning, as opposed to earlier variants? Using a constructive approach, we pinpoint input-dependent input and output processing, as the key features required for in-context learning". I hoped the paper would provide insight into this question. However, it seems that the constructed GD-SSM is quite different from the currently used SSM architectures (e.g. Griffin, Mamba), and performs much better empirically in the experiments. Meanwhile the currently used SSM architectures do not seem to perform regression in-context learning (at least for one layer). So how do the results in this paper answer (or point to answering) the first question in line 052 or support the claim in line 488? This is my main question regarding this paper. I list some additional questions below that are an attempt to clarify some of the presented empirical results below.

3. Why does Griffin do so poorly compared to linear transformers and S5? Shouldn't Griffin be a mix of sliding window attention and SSM (the RG-LRU), making it similar to the GD-SSM formulation? Or is the model referred to as Griffin just RG-LRU?

4. Why does the time-invariant S5 appear to consistently outperform the input-dependent Mamba and Griffin models (even though all fail to converge)? I would have expected the models with input-dependent dynamics to perform ICL better.  Is this difference at all interesting?

5. Does it benefit the other architectures (1 layer linear attention, S5, Griffin, Mamba) to also provide them with the local self-attention preprocessing? Shouldn't they then be able to potentially learn the GD-SSM model if provided this? If not, why not?

6. Can 2 layers of the SSMs solve the task? Note that combining the local self attention with the diagonal SSM is a combination of 2 sequence processors.


Other questions and comments:

7. Where is the 1-step GD in Figure A?

8. In Figure 4B, the 1 or 2 layer GD-SSM formulation always seems to achieve lower loss than the corresponding number of steps GD model. Why is this? What happens if we compare more layers and more steps, e.g. 5 or 10 steps/layers? 

9. Note that the color scheme in Figure 4C etc is hard to read and tell which line corresponds to which model.

10. Note that LSA is first introduced in equation 2, but only later defined in line 119.
soundness: 2
strengths: - To my knowledge, the insight viewing SSMs as a gradient accumulator allowing them to emulate gradient descent on in-context learning tasks is novel and the combination with local self-attention for preprocessing is interesting
- The mathematical theory appears to be sound
- The presentation of the material and step by step walk through of the theory from simple cases to more complex is clear and helpful
- The theoretical findings potentially point to a mechanistic understanding of architectural requirements (for SSMs or other models) that enable types of in-context learning
summary: This paper investigates how state-space models (SSMs) can perform in-context learning through gradient descent. The authors provide both theoretical and empirical evidence that SSMs augmented with local self-attention can emulate gradient descent on implicit regression models. Their key insight is that the diagonal linear recurrent layer in SSMs can act as a gradient accumulator.
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.
   - Note that line 488 says: "These findings not only explain the success of recent SSM variants in in-context learning tasks but
also provide valuable insights for the design of future sequence models." 
   -  Note that Lines 052-053 say, referring to the modern SSM variants : "Which features of these successful models contribute to in-context learning, as opposed to earlier variants? Using a constructive approach, we pinpoint input-dependent input and output processing, as the key features required for in-context learning". 
   - But I do not think the current version of the paper supports either of the claim that this paper sheds light on these questions 
   - The approach constructed seems to be very different from those used in practice. In addition, the methods commonly used in practice do not appear to do well on the regression ICL tasks in this paper. So it is unclear what I should take away from this related to prior SSM methods?
   - I think the empirical results were an opportunity to provide insight here, but didn't seem to fully achieve this. Please see my questions below which may clarify this for me.

- 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.

- 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. 

- 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.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: To the best of my understanding and as mentioned above, the results presented here are a subset of the results presented in Zucchet et al. 2023. Can the authors compare their work to that paper and highlight what their insights are?
soundness: 4
strengths: The paper is overall well written and easy to follow. The theoretical part is sound and experiments convincingly demonstrate the paper's claims in toy settings.
summary: In-context learning (ICL) is one of the surprising capabilities of LLMs at scale. Seminal results have shown that Transformer-based architectures have this ability and more recent ones confirmed that SSMs also do. This work studies shows that SSMs can implement ICL by gradient descent. They provide a constructive proof showing that 1 layer can implement one GD step and confirm empirically on toy tasks that the networks find this solution.
weaknesses: I am concerned by the novelty of this paper. [Zucchet et al. 2023](https://arxiv.org/abs/2309.01775) show that 1 SSM layer can implement any linear self-attention layer. This results implies that any ICL algorithm LSA can implement, an SSM can. This holds for 1 GD step studied in this paper, but also for any other algorithm the community has been studying over the last few years. Additionally, this paper also have very similar experiments to the ones presented here.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: 1. See Weakness.
2. The behavior of GD and 1-D GD-SSM appears different. Could the authors provide an explanation for this discrepancy? Clarifying this would help the readers better understand the distinctions in learning dynamics between these approaches.
3. Figure 3: The font size in Figure 3 is too small and could be increased for readability.
soundness: 3
strengths: 1. The paper demonstrates that the state-space model with LSA can learn in-context learning tasks on linear and nonlinear regression problems. While I do not dispute the claim that state-space models can achieve in-context learning via gradient descent, my concern lies in whether the specific modification introduced warrants the detailed calculations provided. The conclusions, as presented, seem to offer limited insights into how this work might advance research on improving in-context learning capabilities. A clearer connection to meaningful improvements in this area would significantly enhance the paper's contribution.
2. The experiments compare the state-space model with Griffin and Linear Transformer, but they are restricted to shallow architectures of 1 or 2 layers. This setup is inconsistent with typical in-context learning scenarios, where models need to be sufficiently large for emergent phenomena and meaningful in-context learning capabilities to surface. The current experiments do not effectively capture these dynamics, making it difficult to observe such phenomena in shallow sequence models. Expanding the experiments to include deeper architectures would provide a more realistic assessment of in-context learning in state-space models.
summary: This paper imitates Transformers learn in-context by gradient descent(https://arxiv.org/abs/2212.07677). This paper proves that a single structured state-space model layer with local self-attention can reproduce the outputs of an implicit linear model with least square loss. (The task considered is not general)
weaknesses: 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. 
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.
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.

Meta-review: The paper proposes that ssms can perform in-context learning through gradient descent during their forward pass. The authors try to provide theoretical and empirical evidence that SSMs augmented with local self-attention can emulate gradient descent on implicit regression models. Their key insight is that the diagonal linear recurrent layer in ssms can act as a gradient accumulator. 

positive: 
- novel theoretical perspective on ssms as gradient accumulators for in-context learning
- sound mathematical theory

Weaknesses:

- limited practical insights into currently used ssm variants (e.g., Griffin, Mamba)
- restricted to regression-style in-context learning tasks
- unclear connection between theoretical construction and empirical success of modern ssms
- missing related works. 
- more empirical support is needed for the claims about GD-SSM's general applicability.

The paper is borderline. While authors addressed some of reviewers' concerns, they remained skeptical overall. I vote for a rejection but encourage the authors to improve the paper based on reviewers' suggestions.

Decision: Reject