长上下文嵌入模型中的短分块延迟分块技术

博客第一部分 | 博客第二部分 | ArXiv 论文

在许多应用中，将整个文本文档编码为单一的嵌入表示并不实用。许多应用需要检索文本中较小的部分，而基于稠密向量的信息检索系统通常在与较小文本段配合时表现更好，这是因为嵌入向量的信息容量有限。

RAG（检索增强生成）是最著名的需要将文档集合分割成较小文本块的应用之一。这些块通常存储在向量数据库中，并通过文本嵌入模型生成向量表示。在运行时，相同的嵌入模型将查询文本编码为向量表示，用于识别相关的已存储文本块。然后将这些块传递给大型语言模型（LLM），由 LLM 根据检索到的文本综合生成对查询的响应。

上下文问题

这种简单的 RAG 方法并非没有挑战。它特别难以处理长距离的上下文依赖关系，即当相关信息分散在多个块中，且将文本段脱离上下文会导致其无用时。

上图中，一篇维基百科文章被分割成句子长度的块。可以看到像 "its" 和 "the city" 这样的短语引用了仅在第一个句子中提及的 "Berlin"，这使得嵌入模型更难将其与对应的实体关联起来以产生高质量的嵌入表示。

例如，如果我们将一篇维基百科文章分割成句子长度的段（如上面的例子所示），RAG 系统可能无法回答诸如 "柏林的人口是多少？" 这样的查询。城市名称和人口数字从未出现在同一个段中，并且缺乏更大的文档上下文。当其中一个段被呈现给 LLM 时，它无法解析像 "it" 或 "the city" 这样的回指引用。

上下文感知分块

为了解决这个问题，我们利用了像 jina-embeddings-v2-base-en 这样的最新嵌入模型能够处理的长输入序列优势。这些模型支持更长的输入文本，例如 jina-embeddings-v2-base-en 支持 8192 个 token，大约相当于十个标准页面的文本。这种大小的文本段不太可能出现只能靠更大上下文解决的依赖关系。然而，我们仍然需要更小文本块的向量表示，部分原因是受到 LLM 输入尺寸的限制，但主要还是因为短嵌入向量的信息容量有限。

简单的编码方法（如上图左侧所示）在文本处理之前就进行分块，根据句子、段落和最大长度限制先验地分割文本，然后将嵌入模型应用于生成的分块。而延迟分块则首先将嵌入模型中的 Transformer 部分应用于整个文本或其尽可能大的部分。这会为每个 token 生成一个包含整个文本信息的向量表示序列。为了生成文本的单一嵌入，许多嵌入模型对这些 token 表示应用均值池化来输出一个单一向量。延迟分块则是对此 token 向量序列中较小的段应用均值池化，从而为每个考虑整个文本的分块生成嵌入。

上下文感知分块的效果

这对检索产生了立即可测量的具体效果。例如，在维基百科文章中的 "the city" 和 "Berlin" 的情况下，表示 "the city" 的向量包含了将其与先前提到的 "Berlin" 连接起来的信息，使其与涉及该城市名称的查询更加匹配。

你可以在下面的数值结果中看到这一点，该结果比较了字符串 "Berlin" 的嵌入与柏林相关文章中的各种句子的相似度。"传统相似度"列是使用先验分块的相似度值，"延迟分块相似度"列是使用上下文感知分块的结果。

如你所见，第一个包含 "Berlin" 的分块的相似度得分非常接近。对于其他两个分块，得分存在显著差异，因为延迟分块大幅改进了对未显式使用单词 "Berlin" 但包含对其回指引用的句子的匹配。

检索任务评估

为了验证这种方法在几个简单示例之外的有效性，我们使用 BeIR 中的一些检索基准对其进行了测试。这些检索任务包括一个查询集、一个文本文档语料库以及一个存储每个查询相关文档 ID 信息的 QRels 文件。要识别查询的相关文档，可以对文档进行分块，将其编码为嵌入索引，并为每个查询嵌入确定最相似的分块（kNN）。由于每个分块对应一个文档，可以将分块的 kNN 排序转换为文档的 kNN 排序（对于排序中出现多次的文档，只保留第一次出现的）。之后，可以将生成的排序与对应真实 QRels 文件的排序进行比较，并计算检索指标，如 nDCG@10。我们在多个 BeIR 数据集上使用传统分块和我们的新颖延迟分块方法运行了此评估。为了将文本分割成块，我们选择了一种直接的方法，即将测试文本分割成 256 个 token 的字符串。传统分块和延迟分块测试均使用了 jina-embeddings-v2-small-en 模型。

在所有情况下，延迟分块都提高了分数。在某些情况下，它也优于将整个文档编码为单一嵌入，而对于其他数据集，无分块表现最佳。然而，这仅在你不需要对分块进行排序时才有意义。你还可以看到，文档的平均长度与通过延迟分块改进的 nDCG 分数幅度相关。

要复现此评估，你可以使用 pip install . 安装依赖项，然后为 "SciFactChunked"、"TRECCOVIDChunked"、"FiQA2018Chunked"、"NFCorpusChunked" 和 "QuoraChunked" 任务运行以下脚本：

python3 run_chunked_eval.py --task-name {TASK_NAME}

致谢与参考

感谢 Isabelle Mohr(@violenil) 贡献了部分代码，以及 Scott Martens (@scott-martens) 审阅了本 README。

有关评估任务的更多信息，请参阅 MTEB 仓库，有关训练用于长输入文本模型的详细信息，请参阅我们的论文：《Jina embeddings 2: 8192-token general-purpose text embeddings for long documents》。

如果你发现延迟分块在你的研究中很有用，可以引用论文 《Late Chunking: Contextual Chunk Embeddings Using Long-Context Embedding Models》：

@article{gunther2024late,
  title={Late Chunking: Contextual Chunk Embeddings Using Long-Context Embedding Models},
  author={G{\"u}nther, Michael and Mohr, Isabelle and Williams, Daniel J and Wang, Bo and Xiao, Han},
  journal={arXiv preprint arXiv:2409.04701},
  year={2024}
}