Longhorn: State Space Models are Amortized Online Learners

## Abstract

Abstract Modern large language models are built on sequence modeling via next-token prediction.
While the Transformer remains the dominant architecture for sequence modeling, its quadratic decoding complexity in sequence length poses a major limitation. State-space models (SSMs) present a competitive alternative, offering linear decoding efficiency while maintaining parallelism during training. However, most existing SSMs rely on linear recurrence designs that appear somewhat ad hoc.
In this work, we explore SSM design through the lens of online learning, conceptualizing SSMs as meta-modules for specific online learning problems. This approach links SSM design to formulating precise online learning objectives, with state transition rules derived from solving these objectives.
Based on this insight, we introduce a novel deep SSM architecture, Longhorn, whose update resembles the closed-form solution for solving the online associative recall problem. Our experimental results show that Longhorn outperforms state-of-the-art SSMs, including the Mamba model, on standard sequence modeling benchmarks, language modeling, and vision tasks. Specifically, Longhorn achieves a 1.8x improvement in sample efficiency compared to Mamba, and can extrapolate over contexts that are up to 16x longer during inference. Figure 1: (left) The average perplexity on eight downstream datasets for GLA, Mamba, and Longhorn (1.3B model) over seen tokens on SlimPajama. Longhorn leads to a 1.8x speed up in sampling efficiency. (right) Longhorn, pretrained with 2048 context length, extrapolates up to 16x longer context at inference.

## 1 Introduction

The Transformer model has become the go-to architecture for sequence modeling in deep learning (Vaswani et al., ). However, its utility is constrained by the quadratic growth in training and decoding costs with increasing sequence length.
Despite various optimizations such as efficient decoding (e.g., Chen et al., ; Kuperman & Dyke, ), KV-cache compression (e.g., DeepSeek-AI & Dai, ), and memory efficient implementation (e.g., Dao et al., ),
it remains challenging to scale Transformers for autonomous and continual use with an infinite (or very long) context window.

Recent advances in linear attention models (Katharopoulos et al., ) and state-space models (SSMs) (Gu et al., ) have demonstrated their potential. These models are specialized recurrent neural networks capable of efficiently computing outputs in parallel when input tokens are provided simultaneously during training, thus avoiding the inefficiencies of traditional backpropagation through time. During inference, the recurrent form is employed, resulting in linear decoding efficiency.
Initially, these models underperformed compared to Transformers. However, recent SSMs (e.g., Gu & Dao, ; Yang et al., ; Peng et al., ; De et al., ; Beck et al., ) have achieved performance parity with Transformers in language modeling tasks. Despite extensive research into various design aspects of SSMs, a guiding principle for designing SSMs remains elusive.

In this work, we propose one potential principle. We observe that one can view SSMs (or any sequence mixing layers) as “meta modules” that compress the history *online* into a memory state which is then used by later layers in the network for sequence modeling. From this perspective:

*The recurrent form of SSMs can be viewed as solving an online learning problem*.

Figure: Figure 2: (left) Most existing sequence models consist of channel and sequence mixing layers. The sequence mixing layers can be viewed as “meta-modules” that compress history into a state $S_{t}$, which is then passed to later layers for sequence modeling. (middle) Sequence mixing can be seen as an online learning problem, where the SSM state $S_{t}$ optimizes an online objective. The recurrent update of $S_{t}$ is derived either by solving this objective in closed form or via a proximal update. (right) Longhorn’s update solves online associative recall, where the goal is to recover $x\in\mathbb{R}^{d}$ based on a hint $k\in\mathbb{R}^{m}$ from a state matrix $S\in\mathbb{R}^{d\times m}$. Longhorn’s update corresponds to the implicit online learning solution, where $A_{t}=1_{d\times m}-\varepsilon_{t}\otimes k^{\otimes 2}$ and $B_{t}=(\varepsilon\odot x_{t})\otimes k_{t}$, and $\varepsilon_{t}=\beta_{t}/(1+\beta_{t}k_{t}^{\top}k_{t})$. See the details in Section 3 and Algorithm 1.
Refer to caption: x2.png

As a result, we can draw inspiration from online learning and confine the design choices of SSMs to reflect those learning dynamics that solve specific online prediction problems. The aim is that, by optimizing for the right objective, the model can achieve superior performance with fewer parameters or reduced computational costs. Furthermore, this online learning perspective may offer deeper insights into the function of SSM layers in large models. In particular, the recurrent update (i.e., state-transition dynamics) of an SSM can be interpreted as a proximal update step or a closed-form solution to an online learning objective. We outline the corresponding objectives for several existing SSMs in Table 4. One significant advantage of viewing SSMs through the lens of online learning is their ability to adapt post-training during deployment, allowing them to process arbitrarily long data sequences at inference time.

Based on this insight, we propose a simple yet effective architecture (Longhorn), derived from the implicit closed-form update of an online associative recall problem. The closed-form update naturally leads to a stable recurrent form without a manually designed gating mechanism, automatically balancing *forgetting* and *learning*. Thus Longhorn does not need a separately parameterized forget gate, which saves parameters when the state size is large. We demonstrate that Longhorn performs comparably to or better than state-of-the-art SSMs like Mamba (Gu & Dao, ) on synthetic and large-scale sequence modeling tasks. In particular, Longhorn outperforms Mamba at the size of 1.3B-parameter when trained on 100B tokens from the SlimPajama dataset (Soboleva et al., ). To summarize, our contributions are:

1) Theoretical Framework: We propose a novel framework that views SSMs’ recurrent update as solving online learning objectives. As a result, the design of SSMs reduces to the design of the online learning objectives. In particular, we introduce a novel, simple, and effective SSM, named *Longhorn*, that explicitly solves an online associative recall problem. Longhorn’s recurrent update is obtained by the *closed-form* solution to the online learning objective. Consequently, Longhorn does not require a separately parameterized forget gate that appears in most existing SSMs.

2) Empirical Results:
Longhorn demonstrates better performance than existing SSMs including Mamba, across both synthetic associative recall tasks and the large-scale language modeling task. Moreover, it achieves 1.8x improvement in sample efficiency compared to Mamba (See Figure 1 (left)). Longhorn’s training speed is as fast as Mamba, as we only replace the SSM module in the Mamba architecture with Longhorn’s recurrence. So it serves as a drop-in replacement for Mamba. Lastly, Longhorn, trained with 2048 context length can extrapolate to 32K context length at inference time without much perplexity drop (See Figure 1 (right)).

#### Notation

Throughout this work, we use $\odot$ to denote the Hadamard (elementwise) product, and $\otimes$ to denote the Kronecker (or outer) product between two tensors. Uppercase letters $A,B$, etc. denote matrices, while lowercase $k,v$ are in general vectors. $\left\lVert\cdot\right\rVert$ by default refers to the $\ell_{2}$ norm for vectors.

Figure: Algorithm 1 Longhorn’s Single-layer SSM Recurrence (Inference Time)

## 2 Background

In this section, we provide a brief introduction to contemporary deep state space models (deep SSMs).

Modern large language models are sequence-to-sequence models consisting of a stack of layers
$\bm{y}=\Phi_{L}\circ\cdots\circ\Phi_{1}(\bm{x})$ that sequentially processes an input sequence $\bm{x}=\{x_{t}\}_{t=1}^{T}$, where $T$ is the context length.
Specifically, transformers consist of alternative stacks of self-attention (SA) and multi-layer perceptron (MLP) layers that conduct *mixing* (i.e., information aggregation) on the sequence and channel dimensions, respectively.

Deep SSMs replace the SA layers with SSM layers.
Some variants of SSM models leave the MLP layers unchanged (Sun et al., ; Yang et al., ; De et al., ), while others fuse the SSM layer and the MLP layer into a single unified module (Gu & Dao, ).
But in both cases, the sequence mixing is done by the SSM module, and the channel mixing is done by the channel-wise MLP. Taking Mamba as an example (Gu & Dao, ), a Mamba model consists of a stack of homogeneous modules named Mamba block (the $\Phi_{i}(\bm{x})$); we provide a visualization of a single Mamba block in Figure 3 (Gu & Dao, ),
which consists of an SSM block for sequence mixing (red),
and an MLP block for channel mixing (blue).

#### SSM: General Form

The SSM block (in red) plays the crucial role of sequence mixing.
It works by iteratively updating a memory state matrix $S_{t}\in\mathbb{R}^{d\times m}$ with a linear recurrence:

$$ $\displaystyle S_{t}=A(x_{t})*S_{t-1}+B(x_{t}),$ $\displaystyle\forall t\in\{1,\ldots,T\},$ $\displaystyle S_{0}=0,$ (1) $$

where $x_{t}$ is the input at time $t$, $S_{t}$ is the model’s *state*, $A_{t},B_{t}\colon\mathbb{R}^{d}\to\mathbb{R}^{d\times m}$ are some functions of the input and $*$ is a multiplication operation of choice, such as Hadamard product or matrix product.

Given the state $S_{t}$,
SSMs often give the output token at the next layer via a gated linear unit (GLU) (Dauphin et al., ):

$$ $\displaystyle y_{t}=\mathtt{Readout}(S_{t},x_{t})=W_{1}\big{(}o_{t}\odot\sigma (W_{2}x_{t})\big{)},$ $\displaystyle o_{t}=C(x_{t})S_{t},$ $$

where we first get $o_{t}$
via a state-dependent linear projection on $S_{t}$, which is then fed into a subsequent channel mixing gated linear unit (blue in Figure 3),
where $\sigma(\cdot)$ is a non-linear activation function.

A key feature of this design
in Equation 1
is that $S_{t}$ has a linear recurrence, i.e., $S_{t}$ is linear in $S_{t-1}$.

Figure: Figure 3: Mamba Block
Refer to caption: x3.png

Crucially, this allows us to express all $S_{t}$ in an explicit form that can be calculated in parallel:
when all $\bm{x}=\{x_{t}\}_{t}$ are available as in the training phase, $\{S_{t}\}_{t}$ can be written into

$$ $S_{t}=\sum_{t^{\prime}\leq t}(\overline{A}_{t^{\prime}\to t})B(x_{t^{\prime}}) ,~{}~{}~{}~{}~{}~{}\text{where}~{}~{}~{}~{}~{}~{}~{}\overline{A}_{t^{\prime} \to t}=\prod_{t^{\prime}<\tau\leq t}A(x_{\tau}).$ (2) $$

Here $\prod$ denotes the product induced by multiplication operator $*$. The resulting cumulative product $\overline{A}_{t^{\prime}\to t}$ can be implemented efficiently in parallel with the prefix scan algorithm (e.g., Harris et al., ), which only requires $\mathcal{O}(\log T)$ ($T$ is the sequence length) parallel operations. From now on, we will abbreviate $A(x_{t})$ and $B(x_{t})$ as $A_{t}$ and $B_{t}$, respectively.

#### Designs of ( A t subscript 𝐴 𝑡 A_{t} italic_A start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , B t subscript 𝐵 𝑡 B_{t} italic_B start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , ∗ * ∗ )

Existing variants of SSMs
mainly differ in the design choices of the networks $A_{t},B_{t}$, and the associated operator $*$ in the linear recurrence.
A core issue here is that the memory state $S_{t}\in\mathbb{R}^{d\times m}$, designed to be $m$ times the input $x_{t}$ in size,
must be as large as possible to maintain sufficient information during recurrence.
This makes the architecture design of $A_{t},B_{t}$,
both mapping $\mathbb{R}^{d}$ to $\mathbb{R}^{d\times m}$ challenging.
A naive linear mapping would result in $d\times d\times m$ weights, which is prohibitively large.
This makes it necessary to impose certain low-dimensional structures in $A_{t},B_{t}$, which is the main difference from existing designs of SSMs. In Appendix A, we summarize some existing deep SSM models in the form of Equation 1.

## 3 An Online Learning Perspective For Sequence Mixing

As demonstrated in the previous section, designing a state-space model (SSM) depends on the specific selection of $(A_{t},B_{t},*)$, which is intricate and somewhat artisanal. In this section, we propose to streamline SSM design through an online learning perspective. The main idea is to treat the SSM layers as learning modules that learn to compress information along the sequence dimension. From this perspective, the SSM layers are *learning to learn*, such that during the inference time, these layers are still learning (compressing) new information online.

We begin with an overview of online learning and subsequently demonstrate how SSM can be framed as an online learning problem. Finally, we present a straightforward architecture based on the closed-form solution of the implicit online learning algorithm.

### 3.1 SSM as Online Learning

We advocate viewing the recurrence of SSM as solving an online learning problem.
In online learning, the agent picks a state $s_{t}$ at time $t$ and then incurs a loss $\ell_{t}(s_{t})$. The goal is to minimize

$$ $\min_{\{s_{t}\}}\sum_{t}\ell_{t}(s_{t}).$ (3) $$

For instance, consider online linear prediction, where at each step the agent is given an input-label pair $(x_{t},y_{t})$ and $\ell_{t}(s_{t})=\frac{1}{2}||s_{t}^{\top}x_{t}-y_{t}||^{2}$ is the $\ell_{2}$ regression loss, then the problem becomes an online regression problem, and the goal is to successfully predict $y_{t}$ given $x_{t}$ at future time steps, with the key feature that the prediction ($s_{t}$) can change with each new data point.

Online convex programming (OCP) (e.g., Zinkevich, ) yields a principled approach to
solving Equation 3 when $\ell_{t}$ are convex, by trading-off the “stability” and “plasticity” (e.g., Mermillod et al., ). Formally, an online convex programming algorithm updates $s_{t}$ by
solving a regularized cost function:

$$ $s_{t}=\operatorname*{arg\,min}_{s}L_{t}(s),~{}~{}~{}~{}~{}~{}~{}L_{t}(s)= \underbrace{D_{\phi}(s,s_{t-1})}_{\text{stability}}+\underbrace{\beta_{t}\ell_ {t}(s)}_{\text{plasticity}},$ (4) $$

where $\beta_{t}\in\mathbb{R}^{+}$ and
$D_{\phi}$ is a discrepancy measure, often a Bregman divergence induced by the convex function $\phi$ (e.g., when $\phi(x)=\frac{1}{2}\left\lVert x\right\rVert^{2}$, $D_{\phi}(s,s_{t-1})=\frac{1}{2}||s-s_{t-1}||^{2}$).
Here the first term ensures the updated $s$ will be close to the previous $s_{t-1}$, so the agent suffers less from *catastrophic forgetting*, while the second term ensures the agent is incorporating new knowledge from minimizing the new loss $\ell_{t}(s)$. Hence, $\beta_{t}$ controls the trade-off between stability and plasticity.

### 3.2 The Longhorn Architecture

Under the online learning framework, the design of an SSM reduces to the design of $D_{\phi}$ and $\ell_{t}$ in Equation 4.
This provides a unified framework for the existing SSM variants.
We summarize in Table 4 in Appendix B the online learning interpretation of several existing SSM architectures.

In this work, we explore a highly simplified and natural design called *Longhorn* guided by the online principle (see the last row of Table 4).
In particular, we consider $\{(k_{t},x_{t})\}_{t}$ as the input stream, where $k_{t}\in\mathbb{R}^{m}$ and $x_{t}\in\mathbb{R}^{d}$ are the key-value pairs, just as in the Transformer model (Vaswani et al., ). In practice, as in Mamba (Gu & Dao, ), $k_{t}=W_{k}x_{t}\in\mathbb{R}^{m}$, where $W_{k}\in\mathbb{R}^{m\times d}$, is a linear mapping from $x_{t}$.

We want to recurrently update hidden states $\{S_{t}\}_{t}$, where $S_{t}\in\mathbb{R}^{d\times m}$ is a matrix that summarizes the information up to time $t$.
We posit the following OCP objective for updating $S_{t}$:

$$ $S_{t}=\operatorname*{arg\,min}_{S\in\mathbb{R}^{d\times m}}\left\{||S-S_{t-1}| |^{2}_{\text{F}}+||Sk_{t}-x_{t}||^{2}_{\mathrm{diag}(\beta_{t})}\right\}.$ (5) $$

Here, $||\cdot||_{\text{F}}$ denotes the Frobenius norm of a matrix, $\beta_{t}\in\mathbb{R}^{d}$ is a vector controlling how much new information about $x_{t}$ we want the model to incorporate for $S_{t}$. For instance, $\beta_{t,i}=0$ implies $S_{t,i}=S_{t-1,i}$ (i.e., the $i$-th row of $S$ remains unchanged), while a large $\beta_{t,i}$ implies the model empties some part of $S_{i}$ for incorporating $x_{t,i}$.

From a high-level perspective, Equation 5 is solving an online prediction problem of learning a weight matrix $S$ to predict $x_{t}$ given $k_{t}$ with a linear model $x_{t}\approx S^{\top}k_{t}$.
It is a supervised formulation of the
*associative memory* problem of
memorize $(k_{t},x_{t})$ pairs
by learning a mapping from $k_{t}$ to $x_{t}$, such that given a key (input) $k_{t}$ the model can retrieve (predict) its corresponding value (label) $x_{t}$.

The objective in Equation 5 is motivated by the observation that the self-attention layer of the Transformer exhibits a form of online associative recall (often referred to as the induction head property) (Olsson et al., ). This capability has been shown to underpin the model’s ability to perform in-context learning (Brown, ). To explain the connection, in-context learning refers to the model’s ability, during inference, to generalize from a set of provided ${(k,x)}$ (question-answer) pairs and apply this understanding to a new question. This closely parallels associative recall, where the model retrieves relevant information from past interactions to address new inputs.

Fortunately, this simple objective gives a closed-form solution for $S_{t}$, which coincides with the implicit online learning method (e.g., Kulis & Bartlett, ), according to Theorem 6 (We provide the proof in Appendix C):

###### Theorem 3.1 .

The closed form solution for $S_{t}$ for objective in Equation 5 is

$$ $S_{t,i}=(I-\varepsilon_{t,i}k_{t}k_{t}^{\top})S_{t-1,i}+\varepsilon_{t,i}k_{t} x_{t,i},~{}~{}~{}\text{where}~{}\varepsilon_{t,i}=\frac{\beta_{t,i}}{1+\beta_{ t,i}k_{t}^{\top}k_{t}}\in[0,\infty).$ (6) $$

Here, $S_{t,i}$ refers to the $i$-th row of $S_{t}$, $\beta_{t,i}$ refers to the $i$-th element of $\beta_{t}$. As $k_{t}k_{t}^{\top}$ is a matrix, it is hard to compute its cumulative product for conducting a parallel scan. As a result, in practice, we use the diagonal approximation $1_{m}-\varepsilon_{t,i}k_{t}^{\odot 2}$ in place of $I-\varepsilon_{t,i}k_{t}k_{t}^{\top}$, where $a^{\odot 2}=a\odot a$ and $1_{m}$ is the $m$-dimensional all-one vector. Following Mamba (Gu & Dao, ) and Transformer (Vaswani et al., ), we make $k_{t}=W_{k}x_{t}\in\mathbb{R}^{m}$ and $\beta_{t}=\sigma(W_{\beta}x_{t})\in\mathbb{R}^{d}$ (both are functions of $x_{t}$), where the activation $\sigma$ (the Sigmoid function) is to ensure that $\beta_{t}$ is positive and bounded. In summary, the final Longhorn update of $S_{t}$ becomes:

$$ $S_{t}=A_{t}\odot S_{t-1}+B_{t},~{}~{}\text{where}~{}~{}A_{t}=(1_{d\times m}- \varepsilon_{t}\otimes k_{t}^{\odot 2}),~{}~{}B_{t}=(\varepsilon_{t}\odot x_{t })\otimes k_{t}.$ (7) $$

The final architecture of Longhorn follows Mamba strictly (Figure 3), except that we replace the SSM block with Longhorn’s recurrence. We also provide an efficient CUDA kernel for it.
The full inference-time algorithm is provided in Algorithm 1. One can compare Equation 7 to Equation 8 and other SSMs in Appendix A. Longhorn does not introduce an extra “forgetting” gate (hence it has fewer parameters), because the forgetting gate is naturally derived from the key vector, i.e., $(1_{d\times m}-\varepsilon_{t}\otimes k_{t}^{\odot 2})$.

###### Theorem 3.1 .

#### Advantages of Longhorn

- 1.
While we can derive the learning objective for some of the existing SSMs, Longhorn is the first SSM designed for explicitly solving an online regression problem.
- 2.
Longhorn does not require a specific forget gate (e.g., $\alpha_{t}$ in GLA or $A$ matrix in Mamba). The forgetting is naturally linked to the key vector $k_{t}$ through the derivation. This saves about $\mathcal{O}(d\times m)$ parameters per SSM module, where $m$ is the dimension of $k_{t}$, and $d$ is the dimension of $x_{t}$. However, Longhorn demonstrates better performance even with fewer parameters than Mamba (See Figure 1 (left), Table 1, Table 2).
- 3.
The closed-form solution in Equation 6 does not need any specific initialization. In contrast, Mamba requires a special careful initialization of the $A$ and $\varepsilon_{t}$.
- 4.
Unlike DeltaNet (Yang et al., ), which struggles to extrapolate beyond training contexts, Longhorn successfully extrapolates to contexts 16x longer than it was trained for (Figure 1 (right)).

## 4 Related Work

This section provides a summary of recent advances in linear attention and state space models.

#### Linear Attention Models

Several methods reduce the quadratic complexity of Transformers by making attention linear with respect to context length. Linformer projects keys and values into a constant-size matrix, bypassing the scaling with sequence length (Wang et al., ). Linear Transformer replaces the Softmax function with a decomposable similarity function, achieving linear complexity (Katharopoulos et al., ). Performer approximates softmax attention using orthogonal random features (Choromanski et al., ). RetNet adds constant forgetting and rotation (Sun et al., ), while Gated Linear Attention introduces learnable forget gates (Yang et al., ). Linear attention can also be seen as a fast weight network where a slow net adapts a fast network’s parameters online using inputs (Schlag et al., ).

#### State Space Models

State space models (SSMs) focus on parallelizable linear recurrent networks. Initially, a constant state transition matrix $A$ allows recurrence to be computed via convolution (Li et al., ; Gu et al., ). Key models include Diagonal State Space (DSS) (Gupta et al., ), Gated State Space (GSS) (Mehta et al., ), S5 (Smith et al., ), Bidirectional Gated SSM (BiGS) (Wang et al., ), H3 (Fu et al., ), and Mamba (Gu & Dao, ). Efficient recurrent networks often resemble SSMs, such as Deep Linear Recurrent Units (LRUs) (Orvieto et al., ; De et al., ), Hierarchically Gated Linear RNNs (HGRN) (Qin et al., ; ), and RWKV (Peng et al., ; ).

#### Fast Weight Programmer

The idea of networks modifying their own weights in response to inputs dates back to the Fast-weight Programmer (Schmidhuber, ; ; Schlag & Schmidhuber, ; Schlag et al., ). These models update a weight matrix $W\in\mathbb{R}^{d\times m}$ via the outer product of two vectors: $\Delta W=x_{t}\otimes k(x_{t})$, a mechanism similar to Linear Attention. Our framework extends this concept by adapting the weight update process to suit specific online learning objectives, enhancing its use in dynamic learning environments.

#### Concurrent Work

Two concurrent works share similar ideas with ours. Yang et al. () propose a chunk-wise parallel approach to scale DeltaNet (Schlag et al., ) for large-scale language modeling. DeltaNet’s update rule, viewed as a gradient step for an online regression objective, results in a state transition matrix $A(x_{t})=(I-\beta_{t}k_{t}k_{t}^{\top})$, which can have eigenvalues $>$1, leading to instability. To address this, Yang et al. () normalize the key vector $k_{t}$ by its $\ell_{2}$ norm, which can be restrictive. In contrast, Longhorn ensures stability with a closed-form update, using a diagonal approximation ($k_{t}^{\odot 2}$), allowing for both parallel scan (as in Mamba) and chunk-wise parallel training (as in GLA), making it as fast as existing SSMs. Additionally, we provide a parallel scan CUDA kernel, enabling Longhorn to serve as a drop-in replacement for Mamba. Sun et al. () introduce the Test-Time Training framework, where state updates are derived from a gradient step on an online regression objective. To maintain parallelism, they assume each gradient step at $x_{t}$ uses the initial state $s_{0}$, enabling matrix multiplication. In contrast, Longhorn computes the closed-form solution for *every* token, offering greater flexibility.

## 5 Experiments

We validate Longhorn’s performance through the following experiments:

1)  We compare Longhorn against other SSMs on the multi-query associative recall benchmark (Arora et al., ) and find that Longhorn is the only model to achieve near-perfect recall at sequence lengths up to 512 with a hidden dimension of 64.

2)  Using the OpenWebText dataset (Gokaslan & Cohen, ), we assess Longhorn’s performance on language modeling with model sizes of 120M and 350M, and context lengths of 1024 or 4096, showing it consistently outperforms other SSMs in validation perplexity.

3)  We train a 1.3B language model on the SlimPajama dataset (Soboleva et al., ) with 100B tokens and compare its performance across 8 benchmarks, where Longhorn achieves better final performance and $>$1.8x better sample efficiency than Mamba and GLA.

4)  We additional apply Longhorn to vision domain and compare it against the Vision Mamba (ViM) (Zhu et al., ) model (Appendix 5.5), where Longhorn achieves performance comparable (slightly superior) to that of the ViM model.

### 5.1 Multi-Query Associative Recall

We first consider the synthetic benchmark Multi-Query Associative Recall (MQAR) (Arora et al., ). The agent observes a sequence of tokens $\{k_{1},v_{1},k_{2},v_{2},\dots,k_{T},v_{T}\}$, where each consecutive two-tokens become a key-value pair. At test time, the agent is provided with multiple $k\sim\{k_{1},\dots k_{T}\}$, the goal is to “retrieve” the corresponding values. Following the original benchmark, we consider the sequence length $T\in\{64,128,256,512\}$ and model dimension (size of the latent embedding of a token) $d\in\{64,128,256,512\}$. We compare against 1) Transformer model (Attention), 2) Based architecture, which combines an SSM with local-attention, where the SSM is derived from the Taylor approximation of the self-attention (Arora et al., ), 3) Hyena (Poli et al., ), which is a special SSM that adopts long convolution via fast fourier transform, 4) RWKV (Peng et al., ), which can be viewed as the division of two SSMs (i.e., $y=a/b$, where $a,b$ are outputs from two SSMs). The state-transition matrix is a scalar, 5) BaseConv (Arora et al., ), an SSM that combines linear projection with convolution, and 6) Mamba (Gu & Dao, ), the state-of-the-art SSM that has data-dependent $A$ and $B$ (See equation 8). Each experiment individually searches for the best learning rate from $\{10^{-4},4.6\times 10^{-4},2.2\times 10^{-3},10^{-2}\}$. Results are summarized in Figure 4.

Figure: Figure 4: Comparison of Longhorn to state-of-the-art SSMs on the MQAR benchmark. y-axis is the recall rate.
Refer to caption: x4.png

Observation:
From the figure, we can see that Longhorn, which is designed to perform the associative recall task by solving the online prediction objective, outperforms existing SSM variants even at the sequence length of 512 and a small model dimension of 64.

### 5.2 Scaling Law on OpenWebText

In this section, we consider language modeling tasks on models with 120M or 350M parameters with 1024 or 4096 context length. We choose the OpenWebText dataset as it is small and serves as an easily accessible benchmark for quick benchmarks.(^1^11We adapted code from the nanoGPT repository [https://github.com/karpathy/nanoGPT](https://github.com/karpathy/nanoGPT), which is a minimal reproduction of GPT-2 model using PyTorch.) The details about the architecture is provided in Appendix D. We consider the following baseline models: LLaMA (Touvron et al., ), RetNet (Sun et al., ), Mamba (Gu & Dao, ), RWKV (Peng et al., ), and GLA (Yang et al., ). Then we experiment with 1024 or 4096 context length $T$ and model sizes around 120M or 350M. Results are summarized in Table 1 and Figure 5.

Figure: Figure 5: Scaling law with 1024 and 4096 context length on OpenWebText with various SSM models and the LLaMA (strong Transformer) baseline.
Refer to caption: x5.png

**Table 1: Language modeling scaling law against LLaMA (Touvron et al., ), RetNet (Sun et al., ), RWKV (Peng et al., ), and Mamba (Gu & Dao, ). All models are trained on the OpenWebText dataset (Gokaslan & Cohen, ). Models vary from 120-350M parameters and 1024-4096 context length.**
| Model | # Param. (M) | Val. Loss ($\downarrow$) | # Param. (M) | Val. Loss ($\downarrow$) |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| $T=1024$ | $T=4096$ | $T=1024$ | $T=4096$ |  |  |  |
| RetNet | 129.1 | 3.569 | 3.492 | 373.2 | 3.362 | 3.227 |
| GLA | 123.8 | 3.381 | 3.364 | 361.1 | 3.018 | 3.001 |
| RWKV | 124.4 | 3.291 | 3.276 | 354.8 | 2.983 | 2.931 |
| Mamba | 129.2 | 3.238 | 3.231 | 371.5 | 2.902 | 2.868 |
| LLaMA | 124.4 | 3.247 | 3.273 | 357.7 | 2.891 | 2.883 |
| Longhorn | 128.6 | 3.225 | 3.192 | 369.8 | 2.888 | 2.859 |

Observation:
From the figure and table, we can see that Longhorn consistently outperforms baseline SSMs up to 350M and 4096 context length.

### 5.3 Large-scale Language Modeling

For the large-scale language modeling task, we followed the GLA (Yang et al., ) setup, training a 1.3B parameter model on the SlimPajama (Soboleva et al., ) dataset with 100B tokens and a batch size of 2M. We used the AdamW optimizer (Loshchilov & Hutter, ) with a weight decay of 0.01, cosine learning rate decay (peak: $3e-4$, final: $3e-5$), and gradient clipping of 1.0. Comparisons were made against LLaMA, Mamba, and GLA models (context size: 2048). We evaluated on eight standard downstream tasks, including PIQA (Bisk et al., ), HellaSwag (Hella) (Zellers et al., ), WinoGrande (Wino) (Sakaguchi et al., ), ARC-easy (ARC-e) and ARC-challenge (ARC-c) (Clark et al., ), OpenBookQA (OBQA) (Mihaylov et al., ), Social Interaction QA (SIQA) (Sap et al., ), and Boolean questions (BoolQ) (Clark et al., ). We report the average perplexity across the above eight datasets throughout training in Figure 1 (left). Then we summarize the downstream evaluation results in Table 2.

**Table 2: Language modeling results against LLaMA (Touvron et al., ), RetNet (Sun et al., ), and Mamba (Gu & Dao, ). All models are trained on the same subset of the SlimPajama dataset with the Mistral tokenizer. The 340M/1.3B models are trained for 15B/100B tokens respectively. State Size is the effective state size of an SSM per layer. For instance, GLA’s state size (1024K) is computed by $md/h$, where the key and value dimensions are $m=1024$ and $d=2048$, and there are 4 heads $h=4$. The individual task performance is via zero-shot. The last column shows the average value over the results on all benchmarks.**
| Model | State Size | PIQA | Hella | Wino. | ARC-e | ARC-c | OBQA | SIQA | BoolQ | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| acc $\uparrow$ | acc_norm $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc_norm $\uparrow$ | acc $\uparrow$ | acc_norm $\uparrow$ | acc $\uparrow$ |  |  |  |
| LLaMA | 8M | 55.08 | 55.36 | 71.73 | 59.26 | 32.19 | 43.35 | 45.16 | 62.13 | 53.03 |
| GLA | 512K | 55.55 | 49.10 | 71.12 | 58.86 | 28.11 | 41.67 | 44.91 | 59.21 | 51.07 |
| Mamba | 64K | 54.21 | 53.61 | 71.67 | 61.05 | 30.15 | 43.94 | 44.18 | 59.22 | 52.25 |
| Longhorn | 64K | 55.78 | 52.30 | 71.00 | 60.63 | 29.53 | 43.55 | 44.68 | 61.29 | 52.35 |

Observation:
From Figure 1 (left), it is evident that Longhorn not only achieves a lower average perplexity but also improves sampling efficiency by 1.8x compared to Mamba. In other words, Longhorn reaches the same average perplexity with nearly half the training data required by Mamba. From the Table 2, we can see that up to a 1.3B model, Longhorn remains strong among all baseline models and achieves slightly better result than Mamba, even though it has a bit fewer parameters.

### 5.4 Ablation on Length Extrapolation

We evaluate how Longhorn extrapolates to a context length longer than 2048 (training context length) at inference time. In particular, we pick a disjoint validation set from SlimPajama dataset, rearrange it into batches of sequences of length $T\in\{2048,4096,8192,16384,32768\}$, and then evaluate the pretrained model’s perplexity on those sequences. The results are summarized in Figure 1 (right).

Observation:
From the figure, we observe that Longhorn successfully extrapolates to contexts up to 16x longer than those used during training, this contrasts with DeltaNet (Yang et al., ), which highlights a limitation in that the model cannot extrapolate to longer contexts. In contrast, LLaMA, as a Transformer-based model, fails to extrapolate beyond its training context length.

### 5.5 Vision State Space Models

In addition to language tasks, recent works have also applied state space models to the vision domain, leveraging their superior training efficiency. In particular, following the Vision Mamba (ViM) (Zhu et al., ), we conduct experiments on the ImageNet (Deng et al., ) classification task. Similar to ViM, We apply a bi-directional scan with Longhorn SSM (ViL) and compare the results with ViM on both the tiny and small configurations described in the ViM paper.

**Table 3: Top-1 Accuracy on ImageNet classification for Vision Mamba (ViM) and Vision Longhorn (ViL).**
| Model | # Param | Top-1 Accuracy |
| --- | --- | --- |
| ViM-Tiny | 7M | 76.1 |
| ViL-Tiny (ours) | 7M | 76.4 |
| ViM-Small | 26M | 80.5 |
| ViL-Small (ours) | 26M | 80.7 |

Observation:
The results from Table 3 demonstrate that the Vision Longhorn model (ViL) achieves comparable (slightly better) performance to the original ViM. Note that we use the best hyperparameters for ViM without additional tuning, and ViL does not require two additional parameters for the forward and backward $A$ matrices, as they are computed directly based on the key $k$ vector.

## 6 Conclusion and Future Work

This work introduces a novel approach to designing deep state-space models (SSMs) by conceptualizing the recurrence update as solving an online objective. Based on this, we propose Longhorn, a novel SSM model that explicitly solves online associative recall in closed form. Longhorn is parallelizable and achieves state-of-the-art performance among SSMs on MQAR, language modeling, and image classification tasks. One future direction is to explore other online learning objectives. Additionally, recent studies (Ren et al., ) suggest that incorporating sliding-window attention with Mamba improves performance. We anticipate similar benefits for Longhorn.

## Appendix A Prior Deep State Space Models

###### Example A.1 (Linear Attention Variants) .

Linear Attention (LA) (Katharopoulos et al., ), Retention Network (RetNet) (Sun et al., ), and Gated Linear Attention (GLA) (Yang et al., ) all assume $A_{t},B_{t}$ yield rank-1 (or even constant) outputs:

$$ $S_{t}=A_{t}\odot S_{t-1}+v(x_{t})\otimes k(x_{t}),~{}~{}~{}~{}\text{with}~{}~{ }\begin{cases}A_{t}=1&\quad\text{(LA)}\\ A_{t}=c\in[0,1]&\quad\text{(RetNet)}\\ A_{t}=1\otimes\alpha(x_{t})&\quad\text{(GLA)}\end{cases},$ $$

where $S_{t}\in\mathbb{R}^{d\times m}$,
$v(x_{t})\in\mathbb{R}^{d}$, $k(x_{t})\in\mathbb{R}^{m}$ are linear mappings of $x_{t}$, and $\otimes$ denote the outer product. In practice, one can use $h$ heads as in the multi-head attention to save some computation, where the $m$ and $d$ dimensions are divided into $h$ groups and each group performs its own LA variant. The outer product complexity reduces to $\mathcal{O}(h*m/h*d/h=md/h)$. But then the effective size of $S_{t}$ also shrinks to $md/h$.

###### Example A.2 (Mamba (Gu & Dao, 2023 ) ) .

The Mamba architecture is derived by discretizing a continuous linear dynamics. Its discretized update is:

$$ $\begin{split}S_{t}&=A_{t}\odot S_{t-1}+B_{t},~{}~{}~{}~{}\text{where}\\ A_{t}&=\exp(A\odot(\varepsilon(x_{t})\otimes 1)),\qquad B_{t}=(\varepsilon(x_{ t})\odot x_{t})\otimes k(x_{t}).\end{split}$ (8) $$

where $S_{t}\in\mathbb{R}^{d\times m}$ with $m=16$ by default, $\varepsilon(x_{t})\in\mathbb{R}^{d}$, $k(x_{t})\in\mathbb{R}^{m}$ linear mappings of $x_{t}$, and $A\in\mathbb{R}^{d\times m}$ is a data independent (not depending on $x_{t}$ trainable weight matrix.

In Mamba, both $A_{t}$ and $B_{t}$ depend on $\varepsilon(x_{t})$, which represents the step size for the SSM update.

In practice, Mamba does not use multiple heads as in linear attention variants. Perhaps the main reason is that given a fixed $m$ and $d$, the largest memory state will be with $h=1$ (as the effective size of $S_{t}$ is $md/h$). In addition, Mamba’s output is $o_{t}=C(x_{t})S_{t}+D_{t}\odot x_{t}$, which has an additional residual part $D_{t}\odot x_{t}$.

###### Example A.3 (Griffin (De et al., 2024 ) ) .

In Mamba and the linear attention variants, the outer product serves as a critical role in lifting vectors to matrices. The recent Griffin architecture abandons the outer product and performs pure elementwise product:

$$ $s_{t}=a(x_{t})\odot s_{t-1}+\sqrt{1-a(x_{t})}\odot i(x_{t})\odot x_{t},$ $$

where $s_{t},a(x_{t}),i(x_{t})$ are all $\mathbb{R}^{d}$. This yields smaller memory states, but in practice, Griffin is combined with local attention (i.e., the sliding-window self-attention) to strengthen its capability.

###### Example A.4 (RWKV (Peng et al., 2023 ) ) .

The original RWKV also performs elementwise recurrence. It maintains a state of ratio form $s_{t}=u_{t}/z_{t}$, where $u_{t},z_{t}$ are updated separately by two SSMs:

$$ $\displaystyle s_{t}$ $\displaystyle=u_{t}/z_{t}$ $\displaystyle u_{t}$ $\displaystyle=\exp(-w)\cdot u_{t-1}+\exp(k(x_{t}))\odot v(x_{t}),~{}~{}~{}~{}~ {}~{}~{}~{}~{}~{}~{}~{}~{}~{}z_{t}=\exp(-w)\cdot z_{t-1}+\exp(k(x_{t})),$ $$

where all the vectors are of size $\mathbb{R}^{d}$, and
$w>0$ is a *trainable weight* for controlling the forgetting.
In the most recent RWKV version (Peng et al., ), the denominator $z_{t}$ is removed, and the elementwise product is replaced with the outer product, which makes it more similar to an LA variant.

###### Example A.5 (HGRN2 (Qin et al., 2024a ) ) .

The Gated Linear RNNs with State Expansion (HGRN2) model is represented with the following recurrence:

$$ $S_{t}=(1\otimes f(x_{t}))\odot S_{t-1}+i(x_{t})\otimes(1-f(x_{t})).$ $$

Here, $f(x_{t})\in[0,1]$ is the forget gate, $(1-f(x_{t}))$ is the input gate, and $i(x_{t})$ is the input vector. HGRN2 thus resembles an RNN.

###### Example A.1 (Linear Attention Variants) .

###### Example A.2 (Mamba (Gu & Dao, 2023 ) ) .

###### Example A.3 (Griffin (De et al., 2024 ) ) .

###### Example A.4 (RWKV (Peng et al., 2023 ) ) .

###### Example A.5 (HGRN2 (Qin et al., 2024a ) ) .

## Appendix B Existing State Space Models’ Online Objectives

We reverse-engineer some existing deep SSMs’ online learning objectives in Table 4.

**Table 4: Some of the existing SSMs and their corresponding online learning objectives/updates.**
| Method | Online Learning Objective $L_{t}(s)$ (assume $x_{t}\in\mathbb{R}$) | Online Update |
| --- | --- | --- |
| LA | $\left\lVert S-S_{t-1}\right\rVert^{2}_{\text{F}}-2\langle Sk_{t},x_{t}\rangle$ | $S_{t}=S_{t-1}+x_{t}\otimes k_{t}$ |
| RetNet | $\gamma\left\lVert S-S_{t-1}\right\rVert^{2}+(1-\gamma)\left\lVert S\right<br>\rVert^{2}_{\text{F}}-2\langle Sk_{t},x_{t}\rangle$ | $S_{t}=\gamma S_{t-1}+x_{t}\otimes k_{t}$ |
| GLA | $\left\lVert S-S_{t-1}\mathrm{diag}(\alpha_{t})\right\rVert^{2}_{\text{F}}+2<br>\langle Sk_{t},x_{t}\rangle$ | $S_{t}=S_{t-1}\mathrm{diag}(\alpha_{t})+x_{t}\otimes k_{t}$ |
| Griffin | $\left\lVert\sqrt{\alpha_{t}}\odot(s-s_{t-1})\right\rVert^{2}+\left\lVert\sqrt{<br>1-\alpha_{t}}\odot s\right\rVert^{2}-2\sqrt{1-\alpha_{t}}\odot s\odot i_{t}<br>\odot x_{t}$ | $s_{t}=\alpha_{t}\odot s_{t-1}+\sqrt{(1-\alpha_{t})}\odot i_{t}\odot x_{t}$ |
| Longhorn | $\left\lVert S-S_{t-1}\right\rVert^{2}_{\text{F}}+\left\lVert Sk_{t}-x_{t}<br>\right\rVert^{2}_{\mathrm{diag}(\beta_{t})}$ | $S_{t}=(1_{m\times n}-\varepsilon_{t}\otimes k_{t}^{\odot 2})\odot S_{t-1}+$ |
| $(\varepsilon_{t}\odot x_{t})\otimes k_{t},~{}~{}\varepsilon_{t}=\beta_{t}/(1+<br>\beta_{t}k_{t}^{\top}k_{t})$ |  |  |

## Appendix C Proof

This section provides the proof for Theorem 6. Given the Longhorn’s objective $S_{t}=\operatorname*{arg\,min}_{S\in\mathbb{R}^{d\times m}}\left\{||S-S_{t-1}|
|^{2}_{\text{F}}+||Sk_{t}-x_{t}||^{2}_{\mathrm{diag}(\beta_{t})}\right\}$, we have the following theorem:

###### Theorem C.1 .

The closed form solution for $S_{t}$ for objective in Equation 5 is

$$ $S_{t,i}=(I-\varepsilon_{t,i}k_{t}k_{t}^{\top})S_{t-1,i}+\varepsilon_{t,i}k_{t} x_{t,i},~{}~{}~{}\text{where}~{}\varepsilon_{t,i}=\frac{\beta_{t,i}}{1+\beta_{ t,i}k_{t}^{\top}k_{t}}\in[0,\infty).$ (9) $$

###### Proof.

As the objective in equation 5 is in a quadratic form with respect to $s$, there is a unique minimum. Observe that each row of $S$ (e.g., $S_{i}$) optimizes the objective independently, therefore we can solve the solution row-wise. By setting the derivative of $\nabla_{S_{i}}L_{t}=0$, we have:

$$ $\displaystyle\nabla_{S_{i}}L_{t}=0$ $\displaystyle\Longleftrightarrow(S_{i}-S_{t-1,i})+\beta_{t,i}(S_{i}^{\top}k_{t }-x_{t,i})k_{t}=0$ $\displaystyle\Longleftrightarrow(I+\beta_{t,i}k_{t}k_{t}^{\top})S_{i}=S_{t-1,i }+\beta_{t,i}k_{t}x_{t,i}$ $\displaystyle\underbrace{\Longleftrightarrow}_{(3)}S_{i}=\bigg{(}I-\frac{\beta _{t,i}}{1+\beta_{t,i}k_{t}^{\top}k_{t}}k_{t}k_{t}^{\top}\bigg{)}S_{t-1,i}+ \bigg{(}I-\frac{\beta_{t,i}}{I+\beta_{t,i}k_{t}^{\top}k_{t}}k_{t}k_{t}^{\top} \bigg{)}\beta_{t,i}k_{t}x_{t,i}$ $\displaystyle\Longleftrightarrow\bigg{(}I-\frac{\beta_{t,i}}{I+\beta_{t,i}k_{t }^{\top}k_{t}}k_{t}k_{t}^{\top}\bigg{)}S_{t-1,i}+\frac{(I+\beta_{t,i}k_{t}^{ \top}k_{t}-\beta_{t,i}k_{t}k_{t}^{\top})\beta_{t,i}k_{t}x_{t,i}}{I+\beta_{t,i} k_{t}^{\top}k_{t}}$ $\displaystyle\underbrace{\Longleftrightarrow}_{(5)}\bigg{(}I-\frac{\beta_{t,i} }{I+\beta_{t,i}k_{t}^{\top}k_{t}}k_{t}k_{t}^{\top}\bigg{)}S_{t-1,i}+\frac{ \beta_{t,i}k_{t}x_{t,i}}{I+\beta_{t,i}k_{t}^{\top}k_{t}}$ $$

(3) is derived from the fact that $(I+\beta_{t,i}k_{t}k_{t}^{\top})^{-1}=(I-\frac{\beta_{t,i}k_{t}k_{t}^{\top}}{1
+\beta_{t,i}k_{t}^{\top}k_{t}})$ by the Sherman–Morrison formula. (5) is derived by noticing that $k_{t}^{\top}k_{t}k_{t}x_{t,i}-k_{t}k_{t}^{\top}k_{t}x_{t,i}=0$.
∎

###### Theorem C.1 .

###### Proof.

## Appendix D Additional Experiment Details

We provide the architecture detail for conducting the scaling law experiments on OpenWebText in Table 5. The architecture configs follow exactly from the Mamba paper (Gu & Dao, ).

**Table 5: Training details on OpenWebText.**
| Params | n_layers | d_model | n_heads / d_head | Training steps | Learning Rate | Batch Size | Tokens |
| --- | --- | --- | --- | --- | --- | --- | --- |
| 125M | 12 | 768 | 12 / 64 | 4800 | 6e-4 | 0.5M tokens | 2.5B |
| 350M | 24 | 1024 | 16 / 64 | 13500 | 3e-4 | 0.5M tokens | 7B |

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 6
presentation: 4
questions: Are there any actual wall-time comparisons in terms of training & inference?
soundness: 3
strengths: 1. The paper is well-written and easy to follow. 
2. The formulation via online learning for solving in-context associative recall is interesting and elegant. It explains why Longhorn (and also DeltaNet) performs well in MQAR tasks.
3. Empirical results look good.
summary: Longhorn formulates state-space models by solving an online regression problem. By designing various online learning objectives, it can induce different linear recurrent models, providing a unified framework and principled approach for developing new models. This work adopts an objective that encourages the hidden state to remain close to its previous state (i.e., \( ||S_{t+1} - S_t||_F \)) and includes a term that promotes key-input association (i.e., \( ||S k_t - x_t||_{\text{diag}(\beta_t)^2} \)). While the optimal solution is derived, it presents computational challenges. To address this, a diagonal approximation is used, resulting in a model computationally similar to Mamba (and thus similarly efficient) but with improved empirical performance, as demonstrated on synthetic datasets and medium-scale language modeling and image modeling.
weaknesses: 1. The main issue with this work is that the implementation does not fully align with the theory. Using a diagonal matrix to approximate an “identity-plus-low-rank” dense matrix is coarse, and it’s unclear if the theoretical advantage translates to this setting.
2. In Eq5, the norm \(\text{diag}(\beta_t)\) appears unusual and is not well-motivated or empirically validated. Why is a vector-valued \(\epsilon\) necessary? If not, the DeltaNet structure could leverage the kernel from Yang et al. (2024, https://arxiv.org/abs/2406.06484), which would easily scale up the head dimension and likely benefit recall-intensive tasks requiring a large state size. Longhorn, as it stands, cannot be expressed in matmul form, leading to similar challenges as in Mamba. Would Mamba2-like optimization, potentially resulting in a DeltaNet-like model with scalar \(\epsilon\), be preferable?
3. MQAR is a synthetic dataset and insufficient to demonstrate Longhorn’s advantages in recall-intensive tasks. Results on real-world recall-intensive tasks proposed in Arora 2024 [https://arxiv.org/abs/2402.18668] would provide a stronger case. Could you report zero-shot accuracy on these tasks? A table similar to [Table 1, https://arxiv.org/abs/2407.05483] would be very useful and necessary.
4.  This work lacks several ablation studies. For instance, the "value" projection is removed compared to standard models, yet this change is not analyzed. Additionally, the model does not clarify the benefits of parameter tying.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: see weakness.
soundness: 3
strengths: 1. The paper is well-written and easy to follow. The clarity of explanation makes complex ideas accessible, particularly in sections like Appendices A and B, which provide valuable insights into the nuances of different approaches.
2. The novelty of the new formulation for the Longhorn approach is impressive. The retrieval-based perspective is both innovative and elegantly presented, offering a fresh solution that enhances the field.
3. The exploration of SSM structure variances through online learning optimization is also intriguing and adds depth to the paper’s contribution.
summary: This paper introduces Longhorn, a new state-space model (SSM) architecture. By adopting an online learning optimization perspective, the authors unify several popular SSM architectures, bringing clarity to the structural differences between them. The explanation in Appendix A and Table 4 are especially helpful for understanding these nuances. Building on this unified approach, the authors propose a simplified SSM structure through a novel value retrieval mechanism based on key structures, offering insightful explanations of their method. The paper concludes by deriving a closed-form update formula for the state S in the SSM, supported by effective empirical results.
weaknesses: 1. While this paper presents a focused study on architecture, the data and model scale seem limited. Expanding the experimental scale and providing a more comprehensive analysis would significantly enhance the paper's impact.
2. The reduction in perplexity compared to Mamba is notable. However, the results in Table 2 appear mixed, which could benefit from further clarification or exploration.
3. Including additional experiments, such as MMLU, GSM-8K, and more extensive long-context benchmarks, would strengthen the findings and provide a more robust evaluation of the model's capabilities.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. can you provide more details on the sample efficiency experiments? Say, what kinds of hyper-parameters did you try? Can you do an abaltion study? 

2. Echoing the weakness of the paper, it is unclear to me that after using such an appoximation, is the theory framework still well aligned to the experiments?
soundness: 3
strengths: - the online learning framework provides a fair theoretical underpinning for understanding the Linear attention model / SSMs. This approach not only supports the conceptual innovations presented but also enhances the interpretability of SSM behaviors in practical applications.

- emirpical results: Longhorn has good sample efficiency compared to STOA models such as Mamba and GLA. This advantage is critical in scenarios where computational resources are limited.
summary: The paper introduces Longhorn, a novel state-space model (SSM) architecture designed as a meta-module that effectively handles sequence modeling problems. It described a theoretical framework based on online learning principles to derive the closed-form solutions for the online associative recall problem. 

The empirical results convincingly demonstrate that Longhorn surpasses other state-of-the-art SSMs in performance, particularly highlighted by its impressive recall capabilities on the Multi-Query Associative Recall (MQAR) benchmark.
weaknesses: Appoximation: while the diagonal approximation is a key aspect of Longhorn's implementation, its impact on the theoretical framework's alignment with empirical results remains unclear to me. I would expect a deeper exploration into how this approximation influences model performance could bridge the gap between theoretical predictions and observed outcomes.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - There appears to be a discrepancy between the significant improvements in PPL versus the modest gains in downstream task metrics. Could the authors elaborate on this phenomenon?
- Could the authors provide analysis on why DeltaNet struggles with extrapolation, while Longhorn demonstrates superior extrapolation capabilities, especially given their similarities in the update equation?
- Is there an ablation study on the beta parameters in the OCP formula? What guidelines exist for its optimal selection?
soundness: 2
strengths: - Novel theoretical perspective analyzing SSMs through online learning
- Proposed theoretical framework shows strong potential for length generalization
- Comprehensive empirical validation across various settings
summary: This paper presents a novel perspective on SSM models through the lens of online learning, offering a fresh analytical framework. The learning process consists of two losses: one ensuring minimal state updates, and another optimizing the reconstruction of current x from the state using key k. Within this framework, the authors propose Longhorn, which introduces differential importance weighting for various dimensions during reconstruction. Experimental results across different scales, tasks, and sequence lengths demonstrate Longhorn's superior compression rate on PassKey tasks and better length generalization, while maintaining comparable performance with baseline SSMs on other tasks. While the theoretical framework shows significant value, the paper makes an approximation step without thoroughly discussing its impact on the overall theory, which intuitively could lead to substantial differences.
weaknesses: - Insufficient discussion of the approximation's impact, creating a gap between theory and practice
- Limited comparison with contemporary methods (e.g., Mamba2 and DeltaNet); given the concurrent timing with DeltaNet, would appreciate author response on this comparison

Meta-review: This paper offers a novel online-learning perspective on SSM design, introducing a novel architecture called Longhorn and demonstrating improved performance over strong baselines like Mamba. All reviewers appreciated the clarity of exposition, the elegance of the theoretical framework, and the thorough empirical comparisons showing Longhorn’s effectiveness in tasks requiring long-context processing. While the reviewers initially expressed some concerns (e.g., diagonal approximation), they later found the rebuttal and additional experiments convincing.

Decision: Accept (Poster)