The Surprising Soupability of Documents in State Space Models

📄 arXiv: 2505.24033v1 📥 PDF

作者: Yasaman Jafari, Zixian Wang, Leon Bergen, Taylor Berg-Kirkpatrick

分类: cs.CL, cs.CE, cs.LG

发布日期: 2025-05-29

💡 一句话要点

提出文档混合（Document Souping）方法，提升状态空间模型在长文档推理中的性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 状态空间模型 文档混合 长文档推理 多跳问答 Mamba2 模型混合 模块化编码

📋 核心要点

  1. 现有方法在处理长文档时，需要对整个文档进行重复处理，效率较低，难以支持模块化编码和重用。
  2. 论文提出文档混合（Document Souping）方法，独立编码文档，然后将它们的表示池化成一个单一的上下文状态。
  3. 实验表明，该方法在多跳问答、稀疏检索和长文档推理任务上表现出色，在HotpotQA数据集上接近交叉编码器的性能。

📝 摘要（中文）

本文研究了结构化状态空间模型（SSM）中的隐藏状态是否可以在后处理阶段进行合并，以支持下游推理。受到模型混合（model souping）的启发，我们提出了一种策略，即独立编码文档，并通过简单的操作（如平均）将它们的表示池化成一个单一的上下文状态。这种方法称为文档混合（document souping），能够实现模块化编码和重用，而无需为每个查询重新处理完整的输入。我们对Mamba2模型进行微调，使其产生可混合的表示，并发现它们能够以强大的准确性支持多跳问答、稀疏检索和长文档推理。在HotpotQA上，混合十个独立编码的文档几乎与在相同输入上训练的交叉编码器的性能相匹配。

🔬 方法详解

问题定义：论文旨在解决长文档处理中，现有方法需要重复处理整个文档，导致效率低下的问题。现有方法难以支持模块化编码和重用，限制了模型在复杂推理任务中的应用。

核心思路：论文的核心思路是受到模型混合（Model Souping）的启发，提出文档混合（Document Souping）方法。该方法将多个文档独立编码，然后将它们的隐藏状态进行池化（例如，通过平均），形成一个统一的上下文表示。这样，模型可以独立处理每个文档，然后将它们的信息进行整合，而无需重新处理整个输入。

技术框架：整体框架包含以下步骤：1. 使用结构化状态空间模型（SSM），例如Mamba2，独立编码每个文档。2. 将每个文档的隐藏状态提取出来。3. 使用池化操作（例如，平均）将所有文档的隐藏状态合并成一个单一的上下文状态。4. 使用该上下文状态进行下游任务，例如多跳问答、稀疏检索和长文档推理。

关键创新：最重要的技术创新点在于将模型混合的思想应用到文档层面，提出了文档混合（Document Souping）方法。与传统的端到端训练方法相比，该方法能够实现模块化编码和重用，提高了模型的效率和灵活性。此外，该方法允许模型独立处理每个文档，从而更好地捕捉文档之间的关系。

关键设计：论文使用Mamba2模型作为基础编码器，并对其进行微调，使其产生可混合的表示。池化操作采用简单的平均方法，但也可以尝试其他池化方法，例如加权平均或注意力机制。损失函数的设计需要考虑下游任务的具体要求，例如，在多跳问答任务中，可以使用交叉熵损失函数来训练模型预测答案。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，文档混合方法在多跳问答、稀疏检索和长文档推理任务上表现出色。在HotpotQA数据集上，混合十个独立编码的文档几乎与在相同输入上训练的交叉编码器的性能相匹配，证明了该方法的有效性。该方法为长文档处理提供了一种新的思路。

🎯 应用场景

该研究成果可应用于需要处理大量文档的场景，例如信息检索、知识图谱构建、智能问答系统等。通过文档混合，可以高效地整合多个文档的信息，提高系统的推理能力和响应速度。未来，该方法有望在法律、金融、医疗等领域发挥重要作用。

📄 摘要（原文）

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 propose a strategy where documents are encoded independently and their representations are pooled -- via simple operations like averaging -- into a single context state. This approach, which we call document souping, enables modular encoding and reuse without reprocessing the full input for each query. We finetune Mamba2 models to produce soupable representations and find that they support multi-hop QA, sparse retrieval, and long-document reasoning with strong accuracy. On HotpotQA, souping ten independently encoded documents nearly matches the performance of a cross-encoder trained on the same inputs.