Studying the Soupability of Documents in State Space Models
Yasaman Jafari · Zixian Wang · Leon Bergen · Taylor Berg-Kirkpatrick
We empirically investigate document souping, a distributed document-processing approach that merges independently computed hidden states from fine-tuned Mamba models.
摘要
We investigate whether hidden states from Structured State Space Models (SSMs) can be merged post hoc to support downstream reasoning. Inspired by model souping, we study document souping, a strategy where documents are encoded independently, and their representations are pooled, via simple operations like averaging, into a single context state. This approach enables modular encoding and reuse without reprocessing the full input for each query. We demonstrate that finetuned Mamba2 models with souped representations achieve competitive or superior performance across multi-hop QA, sparse retrieval, and long-document reasoning tasks compared to the standard monolithic encoding approach. For example, on the RACE and QuALITY benchmarks for long document question answering, this method substantially outperforms a traditional concatenation approach. Crucially, this modular design scales to hundreds of segments (we test up to 256), while delivering substantial savings in inference cost, unlocking new possibilities for large-scale corpus reasoning.