Weighted State Souping for Efficient Long-Context Retrieval

Researchers William L. Tong, Aryo Lotfi, Emmanuel Abbe, Josh Susskind, Cengiz Pehlevan and Eran Malach proposed weighted state souping, a method that learns query-dependent weights for composing cached state-space-model encodings, in a paper accepted at the PALM workshop at NeurIPS 2026. The authors report the approach substantially improves long-context performance over state souping and vanilla SSMs across single- and multi-hop document question answering, needle-in-a-haystack retrieval, variable tracking and tool use, and a sparse variant yields near-constant encoding cost with respect to context length while preserving or sometimes improving performance. The paper also offers theoretical arguments for why souping works and when it may fail, aimed at agentic workloads built on large collections of documents, tools and memory records.

content type paper https://machinelearning.apple.com/research/ published November 2026 Weighted State Souping for Efficient Long-Context Retrieval AuthorsWilliam L. Tong† , Aryo Lotfi, Emmanuel Abbe, Josh Susskind, Cengiz Pehlevan†, Eran Malach This paper was accepted at the Workshop on Personalized, Aligned, Long-term Memory for AI Systems PALM at NeurIPS 2026. 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. Recursive Language Models Meet Uncertainty: The Surprising Effectiveness of Self-Reflective Program Search for Long Context July 9, 2026 research area Speech and Natural Language Processing https://machinelearning.apple.com/research/?domain=Speech%20and%20Natural%20Language%20Processing Long-context handling remains a core challenge for language models: even with extended context windows, models often fail to reliably extract, reason over, and use the information across long contexts. Recent works like Recursive Language Models RLMs have approached this challenge by agentic way of decomposing long contexts into recursive sub-queries through programmatic interaction at inference. While promising, the success of RLMs critically… To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models March 27, 2026 research area Methods and Algorithms https://machinelearning.apple.com/research/?domain=Methods%20and%20Algorithms , research area Tools, Platforms, Frameworks https://machinelearning.apple.com/research/?domain=Tools%2C%20Platforms%2C%20Frameworks conference ICLR https://machinelearning.apple.com/research/?event=ICLR State Space Models SSMs have become the leading alternative to Transformers for sequence modeling. Their primary advantage is efficiency in long-context and long-form generation, enabled by fixed-size memory and linear scaling of computational complexity. We begin this work by showing a simple theoretical result stating that SSMs cannot accurately solve any “truly long-form” generation problem in a sense we formally define , undermining…