EmbeddingAlign RAG：提升问答系统

检索增强生成（RAG）方法通过在生成答案之前检索相关信息来改进问答系统。本文介绍了EmbeddingAlign RAG，一种改进RAG系统中文档检索的新方法。我们证明，对查询嵌入和分块嵌入应用线性变换可以显著提高检索精度。

我们的方法使用小数据集，例如包含500个分块的单个PDF，一个开箱即用的模型嵌入器（可以是黑盒嵌入模型，如OpenAI嵌入API），并利用现有嵌入集合。该模型在标准硬件的CPU上训练，检索推理时间增加不到10毫秒。我们展示了EmbeddingAlign RAG提高了检索精度，将命中率从0.89提高到0.95。该方法实现简单，非常适合在实际问答系统中直接应用。

RAG和嵌入优化的介绍

现有RAG系统的挑战

RAG 系统结合了检索器和生成器（如GPT-4）。在该系统中，检索器根据用户查询（问题）获取相关文档或文档分块。该系统通常依赖文档嵌入和查询嵌入来检索相关分块。新文档会定期添加到知识库中，以应对不断扩展的使用案例。

RAG 系统的数学表示如下：

设 $QQ$ 为用户查询集，$DD$ 为文档分块集，$EE$ 为嵌入函数。对于查询 $q\in Qq\; \backslash in\; Q$ 和文档分块 $d\in Dd\; \backslash in\; D$，检索分数 $S(q,d)S(q,\; d)$ 通常计算为： $$S(q,d)=\text{similarity}(E(q),E(d))S(q,\; d)\; =\; \backslash text\{similarity\}(E(q),\; E(d))$$ 其中相似度通常是余弦相似度： $\text{similarity}(a,b)=\frac{a\cdot b}{\mathrm{\mid }a\mathrm{\mid }\mathrm{\mid }b\mathrm{\mid }}\backslash text\{similarity\}(a,\; b)\; =\; \backslash frac\{a\; \backslash cdot\; b\}\{|a|\; |b|\}$

然而，随着引入更多文档，检索器可能会拉取不相关分块，导致检索性能不佳。性能下降的一个原因是现有嵌入（无论是查询还是分块的）与新数据未良好对齐。

此外，尝试或训练新的嵌入模型成本高昂，因为大型文档嵌入集合昂贵且从工程角度难以更新。这使得许多 RAG 系统陷入僵局，即检索需要改进以提升性能，但成本又太高而无法实现。

我们的方法：通过线性变换进行嵌入对齐

我们建议训练一个单一的线性变换，将其应用于查询嵌入和文本分块嵌入向量。目标是学习一个变换，使用户查询的嵌入更接近最相关的文档分块，从而提高检索精度。

设 $TT$ 为我们要学习的线性变换矩阵。新的检索分数变为： $$S^{\mathrm{\prime }}(q,d)=\text{similarity}(T\cdot E(q),T\cdot E(d))S\text{'}(q,\; d)\; =\; \backslash text\{similarity\}(T\; \backslash cdot\; E(q),\; T\; \backslash cdot\; E(d))$$

其中

• $T\in {\mathbb{R}}^{N\times N}T\; \backslash in\; \backslash mathbb\{R\}^\{N\; \backslash times\; N\}$
• $E(q),E(d)\in {\mathbb{R}}^{N}E(q),\; E(d)\; \backslash in\; \backslash mathbb\{R\}^N$

在我们的案例中，$N=1536N\; =\; 1536$（OpenAI text-embedding-3-small 模型的嵌入维度）。为了找到一个能提高我们RAG管道检索质量的矩阵 $TT$，我们需要以下数据：

• 查询嵌入（用户问题）
• 良好分块嵌入（文档中的相关分块）
• 干扰分块嵌入（不相关或误导性分块）

在已部署的RAG系统中，这些数据可以通过反馈机制或带注释的数据集轻松收集。例如，用户或注释者对最终结果的反馈（对结果的赞/踩）可用于定义好分块和坏分块。

从推理角度来看，线性层需要在检索步骤中应用于用户查询嵌入，同时也要应用于向量存储中存储的所有嵌入。考虑到延迟和计算效率，我们建议在开始执行检索之前，一次性将线性变换应用于向量存储中的所有向量。

合成数据集生成

在我们的案例中，我们无法访问专有的生产数据。为了克服这一限制，我们通过从现有文档（代表知识库）生成合成数据集来模拟真实场景。这种方法在生产数据缺失或不足的情况下也很有用，因为只需知识库的文本分块即可改进RAG。

我们基于公开文档和模型构建合成数据集：

• 文档：Uber 和 Lyft 的 SEC 10K 表格报告，其中包含关于类似业务活动（出行服务）的丰富文本。
• 模型：GPT-4o 用于生成查询，LlamaIndex 用于生成文档分块。对于每个分块，都会生成一个查询（问题）以形成（查询，分块）对。

以下是数据集中（查询，分块）对的一个示例：

然后，使用 OpenAI 的 *text-embedding-3-small* 模型对查询和分块进行嵌入。在已部署的 RAG 系统中，分块的嵌入已经存在并存储在向量存储中。使用 Uber 文档，我们生成了 686 对嵌入用于训练和验证数据集。使用 Lyft 文档，我们生成了 818 对嵌入，用作测试数据集。生成此数据集的成本适中。

数据准备与增强

为了训练将用户查询与相应文本分块对齐的线性变换模型，我们使用了三元组损失。这种损失使得正确的对（查询和分块）更接近，并将不正确的分块推开。此外，生成三元组几乎免费地显著增加了训练数据集的大小。

三元组包含三个元素：

• 查询（问题嵌入）
• 分块（正确文档分块嵌入）
• 干扰物（不正确文档分块嵌入）

每对（查询，分块）都会与几个干扰物（不正确的分块）匹配。我们称每个对关联的不正确分块的比例为增强因子。当增强因子为0.3时，我们为训练-验证数据集创建了141,100个三元组。通过控制三元组，我们可以惩罚某些希望避免的特定关联，进一步优化检索性能。例如，可以通过用户反馈机制（点赞或点踩按钮）收集这些干扰物。

训练过程

为了提高检索性能，我们对嵌入层进行线性层训练，以最小化三元组损失。三元组损失鼓励模型最小化查询与正确分块之间的距离，同时最大化查询与干扰物之间的距离。

三元组损失公式为：

$$\text{Loss}(q,c,dis)=\max (d(q,c)-d(q,dis)+\text{margin},0)\backslash text\{Loss\}(q,\; c,\; dis)\; =\; \backslash max(d(q,\; c)\; -\; d(q,\; dis)\; +\; \backslash text\{margin\},\; 0)$$

其中：

• $qq$ 是查询嵌入，
• $cc$ 是正确分块嵌入，
• $disdis$ 是干扰分块嵌入，
• $d(x,y)d(x,\; y)$ 代表两个嵌入之间的距离。在我们的案例中，我们使用归一化余弦距离。
• $marginmargin$ 是一个超参数。它是一个正值，用于确保模型在正确和不正确对之间创建足够的区分。

在训练中，我们仅依赖消费级CPU，不使用GPU，因为线性层和余弦距离所需的计算资源很少。

结果

我们比较了我们适应后的用户查询嵌入（使用训练好的线性变换）与简单嵌入在文档嵌入检索方面的表现，使用了两个指标：

MRR（平均倒数排名）：衡量正确答案排名的靠前程度，MRR越高表示排名性能越好。

$$MRR=\frac{1}{card(Q)}\sum _{q\in Q}\frac{1}{ran{k}_{\text{retrieve(q)}}}MRR\; =\; \backslash frac\{1\}\{card(Q)\}\; \backslash sum\_\{q\; \backslash in\; Q\}\; \backslash frac\{1\}\{rank\_\backslash text\{retrieve(q)\}\}$$

其中 $ran{k}_{\text{retrieve(q)}}rank\_\backslash text\{retrieve(q)\}$ 是查询 $qq$ 在 top-k 结果中第一个相关项的位置。

命中率：正确答案出现在 top-k 结果中的查询百分比，反映了整体检索质量。

$$HR=\frac{1}{card(Q)}\sum _{q\in Q}{\mathbb{1}}_c(retrieve(q))HR\; =\; \backslash frac\{1\}\{card(Q)\}\; \backslash sum\_\{q\; \backslash in\; Q\}\; \backslash mathbb\{1\}\_\{c\}\; (retrieve(q))$$

其中 ${\mathbb{1}}_c(retrieve(q))\backslash mathbb\{1\}\_c(retrieve(q))$ 是指示函数，定义为：

换句话说，命中率计算的是至少一个正确答案在 top-k 结果中被检索到的查询的比例。该指标衡量了检索系统在给定查询下，在指定数量的顶层结果中显示相关信息的能力。k 的值通常根据 RAG 系统的具体要求选择，例如向用户显示的结果数量或生成步骤处理的分块数量。k 的常见值包括 1、5 或 10，具体取决于应用。请注意，这些分数的绝对值取决于 k 的选择值。我们主要使用前 4 个结果的命中率（k = 4）和 MRR，因为它们能最好地突出检索精度的改进。

适应性嵌入在两项指标上都显示出显著改进。

对推理时间的影响

在我们的标准硬件上，该模型将检索时间增加了8.6%（不到10毫秒）。在实践中，这种延迟几乎察觉不到，因为RAG中后续的生成步骤可能需要几秒钟。

结论

在本文中，我们介绍了 EmbeddingAlign RAG，这是一种改进检索增强生成（RAG）系统中文档检索的新方法。通过调整嵌入以包含特定领域上下文，我们以较低的计算成本显著提高了检索精度，同时利用了现有嵌入。我们的方法高效、可扩展且易于实现。这使得 EmbeddingAlign RAG 非常适合实际的问答系统，特别是对于较小的数据集，因为它在将检索时间增加不到 10 毫秒的情况下提高了性能。

如果您对此工作有任何疑问，请随时通过 contact@phospho.ai 联系 phospho 团队，或查看我们的工作：phospho.ai 🧪