Longhorn: State Space Models are Amortized Online Learners
Authors: Bo Liu
Rui Wang
Lemeng Wu
Yihao Feng
Peter Stone
qiang liu
Venue: ICLR 2025 Poster (2025)
Keywords: Deep State Space Models
Linear Attention Models
Online Learning
Language Modeling
Primary area: foundation or frontier models, including LLMs
TL;DR: We present a unified framework to design state space models as solving online learning problems, and propose a new architecture that implicitly conducts online associative recall.

Abstract: The most fundamental capability of modern AI methods such as Large Language Models (LLMs) is the ability to predict the next token in a long sequence of tokens, known as “sequence modeling.”  Although the Transformers model is the current dominant approach to sequence modeling, its quadratic computational cost with respect to sequence length is a significant drawback. State-space models (SSMs) offer a promising alternative due to their linear decoding efficiency and high parallelizability during training. However, existing SSMs often rely on seemingly ad hoc linear recurrence designs.
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 optimizing these objectives.
Based on this insight, we introduce a novel deep SSM architecture based on the implicit update for optimizing an online regression objective. Our experimental results show that our models outperform state-of-the-art SSMs, including the Mamba model, on standard sequence modeling benchmarks and language modeling tasks.

Published as a conference paper at ICLR 2025
LONGHORN: STATE SPACE MODELS ARE AMORTIZED ONLINE LEARNERS
Bo Liu Rui Wang', Lemeng Wu', Yihao Feng', Peter Stone Qiang Liu' The University of Texas at Austin, Meta, HHelixon, "Sony Al. (1bo, 1mwu) kmwu)emmeas @meta .omm {yihao, pstone, lqiang] utexas.
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 major limitation. State- saate-spaceeae 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, explore SSM design through the lens of online learning, conceptualizing SSMs as meta-modules for specific online learning problems. This approach links SSM design formulating precise online learning objectives, with state transition rules derived from solving these objectives. Based on this insight, introduce 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 sample efficiency compared to Mamba, and can extrapolate over contexts that are up to 16x longer during inference. The code provided at ht : thub. com/ Craniai al XIX/ Longhorn.
Avg. PPL Downstream Datasets
PPL SlimPajama
30
(Best: 20 86)
0.0
LLaMA (2048)
28
Mamba (Best: 18.861
9.5
Mamba (2665)
GLA (2046)
26
Longhom (Best: 18.05)
Longhoms (2018)
9.0
24
8.5
20
speed up
8.0
18
7.5
10
20
30
40
50
60
70
90
100
2048
4096
8199
16384
32768
Seen Tokeens SlimPajame (Billion)
Context Length
Figure 1: (left)) The average perplexity on eight downstream datasets for GLAA, Mamba, and Longhorn (1.3B model) over seen tokens SlimPajama. Longhorn leads 1.8x speed up in sampling efficiency. (right) Longhorn, pretrained with 2048 context length, extrapolates up to longer context inference.
INTRODUCTION
The Transformer model has become the go-to architecture for sequence modeling in deep learn- ing (Vaswani et al., 2017). However, its utility is constrained by the quadratic growth in training and decoding costs with increasing sequence length. Despite various optimizations such as efficient de- coding (e.g., Chen al., 2023; Kuperman & Dyke, 2011), 1), KV-cache compression (e.g., DeepSeek-AI Dai, 2024), and memory efficient implementation (e.g.. Dao et al., 2022), it remains challenging to scale Transformers for autonomous and continual with infinite (or very long) context window.
Recent advances in linear attention models (Katharopoulos et al., 2020) and state-space models (SSMs) (Gu et al., 2021) have demonstrated their potential. These models specialized recurrent

---

Published as a conference paper at ICLR 2025
Online Learning Objective
Longhorn Objective
(+(5) = DSS, St-1) (s())
||S St-1l + ||Sk; - 2i =ellang(()))
Channel Mixing
(
Sequence Mixing
Online Update
Longhorn Update
St = argmin Lx(S)
S: = A, @ S + B,
Figure 2: (left) Most existing sequence models consist of channel and sequence mixing layers. The sequence mixing layers can viewed as "meta-module" that compress history into state St, which is then passed to later layers for sequence modeling. (middle) Sequence mixing can seen as an online learning problem, where the SSM state St optimizes an online objective. The recurrent update of St is derived either by solving this objective in closed form or via proximal update. (right) Longhom's 's update solves online associative recall, where the goal is recover R based hint & R from state matrix S g Raxm LLnghomn' 's update corresponds to the implicit online learning solution, where A, 1axm - Es @ 8002 and B, (e O xe) @ kg and B./(1 3ck, k.). See the details in Section 3 and Algorithm 1.
neural networks capable of efficiently computing outputs in parallel when input tokens provided simultaneously during training, thus avoiding the inefficiencies of traditional backpropagation through time. During inference, the recurrent form is employed. resulting linear decoding efficiency. Initially, these models underperformed compared to Transformers. However, recent SSMs (c.g.. Gu Dao, 2023; Yang al., 2023; Peng et al., 2024; De et al., 2024; Beck et al., 2024) have achieved performance parity with Transformers in language modeling tasks. Despite extensive research into various design aspects of SSMs, 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.
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 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 closed-form solution to an online learning objective. We outline the corresponding objectives for several existing SSMs in Table 5. 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 of-te-of-the-rr art SSMs like Mamba (Gu & Dao. 2023) on synthetic and large-scale sequence modeling tasks. In particular, Longhorn outperforms Mamba at the size of --B--aarmetter when trained on 100B tokens from the SlimPajama dataset (Soboleva et al., 2023). To summarize, contributions
1) Theoretical Framework:
We propose a novel framework that views SSMs' recurrent update as solving online learning objectives. As 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 online associative recall problem. Longhorn's '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

---

Published as a conference paper at ICLR 2025
Algorithm Longhomns Single-layer SSM Recurrence (Inference Time)
Parameters: Wa E IRm Wi Rmxdd W Reddd where Wa can be low-rank, horizon
2: Initialize the memory state So + 0dxm.
3: for & [ do
4: 1) Receive input 2; E IR..
5: 2) Compute the query qr, key bi and Be:
Q: Wg RM
h Wi R",
B =Sigmoid(Wgx.)) (0, 1)
6: 3) Update the memory state S Rd via
S. (ldxm &, & 02
+ St1 + (ex 16xO0))
2 kr,
&t = 3:/(11++ Bake ki) E (0,
7: 4) Compute the output or Sige Rd
8: end for
9: Note: O elementwise product and is outer product. 2; in practice is preprocessed through
linear projection followed by a Conv ld operation as in Mamba (Gu & Dao, 2023).
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 Mamba (See Figure (left)). Longhorn's 's training speed is as fast as Mamba, as we only replace the SSM module in the Mamba architecture with Longhorn's 's recurrence. So serves as a drop-in replacement for Mamba. Lastly, Longhorn, trained with 2048 context length extrapolate to 32K context length at inference time without much perplexity drop (See Figure (right)).
2) Empirical Results:
Notation Throughout this work, we use @ denote the Hadamard (elementwise) product, and @ to denote the Kronecker (or outer) product between two tensors. Uppercase letters A, B, etc. denote matrices, while lowercase in general vectors. ||-|| by default refers the Eg norm for vectors.
2 BACKGROUND
In this section, we provide a brief introduction to contemporary deep state space models (deep SSMs).
Modern large language models are requence-to-sequeene models consisting of stack of layers = 4r c ... a Φ (x) that sequentially processes an input sequence æ = {xr)k=11 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., 2023; Yang al., 2023; De al., 2024). while others fuse the SSM layer and the MLP layer into a single unified module (Gu & Dao, 2023). But in both cases, the sequence mixing is done by the SSM module, and the channel mixing done by the channel-wise MLP. Taking Mamba as an example (Gu& & Dao, 2023), Mamba model consists of a stack of homogeneous modules named Mamba block (the (x)); we provide a visualization of a single Mamba block in Figure 3 (Gu& & Dao, 2023), which consists of an SSM block for sequence mixing (red), and MLP block for channel mixing (bluc).
The SSM block (in red) plays the crucial role of sequence mixing. It works by iteratively updating a memory state matrix S, E Rdxm with a linear recurrence:
SSM: General Form
S. A(xz) * S + B(xe),
Vt (1....,T).
S. 0.
where 2, is the input at time t, S, is the model's state, A,, B Rd ddxm are some functions of the input and is multiplication operation of choice, such Hadamard product or matrix product.
3

---

Published as a conference paper at ICLR 2025
Given the state S.. SSMs often give the output token at the next layer via gated linear unit (GLU) (Dauphin et al., 2017):
Y/ Readout z4) W1 (or O (WWz)))
= Czz))s,
where we first get ot via state-dependen linear projection St, which is then fed into subsequent channel mixing gated linear unit (blue Figure 3), where a()) is a non- non-linear activation function.
A key feature of this design in Equation is that S₁ has linear recurrence, i.e., St is linear in St-1- Crucially, this allows us to express all St in an explicit form that can be calculated in parallel: when all { are available as in the training phase, {Se)e can be written into
Channel Mixing
Lineer
Sequence Mixing
S = ((r=)))
where
Ar-
II A(x,). PETEE
(2)
Here [ denotes the product induced by multiplication operator *. The resulting cumulative product Aç-.t can implemented efficiently par- allel with the prefix scan algorithm (e.g., Harris et al., 2007), which only requires (log 7) (T' is the sequence length) parallel operations. From
Linear
Linear
now on, we will abbreviate A(xx) and (xx+) as A, and B,, respectively.
Figure 3: Mamba Block
Designs of (At, Bt, *) Existing variants of SSMs mainly differ in the design choices of the networks At. Br, and the associated operator * in the linear recurrence. core issue here is that the memory state S, RRdmm designed to be times the input I₁ in size, must be as large as possible to maintain sufficient information during recurrence. This makes the architecture design of Ac, Br. both mapping R to Rdx challenging, naive linear mapping would result in d x d x m weights, which is prohibitively large. This makes it necessary to impose certain low-dimensional structures in A,, B,, 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 the previous section, designing state-space model (SSM) depends on the specific selection of (Ar, Br, *), 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 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 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 online learning problem. In online learning, the agent picks state 8, at time and then incurs a loss (x:)). The goal is to minimize
min
Σ ((t)) {81]
For instance, consider online linear prediction, where at each step the agent given an input-label pair (xx, v)) and fe(st) l| I, vall is the & regression loss, then the problem becomes an online regression problem, and the goal is to successfully predict y, given 2; at future time steps, with the key feature that the prediction (s,) can change with each data point.
Online convex programming (OCP) (c.g.. Zinkevich, 2003) yields a principled approach to solving Equation 3 when E, are convex, by trading-off the "stability" and "plasticity" (e.g., Mermillod et al.,
4

---

Published as a conference paper at ICLR 2025
2013). Formally, online convex programming algorithm updates 8, by solving regularized cost function:
$1 arg min L((s), Le(s) Do((, + Byle(s).
(4)
stability
plasticity
where 3, R- and D is discrepancy measure, often a Bregman divergence induced by the convex function & (c.g., when ((x) = 2/2 |=|| (s, st-1) = ž ||| st-1ll²). Here the first term ensures the updated will be close to the previous st-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 (s). Hence, Bt controls the trade-off between stability and plasticity.
3.2 THE LONGHORRN ARCHITECTURE
Under the online learning framework, the design of an SSM reduces to the design of De and & in Equation 4. This provides a unified framework for the existing SSM variants. We summarize in Table 5 in Appendix 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 Table 5). particular, we consider xe)) as the input stream, where 22 E R and E R& are the key-value pairs, just as in the Transformer model ( (Vaswani et al., 2017). In practice, as in Mamba (Gu & Dao, 2023), ke = Wa E R", where Wa E Rmxd is linear mapping from
We want to recurrently update hidden states (St) where S, E Rdxm is a matrix that summarizes the information up to time (. We posit the following OCP objective for updating S₁:
S. arg min Seruxm
{1SS- S + ||Skr zzl zelldtag(3)))
(5)
Here, || ||e denotes the Frobenius norm of a matrix, 3r E Rd is vector controlling how much new information about x, we want the model to incorporate for S.. For instance, 3r,; = 0 implies Sr, St-1, (i.e., the i-th row of S remains unchanged), while large Br, implies the model empties some part of S; for incorporating Ft,i-
From high-level perspective, Equation 5 is solving an online prediction problem of learning weight matrix S to predict 2, given K; with a linear model oe se S kg. It is supervised formulation of the associative memory problem of memorize (e)) pairs by learning mapping from ke to Fr. such that given key (input) k the model can retrieve (predict) its corresponding value (label) 2;-
The objective in Equation Transformer exhibits a form of online associative recall (often referred to as the induction head property) (Olsson et al., 2022). This capability has been shown to underpin the model's s ability to perform in-context learning (Brown, 2020). To explain the connection, in-context learning refers to the model's s ability, during inference, to generalize from a set of provided (k; x) (question-answere new question. This closely parallels associative recall, where
motivated by the observation that the self-attention layer of the
pairs and apply this understanding the model retrieves relevant information from past interactions to address new inputs.
Fortunately, this simple objective gives closed-form solution for St, which coincides with the implicit online learning method (c.g., Kulis & Bartlett, 2010), according to Theorem 3.1 (We provide the proof in Appendix C):
Theorem 3.1. The closed form solution for for objective in Equation 5 is
3r.i S₁, (I (t, + Ei,kke where Ei, I+ Br Ke
E [0, co).
(6)
Here, Sr,; refers to the i-th row of St, Br,; refers to the i-th element of 3, As keck is a matrix, it is hard to compute its cumulative product for conducting a parallel scan. As a result, in practice, we use O a and 1m is
the diagonal approximation 1m k02 in place of Et,kee,, where 00 the m-dimenssiona all-one vector. Following Mamba (Gu & Dao, 2023) and Transformer (Vaswani al., 2017), make k, WA e Rh and 8, a (Wa..) E Idd (both are functions of xx). where
5

---

Published as a conference paper at ICLR 2025
the activation (the Sigmoid function) is to ensure that 3r is positive and bounded. In summary, the final Longhorn update of St becomes:
S A + Br, where A, (1axm X 102), B, (ex C G.r ② kr.
(7)
The final architecture of Longhorn follows Mamba strictly (Figure 3), except that we replace the SSM block with Longhorn's 's recurrence. We also provide an efficient CUDA kernel for it. The full inference-tine timriee--mmeeeee algorithm is provided in Algorithm 1. One can compare Equation 7 to Equation and other SSMs in Appendix A. Longhorn does not introduce an extra "forgetting" gate (hence has fewer parameters), because the forgetting gate is naturally derived from the key vector, i.e., (1dxm El k(2).
Advantages of Longhorn
1. While 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 (c.g., ae in GLA or matrix in Mamba). The forgetting is naturally linked to the key vector kt through the derivation. This saves about Odd x m) parameters per SSM module, where m is the dimension of k2, and is the dimension of 24. However, Longhorn demonstrates better performance even with fewer parameters than Mamba (See Figure (left), Table 2, Table 3).
3. The closed-fon form solution in Equation does not need any specific initialization. In contrast, Mamba requires special careful initialization of the A and et.
4. Unlike DeltaNet (Yang et al., 2024), which struggles to extrapolate beyond training contexts, Longhorn successfully extrapolates contexts 16x longer than it trained for (Figure 1 (right)).
4
RELATED WORK
This section provides a summary of recent advances 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 constant-size matrix, bypassing the scaling with sequence length (Wang et al., 2020). Linear Transformer replaces the Softmax function with a decomposable similarity function, achieving linear complexity (Katharopoulos et al.. 2020). Performer approximates softmax attention using orthogonal random features (Choromanski et al., 2020). RetNet adds constant forgetting and rotation (Sun al., 2023), while Gated Linear Attention introduces learnable forget gates (Yang et al., 2023). Linear attention can also be seen as a fast weight network where slow net adapts fast network's parameters online using inputs (Schlag et al., 2021).
State Space Models State space models (SSMs) focus parallelizable linear recurrent networks. Initially, a constant state transition matrix A allows recurrence to be computed via convolution (Li al., 2022; Gu al., 2021). Key models include Diagonal State Space (DSS) (Gupta al., 2022), Gated State Space (GSS) (Mehta et al., 2022), S5 (Smith et al., 2022), Bidirectional Gated SSM (BiGS) (Wang et al., 2022), H3 (Fu et al., 2022), and Mamba (Gu & Dao, 2023). Efficient recurrent
networks often resemble SSMs, such as Deep Linear Recurrent Units (LRUs) (Orvieto et al., 2023; De et al., 2024), Hierarchically Gated Linear RNNs (HGRN) (Qin et al., 2024b;a), and RWKV (Peng et al., 2023; 2024).
Fast Weight Programmer The idea of networks modifying their own weights in response to inputs dates back to the Fast-weight Programmer (Schmidhuber, 1992; 1993; Schlag & Schmidhuber, 2017; Schlag et al., 2021). These models update a weight matrix W dd via the outer product two ② k(xz), 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. More recently, Transformers (self-attention) has also been discovered as conducting mesa opeaiioptini,tion, adjusting the model as it reads new inputs (Von Oswald al., 2023).
vectors: AW

---

Published as a conference paper at ICLR 2025
Concurrent Work Two concurrent works share similar ideas with ours. Yang al. (2024) propose a chunk-wise parallel approach to scale DeltaNet (Schlag et al., 2021) for large-scale language modeling. DeltaNet's update rule, which is viewed a gradient step for an online regression objective, results in a state transition matrix A(x²) = (T - 3,kkk, ), which can have eigenvalues > 1, leading to instability. To address this, Yang et al. (2024) normalize the key vector & by its 42 norm, which can be restrictive. In contrast, Longhorn ensures stability with a closed-form update, using a diagonal approximation (kP²), 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 parallel scan CUDA kernel, enabling Longhorn to serve as drop-in replacement for Mamba. Sun et al. (2024) introduce the Test-Time Training framework, where state updates are derived from gradient step on an online regression objective. To maintain parallelism, they assume that each gradient step at x uses the initial state 8o. enabling matrix multiplication. In contrast, Longhorn computes the closed-form solution for the every token, offering greater flcxibility.
EXPERIMENTS
We validate Longhom's 's performance through the following experiments:
1) We compare Longhorn against other SSMs on the multi-query associative recall benchmark (Arora al., 2023a) and find that Longhorn is the only model to achieve near-perfect recall at sequence lengths up to 512 with hidden dimension of 64. We further compare Longhorn (1B) against Mamba (IB) on real-world recall-intensive tasks (Arora et al., 2024a;c) and find that Longhorn also achieves a better recall rate compared against Mamba.
2) Using the WpenWebText dataset (Gokaslan & Cohen, 2019), 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 1.3B language model on the SlimPajama dataset (Soboleva et al., 2023) 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., 2024) model (Appendix 5.6), 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., 2023a). The agent observes sequence of tokens [k1, 11, k2. u2....,,k. UT). two-tokens become a key-value pair. At test time, the agent is provided with multiple {kI.... kT), consider
where each consecutive
the goal is to "retrieve" the corresponding values. Following the original benchmark, the sequence length T E (64, 128, 256, 512) and model dimension (size of the latent embedding of token) d E (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 al., 2024b), 3) Hyena (Poli al., 2023), which is special SSM that adopts long convolution via fast fourier transform, 4) RWKV (Peng et al., 2023), which can be viewed as the division of two SSMs (i.e.. y = a/b, where a, b are outputs from two SSMs). The combines linear projection with convolution, and 6) Mamba (Gu & Dao, 2023), the state-of-the-an 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, 4.6 x 10-4 2.2 x 10 10 10-2). Results are summarized in Figure 4.
matrix is scalar, 5) BaseConv (Arora et al., 2023a), an SSM that
Observation: From the figure, we see that Longhom, 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 small model dimension of 64.
7

---

Published as a conference paper at ICLR 2025
Attention
Based
Hyena
RWKV
BaseConv
Mamba
Longhorn
Sequence Length:
Sequence Length: 128
Sequence Length: 256
Sequence Length:
1.00
Accuracy
0.75
0.50
0.25
0.00
64
128
256
512
64
128
256
512
64
256 Model dimension
128
512
64
256 Model dimension
128
512
Model dimension
Model dimension
Figure 4: Comparison of Longhorn to statc-of-the-art SSMs the MQAR benchmark, y-axis is the recall rate.
5.2 REAL-WORLD RECALL INTENSITIVE TASK
To further evaluate Longhorn's ability to recall in long real-world sequences, following the experiment setup Arora et al. (2024c), we consider the following six recall-intensive tasks: information extraction tasks like FDA and SWDE (Arora et al., 2024a; 2023b; Wu al., 2021; Deng al., 2022), and
benchmarks like Squad (Rajpurkar et al., 2018), NaturalQuestion al., 2019), TriviaQA (Joshi et al., 2017), and DROP (Dua et al., 2019).
(NQ) (Kwiatkowski Following Arora et al. (2024c), we report both the vanilla question answering accuracy, and the accurayc under the JRT-Prompe format (Arora al., 2024c), where the context is repeated twice before the model conducts the final completion. The zero-shot prompt includes up to 1k tokens in the input and JRT- JRT-Prompt includes up to 2k tokens in the input for all tasks, repeats the context twice. The IB checkpoint of Longhorn is taken from Section 5.4. The results are provided in the following table:
Model
FDA
SWDE
Squad
NQ
TriviaQA
DROP
Average
Mamba-1 1.3B
33.2 / 40.6
35.0 / 36.2
26.6/ 32.6 / 52.7
56.3 / 56.9
20.4/31.5 34.8 41.8 21.3 33.3 36.0 / 45.0
Longhorn-1..
40.2 50.4
33.2 42.3
27.6 33.2
35.0/ 55.0
58.5 55.9
Table 1: Longhomn and Mamba's recall performance across real-world recall-intensive benehmarks.
From the table, it shows that Longhorn 1.3B achieves a 3.4% improvement of recall accuracy under the Mhilla/JRT-Prropl context format, compared against the same size Mamba.
Observation:
5.3 SCALING LAW ON OPENWER
In this section, we consider language modeling tasks on models with 120M or 350M parameters with 1024 or 4096 context length. We choose the Open WebText dataset as it is small and serves as an easily accessible benchmark for quick benchmarks. The details about the architecture is provided in Appendix D. consider following baseline models: LLaMA (Touvron al., 2023), RetNet (Sun et al., 2023), Mamba (Gu & Dao, 2023), RWKV (Peng et al., 2023), and GLA (Yang et al., 2023). Then we experiment with 1024 or 4096 context length T and model sizes around 120M or 350M. Results summarized in Table 2 and Figure 5.
Observation: From the figure and table, SSMs up to 350M and 4096 context length.
that Longhorn consistently outperforms baseline
5.4 LARGE-SCALE LANGUAGE MODELING
For the large-scale language modeling task, we followed the GLA (Yang et al., 2023) setup, training a 1.3B parameter model on the SlimPajama (Soboleva et al., 2023) dataset with 100B tokens and a batch size of 2M. We used the AdamW optimizer (Loshchiloy & Hutter, 2017) with weight decay of 0.01, cosine learning rate decay (peak: 3e 4, final: 3e 5), and gradient clipping of
We adapted code from the nanoGPT repository minimal reproduction of GPT-2 model using PyTorch.
thub.com/eerppattyy/nnorr which
8

---

Published as a conference paper at ICLR 2025
Context Length 1024
Context Length 4096
34
RetNet
Val. Loss
LLaMA
Mamba
Val.
GLA
3.0
5.0
Langharn
2.8
2.8
250
300
300
350
Parameters (M)
Parameters (M)
Figure 5: Scaling law with 1024 and 4096 context length Open WebText with various SSM models and the LLaMA (strong Transformer) baseline.
Val. Loss (L)
Val. Loss (L)
Model
Param. (M)
# Param. (M)
1024
4096
T 1024
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
Table 2: Language modeling scaling law against LLaMA (Touvron et al., 2023), RetNet (Sun al., 2023),
RWKV (Peng et al., 2023), and Mamba (Gu Dao, 2023). All models are trained on Open WebText dataset (Gokaslan & Cohen, 2019). Models vary from 120-350M parameters and 1024-4096 context length.
1.0. Comparisons were made against LLaMA, Mamba, and GLAA models (context size: 2048). We evaluated eight standard downstream tasks, including PIQA (Bisk et al., 2020), HellaSwag (Hella) (Zelless et al., 2019), WinoGrande (Wino) (Sakaguchi et al., 202 I), ARC-easy (ARC-e) and ARC-challenge (ARC-e) (Clark et al., 2018), OpenBookQA (OBQA) (Mihaylov et al., 2018), Social Interaction QA (SIQA) (Sap et al., 2019), and Boolean questions (BoolQ) (Clark et al., 2019). 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 3.
PIQA
Hella
Wino.
ARC-e
ARC-c
OBQA
SIQA
BoolQ
Model
State Size
Avg.
acc 1
acc norm
acc 1
acc
acc norm
acc 1
acc norm 1
acc 1
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
$1.07
Mamba
64K
54.21
53.61
71.67
61.05
30.15
43.94
44.18
59.22
$2.25
Longhorn
64K
55.78
52.30
71.00
60.63
29.53
43.55
44.68
61.29
52.35
Table 3: Language modeling results against LLaMA (Touvron et al., 2023), RetNet (Sun et al., 2023), and Mamba (Gu & Dao, 2023), All models trained the same subset of the SlimPajama dataset with the Mistral tokenizer. The 340M/1.33 models are trained for 15B/100B tokens respectively. State Size the effective state size an SSM per layer. For instance, GL.A's state size (1024K) computed by mn// h, where key 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.
From Figure I (left), it is evident that Longhorn not only achieves 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 3, we can see that up to 1.3B model. Longhorn remains strong among all baseline models and achieves slightly better result than Mamba, even though has a bit fewer parameters.
Observation:

---

Published as a conference paper at ICLR 2025
5.5 ABLATION ON LENGTH EXTRAPOLATION
We evaluate how Longhorn extrapolates to a context length longer than 2048 (training context length) inference time. In particular, we pick disjoint validation set from SlimPajama dataset, rearrange it into batches of sequences of length T E {2048, 4096, 8192, 16384, 32768), and then evaluate the pretrained model's perplexity on those sequences. The results are summarized in Figure (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., 2024), which highlights& a limitation in that the model cannot extrapolate to longer contexts. In contrast, LLaMA, as Transformer-based model, fails to extrapolate beyond its training context length.
5.6 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., 2024), we conduct experiments on the ImageNet (Deng et al., 2009) classification task. Similar ViM, We apply 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.
Model
# Param
Top-1 1 Accuracy
ViM-Tiny
7M
76.1
ViL-Tiny (ours)
7M
76.4
ViM-Small
26M
80.5
ViL--Smll (ours)
26M
80.7
Table 4: Top-1 Accuracy on ImageNet classification for Vision Mamba (ViM) and Vision Longhorn (ViL).
Observation:
The results from Table demonstrate that the Vision Longhorn model (ViL) achieves comparable (slightly better) performance to the original ViM. Note that we use the best hyperparame- ters for ViM without additional tuning, and ViL does not require 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 conceptualiz- ing the recurrence update 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-as 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 al., 2024) suggest that incorporating diding-windon attention with Mamba improves performance. We anticipate similar benefits for Longhorn.
7
ACKNOWLEDGEMENT
This work has taken place in the Learning Agents Research Group (LARG) and the Statistics & AI group at UT Austin. LARG research is supported in part by NSF (FAIN-20199 1844, NRT-2125858), ONR (N00014-15 1000014-8-22444 18-2243), ARO (W91 NFF23-2200004 Agreement HR001 12520004 on Ad Hoe Teamwork) Lockheed Martin, and UT Austin's 's Good Systems grand challenge. Peter Stone serves the Executive Director of Sony AI America and receives financial compensation for this work. The terms of this arrangement have been reviewed and approved by the University of Texas at Austin in accordance with its policy on objectivity in research. The Statistics & Al group receives support in part from NSF CAREER 846421, SenSE2037267, Office of Navy Research, and NSF AI Institute for Foundations of Machine Learning (IFML).
1911NNF17-2-01811 DARPA (Cooperative
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Mcasuring and improving recall in efficient language models. arXiv preprint wwXiv:2312.04927, 2023a.
Simran Arora, Brandon Yang, Sabri Eyuboglu, Avanika Narayan, Andrew Hojel, Immanuel Trummer, and Christopher Ré. Language models enable simple systems for generating structured views of heterogencous data lakes. Proc. VLDB Endow. 17:92-105, 2023b. URL https://aapp
semanti www.nntt.scccllloo.....ooooo Cttp:/orussiiii : 258212828.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zins- ley, James Zou, Atri Rudra, and Christopher R'e. Simple linear attention language mod- balance the recall-throughput tradeoff. ArXiv, abs/2402.18666, 2024a. URL https:
//api. semant cecccccchllaaco wttg/oorpusscceooooooooooo :D
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance tecaul-throughput tradeoff. arXiv preprint arXiv:2402.18666 2024b.
Simran Arora, Aman Timalsina, Aaryan Singhal, Benjamin F. Spector, Sabri Eyuboglu, Xinyi Zhao, Ashish Rao, Atri Rudra, and Christopher R'e. Just read twice: closing the recall gap for recurrent language models. ArXiv, abs/2407.0544 .05483, 2024c. URL https: /apppapp..../// sewanticschooiioooo......cccccoo
org/CozpusID: 271051359.
Maximilian Beck, Korbinian Poppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael K Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Ex- tended long short-term memory. 2024. URL https:/ http:///.a..smaisssssssss.ee
semanticscholar org/
CorpusID 26961433.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005.14165 2020.
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, tean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Christopher Clark, Kenton Lee, Ming-Way Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint 19X.::9905.1004 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv: 1803::8800005535 2018.
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher R'e. Flashattention: Fast and memory-efficient exact attention with io-awareness. ArXiv, abs/2205.14135, 2022. URL https: //api. www.nteistceeeeeceeeeeeeee icccholar. :rgg/CoppusID::200 2449151871.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference machine learning, pp. 933-941. PMLR. 2017.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Murardl Albert Gu, Ruba Haroun, Leonard Berrada. Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 19427, 2024.
11

---

Published as a conference paper at ICLR 2025
DeepSeek-Al and Damai Dai. Deepseek-v2. A strong, economical, and efficient mixture-of- experts language model. 2024. httsttps:// /aaps://appl http://app.....ai.eeeeeaeem wwwant conmanttccchoarr.. org/CorpusID: 269613809.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 255. leee, 2009.
Xiang Deng, Prashant Shiralkar, Colin Lockard, Binxuan Huang, and Huan Sun. Dom-Im: Learning generalizable representations for html documents. ArXiv, 222222011100008 2022. URL https:
//api. swwant wcscholar...cccccccooooooo.a.... wwt.aatttccccooo....oooooooaaaaa
wrg/Corpus///rroroooooooo : 246285527.
Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gardner. Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs. In North American Chapter of the Association for Computational Linguistics, 2019. URL https:
//api. sema ntt cacholas org httg///rpussaie I : 67855846.
Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint 1:X0::2212.14055, 2022.
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus, http::/99k/111/0007 OpenWebTextCorpuu 2019.
io/
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint arXiv:2312.00752, 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.00396 2021.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces, Advances in Neural Information Processing Systems,
2022.
Mark Harris, Shubhabrata Sengupta, and John D Owens. Parallel prefix sum (scan) with cuda. GPU gems, 3(39):851-876 2007.
Mandar Joshi, Funsol Choi, Daniel S. Weld, and Luke Zettlemoyer. Triviaga: A large scale distantly 2017. URL.
supervised challenge dataset for reading comprehension. ArXiv, https://apt semant inteentiicchhoo ar..a..... wwwpoorg//oooo Cotp:/orrps::::20000 25004419
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference Machine Learning, pp. 1156-5165. PMLR, 2020.
Brian Kulis and Peter L Bartlet. Implicit online learning. In Proceedings of the 27th International Conference on Machine Learning (ICML-10), pp. 575-582, 2010.
Victor Kuperman and Julie A. Van Dyke. Individual differences in visual comprehension of morpho- logical complexity. Cognitive Science, 33, 2011. URL https: //ppp/apppipp wwwanticschonttoooo org Tngg/Corre CCrrpss TD: 5555496.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur P. Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc V. Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Trans- actions of the Association for Computational Linguistics, 7:453-466, 2019. URL https: //api. semanteattt nt i csscholrr.
86611921
Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey. What makes convolutional models great on long sequence modeling? arXiv preprint arXiv:2210.09299, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. 17110:1711.05501 2017.
arXiv preprint
12

---

Published as a conference paper at ICLR 2025
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language modeling via gated state spaces. arXiv preprint arXiv:2206. 2022.
Martial Mermillod, Aurélia Bugaiska, and Patrick Bonin. The stabilit--plastiicl dilemma: Investi- gating the continuum from catastrophic forgetting to age-limited learning effects, 2013.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv: arXiiv:809.000 1809.02789, 2018.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et In-context learning and induction heads. arXiv preprini arXiv:2209. 1895, 2022.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pp. 26690-26698. PMLR, 2023.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing rnns for the transformer era. arXiv preprint arXiv:2305. 3048, 2023.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemyslaw Kazienko, et al. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. arXiv preprint arXiv:2404.05892 2024.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pp. 28043-28078. PMLR, 2023.
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong, Hgrn2: Gated linear rans with state expansion. arXiv preprint arXiv:2404.07904 2024a.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. Advances in Neural Information Processing Systems, 36, 2024b.
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't know: Unanswerable questions for squad. ArXiv, abs/1806.00 18011888660222.. 2018. URL https: /ppi. wwt.p..eeaaiiiiioooooooooooooooo semanticscholar org//
CorpusID: 47018994.
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.07522, 2024.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9)999-106, 2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. arXiv preprint /004:: 09728, 2019.
Imanol Schlag and Jürgen Schmidhuber. Gated fast weights for on-the-fly fly neural program generation. In NIPS Metalearning Workshop, 2017.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Jürgen Schmidhuber. Learning to control fast-weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
Jürgen Schmidhuber. Reducing the ratio between learning complexity and number of time varying variables in fully recurrent nets. In ICANN'93: Proceedings of the International Conference Artificial Neural Networks Amsterdan, The Netherlands 13-16 September 1993 3, pp. 460-463. Springer, 1993.
13

---

Published as a conference paper at ICLR 2025
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.0493: 2022.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hes- tness, and Nolan Dey. SlimPajama: plicated version of RedPajama. s1 slimpa jamaajaaaaaa a- - cleaned- and dedupliicated-vversior of reff-addaaaaa,, 627B.
A 627B token cleaned and dedu- https://wwer cerebraseerrrs... /ttp:/.ccrrbbr..ttt/oooo///////
2023. URL https: http///ggggggcceeeeee// co/ /ttaa/aattteeeeerrrr/////
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang. Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, Tatsunori Hashimoto, and Carlos Guestrin. Learning to (lear at test time): Rnns with expressive hidden states. 2024. URL https: /api...... semantiec ccolaclar..... org/CorpusID) 200129711200
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.139777 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Johannes Von Oswald, Maximilian Schlegel, Alexander Meulemans, Seijin Kobayashi, Eyvind Niklasson, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, et al. Uncovering mesa-optimizzttin algorithms in transformers. arXiv preprint arXhv:3309.05858 2023.
Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. arXiv preprint arXiv:2212. 2022.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04770 2020.
Eric Wu, Kevin Wu, Roxana Daneshjou, David Ouyang, Daniel E. Ho, and James Zou. How medical ai devices evaluated: limitations and recommendations from an analysis of fda approvals. Nature Medicine, 27:582 - 584, 2021. URL https: /api. semant i escholar org/ CorpusID: 233037201.
Songlin Yang. Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06633, 2023.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint arXiv:2406.06488 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint
2019.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint arXiv:2401.09417. 2024.
Martin Zinkevich. Online convex programming and generalized infinitesimal gradient ascent. In Proceedings of the 20th international conference on machine learning (icml-03), pp. 928-936, 2003.
14

---

Published as a conference paper at ICLR 2025
A PRIOR DEEP STATE SPACE MODELS
Example A.1 (Linear Attention Variants). Linear Attention (LA) (Katharopoulos et al., 2020).
Retention Network (RetNet) (Sun et al., 2023), and Gated Linear Attention (GLA) (Yang et al., 2023) all assume As, B3 yield rank-l (or even constant) oulputs:
A₁ =
(LA)
= + v(x)) & k(x)),
with
(RetNet). A & a(x)) (GLA)
A, CE [0.1]
where St Rdxm v(21) E R., k(xz) E T are linear mappings of zω, and denote the outer product. In practice, one can use h heads as in the multi-head attention to save some computation, where the and d dimensions The outer product complexity reduces to O(h # m/h * d/h = md/h) h). But then the effective size of St also shrinks to md/h.
divided into h groups and each group performs its own LA variant.
Example A.2 (Mamba (Gu & Dao, 2023)). The Mamba architecture is derived by discretizing continuous linear dynamics. Its discretized update is:
S. A, ④ S Br, where
(8)
A exp(A@ (E(Z))
2 ()))(()))))
where St Raxm with 16 by default, e(xr) E R., k(xz) R linear mappings of xu. and Raxm is a data independent (not depending on 24 trainable weight matrix.
In Mamba, both A, and B, depend on e(x:), 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 fixed and d, the largest memory state will be with h (as the effective size of St is md/ h). In addition, Mamba's 's output is o, S, + D, O T₁, which has an additional residual part D, Nt.
Example A.3 (Griffin (De et al., 2024)). In Mamba and the linear atlention variants, the outer product serves as critical role in lifting vectors to matrices. The recent Griffin architecture abandons the outer product and performs pure elementwise product:
St a(xx)) O st-1 /1- ((z.))((())) 21,
where St, a(xz), i((;) are all Rdd This yields smaller memory states, but in practice, Griffin is combined with local attention (i.e., the sliding- window self- attention) strengthen its capability.
alddingwwindow
Example A.4 (RWKV (Peng al., 2023)). The original RWKV also performs elementwise recurrence. maintains state of ratio form 8, = u/ 2t, where us 21 are updated separately by two SSMs:
= w2/21 Ue exp(p w) u exp(k(x))) 3 Ov(x))
24 = exp(-w) - -2-1-- +
where all the vectors are of size Rd, and 0 is a trainable weight for controlling the forgetting, In the most recent RWKV version (Peng et al., 2024), the denominator 2 is removed, and the elementwise product is replaced with the outer product, which makes it more similar 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 = @1 f( 2 S1-1 i(zx) 2 (1 - f(xz)).
Here, f(xe) [], is the forget gate, (1 f(xx)) is the input gate, and i(zx) is the input vector. HGRN2 thus resembles an RNN.
B EXISTING STATE SPACE MODELS ONLINE OBJECTIVES
We reverse-engineer some existing deep SSMs online learning objectives in Table 5.
15

---

Published as a conference paper at ICLR 2025
Method
Online Learning Ohjective L,(8) (assume 2, e R)
Online Update
LA
[| St Alfi 2.,
S, S. 6
RetNet
1/5 - 8. 7) 1) m()
S 75: 2 &
GLA
|| diddigg(aell 3.)
S. S_,diag(aa)
kz
Griffin
V (s (1--i)l) IVv ofl 2 Or O On
8. or 0 V(I 0.l On O
Longhom
||s - S:-illi [Sk (laing())))
S (1mxn 5x @
(e: (ex) @ lr. 3 P.A kz)
Table 5: Some of the existing SSMs their corresponding online learning ohjectivess/pdatts.
C PROOF
This section provides the proof for Theorem 3.1. Given the Longhorn's objective S. arg g {I|S - S.-1lli + I[Skr - xel xellding(..)) we have following theorem:
Theorem C.1. The closed form solution for St for objective in Equation
S₁ (I (u.-ea + Ee,kketi,s where Er 1 Be,ck Kr
E [0, (0)
(9)
Proof. As the objective in equation 5 is in a quadratic form with respect to there is unique minimum. Observe that each row (e.g., S,) optimizes the objective independently, therefore we solve the solution row-wise. By setting the derivative of V O, we have:
V (S, S1-1.1) + (S lo 21,1)h+ 11))==1 =0
(I + kikk ( )s = S1-1, + Becke
Bes kiki S
Be, I Re ke
S
kski Beckeerr
1 Be,kki Kt
Br,
(I Se, hie
keke
S
ke Kt
I I 3e kr
Br,5
Be,keoas + Be,kk, Ke
keck
I kit lis
(3) is derived from the fact that (I + Be, formula. (5) is derived by noticing that Kj] kekkaaee keki heree 0.
(I
1+ by the Morrison
Beskekk
D ADDITIONAL EXPERIMENT DETAILS
We provide the architecture detail for conducting the scaling law experiments on OpenWeeText in Table The architecture configs follow exactly from the Mamba paper (Gu & Dao, 2023).
Params
n layers
d.model heads / head Training steps Learning Rate
Batch
Tokens
125M
12
768
12/6e 64
4800
6e-4
0.5M tokens
2.5B
350M
1024
64
135500
3e-4
0.5M tokens
Table Training details on OpenWebText.
16


Reviews:
[{"review_id": "TUDspi943d", "reviewer": "Reviewer_dxPc", "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.", "questions": "Are there any actual wall-time comparisons in terms of training & inference?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well-written and easy to follow. \n2. 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.\n3. Empirical results look good.", "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.\n2. 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?\n3. 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.\n4.  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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "n3Sj9xsUAs", "reviewer": "Reviewer_LqBo", "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.", "questions": "see weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n2. 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.\n3. The exploration of SSM structure variances through online learning optimization is also intriguing and adds depth to the paper’s contribution.", "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.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DVuH65bQEN", "reviewer": "Reviewer_FTDh", "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. \n\nThe 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.", "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? \n\n2. 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n\n- 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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2PYcgW8XDj", "reviewer": "Reviewer_DVh4", "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.", "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?\n- 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?\n- Is there an ablation study on the beta parameters in the OCP formula? What guidelines exist for its optimal selection?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Novel theoretical perspective analyzing SSMs through online learning\n- Proposed theoretical framework shows strong potential for length generalization\n- Comprehensive empirical validation across various settings", "weaknesses": "- Insufficient discussion of the approximation's impact, creating a gap between theory and practice\n- Limited comparison with contemporary methods (e.g., Mamba2 and DeltaNet); given the concurrent timing with DeltaNet, would appreciate author response on this comparison", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)