本文证明了信息检索可以通过一个Transformer来完成，其中关于语料库的所有信息都被编码在模型的参数中。为此，我们介绍了可区分搜索索引（Differentiable Search Index，DSI），这是一种新的范式，它学习一种文本到文本的模型，将字符串查询直接映射到相关的文档；换句话说，DSI模型只用它的参数就能直接回答查询，极大地简化了整个检索过程。
我们研究了如何表示文档及其标识符的变化、训练程序的变化以及模型和语料库规模之间的相互作用。实验表明，在适当的设计选择下，DSI明显优于强大的基线，如双编码器模型。此外，DSI展示了强大的泛化能力，在零样本学习的情况下超过了BM25基准的表现。
简介
信息检索（IR）系统将用户查询q∈Q映射到相关文档的排序列表{d1, . . dn}⊆D，通常用整数或短字符串表示，称为文档标识符（docids）。最广泛使用的IR方法是基于流水线式的检索-然后排名策略。对于检索来说，基于倒排索引或最近邻搜索的方法是常见的，而基于反演学习的双编码器是目前最先进的。
本文提出一个替代架构，其中一个序列到序列（seq2seq）学习系统被用来直接映射查询q到相关的文档j∈Y。这个方案在图1的下半部分显示，是一个序列到序列的编码器-解码器架构。
我们把这个提出的架构称为可区分的搜索索引（DSI），并以大型预训练的Transformer模型来实现，建立在最近大型生成语言模型（LM）的成果之上。在这个提出的架构中，语料库的所有信息都被编码在Transformer语言模型的参数中。
在推理时间，训练好的模型将文本查询q作为输入，并输出一个文档j。如果需要，可以使用波束搜索来产生一个潜在相关文档的排序列表。正如我们所展示的那样，如果训练得当，这个过程会有令人惊讶的效果。 在我们的实验中，它可以持续地优于DE基准，有时甚至是大幅度地：对于一个基本大小的T5模型，在最小的语料库上Hits@1提高了20多个点，从DE的12.4%提高到DSI的33.9%；而在一个30倍大的语料库上，性能提高了近7个点。当使用更大的模型时，这些收益会增加：对于一个11B参数的T5模型，Hits@1的性能在小语料库上比DE（双编码器）提高了超过25个点，在大语料库上则超过了15个点。DSI在零样本学习的情况下也表现得非常好，比如说在BM25上提高了14个点。
除了这些数量上的提高，DSI的结构也比DE简单得多（见表1）。DE系统固定了一个搜索程序（MIPS），并学习了内部表征来优化该搜索程序的性能；相比之下，DSI系统不包含特定目标的固定搜索程序，而是使用标准的模型推理来从编码映射到文档。
正如表1所示，机器学习界对DSI特别感兴趣，在检索的各个方面被映射到公认的ML任务。这可能导致解决长期存在的IR问题的新的潜在方法。作为一个例子，由于索引现在是模型训练的一个特例，增量更新索引成为模型更新的一个特例。
在本文中，DSI被应用于中等大小的语料库（从10k到32k文档），所有这些语料库都来自一个具有挑战性的检索任务，我们将DSI扩展到更大的语料库这一重要问题留给未来工作。我们考虑的任务是检索自然问题（NQ）数据集中的问题的支持段落，这对词汇模型来说是一项具有挑战性的任务。
虽然DSI的想法很简单，但有许多方法可以实现，其中一些效果好得令人吃惊，而另一些效果差得令人吃惊。 下面我们探讨了DSI架构的一些变体。
文件表示：我们探讨了几种表示文件的方法，包括使用文件全文的 "朴素 "方法，以及传统IR引擎使用的词袋表示法的变体。
文档表示：我们考虑了几种方式去表示文档。除了将整数表示为文本字符串之外，还考虑了非结构化的原子文档，其中每个文档都被分配了一个唯一的标记，以及一些简单的基准，用于构建结构化语义文档，描述如何通过语料库的分层聚类来导航到一个文档。结构化文档–无论是通过聚类的语义结构化，还是作为标记化整数的主动结构化，都能更好地扩展到大型语料库，因为解码器中使用的词汇量变得更大。
编制索引：一个可训练的IR系统传统上有两个阶段：为语料库建立索引（即记忆每个文档的信息），以及学习如何有效地从索引中检索。在DSI中，索引被存储在模型参数中，而建立索引只是另一种模型训练。图1提出了一种索引语料库的方法：即在将文档dj与docid j配对的例子(1)(dj, j)上进行训练，此外还有将查询q与相关docid j配对的例子(2)(q, j)。在此设置中，类型（1）的示例是“索引”示例。
虽然很明显，单单是第（2）种类型的例子并不能为一个系统提供足够的信息来概括新的检索，但有许多替代第（1）种类型的例子的方法，可以合理地 "教导 "一个关于文档和文档ID之间关联的模型。我们在下文中探讨了其中的一些方法，并表明一些看似合理的技术表现很差。我们还探讨了一些替代的多任务优化和课程学习方案，以结合这些类型的例子。
模型和语料库大小的影响：由于最近的结果表明，大型LM的一些特性只适合于非常大的模型规模，我们探讨了DSI在一系列模型规模和10k、100k和320k文档的语料库规模下的表现。
总结：我们表明，即使是文档和docids的原始表征，加上适当的训练程序来微调现代大型LM，也能有惊人的表现；我们提出了两种改进的docid表征，即非结构化docids和语义结构化docids，它们改进了原始表征的选择。我们表明，不同的索引/训练策略在性能上有很大的差异，而且我们表明，DSI的性能随着模型规模的扩大而显著地、持续地提高。据我们所知，这是第一个生成性索引在一个经过充分研究的文档检索任务中在强大的基准上提高性能的案例。
相关工作
自回归实体链接描述了一个相关的序列到序列系统，文件中提到一个实体–也许是隐含的，例如，通过提出一个问题，该实体是一个答案–被映射到该实体的规范名称。在维基百科的案例中，规范实体名称与页面标题相对应，因此这可以被视为一种文档检索。这种方法已经被改编为其他用途，例如以规范形式生成知识库三要素。
我们考虑的任务与自回归实体链接中考虑的任务不同：我们的目标是检索包含答案的文档，而不是检索标题为答案的文档。更重要的是，在自回归实体链接中，生成目标是一个有语义的名字，而我们允许目标是任意的文档。这使得我们的方法适用于一般的检索任务，但也引发了关于docid表示和索引策略的新问题。
在自回归实体链接中，生成被限制为从一个固定的集合中返回一个输出。将DSI生成的输出约束为有效的docid也是可行的。虽然我们没有使用这种技术，但这在多大程度上可以提高性能是一个值得探讨的问题。
有大量关于检索增强生成的工作，即检索辅助文档以增强语言模型。这些技术对包括问答在内的许多任务都很有用，但依赖于传统的检索方法，如DEs。在这里，我们用生成来代替检索过程，而不是用检索来增强生成过程。
双编码器是一种成熟的检索范式。其关键思想是独立产生查询和文档嵌入，并在向量空间中对所有嵌入对进行相似性检索。查询和候选文档是由一个序列编码器产生的，并使用一种对比性损失的形式进行训练。
在之前的工作中，对大型Transformer模型作为记忆存储的解释已经进行了调查。(Roberts等人，2020年）证明了在闭卷QA任务上的成功，他们通过训练T5模型来检索在预训练期间在模型的参数中编码的事实。 然而，与CBQA不同的是，本文提出的问题是基于文档而不是生成直接答案来检索完整的文档。同时，《Language Models as Knowledge Bases?》也研究了作为知识库的语言模型，发现预训练的LMs可能已经包含了关系知识。《Transformer feed-forward layers arekey-value memories》分析了Transformer前馈层内编码的知识。 还有一些作品证明了Transformers与关联记忆和Hopfield网络的关系，这加强了Transformers在直觉上应该作为一个良好的关联记忆存储或搜索索引的概念。
可区分搜索指数DSI
可区分搜索指数（DSI）的核心思想是在一个单一的神经模型中对传统的多阶段检索-然后排序流程进行全面参数化。要做到这一点，DSI模型必须支持两种基本的操作模式。
1.建立索引：DSI模型应该学会将每个文件dj的内容与其对应的docid j联系起来。本文采用了一种直接的序列到序列（seq2seq）的方法，将文件标记作为输入，并生成标识符作为输出。
2.检索：给定一个输入查询，一个DSI模型应该返回一个候选docids的排序列表。这里，这是通过自回归生成实现的。
在这两项操作之后，一个DSI模型可以被训练来索引一个文档语料库，并且可以选择在一个可用的标记数据集（查询和标记文档）上进行微调，然后用于检索相关的文档–所有这些都在一个统一的模型中。与检索-排序方法相反，这种类型的模型允许简单的端到端训练，并且可以很容易地作为一个更大、更复杂的神经模型的可区分的子组件。
索引策略
我们研究了各种索引策略，这些策略是为了学习文件和它们的标识符之间的关联。我们训练我们的模型来预测给定文档标记序列的docids。这使得我们的模型能够学习哪个标识符属于哪个文档，并且可以被认为是对传统搜索索引的一种可区分的处理。我们考虑了各种替代方案，并在随后的章节中对这些设置进行了删减。最终采用的策略是Inputs2Targets与直接索引。
索引方法
本节讨论我们考虑的索引任务变体。
Inputs2Target：我们把它看作doc_tokens→docid的seq2seq任务。顾名思义，这以一种直接的输入到目标的方式将docid与文件标记结合起来。 这里的好处是，标识符是去噪目标，这使它更接近于损失函数。由于检索任务也与预测标识符有关，这种表述允许网络在序列长度方面遵循类似的输入-目标平衡。一个潜在的不足是，文件标记不是去噪目标，因此没有机会对文档标记进行一般的预训练。
Targets2Inputs：这种表述考虑了与上述相反的情况，即从标识符生成文件标记，即docid→doc_tokens。 直观地讲，这相当于训练一个以docid为条件的自回归语言模型。
Bidirectional：这种提法在相同的训练设置中同时训练Inputs2Targets和Targets2Inputs。前缀标记是预先设置的，它允许模型知道任务在哪个方向执行。
跨度腐败：我们还探索了一种设置，在包含docid标记的情况下，执行基于跨度腐败的去噪。在这种方法中，我们将标识符连接到文件标记上，作为前缀，在跨度腐败目标中可以随机掩盖跨度。这种方法的优点是：（1）在索引期间也进行一般的预训练；（2）实现docids作为去噪目标和输入的良好平衡。
文档表示策略
上一节探讨了 “如何进行索引”。本节研究了 “索引什么？”，即如何最好地表示doc_tokens。我们在这里陈述了我们的选择，并在后面的实验中仔细消减它们。最终的最佳选择是直接索引法。
直接索引：这个策略准确地表示了一个文档。我们取一个文档的前L个标记，保留顺序，并将它们与docid相关联。
集合索引：文档可能包含重复的术语和/或非信息性的词（例如，停止词）。 这个策略使用默认的Python set操作来去除重复的术语，并从文档中去除停止词。过滤后的文档的其余部分以类似于直接索引的方式被传递到模型中。
反转索引：这个策略将分块文档（连续的标记块）而不是整个文档直接映射到docid。我们随机地对一个由k个标记组成的连续块进行子抽样，并将它们与docid相关联。这种方法的主要优点是允许查看前k个标记之外的内容。
为检索表示docids
在基于seq2seq的DSI模型中的检索是通过对输入查询的docids进行解码来完成的。如何以有效的方式进行解码主要取决于docids在模型中的表示方式。本节的其余部分探讨了表示docids的一些可能的方式，以及如何处理其解码。
非结构化原子标识符：表示文档的最朴素方式是为每个文档分配一个任意的（可能是随机的）唯一的整数标识符。我们把这些称为非结构化的原子标识符。
有了这些标识符，一个明显的解码方法是学习标识符的概率分布。在这种情况下，模型被训练成对每个独特的docid（|Ndocuments|）发出一个logit。这类似于标准语言模型的输出层，但扩展到包括docids。为了适应这一点，我们将标准语言模型的输出词汇扩展如下：
其中[；]是行间连接运算符，W∈R和W∈R。h是解码器栈的最后一层隐藏状态（∈R）。 在检索给定查询的前k个文档时，我们只需对输出logits进行排序并返回相应的索引。这也让人联想到标准的列表式学习排名，其中所有的文档都是同时考虑的。
朴素的结构化字符串标识符：我们还考虑了一种表面上很荒谬的方法，把非结构化标识符，即任意的唯一整数，当作可标记的字符串。我们把这些称为真正的结构化标识符。
在这种表述中，检索是通过每次对一个docid字符串按顺序解码来完成的。 这就消除了对非结构化原子标识符所带来的巨大的softmax输出空间的需要。它还消除了为每个单独的 docid 学习嵌入的需要。
在解码时，使用波束搜索来获得预测的最佳docid。有了这个策略，获得前k排名就不那么直接了。 我们可以详尽地梳理整个文档空间，并获得每个文档在查询中的可能性。相反，我们使用部分波束搜索树来构建top-k检索得分。我们发现这种近似方法在实践中是相当高效且有用的。
语义结构化的标识符：迄今为止，所有表示docids的方法都假定标识符是以任意方式分配的。虽然探索任意标识符的局限性是相当有趣的，但直观地看，给文档空间注入语义结构可以带来更好的索引和检索能力。因此，本节探讨了语义结构化的标识符。
具体来说，我们的目标是自动创建满足以下属性的标识符：（1）docid应该捕获一些关于其相关文档语义的信息，（2）docid的结构方式应该是在每个解码步骤之后有效地减少搜索空间。这就导致了语义相似的文件共享标识符前缀。
在这项工作中，我们将其作为一个完全无监督的预处理步骤。然而，作为未来工作的一部分，有可能以一种完全端到端的方式整合并自动学习语义标识符。
为了构建具有这种属性的标识符，我们对文档嵌入采用了一个简单的分层聚类过程，以诱导出一棵十进制树（或者更一般地说，一个字典树）。
给定一个要被索引的语料库，所有的文档被聚类为10个聚类簇。每个文档都被分配一个标识符，其簇的编号为0-9。对于每个包含超过c个文档的聚类，该算法被递归应用，下一级的结果（标识符的剩余后缀）被附加到现有的标识符上。
对于有c个文件或更少的集群，每个元素被分配一个从0到最大c-1的任意数字，同样，它的数字被附加到现有的标识符上。虽然这个具体的过程会诱导出一个十进制树，但使用任何数量的其他合理的策略都有可能诱导出类似类型的字典树。在实践中，我们只是将k-means应用于由一个小型的8层BERT模型产生的嵌入，c=100。
训练和优化
我们训练的DSI模型针对seq2seq交叉熵损失进行了优化，并通过教师强迫进行训练。我们探索了两种训练DSI模型的主要策略。第一个也是更直接的策略是首先训练一个模型来执行索引（记忆），然后是一个微调阶段，用训练好的模型将查询映射到docids（例如检索）。第二种策略是在多任务设置中一起训练它们。为此，我们以类似于T5式协同训练的方式来构建协同训练任务（例如，使用任务提示来区分它们）。后者的表现明显更好，特别是当检索任务实例索引的比例很高时。因此，我们采用多任务学习作为默认策略。
在这里，我们观察到，我们的设置是独特的，与传统的多任务学习或迁移学习不同。在典型的多任务设置中，两个任务有共同点，如果它们一起学习，可以提高两个任务的性能。然而，在我们的设置中，检索任务完全依赖于索引任务。特别是，如果没有索引任务，检索任务所利用的标识符就完全没有意义。因此，为了解决任务B（检索），该模型需要充分学习任务A（索引）。这个问题的设置提出了独特的、在很大程度上未经探索的研究挑战，可能会引起ML社区的兴趣。
实验
本节讨论了实验设置、使用的数据集和基线的比较。我们还讨论了实验结果、发现和本文前面部分讨论的各种策略的效果。由于这是一个相当新的概念，这项工作的目的是提出一个概念证明，并寻求回答研究问题，而不是进行一个 "sotaeesque "的比较。我们把在其他设置和基线上的广泛比较留给未来的工作。
数据集：我们在具有挑战性的自然问题（NQ）数据集上进行了实验。NQ包括307K查询-文档训练对和8K验证对，其中查询是自然语言问题，文档是维基百科文章。给定一个问题，检索任务是确定回答该问题的维基百科文章。为了评估DSI模型在不同规模下的表现，我们从NQ中构建了三个集合来形成我们的测试平台，即NQ10K、NQ100K和NQ320K，表示在合并训练和验证分割中不同数量的总查询-文档对。NQ320K是完整的NQ集，并使用其预先确定的训练和验证分割进行评估。与NQ320K不同，NQ10K和NQ100K构建了随机抽样的验证集。对于所有的数据集，我们使用相同的docid空间/320K tokens预算，用于所有非结构化的原子和朴素结构化的标识符实验。语义结构化的标识符是为每个数据集单独生成的，以防止语义信息从较大的拆分中泄漏到较小的拆分中。 文本是小写的。 请注意，在这些数据集中，存在着比查询-文档对更少的独特文档。
实施细节：
所有的DSI模型都使用标准的预训练T5模型配置进行初始化。配置名称和相应的模型参数数量： Base（0.2B），Large（0.8B），XL（3B）和XXL（11B）。对于非结构化原子标识符的运行，我们随机地将标识符初始化为新参数，只在索引阶段对权重进行微调。 我们使用Jax/T5X实现来做实验。 DSI模型使用128个批处理量，最多可训练100万步。我们根据检索验证性能挑选最佳检查点。我们的训练硬件由128-256个TPUv4芯片组成，用于1B参数以上的模型，其他则由64-128个TPUv3或TPUv4芯片组成。作为估计，高于1B参数的模型通常至少要花费一整天的NQ320K收敛。 我们在{0.001,0.0005}之间调整学习率，在{10K,100K,200K,300K}和/或空中调整线性预热。语义结构的标识符是使用8层BERT模型和scikit-learn中的默认k-means聚类法生成的。基于我们对各种DSI设置的早期消融实验，所展示的主要结果使用直接索引（L=32）和Inputs2Targets索引策略。我们介绍了所有DOCID表示方法的结果。针对主要结果，我们提供了消融研究。
基线
对于基线，我们使用基于T5的双编码器。我们使用gensim包来计算BM25的分数。 对于基于T5的双编码器，我们对NQ对进行对比学习训练，直到收敛（≈10K步），并通过类似于ScaNN的系统获得top-k最近邻。对于零样本检索，我们还与最先进的无监督基线Sentence T5进行了比较，该系统已被专门预训练了相似度学习任务。有两个原因使我们认为Sentence T5是本工作的相关双编码器基线，而不是其他密集的检索工作，如DPR。
首先，我们采用了完全相同的预训练模型，这使得系统地消减所提出的方法，而不混淆其他因素。在科学上，我们认为这种与微调T5的比较是我们提供的最好的比较。其次，微调的T5双编码器被认为在结构和方法上与DPR非常相同（有一些小的区别，如参数共享，但使用相同的批内否定概念）。
实验结果
表2报告了NQ10K、NQ100K和NQ320K的微调检索结果，表3报告了零样本学习的检索结果。对于零样本学习检索，模型只在索引任务上训练，而不是在检索任务上训练，所以模型没有看到标记的查询→docid数据点。
有监督微调结果：我们的结果显示，DSI在所有的数据集大小上都优于DE。在小数据集（NQ10K）上，DSI和DE之间的性能差距很大，例如，最好的DSI变体比DE要好2倍。在NQ100K上，差距变得不那么突出，最好的DSI模型（非结构化的原子标识符）比DE的性能要好+5% Hits@1和Hits@10。在大型数据集（NQ320K）上，最好的DSI模型（结构化语义标识符）超过了最好的DE模型，相对Hits@1和Hits@10分别为+66%和+4.5%。
零样本学习：表3报告了零样本学习检索的结果。回顾一下，零样本学习检索是通过只执行索引而不执行检索任务来进行的。换句话说，该模型没有看到任何注释的查询或文档对。一般来说，在NQ100K和NQ320K上，使用非结构化原子标识符的DSI获得了最佳结果。在所有的NQ数据集上的最佳表现都超过了成熟的无监督检索基线，如BM25。此外，DSI优于无监督的表征学习方法，如SentenceT5，T5被训练为通过对比学习来学习相似性感知表征。我们还注意到，原始T5嵌入的表现极差，在无监督检索的任务中没有产生合理的结果。鉴于无监督的神经方法通常很难超越BM25，我们发现这些早期结果非常令人鼓舞。
文档标识符：本文的一个关键研究问题是如何表示文档标识符的关键选择。一般来说，我们发现，结构化的语义标识符是有帮助的，而且比非结构化的标识符有改进。在比较朴素与语义字符串标识符时，如果可能的话，似乎必须使用语义标识符。这是直观的，因为给目标空间注入语义结构可以促进更容易的优化和额外的无监督表示学习方法作为外部知识。非结构化的原子标识符的竞争力有些参差不齐，我们在优化这种模型时遇到了一些困难。我们假设，这可能是因为新初始化的softmax层，从头开始训练这样的系统会减轻这些问题。然而，我们将这一调查方向推迟到未来的工作中。取而代之的是，非结构化原子标识符的不稳定性和高变异性，在不同的数据集中，其性能并不一致。此外，这些标识符也可能遇到间歇性的不融合，我们将其追溯到与优化有关的怪事。然而，我们也注意到，非结构化的原子标识符在零样本学习检索的设置上表现最好，而且取得的性能往往比波束解码方法的性能高出一倍以上。
索引策略：本节探讨了不同索引方法的效果。我们在NQ100K上用前面描述的不同索引策略进行了实验。模型是用Naive Docid方法训练的。在没有索引的情况下，模型的点击率为0%@1。这是直观的，因为如果没有索引任务，Docids是没有意义的。其次，Inputs2Targets和Bidirectional表现最好，与前者相比，bidirectional方法表现稍差（13.5 vs 13.2）。最后，Targets2Inputs和Span Corrpution与Docids的准确性没有产生有意义的结果（0%的准确性）。这表明，不同的索引策略之间可能存在巨大的差异，有些策略工作得相当好，有些则完全不起作用。
文档表示：在本节中，我们将探讨前文中描述的不同文档表示策略的性能。图5报告了NQ320K上的结果。总的来说，我们发现直接索引的方法效果最好。我们还发现，由于docid被反复暴露在不同的tokens中，所以很难训练出倒置的索引方法。我们还发现，较短的文档长度似乎效果很好，而超过64个标记时，性能似乎大幅下降，这表明当文档标记数量较多时，可能更难优化或有效记忆。最后，我们还发现，对文档标记进行集合处理或停用词预处理并没有额外的优势。
缩放规律：另一个有趣的见解是DSI的缩放规律与双编码器有什么不同。 近年来，了解Transformer的缩放行为引起了人们的极大兴趣。我们发现，在DE中，通过增加模型参数化而获得的检索性能增益似乎相对较小。相反，DSI的缩放特性似乎更为乐观。
图3描绘了三种方法（具有朴素和发散ID的DE和DSI）的缩放行为（对数尺度）。DSI（朴素）强烈地受益于从base到XXL的扩展，并且似乎仍有改进的空间。同时，DSI（语义）开始时与DEbase同样具有竞争力，但随着规模的扩大，表现得更好。遗憾的是，DE模型在较小的参数化时或多或少地出现了停滞。
索引和检索之间的相互作用：我们的早期实验表明，首先学习索引任务，然后以连续的方式学习检索任务，结果表现平平。因此，我们着重于探索使用多任务学习来共同训练索引和检索任务的良好比率r。图4显示了修改索引与检索样本的比例的效果。 我们发现优化过程受到了索引和检索任务之间相互作用的显著影响。r设置太高或太低一般都会导致性能不佳。我们发现，32的比率通常表现良好。
结论
本文提出了可区分搜索索引（DSI），这是一种以统一方式学习端到端搜索系统的新范式，为下一代搜索铺平了道路。我们定义了新颖的索引和检索任务，将术语和文档之间的关系完全编码在Transformer模型的参数中。该论文提出了一些表示文件和文档的不同方式，并探索了不同的模型架构和模型训练策略。在Natural Questions数据集上进行的实验表明，无论是在标准的微调设置中，还是在零样本学习设置中，DSI的表现都优于常见的基准，如BM25和双编码器。
虽然这里提出的模型和结果是有前景的，但在这项工作的基础上，还有大量潜在的未来研究可以探索，以改进这种方法。例如，探索文件和文档表示的替代策略，以及研究专家混合模型来扩展DSI的记忆容量。一个重要的方向是探索如何为动态语料库更新此类模型，因为在动态语料库中，文件可能被添加或删除。最后，进一步研究DSI作为一种无监督的表征学习方法和/或其他语言模型的内存存储，可能也是有趣的。