AUGMENTING ZERO-SHOT DENSE RETRIEVERS WITH PLUG-IN MIXTURE-OF-MEMORIES

问题 (Problem)

本文旨在解决稠密检索（Dense Retrieval）模型在**零样本（Zero-Shot）**场景下的泛化能力差的问题。

传统的稠密检索模型在一个大规模的源领域（如网页搜索）上训练后，直接迁移到新的、未见过的目标领域（如生物医药、金融）时，性能会显著下降。虽然通过不断增大语言模型的参数量可以提升泛化性，但这种方式成本高昂且收益递减，在经济上是不可持续的。

因此，核心问题是：如何不通过暴力增加模型参数，而是通过更高效的方式，提升稠密检索模型在不同领域间的零样本迁移和泛化能力？

方法 (Method)

作者提出了一种名为**混合记忆增强（Mixture-Of-Memory Augmentation, MoMA）**的机制，并将其应用于一个基于T5的强大检索器，构建了名为 MoMA-DR 的新系统。

其核心思想是：在对查询（Query）进行编码表示之前，先从一个由**多个不同信息源（语料库）组成的“混合记忆”**中检索出相关的“增强文档”，将这些文档的信息融入查询中，生成一个内容更丰富、意图更明确的“增强后查询表示”，然后再用这个增强后的查询去目标语料库中检索最终的答案。

模型架构

MoMA-DR 系统包含两个核心组件，如下图所示：

图1：MoMA-DR 模型架构示意图

增强组件 ( $f^a$ ): 这是一个独立的稠密检索器。它的任务是接收原始查询 $q$ ，然后从“混合记忆” $\mathcal{M}$ 中检索出 $K$ 个最相关的增强文档 $D^a = {d_1^a, ..., d_K^a}$ 。
主模型 ( $f^{MoMA}$ ): 这是最终用于检索的稠密检索器。它利用 Fusion-in-Decoder (FiD) 机制来融合原始查询和增强文档的信息。
- 首先，使用 T5 的编码器（Encoder）分别对原始查询 $q$ 和 $K$ 个增强文档 $d_i^a$ 进行编码。
- 然后，T5 的解码器（Decoder）同时关注（attend to）所有这些编码后的向量。
- 最后，解码器输出一个融合了所有信息的、增强后的查询向量 $q^a$ 。这个 $q^a$ 会被用来与目标语料库中的文档向量进行点积计算，以得到最终的排序。

数学公式解读

基础稠密检索 (Dense Retrieval)
模型的打分函数是查询和文档向量的点积：

$f(q,d)=q \cdot d$

其中向量由一个基于 Sentence-T5 的编码器 $g(\cdot)$ 生成，该编码器利用 T5 的 Encoder-Decoder 架构，将文本序列输入 Encoder，并取 Decoder 输出的 [CLS] 位置的表征作为最终的文本向量：

$g(x)=Dec(Enc(x))$
MoMA 的增强过程
1. 从混合记忆中检索增强文档：
  
  $D^{a} = ANN_{f^{a}(q,\circ)}^{\mathcal{M}} \quad ; \quad \mathcal{M}=\{C_{1},...,C_{M}\}$
  
  这里， $f^a$ 是增强组件的检索模型， $\mathcal{M}$ 是由多个语料库 $C_i$ 组成的混合记忆池， $ANN$ 代表高效的近似最近邻搜索。
2. 生成增强查询向量 (Fusion-in-Decoder):
  
  $q^a = g^{MoMA}(q) = Dec(Enc(q),Enc(d_{1}^{a}),...,Enc(d_{K}^{a}))$
  
  这个公式清晰地展示了 FiD 的过程：将查询和所有增强文档分别编码，然后由一个解码器统一处理，输出最终的查询向量 $q^a$ 。
3. 最终检索打分:
  
  $f^{MoMA}(q,d) = q^{a} \cdot d$
  
  使用增强后的查询向量 $q^a$ 和常规的文档向量 $d$ （由 $g(d)$ 生成）进行匹配。

联合学习机制 (Joint Learning)

如何让增强组件 $f^a$ 学会“挑选”有用的增强文档是整个方法的关键，因为我们并没有直接的标注数据。作者设计了一套巧妙的迭代式联合学习流程：

训练主模型 ( $f^{MoMA}$ ):
这一步相对直接。给定增强组件 $f^a$ 检索出的文档，通过标准的排序损失函数（带 ANCE 难负样本）来优化主模型。损失函数的目标是让 $f^{MoMA}(q^s, d^+)$ 的分数高于 $f^{MoMA}(q^s, d^-)$ 。
训练增强组件 ( $f^a$ ):
这是创新的核心。为了给 $f^a$ 提供监督信号，作者利用了主模型解码器的注意力分数作为“伪标签”或“软标签”。
1. 计算文档有用性: 解码器在生成 $q^a$ 时，对每个输入（原始查询和增强文档）的注意力权重总和 $FidAtt(d_i^a)$ ，可以被看作是该文档 $d_i^a$ 的“有用性”得分。
2. 构建正样本集 ( $D^{a+}$ ): 用于训练 $f^a$ $f^{a}$ 的正样本，由两部分组成：
  - 源任务中已知的相关文档 $D^{s+}$ 。
  - 在增强文档中，获得最高注意力分数的前 $N$ 个文档。
    $D^{a+}=D^{s+}\cup \text{Top-N}_{FidAtt(d^{a}),D^{a}}$
3. 构建负样本集 ( $D^{a-}$ ): 同样采用 ANCE 方式，但关键在于，负样本是从整个混合记忆池 $\mathcal{M}$ 中挖掘的。这迫使增强组件 $f^a$ 必须学会在不同风格和领域的语料中进行区分和选择，从而提升其泛化能力。
4. 迭代训练: 整个训练过程是迭代的，在每个 episode 中，先用上一轮的模型挖掘难负样本和增强文档，然后依次更新主模型和增强组件。

即插即用记忆 (Plug-in Memory)

MoMA 机制的一个巨大优势是其灵活性。在推理（测试）阶段，当模型面对一个新领域的任务时，我们可以直接将该任务的目标语料库 $C^t$ “即插即用”地加入到混合记忆池 $\mathcal{M}$ 中（即 $\mathcal{M} = \mathcal{M} \cup C^t$ ），而无需重新训练任何模型参数。

这样，增强组件 $f^a$ 就可以从目标语料库中直接检索到领域相关的上下文信息来增强查询，极大地提升了模型对新领域的适应能力。

图例解读

图 2: 展示了联合学习过程中，模型行为的变化。随着训练的进行（从 Epi-0 到 Epi-2），MoMA-DR 逐渐学会对源数据中的“相关文档”（Marco Rel）给予更高的注意力，并且更多地将它们检索为增强文档。这证明了联合学习机制有效地将最终任务的信号反馈给了增强组件。

图 3: 展示了“即插即用”机制的效果。在测试阶段，即使目标语料库是新加入的，模型依然会给予它最高的关注，并从中检索大量的增强文档。例如，在处理基于维基百科的任务（如 NQ, HotpotQA）时，模型会侧重从维基百科中获取信息；在处理医学任务（如 TREC-COVID）时，则会侧重于从新加入的目标语料库和 MeSH（医学知识库）中获取信息。这表明增强组件具备了出色的零样本泛化能力。

图 4: 显示了在训练过程中，增强组件和主模型（End Retriever）的性能（在源任务 MSMARCO 上的 MRR）都在稳步提升。这说明两者形成了互相促进的良性循环。

Baseline

本文的实验对比了多种类型的检索模型：

稀疏检索:
- BM25: 经典的基于词频的算法。
稠密检索 (单向量基础模型):
- DPR, ANCE: 基于 BERT 的经典稠密检索模型。
- T5-ANCE: 将 ANCE 的骨干网络换成 T5，是 MoMA-DR 最直接的对照组（可以看作是“无记忆增强”的版本）。
- coCondenser: 通过持续预训练提升泛化性的模型。
- GTR: 当时最先进的基于 T5 的稠密检索模型。
使用生成伪标签的稠密检索:
- GenQ: 为每个目标任务生成伪查询-文档对，然后独立训练模型。
作为参考的更强模型:
- ColBERT: 多向量表示的晚期交互模型，索引和计算开销更大。
- GTR-large: 使用了更大参数量 T5-large 的 GTR 模型。

数据集 (Datasets)

源领域 (Source Domain):
- MS MARCO: 一个大规模的网页问答数据集，用于模型的基础训练。
评估/目标领域 (Target Domains):
- BEIR Benchmark: 一个包含 18 个不同数据集的、异构的零样本信息检索评测基准。涵盖的领域非常广泛，包括：
  - 生物医药: TREC-COVID, BioASQ
  - 开放域问答: NQ, HotpotQA
  - 金融: FIQA-2018
  - 新闻: TREC-NEWS
  - 事实核查: SciFact, FEVER
  - …
增强语料库 (Augmenting Corpora):
- 在训练阶段，混合记忆池由 MS MARCO、Wikipedia（英文维基百科）和一个医学知识图谱 (MeSH) 组成。

可复现性 (Reproducibility)

代码: 作者计划在论文被接收后，开源全部代码、增强数据、模型检查点和分析脚本。
算力: 实验使用了 8 块 Nvidia A100 (80GB) GPU 进行 FP16 混合精度训练。完成三个 episode 的联合训练，总耗时约 13 个小时（增强组件 6.6 小时，主模型 6.3 小时）。论文附录中给出了详细的超参数设置（如 Table 7），可复现性较好。

可改进的几个点 (Potential Improvements)

记忆工程 (Memory Engineering): 论文只探索了将源语料库替换为目标语料库。未来可以研究更复杂的记忆组合策略，比如根据查询的意图动态选择或加权不同的记忆库，或者允许用户手动注入特定的知识库来控制模型行为。
效率优化: 增加的增强步骤（检索+多文档编码）无疑会带来额外的计算延迟。虽然 FiD 机制相对高效，但对于在线服务来说，延迟可能仍然是一个挑战。研究如何蒸馏、量化或缓存这个增强过程，以降低延迟，将是重要的方向。
超越查询增强: 目前的工作只对查询端进行了增强。未来可以探索是否也可以对文档端进行增强，或者设计一种查询和文档联合增强的机制。
记忆库的选择: 如何选择最优的初始记忆库组合？除了“大而全”之外，是否存在一些指导原则来构建一个高效且泛化能力强的混合记忆池。
增强文档的数量: 文中将增强文档数量 $K$ 固定为 10。对于不同类型的查询，最优的增强文档数量可能是不同的。可以研究自适应地决定 $K$ 值的方法。

可以被引用的一些结论 (Key Takeaways)

混合记忆是提升泛化性的有效途径: 通过从多个、异构的语料库中进行检索增强，可以显著提升稠密检索模型的零样本泛化能力，这是除了扩大模型参数之外的另一条有效路径。
记忆的多样性至关重要: 实验表明，仅使用单一的外部记忆（即使是源领域的语料库）进行增强，对未见任务的泛化帮助有限，甚至可能产生负面影响。一个内容丰富且多样化的混合记忆池是成功的关键。
联合学习+难负样本是必要条件: 简单的使用注意力分数进行软标签蒸馏不足以训练出强大的增强组件。必须将其与跨混合记忆库的难负样本挖掘相结合，才能让增强组件学会在多样化的信息源中进行有效检索。
“即插即用”机制展现了强大的零样本适应能力: MoMA 机制允许在推理时动态地加入新的、与任务相关的语料库，模型无需重新训练就能立即利用这些领域内信息，这是一种非常实用和高效的领域自适应方法。
MoMA-DR 性能卓越: 在包含 18 个任务的 BEIR 基准测试上，MoMA-DR 的零样本性能优于同等规模的其他稠密检索模型，并能达到与更大参数量模型或更复杂架构（如 ColBERT）相媲美的水平。