One-sentence Summary

Researchers from Shanghai AI Laboratory, Tsinghua University, and Fudan University propose MARCH (Memory-Anchor Routing across Context History), which scales state-space models beyond fixed-size states by periodically caching recurrent-state checkpoints as content-routed state anchors and aggregating over them, and after standard pretraining it improves long-context retrieval and outperforms linear attention variants on commonsense reasoning, LongBench, and in-context retrieval.

Key Contributions

• MARCH (Memory-Anchor Routing across Context History) is introduced as a network architecture that augments recurrent state-space models by periodically caching cumulative recurrent-state checkpoints as state anchors and associating them with compact content-conditioned anchor keys.
• MARCH maintains a memory bank that grows with context length and computes outputs through attention-style aggregation over causally available anchors together with the current state, decoupling memory capacity from dense per-token updates while preserving efficient recurrence.
• After standard pretraining, MARCH consistently outperforms multiple linear attention variants across commonsense reasoning, LongBench, in-context retrieval, and NIAH, with ablations showing a controllable retrieval-efficiency trade-off through checkpoint density and sparse routing.

Introduction

Large language models increasingly need to integrate information across long documents, multi-turn interactions, and in-context examples. Standard self-attention supports fine-grained recall but incurs quadratic training cost and a key-value cache that grows linearly with sequence length. Linear attention and modern recurrent models offer constant-memory decoding, but compress the entire history into a fixed-size state, so earlier associations can be overwritten or weakened and cannot be recovered from the latest state. Prior work expands memory capacity or keeps multiple compressed states, yet query-dependent retrieval over historical states remains underexplored. The authors introduce Memory-Anchor Routing across Context History (MARCH), which periodically saves recurrent-state snapshots as anchors and uses compact learned descriptors to route each token to relevant earlier states while preserving efficient recurrence.

Method

The authors introduce MARCH, a content-routed recurrent memory framework that enables selective retrieval from earlier versions of an evolving recurrent state. Rather than routing over predefined state indices or temporal scales, MARCH matches each query against individual historical states based on their contents.

As shown in the framework diagram, the method operates by periodically checkpointing the cumulative recurrent state as tokens are processed, producing a bank of state anchors. Each checkpoint is paired with an occurrence of a shared learned anchor token, whose hidden representation yields a compact routing key. For each text token, a routing query scores all causally visible anchors alongside a learned null option, allowing the model to use historical memory only when useful. The resulting routing probabilities define a weighted combination of the visible anchor states, which is read using the standard recurrent query of the token. This historical readout is added to the current-state readout, preserving the native recurrent path while introducing a content-dependent route to earlier memory. Together, state anchoring and content-routed retrieval turn the otherwise transient state trajectory into a persistent source of long-range memory.

The anchoring process begins by defining an ordered set of text boundaries. An anchor position is inserted after each boundary using a shared learned anchor embedding. Text positions apply the base recurrent update, allowing the matrix-valued state to evolve continuously across anchor boundaries. Immediately after processing a text token at a boundary, MARCH checkpoints the resulting cumulative state to form a state anchor. The following anchor position does not modify the recurrent state; instead, its hidden representation provides the routing metadata associated with that checkpoint. Because the recurrence is not reset between anchor boundaries, each state anchor encodes the cumulative prefix up to that position, tracing the temporal evolution of a single recurrent memory and preserving earlier versions before subsequent decay and delta updates attenuate their contents.

To make routing explicitly dependent on the content of each state anchor, the anchor position reads only its aligned state checkpoint. The same input representation is projected into a compact routing key. The aligned readout is incorporated into the anchor position through the standard output projection and residual pathway. Consequently, the representation at the next layer depends on the state anchor, and the routing key produced becomes conditioned on the content retained by the aligned state anchor. Although all anchor positions share the same learned input embedding, they acquire distinct, state-dependent representations after the first layer.

For content-based routing, the model projects the normalized hidden state of a text token into a routing query and scores it against the key of each visible anchor. To allow the model to bypass historical memory, the visible anchor set is augmented with a null option whose payload is fixed to zero. The routing probabilities are computed via a softmax over the logits of the visible anchors and the null option. Since the selected routing probabilities directly weight the historical state readouts, their scores remain jointly optimized by the language-modeling objective. The aggregation formulation also admits a sparse variant by restricting aggregation to the highest-scoring visible anchors, which reduces aggregation cost with minimal performance degradation.

Given the routing probabilities, the causally visible state anchors are aggregated into a query-dependent historical state, which is read using the same state-read query as the current state. The resulting historical readout is then added to the current-state readout. This additive formulation preserves the original recurrent path and introduces historical retrieval as an auxiliary residual branch, without modifying the underlying recurrent update. Since the routing probabilities directly affect the layer output, the routing queries and anchor-derived keys are optimized end-to-end.

The authors implement MARCH as a two-stage producer-reader computation. Following a hardware-efficient chunkwise formulation, the producer processes recurrent updates in blocks amenable to tensor-core acceleration, computes the current-state output for each token, and checkpoints the recurrent state at each anchor boundary. The resulting state anchors are consumed by the historical reader. Inspired by I/O-aware principles, the reader jointly tiles query tokens and state anchors, reuses each anchor tile across a block of queries, and fuses routing-score computation, online softmax updates, and the accumulation of weighted state readouts into a streaming reduction. This fused schedule avoids materializing either the dense token-to-anchor routing matrix or the substantially larger tensor of per-anchor candidate readouts, thereby reducing intermediate storage and the associated memory traffic. As demonstrated in the efficiency analysis, despite the cost of historical retrieval, the fused dense implementation exceeds standard attention baselines in throughput at longer sequence lengths and incurs lower core runtime.

Experiment

The paper evaluates MARCH by pretraining models from scratch on 50B tokens with 16K context and comparing it against Gated DeltaNet, a log-linear variant, and Transformer baselines. It tests zero-shot commonsense reasoning, long-context understanding, needle-in-a-haystack retrieval, and in-context retrieval, with ablations on chunk size and routing design. MARCH consistently improves over recurrent baselines, especially on retrieval-intensive and long-context tasks, and remains competitive with full-attention models on short-context understanding. Ablations show that a chunk size of 512 balances accuracy and cost, while sparse routing improves efficiency and the learned null option is beneficial.

Across eight zero-shot commonsense reasoning benchmarks, MARCH improves over both vanilla and log-linear Gated DeltaNet variants on every task and raises average accuracy above both Transformer baselines. The improvement is largest over the vanilla recurrent baseline on OpenBookQA. MARCH outperforms the standard Transformer on six of eight tasks and the 24-layer Transformer on four, indicating competitive short-context language understanding. MARCH consistently outperforms vanilla and log-linear Gated DeltaNet across all eight commonsense benchmarks, with the strongest single-task gain over the vanilla backbone on OpenBookQA. MARCH surpasses the standard Transformer on six of eight tasks and the deeper Transformer on four while achieving a higher average score than both full-attention baselines.

On twelve LongBench tasks, MARCH outperforms both vanilla Gated DeltaNet and its log-linear variant on every task, with a roughly 25 percent average relative gain over Gated DeltaNet. Improvements are especially strong in multi-document QA and summarization, and the average score surpasses the standard Transformer while approaching the 24-layer Transformer. MARCH improves over both Gated DeltaNet variants on all twelve LongBench tasks, spanning single-document QA, multi-document QA, summarization, and few-shot learning. The largest relative gains appear in multi-document QA, where MARCH substantially raises 2WikiMultihopQA and MuSiQue scores relative to the stronger recurrent baseline. Summarization also benefits notably, with QMSum increasing by 32 percent relative to the stronger recurrent baseline. MARCH achieves an average LongBench score higher than the standard Transformer and close to the 24-layer Transformer.

MARCH improves in-context retrieval over both Gated DeltaNet variants across all six benchmarks, yielding relative gains over the stronger recurrent baseline in every task. Average retrieval accuracy rises from 20.5 to 23.3, a 14% relative improvement. Transformer baselines still lead on average, but MARCH narrows the gap while consistently strengthening the recurrent backbone. MARCH outperforms the stronger Gated DeltaNet baseline on all six retrieval tasks, with relative gains ranging from 8% on SQuAD to 23% on TriviaQA. The largest relative improvements over the recurrent baselines occur on TriviaQA, NQ, and FDA. Average retrieval accuracy across the six benchmarks improves by 14% relative to the stronger Gated DeltaNet baseline. MARCH trails the Transformer baselines on average retrieval accuracy but substantially improves over the Gated DeltaNet variants.

Matched training and inference chunk sizes show that a chunk size of 512 offers the best overall balance between retrieval quality and anchor count. Smaller chunks improve some longer-context results but increase memory and routing costs, while larger chunks generally reduce retrieval performance because checkpoints become sparse. Inference-time chunk size changes provide a flexible accuracy-memory tradeoff, and hierarchical Fenwick tree organization remains competitive. A chunk size of 512 provides the best overall retrieval quality relative to anchor count. Smaller chunks create denser anchors and improve some long-context NIAH results but require more memory and routing overhead. Larger chunk sizes reduce anchor density and generally degrade retrieval, especially at longer contexts. At inference, denser anchors tend to improve retrieval at higher cost, while overly sparse anchors cause substantial degradation. Fenwick tree organization of the state bank performs close to the reference scheme, indicating flexibility in anchor organization.

The default dense router with a query-key dimension of 64 and the learned null option provides the best overall trade-off, with the highest commonsense, LongBench, and NIAH averages. A larger router dimension improves retrieval but reduces other aggregate scores, while top-4 routing is close on commonsense and retrieval but weaker on NIAH. Removing the null option lowers all aggregates, indicating that learning to bypass irrelevant historical states is beneficial. The default dense routing configuration with the null option leads on commonsense, LongBench, and NIAH averages, though retrieval is slightly lower than with a wider router. Top-4 routing nearly matches dense routing on commonsense and retrieval but trails on NIAH, and removing the null option consistently degrades aggregate performance.

Across zero-shot commonsense reasoning and LongBench tasks, MARCH consistently improves over both vanilla and log-linear Gated DeltaNet variants, surpasses the standard Transformer on average, and approaches or exceeds a deeper Transformer on several tasks. It also strengthens in-context retrieval over the recurrent baselines, though full-attention Transformers still lead on average. Ablations show that a chunk size of 512 best balances retrieval quality and overhead, while the default dense router with a learned null option provides the strongest overall trade-off across commonsense, long-context, and retrieval benchmarks.