💬 NLP

Studying the Soupability of Documents in State Space Models

This paper introduces "document souping," a modular strategy for Structured State Space Models (SSMs) that merges independently encoded document representations via simple operations like averaging, enabling scalable, cost-efficient long-document reasoning and retrieval that outperforms traditional monolithic encoding approaches.

Original authors: Yasaman Jafari, Zixian Wang, Leon Bergen, Taylor Berg-Kirkpatrick

Published 2026-02-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yasaman Jafari, Zixian Wang, Leon Bergen, Taylor Berg-Kirkpatrick

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a librarian trying to answer a complex question that requires reading through a massive library of books.

The Old Way (The "Monolithic" Approach):
Traditionally, if you wanted to answer a question using 10 different books, you would have to take all 10 books, rip out their pages, and staple them together into one giant, 10,000-page manuscript. Then, you would read this entire giant manuscript from start to finish to find the answer.

The Problem: If you just wanted to swap out one book for a different one, you'd have to rip out the pages again, re-staple the whole thing, and re-read the entire 10,000 pages. It's slow, wasteful, and inflexible.

The New Idea (The "Soup" Approach):
This paper introduces a new method called "Document Souping." Instead of stapling books together, imagine you have a magical blender.

Independent Blending: You take each book and run it through the blender individually. The blender turns the whole book into a single, concentrated "flavor packet" (a hidden state) that captures the essence of that book.
The Soup: When you get a question, you don't need to re-blend the books. You just grab the pre-made flavor packets for the specific books you need, dump them into a pot, and stir them together (this is the "pooling" or "soup" part).
The Result: You now have a single "soup" that contains the combined knowledge of those books, ready to answer your question immediately.

What the Paper Actually Found:

It Works with "Mamba" Models: The researchers tested this on a specific type of AI called Mamba2 (a Structured State Space Model). They found that these models are uniquely good at this. You can blend the "flavor packets" of 256 different documents together, and the AI can still understand the combined story to answer questions.
It Needs Training: You can't just take a pre-trained AI and start blending documents; it won't work well. The AI needs to be "finetuned" (trained specifically) to understand how to taste and mix these blended packets. Once trained, it gets very good at it.
The Best Recipe: The simplest way to mix the soup works best. Just averaging the flavor packets (adding them up and dividing by the number of books) is better than trying to be fancy with complex math.
Speed and Savings: This is the biggest win. Because you only blend each book once and save the "flavor packet," you can reuse it forever.
- Analogy: If you have a question about 10 books, the old way takes 10 hours to read. The new way takes 10 minutes to blend the 10 pre-made packets and 1 minute to answer. If you ask a new question about a different set of 10 books, you don't re-read; you just grab the saved packets and mix them.
Where It Fails:
- Transformers Don't Like It: The researchers tried this same "blending" trick on standard AI models (Transformers, like the ones behind many chatbots). It failed miserably. The "flavor packets" from standard models get confused when mixed. This suggests that the "soup" method only works because of the specific mathematical structure of Mamba models.
- Too Many Books at Once: If you train the AI on small batches (like 4 books) and then suddenly ask it to blend 256 books, it gets confused and fails. However, if you train it on larger batches first, it can learn to handle huge libraries.

In Summary:
The paper proves that for certain types of AI (Mamba), you can process documents separately, save their "essence," and mix them together later to answer complex questions. This is much faster and cheaper than reading everything together every time, but it requires specific training and only works with this specific type of AI architecture.

Technical Summary: Studying the Soupability of Documents in State Space Models

Problem Statement

Many real-world NLP tasks, such as multi-document question answering (QA), scientific summarization, and legal analysis, require reasoning over entire corpora rather than individual long documents. Current Transformer-based models struggle with this due to prohibitive $O(L^2)$ attention costs and the inflexibility of monolithic encoding strategies. In standard approaches, all documents are concatenated into a single sequence before processing; consequently, any modification to the corpus (e.g., adding or removing a single document) necessitates re-encoding the entire input from scratch. This prevents the modular reuse of document representations and creates a bottleneck for scalable, dynamic retrieval systems.

While Structured State Space Models (SSMs) like Mamba and Mamba2 offer linear-time sequence processing and compress inputs into fixed-length hidden states, they are typically deployed using the same monolithic concatenation strategy, inheriting the same re-encoding limitations.

Methodology: Document Souping

The authors propose Document Souping, a strategy inspired by model souping (averaging weights of finetuned checkpoints) but applied to hidden states. The core hypothesis is that the linear recurrence mechanisms in SSMs allow independently encoded document representations to be merged post hoc via simple linear operations (e.g., averaging) to form a composite state suitable for downstream reasoning.

The Architecture

Independent Encoding: Each document $d_i$ in a corpus is encoded independently by a shared SSM encoder ( $SSM_\theta$ ), producing a set of layer-wise hidden states $\{h^{(l)}_i\}$ .
Pooling: These per-document states are aggregated at each layer using a commutative pooling operator (e.g., element-wise average, sum, or max) to create a single "souped" state $h^{(l)}_{soup}$ .
Decoding: The pooled states $\{h^{(1)}_{soup}, \dots, h^{(L)}_{soup}\}$ are injected into the decoder alongside the query $q$ to generate the answer $\hat{y}$ .

Training and Optimization

Finetuning: The paper demonstrates that pretrained SSMs do not natively support souping; full encoder-decoder finetuning is required to teach the model to interpret pooled states effectively.
Memory Efficiency: To handle the memory overhead of backpropagating through multiple independent document encoders, the authors employ document-level activation checkpointing. This recomputes forward activations during the backward pass, allowing memory usage to remain constant regardless of the number of documents ( $k$ ).
Pooling Strategy: Extensive experiments identify element-wise averaging without normalization as the most robust and stable aggregation method, outperforming summation (which suffers from unbounded magnitude growth) and max pooling.

Key Contributions

Systematic Empirical Study: The first comprehensive investigation into "soupability"—the ability of SSMs to merge independently encoded document representations for multi-document reasoning.
Performance Validation: Demonstration that finetuned Mamba2 models with souped representations achieve competitive or superior performance on multi-hop QA (HotpotQA), long-document QA (RACE, QuALITY), and sparse retrieval (NIAH) compared to standard concatenation baselines.
Scalability and Generalization: Evidence that the approach scales to hundreds of documents (tested up to 256) and generalizes to varying context sizes, provided the model is finetuned on a sufficient number of segments.
Architectural Specificity: Proof that this modular approach is uniquely suited to SSMs. An analogous "cache souping" experiment on Transformers (LLaMA 3) failed significantly, highlighting that the fixed-size, recurrent state of SSMs is a critical enabler for this technique.
Inference Efficiency: A new inference workflow where documents are encoded once, cached, and dynamically reassembled, offering substantial latency savings (up to 14.6x speedup for 32k tokens) compared to re-encoding concatenated contexts.

Experimental Results

The authors evaluated Mamba2-8B and Mamba2-2.7B models across several benchmarks:

Multi-Hop QA (HotpotQA): Full encoder-decoder finetuned models with average pooling achieved 47.8 EM on 5 documents, outperforming decoder-only finetuning and matching/exceeding concatenation baselines. The method remained robust as the number of distractors increased.
Long-Document QA (RACE/QuALITY): Souping substantially outperformed traditional concatenation approaches on long-document benchmarks, particularly when the training configuration aligned with the test-time soup size.
Sparse Retrieval (NIAH): Models trained on larger numbers of segments (e.g., 8 segments) maintained high accuracy (71.45% EM) when tested on 16 segments, whereas models trained on fewer segments failed to generalize.
Scalability: While a model trained on 64 documents failed to extrapolate to 256, continued finetuning on 128 documents enabled the model to achieve 37.25 EM on 256 documents, proving the architecture can adapt to vast contexts.
Latency: For a 32k token context, processing a query using cached souped states took 65.8 ms, compared to 963.2 ms for the standard concatenation method.

Significance and Claims

The paper claims that document souping unlocks a new paradigm for modular, scalable corpus reasoning. By decoupling document encoding from query processing, the approach enables:

Efficient Updates: Documents can be encoded once and cached, allowing for dynamic corpus updates and retrieval without reprocessing the full input.
Scalable RAG: The method serves as a drop-in replacement for the "concatenate retrieved passages" step in Retrieval-Augmented Generation (RAG) systems, supporting variable but bounded numbers of documents efficiently.
Architectural Fit: The work establishes that the linear recurrence properties of SSMs make them uniquely capable of supporting state-level composition, a capability that standard Transformers lack due to their reliance on variable-length KV caches.

The authors conclude that while extrapolation to unseen soup sizes can be challenging, appropriate finetuning strategies allow SSMs to effectively handle large-scale, dynamic document sets, offering a lightweight and effective solution for long-context reasoning.