# 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, ‡Helixon
<EMAIL_ADDRESS>

###### 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.111Code is available at https://github.com/Cranial-XIX/longhorn.git.

Figure 1: (left) The average perplexity across 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, successfully extrapolates to longer context (up to
16x of the trained context).

##  1 Introduction

The Transformer model has become the go-to architecture for sequence modeling
in deep learning (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 decoding
(Chen et al., 2023; Kuperman & Dyke, 2011), KV-cache compression (DeepSeek-AI
& Dai, 2024), and memory efficient implementation (Dao et al., 2022), scaling
Transformers for autonomous and continual use with an infinite (or very long)
context window remains a challenge.

Recent advancements in linear attention models (Katharopoulos et al., 2020)
and state-space models (SSMs)(Gu et al., 2021) 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 (Gu & Dao, 2023;
Yang et 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, 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 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 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) One can view sequence mixing as addressing an
online learning problem, where the state of the SSM follows certain dynamics
that optimize an online learning objective. The recurrent update of $s$ can
then be derived by solving the online objective in closed form or by
conducting a gradient step. (right) Longhorn’s dynamics is derived from
solving an online regression problem, where the goal is to recover
$x\in\mathbb{R}$ (if $x\in\mathbb{R}^{d}$ we can run $d$ SSMs in parallel)
based on a hint $k\in\mathbb{R}^{m}$ from a state vector $s\in\mathbb{R}^{d}$.
The closed-form solution to the online regression problem corresponds to the
implicit online learning update, where
$\Delta_{t}=\beta_{t}/(1+\beta_{t}k_{t}^{\top}k_{t})$. Please see the details
in Section 3.

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 hope is that by selecting the right objective,
one can save parameters/computation without sacrificing performance. In
addition, the proposed online learning perspective might lead to a better
understanding of what SSM layers are doing in a large model. Specifically, we
view the recurrent update (a.k.a., the state-transition dynamics) of an SSM
model as either a gradient descent step or a closed-form solution to an online
learning objective. We provide the objectives for some existing SSMs in Table
1.

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. So 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, 2023) on both
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., 2023). To summarize, our
contributions are:

1) Theoretical Framework: We propose a novel framework that views SSM’s
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 enhancement in
sampling 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.

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

1: Parameters: $W_{q}\in\mathbb{R}^{m\times d},W_{k}\in\mathbb{R}^{m\times
d},W_{\beta}\in\mathbb{R}^{d\times d}$, where $W_{\beta}$ can be low-rank,
horizon $T$.

2: Initialize the memory state $S_{0}\leftarrow 0^{d\times m}$.

3: for $t\in\\{1,\dots,T\\}$ do

4: 1) Receive input $x_{t}\in\mathbb{R}^{d}$.

5: 2) Compute the query $q_{t}$, key $k_{t}$ and $\beta_{t}$ (as in objective
$||s-s_{t-1}||^{2}+\beta_{t}||s^{\top}k_{t}-x_{t}||^{2}$):  $\displaystyle
q_{t}=W_{q}x_{t}\in\mathbb{R}^{m},$ $\displaystyle
k_{t}=W_{k}x_{t}\in\mathbb{R}^{m},$
$\displaystyle\beta_{t}=\text{Sigmoid}(W_{\beta}x_{t})\in(0,1)^{d}.$

6: 3) Update the memory state $S_{t}\in\mathbb{R}^{d\times m}$ via
$\displaystyle S_{t}=\big{(}1-\Delta_{t}\otimes k_{t}^{\odot 2}\big{)}\odot
S_{t-1}+\big{(}\Delta_{t}\odot x_{t}\big{)}\otimes k_{t},$ where $\Delta_{t}$
is the step size:
$\Delta_{t}=\beta_{t}/(1+\beta_{t}k_{t}^{\top}k_{t})\in(0,1)^{d}.$

7: 4) Compute the output $o_{t}=S_{t}q_{t}\in\mathbb{R}^{d}$.

8: end for

9: Note: $\odot$ elementwise product and $\otimes$ is outer product. $x_{t}$
in practice is preprocessed through a linear projection followed by a Conv1d
operation as in Mamba (Gu & Dao, 2023).

##  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., 2023; Yang et al., 2023; De et
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 is done by the channel-wise
MLP. Taking Mamba as an example (Gu & Dao, 2023), a Mamba model consists of a
stack of homogeneous modules named Mamba block (the $\Phi(\bm{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 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
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., 2017):

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

Figure 3: Mamba Block

A key feature of this design in Equation 1 is that the update of $S_{t}$ is a
_linear_ recurrence relation. That is, $S_{t}$ is a linear function of
$S_{t-1}$. 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 (Harris et
al., 2007), which only requires a sublinear complexity in terms of sequence
length (e.g., $\mathcal{O}(\log t)$). From now on, we will abbreviate
$A(x_{t})$ and $B(x_{t})$ as $A_{t}$ and $B_{t}$, respectively.

#### Designs of ($A_{t}$, $B_{t}$, $*$)

Existing variants of SSMs mainly differs 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 be $m$ times of 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 lay 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 what differentiates the existing designs of SSM. In the following, we
provide some examples in the form of Equation 1 for some existing SSM models.

######  Example 2.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 $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 2.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(\Delta(x_{t})\otimes 1)),\qquad B_{t}=(\Delta(x_{t})\odot
x_{t})\otimes k(x_{t}).\end{split}$ (3)

where $S_{t}\in\mathbb{R}^{d\times m}$ with $m=16$ by default,
$\Delta(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 $\Delta(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 2.3 (Griffin (De et al., 2024)).

In Mamba and the linear attention variants, the outer product serves as a
critical rule of 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 2.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., 2024), 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 2.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.

##  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 (the algorithm) 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}).$ (4)

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.

Online convex optimization (OCP) yields a principled approach to solving
Equation 4 when $\ell_{t}$ are convex, by trading-off the “stability” and
“plasiticity” (Mermillod et al., 2013). 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}},$
(5)

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 SSM as Online Learning and The Longhorn Architecture

Under the online learning framework, the design of SSM reduces to the design
of $D_{\phi}$ and $\ell_{t}$ in Equation 5. This provides a unified framework
for the existing SSM variants. We summarize in Table 1 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 1). In
particular, we consider $\\{x_{t}\\}_{t}$ as the input stream, where
$x_{t}\in\mathbb{R}$ is a scalar (as in Mamba, one can consider there are $d$
parallel SSMs if $x_{t}\in\mathbb{R}^{d}$ is a vector). We want to recurrently
update hidden states $\\{s_{t}\\}_{t}$, where $s_{t}\in\mathbb{R}^{m}$ is a
vector that aims to summarize 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}^{m}}||s-s_{t-1}||^{2}+\beta_{t}(x_{t})||s^{\top}k_{t}(x_{t})-x_{t}||^{2}.$
(6)

Here, $k_{t}(x_{t})\in\mathbb{R}^{m}$ is the linear transformation of the
input token $x_{t}\in\mathbb{R}$, which represents the key vector in a
standard transformer. $\beta_{t}(x_{t})$ controls how much new information
about $x_{t}$ we want the model to incorporate. For instance,
$\beta_{t}(x_{t})=0$ implies $s_{t}=s_{t-1}$, while a large $\beta_{t}(x_{t})$
implies the model sacrifices stability for plasticity. For simplicity, we will
write $k_{t}=k_{t}(x_{t})$ and $\beta_{t}=\beta_{t}(x_{t})$ from now on.

From a high-level, Equation 6 is solving an online prediction problem of
learning a weight $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}$.

Fortunately, this simple objective gives a closed-form solution for $s_{t}$,
which coincides with the implicit online learning method (Kulis & Bartlett,
2010).

######  Theorem 3.1.

The closed form solution for $s_{t}$ for objective in Equation 6 is

$s_{t}=(I-\Delta_{t}k_{t}k_{t}^{\top})s_{t-1}+\Delta_{t}k_{t}x_{t},~{}~{}~{}\text{where}~{}\Delta_{t}=\frac{\beta_{t}}{1+\beta_{t}k_{t}^{\top}k_{t}}\in[0,\infty).$
(7)

###### Proof.

As the objective in Equation 6 is in a quadratic form with respect to $s$,
there is a unique minimum:

$\displaystyle\nabla L_{t}(s)=0$
$\displaystyle\Longleftrightarrow(s-s_{t-1})+\beta_{t}(s^{\top}k_{t}-x_{t})k_{t}=0$
$\displaystyle\Longleftrightarrow(I+\beta_{t}k_{t}k_{t}^{\top})s=s_{t-1}+\beta_{t}k_{t}x_{t}$
$\displaystyle\underbrace{\Longleftrightarrow}_{(3)}s=\bigg{(}I-\frac{\beta_{t}}{1+\beta_{t}k_{t}^{\top}k_{t}}k_{t}k_{t}^{\top}\bigg{)}s_{t-1}+\bigg{(}I-\frac{\beta_{t}}{I+\beta_{t}k_{t}^{\top}k_{t}}k_{t}k_{t}^{\top}\bigg{)}\beta_{t}k_{t}x_{t}$
$\displaystyle\Longleftrightarrow\bigg{(}I-\frac{\beta_{t}}{I+\beta_{t}k_{t}^{\top}k_{t}}k_{t}k_{t}^{\top}\bigg{)}s_{t-1}+\frac{(I+\beta_{t}k_{t}^{\top}k_{t}-\beta_{t}k_{t}k_{t}^{\top})\beta_{t}k_{t}x_{t}}{I+\beta_{t}k_{t}^{\top}k_{t}}$
$\displaystyle\underbrace{\Longleftrightarrow}_{(5)}\bigg{(}I-\frac{\beta_{t}}{I+\beta_{t}k_{t}^{\top}k_{t}}k_{t}k_{t}^{\top}\bigg{)}s_{t-1}+\frac{\beta_{t}k_{t}x_{t}}{I+\beta_{t}k_{t}^{\top}k_{t}}$

(3) is derived from the fact that
$(I+\beta_{t}k_{t}k_{t}^{\top})^{-1}=(I-\frac{\beta_{t}k_{t}k_{t}^{\top}}{I+\beta
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}-k_{t}k_{t}^{\top}k_{t}x_{t}=0$. ∎

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}-\Delta_{t}k_{t}^{\odot 2}$ in place of
$I-\Delta_{t}k_{t}k_{t}^{\top}$, where $x^{\odot 2}=x\odot x$ and $1_{m}$ is
the $m$-dimensional vector of all ones. The final SSM recurrence therefore
becomes:

$s_{t}=\underbrace{(1-\Delta_{t}k_{t}^{\odot 2})}_{A_{t}(x_{t})}\odot
s_{t-1}+\underbrace{\Delta_{t}k_{t}x_{t}}_{B_{t}(x_{t})}.$ (8)

Following the Mamba design, when $x_{t}\in\mathbb{R}^{d}$ is a vector, we use
$d$ parallel SSMs. As a result,
$x_{t},\beta_{t},\Delta_{t}\in\mathbb{R}^{d},k_{t}\in\mathbb{R}^{m},S_{t}\in\mathbb{R}^{d\times
m}$, Longhorn’s update becomes:

$S_{t}=A_{t}\odot
S_{t-1}+B_{t},~{}~{}\text{where}~{}~{}A_{t}=(1-\Delta_{t}\otimes k_{t}^{\odot
2}),~{}~{}B_{t}=(\Delta_{t}\odot x_{t})\otimes k_{t}.$ (9)

The final architecture of Longhorn follows Mamba strictly (Figure 3), and we
only replace the SSM block with the 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 9 to Equation 3 and other SSM in Section
2. Longhorn does not introduce an extra “forgetting” gate (hence it has fewer
parameters), because the forget gate is naturally derived from the key vector,
i.e., $(1-k^{\odot 2}\otimes\Delta_{t})$.

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 3, Table 4). 3. The closed-form solution in
Equation 7 does not need any specific initialization. In contrast, Mamba
requires careful initialization of the $A$ and $\Delta_{t}$. Compared to
DeltaNet (Schlag et al., 2021; Yang et al., 2024), we restrict the state-
transition matrix to be diagonal, hence Longhorn can leverage both parallel
scan (as in Mamba) and the chunk-wise parallelization (as in GLA). Therefore,
it is as fast as existing SSMs. 4. Unlike DeltaNet, which, as noted in Yang et
al. (2024), cannot extrapolate beyond the training context, Longhorn
demonstrates the capability to successfully extrapolate to contexts up to 16x
longer at inference time (See Figure 1 (right)).

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}-2\langle s^{\top}k_{t},x_{t}\rangle$ | $s_{t}=s_{t-1}+k_{t}x_{t}$
RetNet | $\gamma\left\lVert s-s_{t-1}\right\rVert^{2}+(1-\gamma)\left\lVert s\right\rVert^{2}-2\langle s^{\top}k_{t},x_{t}\rangle$ | $s_{t}=\gamma s_{t-1}+k_{t}x_{t}$
GLA | $(s-s_{t-1})^{\top}\mathrm{diag}(\alpha_{t})(s-s_{t-1})+s^{\top}\mathrm{diag}(1-\alpha_{t})s-2\langle s^{\top}k_{t},x_{t}\rangle$ | $s_{t}=\mathrm{diag}(\alpha_{t})s_{t-1}+k_{t}x_{t}$
Griffin | $\left\lVert\sqrt{\alpha_{t}}\odot(s-s_{t-1})\right\rVert^{2}+\left\lVert\sqrt{1-\alpha_{t}}\odot s\right\rVert^{2}-2\sqrt{1-\alpha_{t}}\odot s\odot i_{t}\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}+\beta_{t}\left\lVert s^{\top}k_{t}-x_{t}\right\rVert^{2}$ | $s_{t}=(1-\Delta_{t}k_{t}^{\odot 2})s_{t-1}+\Delta_{t}k_{t}x_{t}$
$\Delta_{t}=\beta_{t}/(1+\beta_{t}k_{t}^{\top}k_{t})$

Table 1: Some of the existing SSMs and their corresponding online learning
objectives/updates.

##  4 Related Work

This section provides a brief summary of recent advances in linear attention
(LA) and deep state space models (SSMs).

#### Linear Attention Models

Several methods have been developed to address the quadratic complexity of the
Transformer by making attention linear with respect to context length. In
particular, Linformer uses a low-rank approximation of the self-attention by
projecting the keys and values into a constant size matrix (instead of scaling
with the sequence length)(Wang et al., 2020). Realizing that the main
bottleneck causing the quadratic cost in Transformers is the Softmax function,
the Linear Transformer replaces the Softmax function with a decomposable
similarity function analogous to kernel methods, thereby making the
computation linear with respect to sequence length(Katharopoulos et al.,
2020). Performer approximates the softmax attention using positive orthogonal
random features (Choromanski et al., 2020). More recently, based on the Linear
Transformer, the Retentive Network (RetNet) adds additional constant
forgetting and rotation (Sun et al., 2023). Gated Linear Attention (GLA)
further experiments with learnable forget gates (Yang et al., 2023). Notably,
linear attention can be viewed as a fast weight network where a slow net (the
model’s parameters) learns to program a changing fast network (e.g., a linear
predictor) by adapting its parameters (e.g., the $s_{t}$) online using inputs
(Schlag et al., 2021).

#### State Space Models

Instead of trying to linearize transformers, state space models (SSMs) start
with parallelizable linear recurrent networks directly. Initially, the state
transition matrix $A$ is assumed to be constant so that the recurrence can be
computed in parallel using a convolution (Li et al., 2022; Gu et al., 2021).
Subsequent developments include the Diagonal State Space (DSS) model (Gupta et
al., 2022), Gated State Space (GSS) models (Mehta et al., 2022), S5 model
(Smith et al., 2022), Bidirectional Gated SSM (BiGS)(Wang et al., 2022), H3
model(Fu et al., 2022), and Mamba (Gu & Dao, 2023). In addition, there are
also works directly trying to make recurrent networks efficient, which often
results in a particular form of SSM as well. This includes Deep Linear
Recurrent Units (LRUs)(Orvieto et al., 2023; De et al., 2024), Hierarchically
Gated Linear RNN (HGRN)(Qin et al., 2024b; a), and RWKV (Peng et al., 2023;
2024).

#### Fast Weight Programmer

The concept of a network modifying its own weights in response to inputs is
not novel and has historical roots in the Fast-weight Programmer (Schmidhuber,
1992; 1993; Schlag & Schmidhuber, 2017; Schlag et al., 2021). Specifically,
these models propose updating a weight matrix $W\in\mathbb{R}^{d\times m}$
using the outer product of two vectors, expressed as $\Delta W=v(x_{t})\otimes
k(x_{t})$. This mechanism closely aligns with the principles underlying the
Linear Attention model. Our framework builds upon the Fast Weight concept by
tailoring the weight update process to address a specific online learning
objective, thereby extending its applicability and effectiveness in dynamic
learning environments.

#### Concurrent Work

There are two concurrent works that share similar ideas to ours. Yang et al.
(2024) propose a chunk-wise parallel form for scaling the DeltaNet (Schlag et
al., 2021) to large-scale language modeling tasks (Yang et al., 2024). The
update rule in DeltaNet can be viewed as conducting the gradient step instead
of solving for the closed form of the online regression objective, hence the
state transition matrix is $A(x_{t})=(I-\beta_{t}k_{t}k_{t}^{\top})$, which is
not guaranteed to have eigenvalues $\leq 1$ (the forgetting can even be
negative). Therefore, in practice Yang et al. (2024) normalize the key vector
$k_{t}$ by its $\ell_{2}$ norm (Schlag et al., 2021), which can be
restrictive. In contrast, the closed-form update in Longhorn ensures that the
state transition matrix is always stable. In addition, $A(x_{t})$ is a matrix
in DeltaNet while we apply a diagonal approximation ($k_{t}^{\odot 2}$).
Therefore, Longhorn can leverage either parallel scan (as in Mamba) and chunk-
wise parallel training (as in GLA), making it as fast as any existing SSM. In
practice, we provide a parallel scan CUDA kernel, so Longhorn serves as a
drop-in replacement for Mamba. On the other hand, Sun et al. (2024) propose
the Test-Time Training framework, where the update of the state is derived
from applying one gradient step of an online regression objective (Sun et al.,
2024). To make the model parallelizable, the authors assume that each gradient
step at $x_{t}$ ($t\geq 0$) is conducted at state $s_{0}$, hence deriving a
parallel form that utilizes matrix multiplication. In contrast, Longhorn
solves for the closed-form solution on _every_ token, which potentially leads
to a more flexible model.

##  5 Experiments

We conduct experiments on both synthetic benchmarks and the language modeling
task to validate the performance of Longhorn:

  1. 1.

To test Longhorn’s performance on memorizing associations, we compare it
against other SSMs on the synthetic multi-query associative recall benchmark
(Arora et al., 2023). We find that Longhorn is the only model that achieves
near perfect recall rate when the sequence length scales up to 512 with a
hidden dimension of 64).

  2. 2.

To see how Longhorn performs on language modeling tasks, we first conduct the
scaling law experiments using the OpenWebText dataset (Gokaslan & Cohen,
2019), with model sizes of 120M and 350M and context length 1024 or 4096, and
then compare the validation perplexity of Longhorn against that of the other
SSM models. We find that Longhorn consistently outperforms baseline deep SSM
models across context lengths and model sizes.

  3. 3.

In the end, to see how Longhorn performs on downstream language evaluation
tasks, we train a 1.3B language model on the SlimPajama dataset (Soboleva et
al., 2023) with 100B tokens, and then compare its downstream evaluation
performance across 5 standard benchmarks. We find that Longhorn achieves 1.8x
speedup sampling efficiency compared to Mamba, and also outperforms GLA and
Mamba on average in downstream evaluation tasks.

###  5.1 Multi-Query Associative Recall

We first consider the synthetic benchmark Multi-Query Associative Recal
(MQAR)l (Arora et al., 2023). The task is the agent will see 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 for it to
“retrieve” the corresponding values, hence the name multi-query associative
recall. 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., 2024), 3) Hyena (Poli et al., 2023), which is a
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 state-transition matrix
is a scalar, 5) BaseConv (Arora et al., 2023), an SSM that combines linear
projection with convolution, and 6) Mamba (Gu & Dao, 2023), the state-of-the-
art SSM that has data-dependent $A$ and $B$ (See Equation 3). Each experiment
individually searches for the best learning rate from
$\\{1e-4,4.6e-4,2.2e-3,1e-02\\}$ (evenly separated in the log space). Results
are summarized in Figure 4.

Figure 4: Comparison of Longhorn against state-of-the-art SSM variants on the
MQAR benchmark. The y-axis is the recall rate.

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.222We adapt the code from the nanoGPT repo, a minimal reproduction
of the GPT-2 model. See https://github.com/karpathy/nanoGPT/tree/master.
Following the Mamba (Gu & Dao, 2023) experiment setting, the training details
are provided in Table 2:

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
Table 2: Training details on OpenWebText.

We consider the following baseline models: LLaMA (Touvron et 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$, which model sizes around 120M or 350M. Results are summarized in
Figure 5 and Table 3.

Figure 5: Scaling law with 1024 and 4096 context length on OpenWebText with various SSM models and the LLaMA (strong Transformer) baseline.  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
Table 3: Language modeling scaling law against LLaMA (Touvron et al., 2023),
RetNet (Sun et al., 2023), RWKV (Peng et al., 2023), and Mamba (Gu & Dao,
2023). All models are trained on the OpenWebText dataset (Gokaslan & Cohen,
2019). Models vary from around 120M and 350M parameters, with context length
1024 or 4096.

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 experiment setting
in GLA (Yang et al., 2023), where a model with around 1.3B parameters is
trained on the SlimPajama (Soboleva et al., 2023) dataset over 100B tokens
with a batch size of 2M tokens. We use the AdamW optimizer (Loshchilov &
Hutter, 2017) with weight decay of 0.01, and apply cosine learning rate decay
with a peak learning rate of $3e-4$ and the final learning rate of $3e-5$. In
addition, we use gradient clipping of 1.0. Due to the large size of the
experiment, we only compare against LLaMA, Mamba and GLA models, all with a
context size of 2048.333Our result is in general better than those reported in
the GLA paper, this might due to the fact that we are using a different 100B
subset of SlimPajama since the original split of the SlimPajama is not
provided.

We evaluate the trained models on a wide range of downstream tasks covering
both common-sense reasoning and question-answering. In particular, we consider
PIQA (Bisk et al., 2020), HellaSwag (Hella) (Zellers et al., 2019), WinoGrande
(Wino) (Sakaguchi et al., 2021), ARC-easy (ARC-e) and ARC-challenge (ARC-c)
(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 4.

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 | 1024K | 55.08 | 55.36 | 71.73 | 59.26 | 32.19 | 43.35 | 45.16 | 62.13 | 53.03
GLA | 1024K | 55.55 | 49.10 | 71.12 | 58.86 | 28.11 | 41.67 | 44.91 | 59.21 | 51.07
Mamba | 32K | 54.21 | 53.61 | 71.67 | 61.05 | 30.15 | 43.94 | 44.18 | 59.22 | 52.25
Longhorn | 32K | 55.78 | 52.30 | 71.00 | 60.63 | 29.53 | 43.55 | 44.68 | 61.29 | 52.35

Table 4: Language modeling results against LLaMA (Touvron et al., 2023),
RetNet (Sun et al., 2023), and Mamba (Gu & Dao, 2023). 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.

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 4, 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

Moreover, 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., 2024), which highlights a
limitation in that the model cannot extrapolate to longer contexts. Moreover,
it is expected that LLaMA, which is a Transformer-based model, fails to
extrapolates beyond its training context length.

##  6 Conclusion and Future Work

In this work, we introduce a novel approach to designing deep state-space
models (SSMs) by conceptualizing the recurrence update as solving an online
objective. We propose a straightforward online regression objective, adopting
its implicit closed-form update to define our model, which we refer to as
Longhorn. Notably, Longhorn is highly parallelizable and achieves state-of-
the-art performance in both synthetic sequence modeling and language modeling
tasks. For future research, an intriguing avenue would be exploring other
online learning objectives that can be efficiently implemented on modern
hardware. Additionally, while the current implementation of Longhorn closely
aligns with Mamba, recent studies (Ren et al., 2024) suggest that
incorporating sliding-window attention with Mamba improves performance. We
anticipate similar benefits for Longhorn.

## References

  * Arora et al. (2023) Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré.  Zoology: Measuring and improving recall in efficient language models.  _arXiv preprint arXiv:2312.04927_ , 2023.
  * Arora et al. (2024) 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 the recall-throughput tradeoff.  _arXiv preprint arXiv:2402.18668_ , 2024.
  * Beck et al. (2024) Maximilian Beck, Korbinian Poppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael K Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter.  xlstm: Extended long short-term memory.  2024\.  URL https://api.semanticscholar.org/CorpusID:269614336.
  * Bisk et al. (2020) Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al.  Piqa: Reasoning about physical commonsense in natural language.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 34, pp. 7432–7439, 2020.
  * Chen et al. (2023) Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper.  Accelerating large language model decoding with speculative sampling.  _arXiv preprint arXiv:2302.01318_ , 2023.
  * Choromanski et al. (2020) 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.
  * Clark et al. (2019) Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova.  Boolq: Exploring the surprising difficulty of natural yes/no questions.  _arXiv preprint arXiv:1905.10044_ , 2019.
  * Clark et al. (2018) 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.05457_ , 2018.
  * Dao et al. (2022) Tri Dao, Daniel Y. 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.semanticscholar.org/CorpusID:249151871.
  * Dauphin et al. (2017) Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier.  Language modeling with gated convolutional networks.  In _International conference on machine learning_ , pp. 933–941. PMLR, 2017.
  * De et al. (2024) Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al.  Griffin: Mixing gated linear recurrences with local attention for efficient language models.  _arXiv preprint arXiv:2402.19427_ , 2024.
  * DeepSeek-AI & Dai (2024) DeepSeek-AI and Damai Dai.  Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.  2024\.  URL https://api.semanticscholar.org/CorpusID:269613809.
  * Fu et al. (2022) 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 arXiv:2212.14052_ , 2022.
  * Gokaslan & Cohen (2019) Aaron Gokaslan and Vanya Cohen.  Openwebtext corpus.  http://Skylion007.github.io/OpenWebTextCorpus, 2019.
  * Gu & Dao (2023) Albert Gu and Tri Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  _arXiv preprint arXiv:2312.00752_ , 2023.
  * Gu et al. (2021) Albert Gu, Karan Goel, and Christopher Ré.  Efficiently modeling long sequences with structured state spaces.  _arXiv preprint arXiv:2111.00396_ , 2021.
  * Gupta et al. (2022) Ankit Gupta, Albert Gu, and Jonathan Berant.  Diagonal state spaces are as effective as structured state spaces.  _Advances in Neural Information Processing Systems_ , 35:22982–22994, 2022.
  * Harris et al. (2007) Mark Harris, Shubhabrata Sengupta, and John D Owens.  Parallel prefix sum (scan) with cuda.  _GPU gems_ , 3(39):851–876, 2007.
  * Katharopoulos et al. (2020) Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret.  Transformers are RNNs: Fast autoregressive transformers with linear attention.  In _International Conference on Machine Learning_ , pp. 5156–5165. PMLR, 2020.
  * Kulis & Bartlett (2010) Brian Kulis and Peter L Bartlett.  Implicit online learning.  In _Proceedings of the 27th International Conference on Machine Learning (ICML-10)_ , pp. 575–582, 2010.
  * Kuperman & Dyke (2011) Victor Kuperman and Julie A. Van Dyke.  Individual differences in visual comprehension of morphological complexity.  _Cognitive Science_ , 33, 2011.  URL https://api.semanticscholar.org/CorpusID:5555496.
  * Li et al. (2022) Yuhong Li, Tianle Cai, Yi Zhang, Deming Chen, and Debadeepta Dey.  What makes convolutional models great on long sequence modeling?  _arXiv preprint arXiv:2210.09298_ , 2022.
  * Loshchilov & Hutter (2017) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * Mehta et al. (2022) Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur.  Long range language modeling via gated state spaces.  _arXiv preprint arXiv:2206.13947_ , 2022.
  * Mermillod et al. (2013) Martial Mermillod, Aurélia Bugaiska, and Patrick Bonin.  The stability-plasticity dilemma: Investigating the continuum from catastrophic forgetting to age-limited learning effects, 2013.
  * Mihaylov et al. (2018) 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:1809.02789_ , 2018.
  * Orvieto et al. (2023) 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. 26670–26698. PMLR, 2023.
  * Peng et al. (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.13048_ , 2023.
  * Peng et al. (2024) Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemysław Kazienko, et al.  Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence.  _arXiv preprint arXiv:2404.05892_ , 2024.
  * Poli et al. (2023) 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.
  * Qin et al. (2024a) Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong.  Hgrn2: Gated linear rnns with state expansion.  _arXiv preprint arXiv:2404.07904_ , 2024a.
  * Qin et al. (2024b) Zhen Qin, Songlin Yang, and Yiran Zhong.  Hierarchically gated recurrent neural network for sequence modeling.  _Advances in Neural Information Processing Systems_ , 36, 2024b.
  * Ren et al. (2024) 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.
  * Sakaguchi et al. (2021) Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi.  Winogrande: An adversarial winograd schema challenge at scale.  _Communications of the ACM_ , 64(9):99–106, 2021.
  * Sap et al. (2019) Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi.  Socialiqa: Commonsense reasoning about social interactions.  _arXiv preprint arXiv:1904.09728_ , 2019.
  * Schlag & Schmidhuber (2017) Imanol Schlag and Jürgen Schmidhuber.  Gated fast weights for on-the-fly neural program generation.  In _NIPS Metalearning Workshop_ , 2017.
  * Schlag et al. (2021) 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.
  * Schmidhuber (1992) Jürgen Schmidhuber.  Learning to control fast-weight memories: An alternative to dynamic recurrent networks.  _Neural Computation_ , 4(1):131–139, 1992.
  * Schmidhuber (1993) 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 on Artificial Neural Networks Amsterdam, The Netherlands 13–16 September 1993 3_ , pp. 460–463. Springer, 1993.
  * Smith et al. (2022) Jimmy TH Smith, Andrew Warrington, and Scott W Linderman.  Simplified state space layers for sequence modeling.  _arXiv preprint arXiv:2208.04933_ , 2022.
  * Soboleva et al. (2023) Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey.  SlimPajama: A 627B token cleaned and deduplicated version of RedPajama.  https://www.cerebras.net/blog/slimpajama-a-627b-token-cleaned-and-deduplicated-version-of-redpajama, 2023.  URL https://huggingface.co/datasets/cerebras/SlimPajama-627B.
  * Sun et al. (2024) 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 (learn at test time): Rnns with expressive hidden states.  2024\.  URL https://api.semanticscholar.org/CorpusID:271039606.
  * Sun et al. (2023) 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.
  * Touvron et al. (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.13971_ , 2023.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wang et al. (2022) Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush.  Pretraining without attention.  _arXiv preprint arXiv:2212.10544_ , 2022.
  * Wang et al. (2020) Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma.  Linformer: Self-attention with linear complexity.  _arXiv preprint arXiv:2006.04768_ , 2020.
  * Yang et al. (2023) Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim.  Gated linear attention transformers with hardware-efficient training.  _arXiv preprint arXiv:2312.06635_ , 2023.
  * Yang et al. (2024) 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.06484_ , 2024.
  * Zellers et al. (2019) Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi.  Hellaswag: Can a machine really finish your sentence?  _arXiv preprint arXiv:1905.07830_ , 2019.