长文本 Embedding 模型中的“迟分”策略

来自：Jina AI

大约一年前，2023 年 10 月，我们推出了全球首个支持 8K 上下文长度的开源 Embedding 模型 —— jina-embeddings-v2-base-en。自此，长文本在 Embedding 模型中的应用引发了广泛讨论和争议。

1. 信息压缩问题 ：将数千字的长文本编码为单一 Embedding 表示会导致语义信息的"过度压缩"，使得检索系统难以准确定位特定信息。

2. 检索粒度不足 ：许多应用，尤其是检索增强生成（RAG）系统，需要检索文档中的较小片段，而非整个长文档。

3. 短文本检索优势 ：基于密集向量的检索系统在处理短文本时通常表现更好，因为短文本的语义信息更容易被准确编码和检索。

那么，如果行业只需要具有 512 上下文长度的 Embedding 模型，那么训练 8192 上下文长度的模型又有什么意义呢？

在本文中，我们通过探讨 RAG 中传统分块 -> Embeddings 流程的局限性，来重新审视这个问题。同时，我们还引入了一种新策略，称为 迟 分(Late Chunking) ，能够在保留长文本 Embedding 模型优势的同时，也能满足精细粒度检索的需求。

上下文丢失问题

传统的分块 - Embedding - 检索 - 生成流程在处理长文档时可能会丢失长距离的上下文依赖关系，这对于信息检索和理解是一大隐患。换句话说，当关键信息散落在多个文本块中，脱离上下文的文本片段很可能失去其原有的意义，导致处理效果大打折扣。

以维基百科上的一篇关于柏林的文章为例，若将其分割为句子块，不难发现诸如“其”和“这座城市”等指代表达，实际上指向的是文章开头提到的“柏林”。 这种情况下，Embedding 模型很难将这些指代正确链接到实体，进而产生质量低下的向量表示。

这意味着，若我们将一篇长文按照句子长度进行分块处理，RAG 系统在回答诸如“柏林的总人口是多少？”之类的查询时，可能会遇到困难。因为城市名称和人口数据从未出现在同一文本块中，且缺乏更大的文档上下文，处理这些块的 LLM 难以解决这样的指代问题。

尽管有一些启发式算法试图缓解这一问题，如滑动窗口重新采样、多种上下文窗口长度及多次文档扫描等，然而，像所有启发式算法一样，这些方法时灵时不灵；它们可能在某些情况下有效，但没有理论上的保证。

迟分让 Embedding 更懂上下文

在传统的文本编码领域，我们常常采用一种“预先分块”的策略，即先分块，再过 Embedding 模型。首先依据句子、段落或预设的最大长度对文本进行切割。随后，Embedding 模型会对这些切割好的文本块逐一进行处理，通过平均池化等方法，将 token 级的 Embedding 聚合成单一的块 Embedding 向量。如下图左侧所示：

左侧为传统分块策略，右侧为迟分策略的图解。

然而，本文提出的“迟分”策略，则是先过 Embedding 模型再分块，我们先将 Embedding 模型的 transformer 层应用到整个文本或尽可能多的连续文本，为每个 token 生成一个包含丰富上下文信息的向量表示序列。随后， 再对这些 token 向量 序列进行平均池化，进而得到考虑了整个文本上下文的块 Embedding。

与传统的独立同分布（i.i.d.）块 Embedding 不同， 迟分生成的块 Embedding 是“有条件的”，这意味着每个块都编码了更多的上下文信息 ，从而大大提升了编码的质量和准确性。

值得注意的是，为了充分发挥迟分的优势，我们需要借助支持长上下文的 Embedding 模型，如 jina-embeddings-v2-base-en，它能够处理长达 8192 个 token 的文本(相当于 10 页 A4 纸)，基本满足了大多数长文本的上下文需求。

当然，迟分也并非完美，它同样需要边界线索来辅助处理。但与传统方法不同的是，这些边界线索是在获取到 token 级 Embedding 之后才使用的，这也是“迟分”名称的由来。

迟分效果一览，定性评估案例

Google Colab: https://colab.research.google.com/drive/15vNZb6AsU7byjYoaEtXuNu567JWNzXOz

迟分的实现可以在上述链接的 Google Colab 中找到。在这里，我们利用了最近在 Tokenizer API 中的最新功能，借助所有可能的边界线索，将长文档分割成有意义的块。关于此功能背后的算法的更多讨论可以在 twitter@JinaAI_ 上找到。

Tokenizer API: https://jina.ai/tokenizer/

当将迟分应用于上述柏林的维基百科示例时，可以立即看到语义相似性的改善。例如，以“这座城市”和“柏林”的指代关系为例，采用迟分后，“这座城市”的向量表示成功捕捉到了与“柏林”的关联信息，在与涉及该城市名称的查询匹配时表现得更为出色。

以下是具体的数值对比结果，展示了“柏林”一词的 Embedding 与柏林维基百科文章中各个句子的余弦相似性。从表格中可以清晰地看出，与传统分块相比，迟分在提升语义相似性方面取得了显著成效。

BEIR 实验验证，迟分效果显著

为了全面验证迟分在各种场景下的有效性，我们借助 BEIR 中的检索基准进行了一系列严谨的测试。这些测试涵盖了查询集、文本文档语料库以及存储查询与文档关联信息的 QRels 文件。

在评估过程中，我们将文档分块并编码至 Embedding 索引，通过 k 近邻（kNN）算法寻找与查询 Embedding 最相似的块。随后，将块的 kNN 排名转换为文档排名，并与真实排名进行对比，计算 nDCG@10 等检索指标。

整个评估流程的脚本已开源，可在此处获取：https://github.com/jina-ai/late-chunking

我们在多个 BeIR 数据集上进行了传统分块与迟分的对比测试。为提取边界线索，我们采用正则表达式将文本切分为约 256 个 token 的字符串。同时，选用 jina-embeddings-v2-small-en 作为 Embedding 模型，虽为较小版本的模型，但仍具备 8192-token 的处理能力。

以下是各数据集上的测试结果：