Longhorn: State Space Models are Amortized Online Learners

1 INTRODUCTION

• Transformer utility is constrained by quadratic training and decoding costs as sequence length grows, making very long context difficult to scale.The introduction notes this remains challenging despite several efficiency optimizations.Source: §1 · p. 1
• The paper proposes viewing sequence-mixing layers as meta-modules that compress history into a state used by later layers.This perspective treats the recurrent state as an online memory.Source: §1 · p. 2
• SSM recurrence can be viewed as solving an online-learning problem, with updates derived as proximal steps or closed-form solutions to online objectives.This frames SSM design around objectives for specific online prediction problems.Source: §1 · p. 2
• Longhorn is derived from the implicit closed-form update for online associative recall and does not require a separately parameterized forget gate.The update naturally balances forgetting and learning, saving parameters when the state size is large.Source: §1 · p. 2
• 1.8x improvement in sample efficiency compared to Mamba; Longhorn also outperforms existing SSMs on synthetic associative-recall and large-scale language-modeling tasks.Longhorn trained with 2048 context length extrapolates to 32K at inference without much perplexity drop.Source: §1 · pp. 2–3

2 BACKGROUND

• Deep SSMs replace self-attention layers with SSM sequence mixers, while channel mixing is handled by MLPs or fused modules.Mamba combines sequence and channel mixing in homogeneous Mamba blocks.Source: §2 · p. 3
• An SSM updates its memory state through a recurrence in which input-dependent functions transform the previous state and add a new contribution.The recurrence begins from an all-zero state, and its operator can be Hadamard or matrix multiplication.Source: §2 · p. 3
• Because the state recurrence is linear in the previous state, all states can be expressed explicitly and computed in parallel during training.The cumulative product can use a prefix scan requiring O(log T) parallel operations.Source: §2 · p. 4
• The state must preserve substantial information, but naive mappings for its input-dependent update functions require prohibitively many weights.This motivates low-dimensional structure in the update functions, a key difference among SSM designs.Source: §2 · p. 4

3 AN ONLINE LEARNING PERSPECTIVE FOR SEQUENCE MIXING

• The framework treats SSM layers as online learning modules that compress sequence information during inference, and frames SSM design as choosing an objective.This perspective views recurrent updates as solutions to online learning problems.Source: §3 · p. 4, §3.2 · p. 5

3.1 SSM AS ONLINE LEARNING

• Online regression updates predictions as new input-label pairs arrive, while online convex programming balances stability against incorporating each new loss.The regularization term keeps the state near its previous value; β_t controls the stability–plasticity trade-off.Source: §3.1 · pp. 4–5

3.2 THE LONGHORN ARCHITECTURE

• Longhorn learns a state matrix to predict each value from its key, treating associative recall as online regression motivated by attention’s recall behavior and its connection to in-context learning.Its closed-form update incorporates new information according to β_t, with zero β_t,i leaving the corresponding state row unchanged.Source: §3.2 · p. 5
• Longhorn performs better than Mamba even with fewer parameters and extrapolates to contexts 16x longer than its training contexts.The section also reports that Longhorn extrapolates beyond training contexts where DeltaNet struggles.Source: §3.2 · p. 6
• The closed-form update derives forgetting from the key vector, eliminating a separate forget gate and saving about O(d × m) parameters per module; it also needs no specific initialization.The section identifies the missing gate as a parameter-saving difference from Mamba and GLA.Source: §3.2 · pp. 5–6

Linear Attention Models

• Linear attention methods reduce Transformer attention’s quadratic context cost through key-value projection, decomposable similarity, random-feature approximation, or forgetting and rotation.Gated Linear Attention uses learnable forget gates (Yang et al., 2023), while linear attention is also framed as online fast-weight adaptation (Schlag et al., 2021).Source: §4 · p. 6

State Space Models

• SSMs are parallelizable linear recurrent networks, spanning constant-transition convolutional models and later architectures including Mamba.The survey also lists efficient recurrent alternatives such as LRUs, HGRN, and RWKV.Source: §4 · p. 6

Fast Weight Programmer

• Fast-weight programmers update a weight matrix by an input-dependent outer product, and this framework adapts weight updates to specific online-learning objectives.Source: §4 · p. 6

Concurrent Work

• Concurrent DeltaNet work uses a gradient step for online regression and key normalization to address possible instability, whereas Longhorn uses a stable closed-form update and supports parallel scan and chunk-wise training.Concurrent test-time training instead applies each gradient step from the initial state; Longhorn computes a closed-form solution at every token.Source: §4 · p. 6

5 EXPERIMENTS

• The experiments compare Longhorn with SSMs and other baselines across recall, language modeling, and vision tasks.The large-scale language-modeling setup trains a 1.3B model on SlimPajama with 100B tokens and evaluates eight downstream tasks.Source: §5 · p. 7, §5.3 · p. 8, §5.5 · p. 9

5.1 MULTI-QUERY ASSOCIATIVE RECALL

• Longhorn is the only compared model reported to achieve near-perfect MQAR recall at sequence lengths up to 512 with hidden dimension 64.The benchmark varies sequence length and model dimension; Figure 4 plots recall rate.Source: §5 · p. 7, §5.1 · p. 7, Figure 4 · p. 7

5.3 LARGE-SCALE LANGUAGE MODELING

• 1.8x better sample efficiency than Mamba and GLA, with better final performance, is reported for the 1.3B SlimPajama language model.The model was trained on 100B tokens and evaluated across eight benchmarks.Source: §5 · p. 7

5.2 SCALING LAW ON OPENWEBTEXT

• In the reported comparisons, Longhorn consistently outperforms baseline SSMs, with experiments covering models up to 350M parameters and context lengths up to 4096.Experiments used models around 120M or 350M parameters and context lengths of 1024 or 4096.Source: §5.2 · p. 8, Table 1 · p. 8

5.4 ABLATION ON LENGTH EXTRAPOLATION

• After training with context length 2048, Longhorn successfully extrapolates at inference to contexts up to 16x longer; pretrained-model perplexity was evaluated on disjoint SlimPajama validation sequences of lengths 2048, 4096, 8192, 16384, and 32768.Source: §5.4 · p. 9

5.5 VISION STATE SPACE MODELS

• Vision Longhorn achieves comparable, slightly better ImageNet classification performance than Vision Mamba.ViL uses a bidirectional scan and is compared with ViM on TINY and SMALL configurations.Source: §5.5 · p. 9, Table 3 · p. 9

Conclusion

• Longhorn is reported to achieve state-of-the-art performance among SSMs on MQAR, language modeling, and image classification.The conclusion describes its recurrence as explicitly solving online associative recall in closed form.Source: §6 · p. 9
• The paper frames SSM design as choosing an online objective that the recurrence update solves.Longhorn is the paper’s example, based on online associative recall.Source: §6 · p. 9

Future Work

• Future directions include exploring other online learning objectives and adding sliding-window attention, which the authors anticipate could similarly benefit Longhorn.Source: §6 · p. 9

Linear Attention Variants

• Linear Attention, RetNet, and GLA assume rank-1 or constant outputs for their A_t and B_t terms.The appendix also describes a linear-attention recurrence using an outer product to update its matrix state.Source: A PRIOR DEEP STATE SPACE MODELS · p. 13

Mamba

• Mamba derives its update by discretizing continuous linear dynamics, with both A_t and B_t depending on the input-conditioned step size.Its state update combines A_t times the previous state with B_t.Source: A PRIOR DEEP STATE SPACE MODELS · p. 13

Griffin

• Griffin abandons the outer product used by Mamba and linear-attention variants in favor of elementwise recurrence.The supplied passage also notes Griffin is combined with local attention in practice.Source: A PRIOR DEEP STATE SPACE MODELS · p. 13

RWKV

• RWKV maintains a ratio state formed from separately updated numerator and denominator SSMs, with a trainable weight controlling forgetting.A later RWKV version removes the denominator and replaces the elementwise product with an outer product.Source: A PRIOR DEEP STATE SPACE MODELS · p. 13

HGRN2

• HGRN2 uses a forget gate and complementary input gate in an elementwise recurrence, making it resemble an RNN.Source: A PRIOR DEEP STATE SPACE MODELS · p. 13

• Table 4 collects the paper’s reverse-engineering of existing SSMs’ online learning objectives.Source: B EXISTING STATE SPACE MODELS’ ONLINE… · p. 14

C PROOF

• Theorem C.1 gives a closed-form solution for S_t under the objective in Equation 5.Source: C PROOF · p. 14
• The objective is quadratic in S, so it has a unique minimum.Source: C PROOF · p. 14
• Because each row of S optimizes independently, the proof solves row-wise by setting its derivative to zero.Source: C PROOF · p. 14

D ADDITIONAL EXPERIMENT DETAILS

• The architecture configurations follow exactly from the Mamba paper.Source: D ADDITIONAL EXPERIMENT DETAILS · p. 14
• The architecture details are for conducting scaling-law experiments on OpenWebText.Source: D ADDITIONAL EXPERIMENT DETAILS · p. 14
• Table 5 reports training details for OpenWebText.Source: Table 5 · p. 14