Weighted State Souping for Efficient Long-Context Retrieval

William L. Tong; Aryo Lotfi; Emmanuel Abbe; Josh Susskind; Cengiz Pehlevan; Eran Malach

cs.LGOct 19, 2026

Weighted State Souping for Efficient Long-Context Retrieval

Authors: William L. Tong, Aryo Lotfi, Emmanuel Abbe, Josh Susskind, Cengiz Pehlevan, Eran Malach

Organizations: Harvard University · Apple

Abstract

State space models (SSMs) combine fast inference and fixed state size with competitive language-modeling performance. Recently, state souping has been proposed as a procedure whereby an SSM independently encodes chunks of tokens like documents, tool definitions, or memory records. The resulting state encodings may be cached and recombined to cheaply construct a context for downstream queries. However, both state souping and vanilla SSMs deteriorate on long, crowded contexts, limiting their utility for long-horizon domains. To address this issue, we propose weighted state souping, a simple method that learns query-dependent weights for composing cached states. The resulting models improve long-context performance substantially across both single- and multi-hop document question answering, needle-in-a-haystack retrieval, variable tracking, and tool use. Our method can be further adapted into an efficient sparse variant, yielding near-constant encoding cost with respect to context length while preserving (and sometimes improving) performance. We also present simple theoretical arguments that illustrate why souping works and when it may fail, and discuss informally the advantages of weighted souping. Weighted state souping thus combines reusable context representations with selective retrieval, offering an efficient approach to language modeling over large collections of documents, tools, and other forms of structured context that increasingly dominate modern agentic workloads.

Figures & tables

Figure 1: Sequential encoding versus state souping. (a) A typical SSM operates by concatenating input chunks together with the user query , and performs a single, long autoregressive pass to produce the answer . (b) A state souping SSM encodes each input chunk independently, then soups the resulting state encodings into a single encoding used to initialize the model as it processes the user query . (c) In our weighted state souping method, we augment the SSM with an additional linear readout that produces key vectors from the input chunks and a query vector from the user query . The soup is constructed as a weighted sum with weighting function

\sigma

Figure 2: Weighted state souping benchmark results. Vertical gray dashed lines indicate the max pool size or context length observed during training. (a - c) Single-hop (SQuAD) and multi-hop (HotpotQA, MuSiQue) benchmarks. We plot accuracy as an exact match between the gold answer span and the model prediction. Dotted lines represent closed book accuracy, where no documents are prompted in-context. Sparse souping keeps top

N_{top}=4

chunks. (d,e) NIAH subsets of RULER. In the easy version, a single needle consisting of a random number is hidden among a natural language haystack of variable size. In the hard version, the haystack consists entirely of needle-like distractors, requiring precise retrieval. Sparse souping keeps top

N_{top}=20

chunks. (e) Variable tracking subset of RULER. The input consists of random variable assignments embedded in random noise sentences, and the model must infer the correct value of a randomly chosen variable. Sparse souping keeps top

N_{top}=20

chunks. (f) Simple Python subset of BFCL. The model is prompted with a list of Python functions and descriptions, and must execute a correct function call given a user query. Sparse souping keeps top

N_{top}=4

chunks. (g) Parallel subset of BFCL. The setup is the same as in (f) , but the model needs to execute the function multiple times.

Figure 3: Computational cost of souping. We measure the time it takes to encode a set of input chunks on an AWS H100 node. For the souping models, we include the time it takes to encode a 128 token query and compute the softmax weights. The state encodings are retrieved from either GPU high bandwidth memory (HBM), system RAM, or system disk. We plot a subset of timings here, with the full results in Figure 7 . We include two sparse souping variants, one that keeps the top

N_{top}=4

documents and another that keeps the top

N_{top}=32

. ( a ) We hold the chunk length fixed at

L=512

tokens and increase the total number of chunks. Encoding time therefore grows linearly with increasing chunks for all models. For the HBM model, at low

K

encoding the query dominates the encoding time, resulting in a flat trend. ( b ) We hold the number of chunks fixed at

K=8

and increase the length of each chunk. Since the number of chunks is constant, souping retrieves the same number of chunks for each query, resulting in a flat trend across all

L

where the height is controlled by retrieval speed.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Squad-n benchmark results. As a further test of souping when the task requires attending to multiple documents, we construct a task based on SQuAD called SQuAD-n . Each example consists of

n

individual SQuAD document-question-answer sets. The model is prompted with all

n

documents (mixed among

\text{max}(k-n,0)

distractors) followed by the questions, and must produce the answer spans corresponding to each question. We plot both per-question exact match accuracy in ( a, c ) as well as joint accuracy in ( b, d ) for

n=3

and

n=5

. Souping generally performs competitively with vanilla Mamba, and sparse souping in particular maintains long context generalization better than full souping. However, both Mamba and souping underperform Transformers on in-context

K

in this setting. This may be in part due to the Transformer’s sharper ability to track its context, and determine which questions have already been answered and which require an answer still.

Figure 5: Ruler benchmark, additional results. In the main text, we present results on needle-in-a-haystack (a,b) and variable tracking (c) . We additionally plot a sweep over the number of chunks

N_{top}

kept by the sparse souping model in (d - f) . For the NIAH tasks, where only a single chunk is relevant, the value of

N_{top}

is irrelevant. However, for the variable tracking task, relevant variable assignments are distributed across many more chunks, so below a certain threshold, the model fails to learn the task well.

Figure 6: BFCL benchmark, additional results. In the main text, we present results on the Python Simple split of BFCL ( a ), in which the model must produce a single tool call given a list of tool definitions and a user query. Here, we also plot performance on the Parallel split (the model must call a single tool multiple times) and the Parallel Multiple split (the model must call multiple tools multiple times). We measure accuracy by BFCL’s built-in abstract syntax tree (AST)-based checker ( a - c ) as well as a match based on function name alone ( d - f ). In general, weighted souping performs well for both Python Simple and Parallel, but a substantial gap opens on small, in-distribution

K

for Parallel Multiple. Mean souping performs better than weighted souping at

K=1

(where only gold tool definitions are prompted), indicating that souping may still work well, but the model is struggling to simultaneously learn appropriate weights and use a souped state initialization.

Figure 7: Additional measurements on encoding time across models. We plot additional sets of timings beyond those showin in Figure 3 . In particular, we show full timings measured for full and sparse souping models retrieving from HBM ( a, b ), system RAM ( c, d ), and system disk ( e, f ). The timings for Transformer and sequential Mamba models in the main text are computed with batch size of 1. To give a sense for what the compute-bound timings look like, we plot the per-example timing for batch size 32 for Transformer and sequential Mamba models (dashed lines).

State Space Language Models

See all 19 papers

Sessa: Selective State Space Attention
Apr 20, 2026Liubomyr HorbatkoAttention Mechanisms Self-Attention
Structured Memory for Edge Language Models: Persistent Context and Corpus Retrieval via O(1) SSM State Injection
Aug 3, 2026Anusha Madan Gopal, Aras Pirbadian, Kristofor D. Carlson +2Retrieval-Augmented Generation On-Device Language Model Inference
BrainChip Inc. 23041 Avenida de la Carlota, Laguna Hills, CA
Contextual Representation Flow Across Depth in State-Space and Transformer Architectures
Oct 8, 2025Nhat M. Hoang, Do Xuan Long, Cong-Duy Nguyen +3State Space Models State Space Language Models
Nanyang Technological University, Singapore · National University of Singapore

Apr 20, 2026J/K move · Enter open · S save

Sessa: Selective State Space Attention

Liubomyr Horbatko

Attention Mechanisms Self-Attention Long-Context Language Modeling Persistent Memory for Language Models State Space Models Recurrent Transformers+3

Open paper AskPDF arXiv

Modern sequence modeling is dominated by two families: Transformers, whose self-attention can access arbitrary elements of the visible sequence, and structured state-space models, which propagate information through an explicit recurrent state. These mechanisms face different limitations on long contexts: when attention is diffuse, the influence of individual tokens is diluted across the effective support, while recurrent state propagation can lose long-range sensitivity unless information is actively preserved. As a result, both mechanisms face challenges in preserving and selectively retrieving information over long contexts. We propose Sessa, a decoder that places attention inside a recurrent feedback path. This creates many attention-based paths through which past tokens can influence future states, rather than relying on a single attention read or a single recurrent chain. We prove that, under explicit assumptions and matched regimes, Sessa admits power-law memory tails

O(\ell^{-β})

for

0 < β< 1

, with slower decay than in the corresponding Transformer and Mamba-style baselines. We further give an explicit construction that achieves this power-law rate. Under the same assumptions, Sessa is the only model class among those considered that realizes flexible selective retrieval, including profiles whose influence does not decay with distance. Consistent with this theoretical advantage, across matched experiments, Sessa achieves the strongest performance on long-context benchmarks while remaining competitive with Transformer and Mamba-style baselines on short-context language modeling.

Abstract

State Space Language Models

Explore similar work

Sessa: Selective State Space Attention

Structured Memory for Edge Language Models: Persistent Context and Corpus Retrieval via O(1) SSM State Injection

Contextual Representation Flow Across Depth in State-Space and Transformer Architectures