Weighted State Souping for Efficient Long-Context Retrieval

Introduction

Many applications of large language models (LLMs) require long-lived context organized into discrete chunks of tokens: tool definitions for an agent ([36, 30]), documents for question answering or retrieval-augmented generation ([32, 52, 45, 20, 11]), and memory records or chat histories that persist across interactions ([28, 34] (Wu et al. 2024)). A useful memory architecture must retrieve relevant chunks from an increasingly crowded store without repeatedly processing the entire history.

In these settings, common practice concatenates the chunks of tokens and processes them as a single flat sequence. With the rising popularity of state space models (SSMs) and related fixed-state architectures ([10, 5, 18, 37, 42]), an alternative approach becomes possible. Rather than perform a single autoregressive pass on a concatenated sequence, each chunk can be independently encoded into a state. The resulting collection of states are then point-wise summed to initialize the model before it processes a query (Figure 1). By analogy to model souping ([49]), this procedure is called state souping ([31, 15, 21]): the quantities combined are input-dependent recurrent states, rather than model parameters.

State souping has an obvious efficiency advantage. State encodings can be computed in parallel and, for slowly changing chunks like memory records or tool definitions, cached and reused. This ability to cache representations makes souping naturally suited for handling the structured context that increasingly dominates modern agentic workloads: inference can depend on selected state encodings without repopulating the context with large memory artifacts, tool definitions, and documents. However, finite-state SSMs struggle to copy or retrieve from long contexts [16, 48, 2], and state souping performance similarly decays as more documents are added [15]. This loss of selectivity limits souping precisely where persistent memory is most useful.

In this work, we introduce weighted state souping, which learns attention-like weights over state encodings. Across document QA, needle-in-a-haystack retrieval, variable tracking, and tool use, weighting substantially improves robustness as contexts grow beyond training lengths. Our method can be further adapted into a sparse variant that permits near-constant encoding cost while preserving — and sometimes improving — performance. Our main contributions are:

• a novel weighted state souping procedure that permits robust selectivity, strong generalization on long contexts, and substantial efficiency gains over both Transformers and traditional SSMs,

• a comprehensive evaluation over diverse benchmarks including multi-document question answering, needle-in-a-haystack retrieval, variable tracking, and agentic tool use,

• a simple theoretical account of how souping successfully retrieves and composes across chunks, and its potential weaknesses compared against sequential models. We also discuss informally the advantages of weighted state souping.

Related work

State souping is a relatively recent method. Pióro et al. [31] introduced the term, and explored linear combinations of recurrent states for in-context skill transfer. Jafari et al. [15] study independently encoded document states, showing that finetuning enables simple pooling operators to support downstream QA and retrieval. PICASO [21] derives state-composition rules similar to souping for permutation-invariant inputs. We build on these works by showing how query-dependent weighting permits both significantly stronger length generalization and substantial efficiency wins.

State souping can be seen as an approximate form of state caching, a procedure whereby states generated from common prefixes can be cached and reused exactly for downstream requests. In Transformers, the procedure is more often called KV caching [38, 19]. In SSMs, state caching has been explored particularly in the context of hybrid models, where common KV caching methods for attention blocks may not transfer directly to SSM layers. In particular, Pan et al. [29] explore SSM-specific cache eviction policies to maintain an efficient prefix cache for SBMs. Shirokikh & Nikolenko [39] explore a complementary idea in which cache checkpoints are saved along the sequence; future generation resumes from the deepest checkpoint that matches the current prefix. Gopal et al. [9] investigate state caching in pure SSMs for retrieval augmented generation, and highlights its efficiency for edge-computing applications.

Souping is also related to state tuning, the process of learning an initial state for a recurrent model. This process has been studied in the context of parameter-efficient fine tuning (PEFT) for state space models [17, 53] (Galin et al. 2024). Kang et al. [17] in particular find that state offset tuning, in which a learned offset is added to every recurrent state, outperforms initial state tuning and other PEFT methods like LoRA. State tuning has also been studied in the RWKV community [51], and historically in early RNNs [7].

Our weighted approach to state souping is similar to methods that attempt to introduce attention-like features to state space models, interpolating between Transformers and SSMs. In particular, Akhauri et al. [1] introduce a method in which SSMs produce compressed key-value representations before applying attention. Similarly, Infini-attention [25] combines local attention with a compressive recurrent memory. Memory Caching [3] explores storing and combining recurrent memory checkpoints through attention-like weights. Our approach focuses on structured contexts that permit state souping, and our weighting does not retain token-by-token attention to cached states during generation, ultimately yielding both substantial efficiency benefits and performance gains.

Methods analogous to state souping have also been investigated in Transformers. Notably, Eyuboglu et al. [6] introduce a method for directly training KV caches to encode documents. The resulting “cartridges” can be provided in-context to approximate the behavior of the model had it been prompted with the original documents, similar to how SSM state encodings are used in souping. However, cartridges require offline tuning specific to each document, and multiple cartridges do not compose in-context without significant additional training [12]. Our weighted state souping approach is also similar in motivation to block attention architectures that use block representations to select relevant blocks for token-level attention [24, 22] or use query-dependent scores to load KV cache blocks during decoding [43].

Weighted State Souping

State souping operates in the following way. Let an input contain $k$ token chunks $\mathbf{c}_1,\ldots,\mathbf{c}_k$ and query $\mathbf{q}$. The SSM independently encodes each chunk into its final state $\mathbf{s}_i$, where the state encompasses the collection of recurrent states across all layers. We then perform a point-wise average over the $k$ encodings to make the “soup” $\mathbf{s}_0 = \frac{1}{k}\sum_{i=1}^{k}\mathbf{s}_i$. $\mathbf{s}_0$ becomes the initialization for the SSM as it processes the user query Figure 1. We use the same SSM for encoding and generation, but these models can be different. Souping is not an inference-time only intervention, and requires either learning from scratch or additional finetuning on a pretrained SSM to teach the model to used this composed initialization [15].

Crucially, the only channel through which the input chunks influence the query is the souped initialization. Since SSMs perform simple linear state updates, a soup constructed through an average of states can be viewed as a naive approximation of the state construct from a concatenated, sequential input. Such a mean soup is the approach used in prior work [31, 15], and performs comparably with a sequential SSM. However, both mean souping and sequential SSMs struggle when operating over long contexts containing many input chunks [15, 4, 2]. As the number of chunks increases, a mean soup or finite SSM state becomes increasingly noisy, degrading performance.

Query-dependent weighting may instead concentrate the souped state on relevant chunks, the inspiration for weighted state souping. Our method is related to other schemes that weight input tokens in sequential SSMs [3] (Akhouri et al. 2024), though our focus is on a souping architecture for chunked inputs. To compute these weightings, we add a learnable linear readout that maps the final activation to a key and query vector. Each chunk produces key $\mathbf{k}_i$. Starting from a zero-initialized state, a pass over the user query produces $\mathbf{u}$. We then compute the souped initialization as $\mathbf{s}_0 = \sum_{i=1}^{k}\sigma(\mathbf{k}_i\cdot\mathbf{u})\mathbf{s}_i$, with weighting function $\sigma(\cdot)$. The linear readout and base model are jointly finetuned. See Figure 1c for an illustration of this¹. Unlike a Transformer’s token-level attention, this procedure assigns one scalar weight per chunk, shared across all layers and heads and computed once from the original user query. The weights remain fixed during query processing and generation.

For $\sigma$, we use the softmax function with cosine-normalized scores (see Appendix § 11 for details). We also explore sparse souping, an efficient variant in which we mask at inference time all but the top-$N_{\mathrm{top}}$ logits before softmax, retaining normalized weights while avoiding retrieving zero-weight states. Masked states do not need to be retrieved, providing significant efficiency advantages for sparse souping (§ 4.2). Other reasonable choices of weighting functions include sigmoidal and linear weights, though we find that softmax performs best in our experiments.

Experiments

We use a 1.5B Mamba-2 SSM (Dao & Gu, [5]) and compare softmax-weighted, sparse, and mean souping against sequential Mamba-2 and a matched Transformer with rotary positional encodings (Su et al., [41]).

We evaluate our models across a diverse range of tasks. For document QA, we use SQuAD (Rajpurkar et al., [32]) for single-hop; HotpotQA (Yang et al., [52]) and MuSiQue (Trivedi et al., [45]) for multi-hop. For NIAH and variable tracking, we use RULER (Hsieh et al., [13]). For tool use, we use the Berkeley Function Calling Leaderboard (BFCL) (Patil et al., [30]). To probe long-context performance on document QA and BFCL, we augment the input with additional distractors and plot performance for varying pool size $k$. If $k$ is smaller than the number of gold documents or gold tool definitions (relevant only for the $k = 1$ datapoints), then the input consists only of gold samples. This suite tests selective retrieval and reasoning as the available context grows, crowding the relevant chunks with increasingly more distractors. Our models are finetuned to each task, using either the task-provided official training split (for document QA), a specialized post-training corpus (for tool use), or a disjoint subset of examples (for RULER). See Appendix § 11 for additional details on training, evaluation, and our benchmarks.

Benchmark results

Figure 2 summarizes the main results; Appendix § 9 reports on additional metrics and task subsets.

In general, weighted state souping performs strongly even at large pool size $k$. While mean souping and sequential Mamba tend to deteriorate as the number of chunks increases, weighted souping matches their performance at small $k$, and persists in near-constant performance as $k$ increases to several times beyond the max training size. This phenomenon persists across a diversity of tasks and benchmarks, including in multi-hop settings like HotpotQA, MuSiQue, and RULER variable tracking where the model must link information across multiple chunks. Similar to other works ([23, 56]), we observe that Transformers do not extrapolate well to sequences longer than the training data, though their in-distribution performance tends to be strong.

Interestingly, sparse souping matches or improves on full softmax souping, particularly on multi-hop QA (HotpotQA, MuSiQue). Softmax mass may diffuse across many documents as $k$ increases. Sparse souping instead may beneficially concentrate attention in the most relevant indices. An analogous benefit has been observed in sparse attention for Transformers ([47, 54] (Vaslyenko et al. 2025)). One weakness of sparse souping is the need to specify the number of chunks $N_{\mathrm{top}}$ to keep. If $N_{\mathrm{top}}$ is insufficiently large, critical chunks may be dropped and the model’s performance suffers (Figure 5f). Designing a weighting function with adaptive sparsity represents an interesting avenue for further exploration.

While the performance of weighted souping is generally strong, it underperforms at small $k$ in the parallel multiple subset of BFCL, despite strong performance on python simple and parallel (Figure 6). Whether this reflects an optimization issue or a fundamental limitation remains open.

Computational efficiency

The primary computational advantage of souping is in its ability to cheaply retrieve a state encoding from memory, rather than autoregressively re-encoding every input. Since inputs like memory artifacts and tool definitions change slowly, their state encodings may be routinely cached for future retrieval. Hence, a souped model may be considerably faster than non-souped alternatives, particularly as the context length grows.

Theory: retrieval and composition

Figure 3 compares the measured time to encode an input between souping, Mamba, and Transformer models. See Figure 7 for additional timings.

There are a number of further considerations for souping. 1) Weighted souping must also factor in the time to perform an initial pass over the user query to generate a query vector, and the time it takes to compute the weightings. For a fixed key dimension, scoring and normalizing the weights costs $O(k)$ per query. Once the soup is constructed, the cost per generated token is constant with respect to $k$, since generation does not access the cached chunks again. In general, user queries average under 100 tokens ([55]) and computing softmax weights is extremely fast, so the total added cost of computing the weights is typically negligible and amortized over the number of input chunks. 2) The retrieval time is heavily influenced by the specific memory holding the cached state encodings. GPU HBM is significantly faster than system RAM, which in turn is faster than system disk. Depending on where the encodings live, retrieval may in the worst case be slower than encoding from scratch. 3) Sparse souping improves efficiency further, since chunks assigned zero weight do not need to be retrieved. Hence, if the number of nonzero weights is constant with respect to the context length (which is true in our setting), then the encoding performance of sparse souping is effectively constant as context length grows$^{2}$.

Theory: retrieval and composition

The experiments in § 4 provide the main evidence for the effectiveness of weighted state souping. To complement these empirics, we present a simple theoretical analysis that begins to answer the questions of why souping works, and what potential strengths and weaknesses it may have compared with sequential models. Our arguments are primarily expressivity-based: we provide constructions showing that both souping and sequential models can implement an induction-head retrieval task representative of our empirical benchmarks. We also discuss possible weaknesses that independent chunk encoding introduces, namely a sensitivity to chunk boundaries and the inability to encode cross-chunk representations. We conclude with a discussion of the benefits weighted souping has for retrieval over long, crowded contexts. Our analysis is by no means a complete theoretical account of state souping. Instead, we introduce simple intuitions that build towards a principled understanding of the method.

Task

We study a key-value version of the induction-head retrieval task (Olsson et al., [27]) and its multi-hop extension (Sanford et al., [35]).

Let $[N] = \{1, \ldots, N\}$, and introduce disjoint key and value vocabularies $\mathcal{K} = \{k_a : a \in[N]\}$ and $\mathcal{V} = \{v_a : a \in[N]\}$. An input context defines a partial function $f : [N] \to[N]$ by listing pairs $(k_a, v_{f(a)})$, with each key appearing at most once. Given a query $k_q$, the $n$-hop induction target is

whenever $f(q), \ldots, f^n(q)$ are defined. Thus one hop retrieves the value paired with $k_q$; each subsequent hop uses the index of the retrieved value as the next key. The input context is presented to the model as concatenation of all pairs followed by the query,

A sequential model reads this sequence in one pass. A souping model encodes each pair independently. A chain of length $n$ is a sequence of distinct indices $a_0,\ldots,a_n$ with $f(a_i)=a_{i+1}$. Multiple chains and irrelevant pairs may appear in the context. Pair order is arbitrary except in the explicitly ordered setting below.

Model

We consider a stripped-down linear attention SSM with Mamba-like history decay and augmented residuals. Let $z_1,\ldots,z_T$ be the tokens and let $e(z_t) \in\mathbb{R}^{d}$ be orthonormal vocabulary embeddings, with $d \ge2N$. Initialize the residual stream as $r_t^{(0)}=[e(z_t);0]\in\mathbb{R}^{2d}$. At layer $\ell\in[D]$, a causal width-two convolution copies the previous and current residual vectors into separate coordinates:

This is the copying role of the convolution identified by Huang et al. [14]. We then perform the SSM state update as

where $H_t^{(\ell)}\in\mathbb{R}^{d\times d}$ is the SSM state; $W_k^{(\ell)},W_v^{(\ell)},W_q^{(\ell)}\in\mathbb{R}^{d\times4d}$ are the key, value, and query projections; and $\alpha\in(0,1]$ is the history decay. We omit bias terms.

The residuals are updated as

In this way, $r_t^{(\ell)}=[e(z_t);u_t^{(\ell)}]$ preserves the token embedding in the first block and accumulates updates in the second block.

The final output logits produced by the model are $Uu_T^{(D)}$, where $U\in\mathbb{R}^{N\times d}$ is the readout matrix. To extract a concrete next-token prediction, we take the argmax over these logits.

Sequential encoding starts from zero-initialized recurrent states and a zero convolution buffer. For souping, each chunk $c_i$ is encoded independently from zero at every layer, giving states $H_i^{(\ell)}$. The query then starts from the weighted, souped initialization $H_{\mathrm{init}}^{(\ell)}=\sum_{i=1}^{K}w_iH_i^{(\ell)}$, with its convolution buffer reset to zero. The weights $w_i$ are shared across layers and fixed during query processing. Unless stated otherwise, chunk encoding and weighting have no access to the chunks’ original positions: permuting the chunks permutes their weights and leaves the soup unchanged. This assumption is accurate of mean and content-based softmax souping. Our discussion below focuses on expressivity rather than strict learnability, and assumes exact arithmetic and suitably chosen positive weights.

Remark 1. Linear states permit souping. For a single layer with matching convolution resets, sequential encoding equals a weighted sum of independently encoded chunk states.

There is a simple algebraic reason that states from this model can be souped. Split a context into chunks $c_1,\ldots,c_K$ of lengths $m_1,\ldots,m_K$, and let $H(c_i)$ be the final state obtained by encoding $c_i$ from zero. Let $H_{\mathrm{seq}}(c_1\Vert\cdots\Vert c_K)$ be the final state obtained by encoding the concatenated sequence of chunks. If we reset the convolution at the same boundaries in both schemes, then unrolling (2) gives

Thus a sequential linear state is itself a weighted soup of that layer’s chunk states, an identity also observed in prior state-composition work [21]. With $\alpha=1$ and a one layer model, sequential and souping are functionally equivalent.

However, there are a few important differences to keep in mind. On one hand, learned state souping replaces the fixed recency weights $\alpha^{\sum_{j>i}m_j}$ in (3) with content-dependent weights, permitting selection based on relevance rather than position. On the other hand, independently encoded chunks cannot incorporate information from other chunks during encoding. Particularly when a model has two or more layers, its sequential upper-layer representations may therefore differ from its independently encoded representations, potentially impacting efficiency when cross-chunk representations would be advantageous (see Remark 5 below).

One-hop induction

We begin by considering the one-hop induction-head task, reflecting the setting in Olsson et al. [27]. We observe that both a sequential and souping model can solve the task, but chunk boundaries that separate the task’s key-value pairs can make souping fail.

See Section B.1 for the construction proving this statement. Proposition 2 demonstrates that even a single-layer SSM is sufficiently expressive to perform one-hop induction. Together with our earlier discussion in Remark 1, we observe that both souping and sequential models can readily perform in-context retrieval. However, souping is sensitive to chunk boundaries, as we establish in the next proposition.

Proof. Choose distinct indices $a$, $b$, $u$, $v$. Consider the two token sequences

After splitting into singleton chunks, both inputs have the same multiset $\{k_a, k_b, v_u, v_v\}$. For the same query, souping therefore produces the same state for $C$ and $C'$. Hence, the model must return the same output on both inputs, though their labels differ, failing the task on at least one input. □

The setting in Proposition 3 is extreme, but underlines the sensitivity of souping to correctly formed chunks. However, this does not mean that souping cannot link information between chunks — our multi-hop document QA and variable tracking experiments demonstrate that souping models can robustly link information across multiple chunks. In the next section, we observe that souping can readily link disparate chunks when performing the multi-hop generalization of our task.

Multi-hop induction

We consider the multi-hop generalization of our induction-head task. We discuss how souping can continue linking information between chunks to solve the task, though under certain circumstances, it may be less efficient that a sequential input model.

See Section B.2 for the construction proving this statement. In this way, we establish the basic competency of souping models to solve these classic in-context retrieval tasks. However, one advantage sequential models retain over souping is their ability to learn cross-chunk representations in deeper layers. Because souping encodes chunks independently, it remains impossible to directly encode information between chunks in deeper layers. This expressivity difference may lead to capability separations between sequential and souping models.

For instance, suppose the inputs to our multi-hop induction task are favorably ordered. Consider $C$ node-disjoint chains presented consecutively in forward order, with each chain taking the form

Suppose the query always names a chain’s beginning and asks for its $n$-hop endpoint. Then we observe the following.

The construction is given in Section B.3. Each of the first $L$ layers stores reverse associations whose hop length doubles with depth, allowing later tokens to identify earlier nodes in their chain. The final layer then writes

so a query at a chain start directly retrieves its endpoint. Independently encoding a single pair cannot perform this cross-chunk preprocessing, though souping can still compose the original relations at query time as in Proposition 4.

A simple counting argument clarifies why the number of retained associations matters for storage.

Proof. There are $N^{m}$ possible assignments. Distinct assignments must yield distinct states, since they disagree on at least one query. Thus $2^{B} \ge N^{m}$. Taking the log of both sides results in the desired bound $B \ge\lceil m \log_{2} N \rceil$, for $N \ge2$.

For the ordered task, the final summary retains $C$ endpoint associations instead of all $Cn$ edges. Hence, Proposition 6 suggests that the SSM state for a sequential model could be smaller than that of a souping model to learn a task with a fixed number of chains $C$ and hops $n$. Indeed, to learn $Cn$ associations, a souping model would require $B = Cn \log_{2} N$ bits, whereas a sequential model needs a factor of $n$ fewer with just $B = C \log_{2} N$ bits.

More generally, sequential encoding may permit useful composed representations in deeper layers. The doubling construction is similar in spirit to the pointer-doubling strategy of Sanford et al. [35], and gives a logarithmic-depth sequential solution in this ordered setting. We conjecture that this may yield a depth separation from souping under suitable width and precision constraints.

Weighted state souping permits selective retrieval

While souping may be sensitive to chunk boundaries and exhibit less efficiency for favorably ordered inputs, weighted souping nonetheless retains a decisive advantage over sequential models for long-context retrieval. Poor long-context performance can be understood as a consequence of gating and noisy contexts. In a sequential model with Mamba-like history decay, a contribution at distance $\Delta$ is exponentially decayed by a factor $\alpha^{\Delta}$. Tong et al. [44] have formalized the resulting tradeoff: reducing decay preserves distant signals but permits distractor variance to accumulate, whereas stronger decay controls that variance at the cost of attenuating earlier tokens. Hence, no setting permits a sequential model to retrieve from every token in a long context with even reliability.

Weighted souping bypasses this tradeoff by permitting direct attention on the relevant chunks. In this scheme, distractor variance is determined only by the selected set of chunks rather than the full context size. Earlier chunks are not attenuated compared with later chunks, so weighting is computed evenly across the whole context. Provided that the weights remain concentrated as the number of distractors increases, weighted souping may maintain strong long-context performance when sequential and mean-souped variants fail. In some sense, weighted souping interpolates between recurrent accumulation and attention-like selection. It is not equivalent to full attention (souping weights are scalar, chunk-level, and computed once from the original query). Nonetheless, our experiments support the strong retrieval performance of weighted souping models. Together with their efficiency benefits, weighted state souping presents an attractive paradigm for future models.

Discussion

Weighted state souping augments reusable context encoding with query-dependent selection. A lightweight routing head combines cached chunk states into a fixed-size recurrent initialization, improving long-context performance across document QA, retrieval, variable tracking, and tool use benchmarks. The results show that a fixed state need not imply uniformly compressing the entire available collection: each query determines which information enters the working state. Inference-time sparsification often improves both accuracy and state-transfer cost, particularly beyond the training pool size. The theory complements these results by demonstrating retrieval and cross-chunk composition, while identifying information lost at chunk boundaries and discussing the potential advantages of sequential processing.

There remain many avenues for further exploration. Our focus in these experiments has been on single-turn interactions, but generalizing souping to multi-turn interactions featuring several rounds of back-and-forth between user and assistant remains an unexplored direction. There are also many opportunities for improving weighted state souping, for instance by replacing fixed-$N$ sparse souping with adaptive sparsity that adjusts to the problem, exploring alternative weighting schemes, or using thinking tokens in conjunction with souping. Our theoretical discussion offers initial intuitions for why state souping works and when it may fail, but there remain many opportunities for investigating potential memory and depth separations between souping and sequential models.

Acknowledgments

This work was performed while WLT was at Apple. WLT is supported by a Kempner Graduate Fellowship. CP is supported by an NSF CAREER Award (IIS-2239780), DARPA grants DIAL-FP-038 and AIQ-HR00112520041, the Simons Collaboration on the Physics of Learning and Neural Computation, and the William F. Milton Fund from Harvard University. This work has been made possible in part by a gift from the Chan Zuckerberg Initiative Foundation to establish the Kempner Institute for the Study of Natural and Artificial Intelligence. LLMs were used in implementing the experiments, assisting with proofs, and polishing the writing.

Additional figures

Constructions

We use the simplified SSM model in § 6 throughout. In the constructions below, each block matrix column has width $d$, in the order $[\mathbf{e}(z_{t-1});\mathbf{u}_{t-1}^{(\ell-1)};\mathbf{e}(z_t);\mathbf{u}_{t}^{(\ell-1)}]$. At a reset boundary, the previous residual vector is zero. Recurrent states are read after the current token’s write, as specified in (2).

One-hop retrieval

Proof. Let $E_K=[\mathbf{e}(k_1)\ \cdots\ \mathbf{e}(k_N)]$ and $E_V=[\mathbf{e}(v_1)\ \cdots\ \mathbf{e}(v_N)]$. Orthogonality gives $E_K^\top E_K=E_V^\top E_V=I_N$ and $E_K^\top E_V=0$. Define $A_K=E_KE_K^\top$ and $A_V=E_VE_V^\top$, and use depth $D=1$. The four block columns correspond to $[\mathbf{e}(z_{t-1});\mathbf{u}_{t-1};\mathbf{e}(z_t);\mathbf{u}_t]$. Choose

Together with the stated embeddings, convolution, residual map $W_o$, zero initial state, zero biases, and any fixed $\alpha\in(0,1]$, these choices specify all model parameters. We suppress the layer superscript below.

At a value token $v_{f(a)}$, the previous token is $k_a$, so the state increment is $\mathbf{e}(k_a)\mathbf{e}(v_{f(a)})^\top$. At a key token, including the query, $W_v\mathbf{x}_t=0$, so there is no write. At the query the state is therefore

For sequential encoding, $\gamma_a=\alpha^{T-t_a}$, where $t_a$ is the position of the value token and $T$ is the query position. For souping, if this pair lies in chunk $i(a)$ and its value has local position $t_{a,i(a)}$, then $\gamma_a=w_{i(a)}\alpha^{m_{i(a)}-t_{a,i(a)}+1}$. These coefficients include the decay at the query token. Orthogonality now gives $\mathbf{h}_q=\gamma_q\mathbf{e}(v_{f(q)})$. The residual update gives $\mathbf{u}_q^{(1)}=\mathbf{h}_q$ because $\mathbf{u}_q^{(0)}=0$, so $U\mathbf{u}_q^{(1)}$ has its only nonzero entry, $\gamma_q>0$, at coordinate $f(q)$. Argmax returns the correct value. In the sequential model, continuous convolution across pair boundaries gives the same writes, since keys have zero value projection.

Multi-hop retrieval with additive residuals

Proof. Use the embedding matrices and $A_K,A_V$ from the one-hop proof, and let $F=E_KE_V^\top$, so $F\mathbf{e}(v_a)=\mathbf{e}(k_a)$. The four block columns below correspond to $[\mathbf{e}(z_{t-1});\mathbf{u}_{t-1}^{(\ell-1)};\mathbf{e}(z_t);\mathbf{u}_{t}^{(\ell-1)}]$. In the first layer, use

Since $\mathbf{u}_t^{(0)}=0$, this is the one-hop construction. Every context output is zero: a context key has not yet been written, and a value has zero query projection. At the final query, $\mathbf{u}_q^{(1)}=\gamma_q\mathbf{e}(v_{f(q)})$.

Subsequent layers must add the next retrieval and subtract the previous working vector. This can be implemented by the attention block itself. Choose a fixed unit vector $\mathbf{b}=\mathbf{e}(v_1)$, orthogonal to all key embeddings, and let $\mathbf{c}=\sum_{a=1}^{N}\mathbf{e}(k_a)$ and $B=\mathbf{b}\mathbf{c}^\top$. Thus $B\mathbf{e}(k_a)=\mathbf{b}$ and $B\mathbf{e}(v_a)=0$. For $\ell\ge2$, set

At context positions the working vectors remain zero: a value writes the original association and has zero query projection; a key has zero value projection and queries $\mathbf{b}$, orthogonal to all association keys. Consequently, every layer’s context state contains precisely the original associations, for both sequential encoding and independently encoded chunks.

At the final query, write $\mathbf{u}=\mathbf{u}_{q}^{(\ell-1)}$ and let

be the context contribution after multiplying the initialized context state by $\alpha$ when processing the query token, with the same coefficients as in the one-hop proof. The query’s key, value, and query projections are respectively $\mathbf{b}$, $-\mathbf{u}$, and $F\mathbf{u}+\mathbf{b}$. Since the recurrence reads after the current write,

Here $M^{\top}\mathbf{b}=0$, $\mathbf{b}^{\top}F\mathbf{u}=0$, and $\lVert\mathbf{b}\rVert=1$. The additive residual therefore gives $\mathbf{u}_{q}^{(\ell)}=M^{\top}F\mathbf{u}_{q}^{(\ell-1)}$. Inductively,

Use $U=E_{V}^{\top}$ and the same $\alpha$ at every layer. Every factor is positive, so the final readout identifies $v_{f_{n}(q)}$. $\square$

Ordered-chain construction

We give the explicit construction for Remark 5. Let $n=2^{L}$, and let the $C$ chains in (4) have mutually disjoint node sets, so $C(n+1)\le N$. Set $\alpha=1$ and use the orthogonal embeddings, residual path, and convolution from § 6. Convolution continues across pair boundaries in the sequential model. All recurrent states start at zero and all biases are zero. Write $A_{K}=E_{K}E_{K}^{\top}$, $A_{V}=E_{V}E_{V}^{\top}$, and $F'=E_{K}E_{V}^{\top}$, as in the retrieval proofs.

The first $L$ layers store reverse associations. Using the same block columns as in the proof of Proposition 4, set

For $2\le\ell\le L$, let $G=F^{\top}$ and set

Working vectors at value tokens remain zero, whereas those at key tokens lie in the key-embedding span. In the first layer, each value token writes its reverse one-hop association using the preceding key’s original embedding. For $\ell\ge2$, a value token $v_{a,c,j}$ writes the reverse association with key $e(k_{a,c,j})$ and value equal to the preceding key’s working vector. A key token $k_a$ instead writes $-e(v_a)(\mathbf{u}_{t}^{(\ell-1)})^{\top}$ and queries $\mathbf{u}_{t}^{(\ell-1)}+e(v_a)$. The value-embedding direction $e(v_a)$ is orthogonal to all reverse-association keys and to the cancellation keys at earlier context positions, since each context key is unique. Its current write therefore subtracts exactly $\mathbf{u}_{t}^{(\ell-1)}$ from the attention output. After residual addition, the working vector contains only the newly retrieved reverse jump. These fixed key and value spans suffice at every layer.

Index edges by $j=1,\ldots,n$, so edge $j$ is $(k_{a,c,j-1},v_{a,c,j})$. Let $\mathbf{u}_{c,j}^{(\ell)}$ denote the working vector at the key token of edge $j$ in chain $c$, for $\ell\ge1$. We claim

For the base case $\ell=1$, the key token of edge $j$ is $k_{a,c,j-1}$. For $j\ge2$, the preceding edge has already written the reverse association from $k_{a,c,j-1}$ to $k_{a,c,j-2}$; for $j=1$ no such association exists. This proves the stated base case. For $\ell\ge2$, suppose the claim holds for layer $\ell- 1$. Put $s = 2^{\ell-1}$. At the value token of edge $j$, layer $\ell$ writes

The component of the state addressed by key embeddings therefore stores reverse jumps of length $s$. At the key token of edge $j$, the corresponding query component is $\mathrm{e}(k_{ac,j-s})$ when $j \ge s$. If $j \ge2s$, edge $j-s$ has already been processed, so its reverse association returns $\mathrm{e}(k_{ac,j-2s})$. For $s \le j < 2s$, no matching association has been written; for $j < s$, this query component is zero. Both cases give zero retrieval. In layers $\ell\ge2$, the additional query component cancels the incoming working vector as described above. Orthogonality and disjoint chain nodes exclude all other matches, proving (7) after residual addition.

For $L \ge1$, the final layer writes forward endpoint associations. Set

By (7), the preceding key’s working vector is nonzero only for $j \ge2^L = n$. Thus only the final edge of each chain writes to this layer, producing

At a chain-start query, every reverse layer returns zero: no reverse association points back from a chain start, and the earlier occurrence of this start wrote no cancellation term. Hence the query enters the final layer with zero working vector, writes zero, and uses its original key embedding to retrieve the appropriate endpoint. The readout $U$ identifies that value. For example, four-hop retrieval uses two reverse layers, storing jumps of lengths one and two, followed by a final layer associating each chain start with its four-hop endpoint.

Implementation details

Tasks

Across the document-QA and tool-use benchmarks we probe long-context robustness with a shared pooling construction. Each example carries one or more gold items — the documents or tool definitions required to answer the query — which we always retain; we then add sampled distractors until the pool holds $k$ items in total and shuffle so the gold position is randomized. Thus $k$ counts the total number of in-context chunks (not the number of distractors), and $k = 1$ denotes a gold-only pool. Distractors are drawn without replacement from a split-local pool that excludes the example’s own gold items, so training and evaluation distractors never cross split boundaries. During finetuning the pool size is sampled per example as $k \sim\operatorname{Uniform}[1, k_{\max}]$; at evaluation we sweep $k$ well beyond $k_{\max}$ to test generalization. All tasks share the Nemotron-Nano tokenizer ([26]) (vocabulary size 131,072) and are evaluated on the public validation split (the official test splits are hidden). RULER is handled differently, as described below.

Berkeley Function Calling Leaderboard (BFCL). A standard tool-use benchmark ([30]): the model is prompted with a list of tool definitions and must emit a tool call that solves the user query. Before evaluation we finetune each model for tool calling on a filtered subset of NVIDIA’s Nemotron agentic SFT data ([26]), keeping the $\approx228{,}700$ examples whose first assistant turn is a tool call. We evaluate on BFCL, reporting the simple python (one function, one call; 400 examples), parallel (one function, multiple calls; 200 examples), and parallel multiple (multiple functions, multiple calls; 200 examples) splits. Tool pools are built with a global BM25 retriever ([33]) over the $\approx14{,}700$-tool non-live vocabulary: candidate tools are scored by BM25, truncated to the top $k$, and the gold tool(s) are then force-included and randomly positioned. Training samples $k \sim\operatorname{Uniform}[1, 24]$; the headline evaluation sweeps $k \in\{1, 12, 24, 36, 48, 64, 96, 128\}$. Correctness is scored by BFCL’s abstract-syntax-tree checker, which validates the function name and every argument value against the reference-answer set. We additionally report an order-independent function-name-only match.

Stanford Question Answering Dataset (SQuAD). A single-document extractive QA dataset ([32]): given a Wikipedia paragraph, the model quotes a text span answering the query. We use SQuAD v1.1 and draw distractor paragraphs from the same split. Training uses $k \sim\operatorname{Uniform}[1, 24]$; evaluation sweeps $k \in\{1, 12, 24, 36, 48, 64, 96, 128\}$ over a fixed subset of 500 validation questions (shared across all arms via a common seed). We report exact match (EM) accuracy under SQuAD normalization (lowercasing, and removal of punctuation, articles, and extra whitespace). Dotted lines in the figures report closed-book accuracy ($k = 0$, no documents in context) as a memorization floor.

HotpotQA. A multi-hop QA dataset ([52]). Each context paragraph becomes one document, and the annotated supporting paragraphs are marked gold (a multi-gold example). Pooling, evaluation $k$-sweep, example count, and metrics match SQuAD.

MuSiQue. A multi-hop QA dataset engineered to require genuine reasoning over many hops and to resist single-document shortcuts ([45]); it is a harder version of HotpotQA. We use the answerable variant (musique_ans), treating the annotated supporting paragraphs as gold and scoring against the gold answer and its aliases. Pooling, evaluation $k$-sweep, example count, and metrics match SQuAD.

SQuAD-$n$. To further stress multi-document attention, we construct SQuAD-n: we merge $n$ SQuAD examples with distinct gold paragraphs into a single example whose prompt contains all $n$ gold paragraphs (mixed among $\max(k-n, 0)$ distractors) followed by the $n$ questions, labeled (1) … (n); the model must return the $n$ answer spans in order. We use $n \in\{3, 5\}$. We report micro-averaged EM over all (example, sub-question) pairs, and a strict joint EM that credits an example only when every sub-answer is an exact match.

RULER. A synthetic long-context benchmark ([13]). We use an in-repository reimplementation of its tasks and focus on three subtasks: single-needle-in-a-haystack (a magic-number “needle” line hidden among Paul Graham essays), multi-key NIAH (the needle is hidden among confusable distractor needles), and variable tracking (chains of variable assignments embedded in filler text, where the model resolves the value of a queried variable by following a four-hop chain). Unlike the QA and tool-use tasks, RULER has no natural chunks and no distractor pool; following Jafari et al. [15] we chunk each input into consecutive 256-token blocks, so the “pool size” is $\lceil\text{context length}/256\rceil$ and grows with the context. A block is treated as gold if its token span overlaps a needle sentence (NIAH) or an assignment in the queried chain (variable tracking). We cap the training context at 2048 tokens (8 chunks) and evaluate at context lengths $\{1, 2, 4, 8, 16, 32, 64, 128\}$K tokens over 200 examples per (subtask, length).

Models

Architectures. Our SSM is a 1.5B-parameter Mamba-2 model (Dao & Gu [5]): 32 Mamba-2 mixer layers (with no interleaved attention or MLP blocks), hidden size 2048, RMSNorm, and no positional embedding. Each mixer uses a state-expansion factor of 2 (inner dimension 4096) partitioned into 64 heads of head dimension 64, state dimension 128, 8 groups, and a convolution kernel of width 4. The recurrent state is therefore $4096 \times128 = 524{,}288$ elements per layer, or $\approx16.8$M elements (33.6 MB in bf16) across all 32 layers—this is the size of a single cached state encoding. Our Transformer is a 1.5B-parameter LLaMA-style model: 20 layers, hidden size 2048, SwiGLU feed-forward width 5504, 16 attention heads with full multi-head attention, rotary position embeddings, and RMSNorm. The two models are matched in parameter count ($\approx1.5$B). Both are pretrained on the Nemotron-CC corpus (Su et al. [40]) (on the order of 300B tokens) with the same optimizer schedule (AdamW, peak learning rate $3 \times10^{-4}$, cosine decay), share the Nemotron-Nano tokenizer, and run in bf16. We deliberately give the Transformer rotary embeddings so that it can process contexts longer than those seen in training, so that the sequential-Transformer baseline is given a fair chance at length generalization.

Weighted-souping read-out. We augment the SSM with a lightweight read-out head that produces the key and query vectors of § 4. Concretely, we read the last-token residual activation at the final decoder block and pass it through two learned linear projections $W_k, W_q \in\mathbb{R}^{d \times d'}$ with $d = 2048$ and $d' = 128$ (equivalently, a single projection to $\mathbb{R}^{2d'}$ split into a key half and a query half). The two projections are separate parameters, initialized to identical values ($W_k = W_q$) so the head begins as the parameter-free cosine ranker over the base model’s activations; they are untied thereafter and trained independently. The head is the only component initialized from scratch; it is trained jointly with the base model but placed in its own optimizer group at learning rate base LR $\times$ lr_scale, with lr_scale = 10.

Cosine normalization. The weight assigned to chunk $i$ is $\sigma(\operatorname{score}_i)$, where the score is a dot-product between the chunk key $\mathbf{k}_i$ and the query vector $\mathbf{u}$. Rather than the usual $1/\sqrt{d}$ dot-product scaling, we $\ell_2$-normalize the key and query and scale their cosine similarity by a learnable inverse temperature $\gamma$: $\operatorname{score}_i = \gamma\frac{\mathbf{k}_i \cdot\mathbf{u}}{\lVert\mathbf{k}_i \rVert\lVert\mathbf{u} \rVert}$. We initialize $\gamma= 14$ and cap it at 100. Because the pre-nonlinearity logits are bounded to $[-\gamma,\gamma]$ regardless of pool size, the weighting stays well-conditioned as $k$ grows, which we found important for generalizing to pool sizes far beyond the training distribution. Our headline weighting function $\sigma$ is the softmax.

Sparse souping. Sparse souping keeps only the top-$N$ weights and zeroes the rest; with softmax we apply the mask before the nonlinearity so the surviving weights still sum to one. It is applied at inference time and reuses the softmax checkpoints unchanged. We use $N_{\mathrm{top}} = 4$ for the document-QA and BFCL figures and $N_{\mathrm{top}} = 20$ for RULER; the larger $N_{\mathrm{top}}$ is needed for variable tracking, where gold information is spread across many chunks (we sweep $N_{\mathrm{top}}$ in Figure 5).

Auxiliary routing loss. We optionally supervise the weights toward the known gold chunks with an auxiliary loss, $\mathcal{L} = \mathcal{L}_{\mathrm{SFT}} + \lambda\mathcal{L}_{\mathrm{aux}}$, where $\mathcal{L}_{\mathrm{aux}}$ is a cross-entropy over the $k$ chunks under softmax weighting, computed from the same scores used to combine the states. We set $\lambda= 1$ for RULER and tool use, where the gold chunks are reliably identifiable, and $\lambda= 0$ for document QA, where the head learns the weighting purely from the downstream SFT signal.

Souping finetuning. Finetuning is necessary to learn the souping format (Jafari et al., [15]). All souping models are finetuned from the pretrained base with the softmax weighting, residual read-out, and cosine normalization above. For document QA we train for 3 epochs (SQuAD, HotpotQA) or 5 epochs (MuSiQue) at base learning rate $5 \times10^{-6}$ with no auxiliary loss, from the train splits included with the benchmarks. For RULER we train for 6000 steps at base learning rate $10^{-5}$ with $\lambda= 1$, chunked at 256 tokens. For tool use we warm-start from a tool-calling-SFT checkpoint already finetuned on the Nemotron Agentic SFT split [26], and train for an additional one epoch at base learning rate $10^{-5}$ with $\lambda= 1$ to learn the souping format. In every case the read-out head uses the 10× learning-rate multiplier described above. The sequential baselines are finetuned on the concatenated input for matched steps at learning rate $10^{-5}$ (QA/RULER) or $5 \times10^{-5}$ (tool use), so that all models see matched data and compute.

Decoding and evaluation. All models decode greedily. For tool use, decoding is constrained by a tokenizer-level logit mask that enforces the tool-call envelope (<TOOLCALL>[ ... ]</TOOLCALL>); the constraint is applied identically to the souping arms and the sequential baselines. The QA and RULER tasks decode without constraints. Evaluation pool sizes, context lengths, and example counts are as given in Appendix § 11.§ 1.

Computational benchmark

The timings in Figure 3 are measured on a single AWS p5.48xlarge node ($8 \times$ H100). We compare four inference paths: a sequential Transformer, a sequential Mamba, full weighted souping, and sparse souping. For the souping paths, we time the encoding of a 128-token query and the computation of the softmax weights, while the chunk state encodings themselves are retrieved rather than recomputed, from one of three memory tiers: GPU high-bandwidth memory (HBM), system RAM, or system disk. Each cached Mamba-2 encoding is the 33.6 MB bf16 state described above, so retrieval cost is dominated by the bandwidth of the tier in which the encodings live; retrieval from HBM is far faster than from RAM, which in turn is faster than from disk, and in the worst case (short chunks retrieved from disk) retrieval can be slower than re-encoding.