EdgeRAG: 检索增强生成的新篇章，开启边缘智能新时代

论文链接: https://arxiv.org/pdf/2412.21023

简介

随着大型语言模型 (LLM) 和检索增强生成 (RAG) 技术的兴起，将它们部署在资源受限的边缘设备上成为一项挑战，因为边缘设备的内存和计算能力有限。传统的 RAG 系统需要将整个嵌入向量数据库加载到内存中，这在边缘设备上是不现实的，会导致内存抖动和性能下降。

为了解决这个问题，EdgeRAG 提出了一种内存高效的 RAG 系统，通过选择性存储，不是所有向量都存储，只有在检索过程中真正需要的嵌入向量才会被生成和存储，此外，还采用自适应缓存策略，以减少冗余计算并进一步优化延迟。

实验结果表明，EdgeRAG 在保持检索和生成质量的同时，显著降低了检索延迟，并能够支持比内存容量更大的数据集。

方法

EdgeRAG索引是一种高效的二级索引系统，设计旨在兼顾内存使用效率和在线计算能力。它以传统的二级倒排文件（IVF）索引（如下图所示）为基础。索引的第一层始终驻留在内存中，主要存储集群质心及其到第二层索引的引用；第二层则包含文本块的引用以及嵌入生成的延迟信息。

与传统方法不同，EdgeRAG并未存储所有文本块的嵌入，而是通过修剪嵌入并在检索过程中动态生成，仅存储计算成本较高的集群索引，以此优化性能并降低延迟。为了进一步提升效率，EdgeRAG采用选择性缓存策略：对于检索中生成的嵌入，系统优先缓存计算成本较高的嵌入，利用缓存命中显著提高性能；而对于生成成本较低、不会影响服务水平目标（SLOs）的嵌入，则避免缓存，为高成本嵌入腾出更多空间。

以下是EdgeRAG方法的详细流程总结，包括索引的构建、检索、插入和删除的全过程：

索引构建（EdgeRAG Indexing）

EdgeRAG 基于传统的倒排文件（IVF）索引，但进行了多方面优化。上图展示了EdgeRAG的索引构建过程，索引构建的具体步骤如下：

1. 文本分块
   将文本语料库分割成较小的数据块，便于管理和处理。

2. 生成嵌入
   为每个数据块生成嵌入向量，用于后续的聚类和检索。

3. 聚类嵌入
   将生成的嵌入向量进行聚类，以减少索引的复杂度和查询时间。

4. 存储质心
   将集群的质心嵌入存储到第一层索引中，并记录对第二层索引的引用。

5. 分配嵌入到集群
   将每个数据块的嵌入分配到其所属的集群，并存储数据块的引用。

6. 计算生成成本
   对每个数据块嵌入的生成成本进行计算，判断是否超过预定义的服务等级目标 (SLO)。

7. 优化存储

• 高成本嵌入：如果嵌入生成成本超过 SLO，则直接存储这些嵌入以减少未来的计算开销。
• 低成本嵌入：如果嵌入生成成本低于 SLO，则丢弃这些嵌入，节约存储空间。

检索过程（EdgeRAG Retrieval）

EdgeRAG 结合了高效的嵌入加载与智能缓存策略，实现了快速检索，其过程如上图所示。具体步骤如下：

1. 查询最相似的质心
   根据查询嵌入，找到与其最相似的集群质心。

2. 检查预存嵌入
   检查该集群是否已有预计算的嵌入。

• 如果存在，则直接加载这些嵌入。
• 如果不存在，则进入下一步。

• 查找嵌入缓存
  查看缓存中是否存储了相关嵌入。

• 缓存命中：从缓存中加载嵌入，跳过生成步骤。
• 缓存未命中：进入下一步。

• 重新生成嵌入
  如果缓存未命中且没有预计算的嵌入，则动态生成嵌入向量，并将其存储在缓存中，以便后续查询使用。

• 加载嵌入并检索数据块
  加载嵌入后，查找最匹配的嵌入，并检索相关的文本数据块。

插入和删除

插入

• 将新的数据块的嵌入向量添加到最相似的簇中。
• 如果添加后该簇的嵌入生成延迟超过 SLO，则重新生成并存储该簇中所有数据块的嵌入向量。

删除

• 从其所属的簇中移除要删除的数据块的嵌入向量，并更新簇索引。
• 如果移除后该簇的嵌入生成延迟低于 SLO，则可以删除该簇中所有数据块的嵌入向量。

总结

本文提出了EdgeRAG，一种新颖的RAG系统，旨在解决边缘平台的内存限制问题。EdgeRAG通过修剪不必要的第二级嵌入、在执行过程中有选择地存储或重新生成嵌入，并通过缓存生成的嵌入来最大限度地减少冗余计算，从而优化两级IVF索引。这种方法使得在超出可用内存的数据集上，RAG应用程序也能高效运行，推动了在边缘计算场景中应用RAG技术的可行性。