PaperlayerView exact revision on arXiv

Source-linked AI summary

Raven: High-Recall Sequence Modeling with Sparse Memory Routing

Arshia Afzal, Aviv Bick, Eric P. Xing, Volkan Cevher, Albert Gu

arXiv:2607.25357v1cs.LGcs.AI

TL;DR

Dense state updates can interfere with long-range recall, while sliding-window attention loses tokens when they leave its fixed window. Raven uses sparse, input-dependent routing to update and decay selected memory slots. It performs strongly on recall-intensive benchmarks and extrapolates to contexts more than 16× longer than its training contexts.

  • Problem

    Dense writes in state-based linear models can interfere with preserving individual past tokens, while sliding-window attention loses older tokens when they leave its fixed window.

    Source: §1 · p. 1

  • Method

    Raven maintains fixed memory slots and uses sparse, input-dependent routing to update and decay only selected slots.

    Source: §1 · p. 1

  • Results

    Raven consistently outperforms dense SSMs on recall-intensive benchmarks, remains competitive on general language modeling, and carries gains to hybrids with full attention or SWA.

    Source: §9 · p. 18

  • Takeaways & Limitations

    Raven isolates retrieval-critical tokens in dedicated slots and leaves unselected slots unchanged, allowing those retrieval slots to persist with less interference.

    Source: §9 · p. 18

Abstract

from arXiv · show

Long-context recall in linear-time sequence models highlights a tradeoff in how they write to memory. State-based linear models, such as state-space models (SSMs) and linear Transformers, write densely, updating the entire state for each newly arrived token, which leads to interference and makes specific past tokens hard to recover. Sliding-window attention (SWA) exhibits the opposite behavior: it writes sparsely by storing explicit token representations, but only within a fixed window, so recall drops once the relevant token is evicted. Interpolating between these models, we introduce Raven, a linear-time sequence model that maintains a fixed set of memory slots and, at each step, decays and updates only a selected subset via learned, input-dependent routing. This lets Raven mitigate SWA's position-based overwriting and hard eviction while reducing interference from dense state updates in SSMs, thereby preserving long-range content much more effectively. Across recall-intensive benchmarks, Raven is competitive with or outperforms prior linear-time baselines, achieving strong long-context recall where both SWA and SSMs sharply degrade. It remains effective when extrapolating to context lengths as large as 16x its training length, with similar gains in hybrid architectures.

1 Introduction

Fixed-size memory models face a tradeoff: dense writes can interfere with specific-token recovery, while sliding-window attention (SWA) loses tokens at window eviction. Raven uses sparse, input-dependent routing with selective decay to update chosen memory slots.

1 Introduction

  • Dense writes can interfere with preserving individual past items, whereas SWA retains explicit token representations only within a fixed recent window.SSMs and linear Transformer variants can in principle carry information indefinitely, but SWA loses persistence when older tokens are dropped (Bick et al., 2025c; Dao and Gu, 2024; Gu and Dao, 2024; Katharopoulos et al., 2020; Yang et al., 2023, 2024b; Beltagy et al., 2020).Source: §1 · p. 1
  • Routing Slot Memories maintain fixed memory slots, route new content into them, and explicitly decay updated slots.The framework presents SSMs as dense writers with gradual forgetting and SWA as sparse, deterministic writing with hard eviction.Source: §1 · p. 1
  • Raven instantiates this framework with sparse, input-dependent routing, updating and decaying only a selected subset of slots.The model is designed to keep specific tokens recoverable as context grows beyond training length, without convolutional or sliding-window attention layers for effective retrieval (Yang et al., 2023).Source: §1 · pp. 1–2
  • Raven achieves perfect recall on Needle-in-a-Haystack benchmarks where SWA performs near zero, and improves recall-heavy tasks in standalone and hybrid settings.It is also competitive with or outperforms recent SSMs and linear models, including when extrapolating 16× beyond training lengths.Source: §1 · p. 2

2 Background

The background formulates linear sequence models through memory states and contrasts dense state updates with SWA’s finite token cache. Both can be described as recurrences, but their memory updates differ.

2.1 State-Space Models

  • SSMs update a fixed-size state using a decay gate and a write term, with the state read by multiplying it by the query.SSM variants differ in decay structure: GLA uses a diagonal gate, Mamba-2 a scalar gate, and DeltaNet a delta-rule update (Yang et al., 2023; Yang et al., 2024b).Source: §2.1 · p. 3

2.2 Sliding-Window Attention as a Dual-State SSM

  • SWA stores key/value representations for the most recent M tokens and restricts attention to that window.It updates a FIFO cache by overwriting one slot per step; viewed as a state-space model, it applies binary slot-wise decay and removes the earliest token (Beltagy et al., 2020).Source: §2.2 · p. 3

3 Routing Slot Memories

Routing Slot Memories make write location explicit in fixed-size, slot-separable state, separating where information is written from how it is forgotten. The framework describes existing models through their routing and decay choices.

3 Routing Slot Memories

  • RSMs frame finite-state memory around where information is written and which slots are preserved, rather than whether the model forgets.They formalize write location with a routing vector.Source: §3 · p. 4

3.1 Memory Slots and Routing

  • The state is organized as independent memory slots, and a router specifies each slot’s write intensity based on the input or time.Slot separability constrains updates so each slot changes independently of the others.Source: §3.1 · p. 4

3.2 The Spectrum of Routing Slot Memories

  • An RSM combines slot-wise decay, feature mixing, write content, and routing; a slot with zero routing is frozen until selected again.Intermediate routing values interpolate between a frozen slot and a fully updated one.Source: §3.2 · p. 4

3.3 Dual-State Models as RSMs

  • GSA uses input-dependent routing but typically sends a write signal to every slot, while ABC uses dense input-dependent routing without decay (Peng et al., 2021).SWA instead uses a time-dependent one-hot router that overwrites the oldest slot, producing hard deletion (Beltagy et al., 2020).Source: §3.3 · p. 5

3.4 State-Space Models as RSMs

  • Standard diagonal SSMs correspond to dense routing, updating every slot and relying on decay for selective retention; Delta Networks also route densely but use a rank-one forget gate (Yang et al., 2024b).This exposes a design gap: dense updates can interfere with slot protection, while SWA’s sparse writes are input-blind; Raven combines sparse, input-dependent routing with explicit decay.Source: §3.4 · pp. 5–6

4 Raven: Persistent Memory with Sparse Routing

Raven fills a gap between dense state updates and sparse but input-blind writes by routing updates to selected memory slots and controlling their persistence with decay. Its sparse, input-dependent, slot-wise update combines selective overwriting with explicit decay.

4 Raven: Persistent Memory with Sparse Routing

  • Raven combines sparse, input-dependent routing with explicit decay, filling the gap between dense state updates and sparse, input-blind writes.It maintains separate key/value states and uses decay to control how long selected writes persist.Source: §4 · p. 6

4.1 Raven Memory Update

  • Selected slots are decayed and overwritten, while slots with near-zero routing weights remain untouched.The update decouples where information is stored from how long it persists.Source: §4.1 · p. 6
  • The router sparsifies input-dependent scores by keeping only the top-K entries, then normalizes the selected routing weights.The normalized weights sum to one across at most K selected slots.Source: §4.1 · p. 7

4.2 Raven Block Design

  • Raven omits short-range convolutions over queries, keys, and values, simplifying the block without sacrificing performance.The block uses a channel mixer, nonlinearities, and normalization.Source: §4.2 · p. 7

4.3 Unifying Sparse and Selective Writes

  • Raven’s update uses a sparse, input-dependent diagonal overwrite matrix, between SWA’s one-hot write and DeltaNet’s dense distributed update.This supports slot-wise selective overwrites while combining SWA-like sparsity with decay.Source: §4.3 · p. 8

5 Empirical Validation

Raven is evaluated on retrieval, general language modeling, and hybrid performance. It sustains long-context recall across these settings while remaining competitive on language modeling.

5.1 Retrieval Abilities

  • At 800M, Raven retains 91.0% on NIAH-1 at 64K, versus 45.2% for next-best GDN; at 400M, it maintains ≥99% through 16K and >91% at 32K.The 32K context is 16× the 400M training length; the 800M models were trained at 4K context.Source: §5.1 · p. 9
  • Raven outperforms linear-time baselines on SQuAD, SWDE, and FDA; on SWDE it reaches 34.1% accuracy versus 25.7% for Mamba-2 and 29.0% for GLA.On FDA, Raven is the only linear model to exceed 22%, improving by approximately 8 points over the best SSM baseline.Source: §5.1 · p. 9
  • Raven achieves state-of-the-art retrieval without short-range convolutions, with routing controlling where information is written and per-slot decay controlling how long it persists.The study reports that Raven maintains strong recall without convolutions, unlike prior linear models.Source: §5.1 · p. 9

5.2 Language Modeling

  • At 400M, Raven matches or exceeds Mamba-2 and GLA on average accuracy and has the best Lambada accuracy; at 800M, it trails GDN on average accuracy, though GDN uses 892M versus Raven’s 792M parameters.Across both scales, the authors report no measurable degradation in general language modeling from sparse routing and input-dependent forgetting.Source: §5.2 · pp. 9–10

5.3 Hybrid Raven

  • With full NoPE attention, Raven hybrids maintain strong accuracy to 32K on NIAH-1 and 16K on NIAH-2, while GDN and Mamba-2 hybrids degrade sharply on long-context NIAH.GDN+Attn degrades on NIAH-2 and NIAH-3 beyond 4K; Mamba-2+Attn nearly collapses on NIAH-1 past 2K.Source: §5.3 · p. 10
  • Hybrid models match or outperform strong Transformer baselines such as FoX on recall-heavy tasks, particularly beyond training context length; Raven with NoPE attention performs best in this regime.On long-sequence NIAH-2 tasks at 32K and 16K, Raven attains over 80% accuracy, while other hybrid approaches fail to achieve meaningful accuracy.Source: §5.3 · p. 10
  • Raven hybrids consistently outperform SWA-RoPE hybrids on recall-intensive benchmarks and NIAH, with substantially improved length generalization.The Table 4 caption specifies full attention with NoPE positional embeddings and reports NIAH-1/2/3 plus SWDE, FDA, and SQuAD.Source: §5.3 · p. 10, Table 4 · p. 11

6 Ablations

The ablations examine architectural components, router design, memory shape, and routing sparsity. They support selective routing and a larger number of memory slots as key design choices.

6.1 Block Components

  • Output gating consistently improves retrieval; Raven’s decay outperforms temperature-scaled sigmoid decay, while convolutions yield only marginal gains and are omitted for simplicity.The decay comparison suggests Mamba-2-style input-dependent scalar decay is better suited to the routed setting.Source: §6.1 · p. 11

6.2 Router Design

  • The router ablations vary scoring, training-time exploration, and projection expressiveness on recall-heavy benchmarks.The router is Raven’s novel input-dependent slot-selection component relative to standard SSMs and linear Transformers.Source: §6.2 · p. 11

6.2.1 Softmax vs. Sigmoid Routing

  • Sigmoid routing scores slots independently and outperforms softmax across most tasks, so the authors adopt it as the default.Its scores are naturally sparse, and Top-K selection mostly confirms the router’s prior preferences.Source: §6.2.1 · p. 12

6.2.2 Gumbel Noise During Training

  • Gumbel noise encourages training-time exploration across all slots and yields consistent gains; it is removed at inference for deterministic routing.Without intervention, early router collapse can leave slots uninformative and reduce effective memory capacity.Source: §6.2.2 · p. 12

6.2.3 Linear vs. MLP Projection

  • The MLP router helps on some benchmarks, notably FDA, but underperforms the linear router on others and adds parameters, so gains are inconsistent.The final design selects the best-performing configuration per Table 6.Source: §6.2.3 · p. 12

6.3 Memory Shape: Head Dimension vs. Number of Slots

  • Under a fixed memory budget, allocating capacity to more slots improves recall even as per-slot dimension shrinks; the authors set M = 256 as default.The study says increased slot granularity reduces the chance that retrieval-critical tokens share a slot with unrelated content and are overwritten.Source: §6.3 · p. 13

6.4 Top-K Sparsity

  • K = 32 of M = 256 slots (12.5% occupancy) provides the best balance, preserving lightly used traces while allowing general content to spread across slots.Smaller K improves SWDE recall; larger K moves Raven toward dense, SSM-like updates.Source: §6.4 · p. 13

7 Analysis: Learned Memory Allocation

The analysis examines whether Raven learns structured, content-aware memory allocation by measuring how much slots process and what they store. It connects non-uniform allocation to variable effective sequence lengths and retrieval specialization.

7.1 Effective Sequence Length in Routed Memory

  • Raven slots span short, intermediate, and long effective sequence lengths, unlike SWA’s uniform one-token slots and SSMs’ slots that each see every token.The effective-sequence-length spread is consistent across layers and heads.Source: Figure 5 · p. 14, §7.1 · p. 14
  • Non-uniform effective sequence lengths implicitly expose Raven to variable-length dynamics within each training example, without explicit intervention.The paper contrasts this with state passing, which varies effective lengths during training.Source: §7.1 · p. 14

7.2 Selective Allocation in a Retrieval Task

  • Different slots specialize for retrieval-critical passkeys or non-retrieval content, with shared slots rare in heads showing strong retrieval behavior.The paper relates dedicated retrieval slots to high retrieval accuracy up to 16× the training length.Source: §7.2 · pp. 14–15, Figure 6 · p. 15

7.3 Intentional Memory Imbalance

  • Raven omits a load-balancing loss so memory usage can specialize, while Gumbel noise encourages exploration and prevents collapse to only a few slots.Uniform slot usage could prevent retrieval-slot formation and expose stored tokens to unnecessary overwrite.Source: §7.3 · p. 15

8 Related Work

Related work situates Raven among linear sequence mixers, retrieval and length-generalization research, and routing-based architectures. It distinguishes Raven’s within-memory-slot routing from approaches that route across channels or heads.

8.1 Linear Sequence Mixers

  • Raven combines diagonal decay at individual memory elements with row-level updates, selectively updating slots rather than rewriting the full memory as GSA does.This design is presented as a way to retain fast training while supporting persistence across timesteps.Source: §8.1 · p. 16
  • Prior work links linear-model retrieval limits to fixed-size state and identifies unexplored states as one source of length-generalization failures; Raven connects sparse routing to variable write horizons.The evaluation benchmarks cover real-world retrieval, synthetic key–value retrieval, and language modeling tasks.Source: §8.1 · p. 16

8.2 Evaluation Benchmarks

  • The evaluation distinguishes retrieval-heavy tasks requiring contextual evidence from language-modeling benchmarks focused more on factual knowledge and commonsense reasoning.Retrieval includes SQuAD, SWDE, FDA, and RULER’s S-NIAH tasks; language modeling includes PIQA, Winogrande, HellaSwag, and ARC.Source: §8.2 · p. 17

8.3 Routing and Mixture of Experts (MoE)

  • Unlike MoE channel routing and MoM or MoH head-level routing, Raven routes within each head’s memory allocation, and head routing could be combined with it.The cited prior work also includes SSE, which sparsely selects components of the key vector.Source: §8.3 · p. 17

9 Conclusion

The conclusion presents Routing Slot Memories as a framework for separating where information is written from how long it persists. Raven uses that separation for selective memory allocation, and the paper reports long-context recall gains across benchmarks and hybrid architectures.

9 Conclusion

  • Routing Slot Memories formalize memory-slot selection, while Raven decouples where information is written from how long it persists.The paper frames sparse input-dependent routing with explicit decay as a gap in prior models.Source: §9 · p. 18
  • Raven isolates retrieval-critical tokens in dedicated slots and induces variable effective sequence lengths, leaving unselected slots unchanged to reduce interference.Compared with SWA, it uses content-dependent allocation and selective decay instead of position-based eviction.Source: §9 · p. 18
  • Raven extrapolates to contexts more than 16× longer than training and consistently outperforms dense SSMs on recall-intensive benchmarks while remaining competitive in general language modeling.The reported gains also carry over to hybrids paired with full attention or SWA, without relying on short-range convolutions.Source: §9 · p. 18
  • The RSM framework treats routing and forgetting as orthogonal design axes that existing models have largely conflated.Source: §9 · p. 18

A.1 Notation

Table 8 summarizes the paper’s notation.

A.1 Notation

  • Table 8 is a summary of notations used throughout the paper.Source: Table 8 · p. 23

A.2 Slot Separability of Linear Updates

This appendix establishes when linear updates keep memory slots separate, describes Raven’s training and design, and interprets its sparse updates and memory behavior. Raven combines selective slot updates with efficient training and shows long-range retrieval strengths in hybrid settings.

A.2 Slot Separability of Linear Updates

  • A linear update is row-wise if and only if its decay matrix Dt is diagonal, so each updated slot depends only on its corresponding prior slot.The proposition states the equivalence for the update defined using Dt, At, and Ut.Source: §A.2 · p. 23

A.3.1 Training Recipe

  • The 400M models train for 15B tokens at sequence length 2048, and the 800M models for 32B tokens at length 4096, using SlimPajama-627B.Both scales use batch size 0.5M; training follows the stated standard pipeline.Source: §A.3.1 · pp. 23–24

A.3.2 Models Configuration

  • For fair retrieval comparisons, single-state SSM baselines GLA and GDN are expanded to match the other models’ total memory elements.The configuration includes GLA short-range convolution and uses 24 layers except for 48-layer Mamba-2; Transformer baselines use 24 layers and 16 heads.Source: §A.3.2 · p. 24

A.5 Efficient Training using Gated Linear Attention

  • Raven training adapts two-pass GLA/GSA-style computation and flash-linear-attention kernels; chunkwise computation addresses numerical instability at large sequence lengths.The chunkwise method splits T tokens into n chunks of length C, with T = n × C.Source: §A.5 · pp. 24–25

A.6 Raven Design Details

  • Raven’s standalone block uses four heads, separate key and value states, a channel mixer, skip connections, QK normalization, and output gating.In the hybrid setup, Raven blocks are interleaved with SWA and softmax attention mixers; RoPE length generalization failed for those mixers in the evaluated setup.Source: §A.6 · p. 25
  • Raven +SWA with RoPE achieves near-perfect NIAH-1 accuracy across 1K–32K, while GDN+SWA degrades steadily and Mamba-2+SWA collapses near zero beyond 4K.Raven also leads on SWDE, FDA, and SQuAD in this configuration; the reported results are in Table 9.Source: §A.6 · p. 26, Table 9 · p. 26

A.7 Test-Time Training Point of View

  • Raven’s router can be viewed as selective dropout in an online learning objective: unselected slots incur no weight decay and preserve their contents exactly.The appendix frames Raven’s sparse state updates as sparse test-time training.Source: §A.7 · p. 26

A.8 Memory Dynamic Visualizations

  • Different Raven heads have different effective sequence lengths, ranging from a few active slots to many slots storing tokens more diversely.Some retrieval heads store passkey tokens in dedicated slots, while others share slots between passkey and ordinary tokens; non-retrieval heads store tokens equally in memory.Source: §A.8 · pp. 27–28

Related papers

Loading 2607.25357v1…