优化向量数据库，增强RAG驱动的生成式AI
作者：Cathy Zhang 和 Dr. Malini Bhandaru
生成式AI（GenAI）模型在我们日常生活中的应用正呈指数级增长。通过检索增强生成（RAG）技术，普通的大型语言模型（LLM）可以利用外部数据源获取事实，从而提高响应的准确性和可靠性。RAG通过利用存储为向量的非结构化数据帮助LLM理解上下文并减少幻觉。
RAG有助于检索更多上下文信息，从而生成更好的响应，但它依赖的向量数据库为了提供丰富的内容而变得越来越大。随着参数达到数万亿的LLM即将出现，拥有数十亿向量的向量数据库也不再遥远。作为优化工程师，我们好奇是否可以使向量数据库更高效地加载数据、更快地创建索引，以确保在添加新数据时的检索速度。这样做不仅可以减少用户等待时间，还可以使基于RAG的AI解决方案更加可持续。
本文将探讨向量数据库及其基准测试框架、不同方面的数据集以及用于性能分析的工具。我们还将分享在两个流行的向量数据库解决方案上的优化成果，激励你在性能和可持续性方面的优化旅程。
理解向量数据库 与传统的关系或非关系数据库以结构化方式存储数据不同，向量数据库包含通过嵌入或变换函数构建的单个数据项的数学表示，即向量。向量通常表示特征或语义含义，可以是短的或长的。向量数据库通过使用距离度量（如欧氏距离、点积或余弦相似度）进行相似性检索。
为了加速检索过程，向量数据使用索引机制进行组织。这些组织方法包括平面结构、倒排文件（IVF）、层次可导航小世界（HNSW）和局部敏感哈希（LSH）等。每种方法都有助于在需要时高效地检索相似向量。
我们来看看在GenAI系统中如何使用向量数据库。图1展示了将数据加载到向量数据库中的过程，以及在GenAI应用中使用它的过程。当你输入提示时，它会经历与生成数据库中向量相同的变换过程。然后使用该变换后的向量提示从向量数据库中检索相似向量。这些检索到的项本质上作为对话记忆，为提示提供上下文历史，类似于LLM的操作方式。这一特性在自然语言处理、计算机视觉、推荐系统和其他需要语义理解和数据匹配的领域特别有用。你的初始提示随后与检索到的元素“合并”，提供上下文，帮助LLM基于提供的上下文生成响应，而不仅仅依赖于其原始训练数据。
图1：RAG应用架构
向量存储和索引以实现快速检索。向量数据库有两种主要类型：传统数据库扩展以存储向量，和专门构建的向量数据库。一些提供向量支持的传统数据库示例包括Redis、pgvector、Elasticsearch和OpenSearch。专门构建的向量数据库示例包括Zilliz和Pinecone等专有解决方案，以及Milvus、Weaviate、Qdrant、Faiss和Chroma等开源项目。你可以在GitHub上的LangChain和OpenAI Cookbook中了解更多关于向量数据库的信息。
我们将详细介绍两种类型的向量数据库，分别是Milvus和Redis。
提高性能 在深入了解优化之前，让我们回顾一下如何评估向量数据库、一些评估框架和可用的性能分析工具。
性能指标 让我们看看一些关键指标，可以帮助你衡量向量数据库的性能。
加载延迟：衡量将数据加载到向量数据库内存并构建索引所需的时间。索引是一种数据结构，用于根据相似度或距离高效地组织和检索向量数据。内存中的索引类型包括平面索引、IVF_FLAT、IVF_PQ、HNSW、可扩展最近邻（ScaNN）和DiskANN。 召回率：搜索算法检索的Top K结果中找到的真实匹配或相关项的比例。较高的召回率表示相关项的检索更好。 每秒查询数（QPS）：向量数据库处理传入查询的速率。较高的QPS值表示更好的查询处理能力和系统吞吐量。 基准测试框架 基准测试向量数据库需要一个向量数据库服务器和客户端。在我们的性能测试中，我们使用了两个流行的开源工具。
VectorDBBench：由Zilliz开发和开源，VectorDBBench帮助测试不同索引类型的不同向量数据库，并提供方便的Web界面。 vector-db-benchmark：由Qdrant开发和开源，vector-db-benchmark帮助测试多种典型向量数据库的HNSW索引类型。它通过命令行运行测试，并提供Docker Compose文件简化服务器组件的启动。
图2：向量数据库基准测试框架
但基准框架只是方程的一部分。我们需要能够处理大量数据、各种向量大小和检索速度的数据集。让我们看看一些可用的公共数据集。
用于测试向量数据库的开放数据集 大型数据集是测试加载延迟和资源分配的好候选项。一些数据集具有高维数据，非常适合测试计算相似度的速度。
数据集的维度范围从25到2048。LAION数据集，一个开放的图像集合，已被用于训练非常大的视觉和语言深度神经模型，如稳定扩散生成模型。OpenAI的5M向量数据集，每个向量有1536个维度，由VectorDBBench通过在原始数据上运行OpenAI创建。由于每个向量元素都是FLOAT类型，仅保存这些向量大约需要29GB（5M * 1536 * 4）的内存，加上用于保存索引和其他元数据的额外内存，总共需要约58GB的内存用于测试。使用vector-db-benchmark工具时，请确保有足够的磁盘存储空间来保存结果。
要测试加载延迟，我们需要大量的向量集合，deep-image-96-angular提供了这样的数据集。为了测试索引生成和相似度计算的性能，高维向量提供了更多的压力。为此，我们选择了500K个维度为1536的向量数据集。
性能工具 我们已经介绍了如何通过施加压力来识别感兴趣的指标，但让我们看看底层发生了什么：计算单元有多忙，内存消耗情况，锁等待等。这些提供了数据库行为的线索，对于识别问题区域特别有用。
Linux的top工具提供了系统性能信息。然而，Linux中的perf工具提供了更深层次的见解。要了解更多信息，我们还推荐阅读Linux perf示例和Intel顶层微架构分析方法。另一个工具是Intel® vTune™ Profiler，它在优化应用程序和系统性能和配置时非常有用，涵盖HPC、云、物联网、媒体、存储等各种工作负载。
Milvus向量数据库优化 让我们看看我们如何尝试提高Milvus向量数据库的性能。
降低数据节点缓冲写入的内存移动开销 Milvus的写路径代理将数据写入日志代理，通过MsgStream传输。数据节点然后消费数据，将其转换并存储到段中。段会合并新插入的数据。合并逻辑分配一个新缓冲区来保存/移动旧数据和要插入的新数据，然后将新缓冲区作为旧数据返回进行下一次数据合并。这导致旧数据逐渐变大，从而使数据移动变慢。性能分析显示这一逻辑的开销很高。
图3：向量数据库中数据合并和移动产生了高性能开销
我们更改了合并缓冲区逻辑，直接将要插入的新数据追加到旧数据中，避免分配新缓冲区和移动大块旧数据。性能分析证实这一逻辑没有开销。微代码指标metric_CPU操作频率和metric_CPU利用率表明，系统不再需要等待长时间的内存移动，性能得到改善。加载延迟提高了超过60％。在GitHub上捕捉到了这一改进。
图4：减少复制后加载延迟提高了超过50％
通过减少内存分配开销来构建倒排索引 Milvus搜索引擎Knowhere使用Elkan k-means算法训练簇数据以创建倒排文件（IVF）索引。每轮数据训练定义一个迭代次数，次数越大，训练结果越好。然而，这也意味着Elkan算法将被更频繁地调用。
Elkan算法每次执行时处理内存分配和释放。具体来说，它分配内存来存储对称矩阵数据的一半大小，不包括对角元素。在Knowhere中，Elkan算法使用的对称矩阵维度设置为1024，结果内存大小约为2MB。这意味着每轮训练中，Elkan重复分配和释放2MB内存。
性能分析数据表明频繁的大内存分配活动。事实上，它触发了虚拟内存区域（VMA）分配、物理页面分配、页面映射设置和内核中内存cgroup统计更新。这种模式的大内存分配/释放活动在某些情况下还会加剧内存碎片化，这是一个显著的负担。
IndexFlatElkan结构专门设计和构建以支持Elkan算法。每个数据训练过程将初始化一个IndexFlatElkan实例。为了缓解Elkan算法频繁内存分配和释放带来的性能影响，我们重构了代码逻辑，将内存管理从Elkan算法函数移到IndexFlatElkan的构建过程中。这使得内存分配仅在初始化阶段进行一次，同时为当前数据训练过程中的所有后续Elkan算法函数调用服务，并有助于将加载延迟提高约3％。在这里可以找到Knowhere补丁。
通过软件预取加速Redis向量搜索 Redis是一个流行的传统内存键值数据存储，最近开始支持向量搜索。为了超越典型的键值存储，它提供了扩展模块；RediSearch模块使得可以直接在Redis中存储和搜索向量。
对于向量相似度搜索，Redis支持两种算法，即蛮力和HNSW。HNSW算法专门用于高效定位高维空间中的近似最近邻。它使用一个名为candidate_set的优先队列来管理所有向量候选者进行距离计算。
每个向量候选者包含大量元数据以及向量数据。因此，从内存中加载候选者时可能会导致数据缓存未命中，从而产生处理延迟。我们的优化引入了软件预取，以在处理当前候选者时主动加载下一个候选者。这一增强使向量相似度搜索的吞吐量提高了2％到3％。补丁正在上游处理中。
通过防止混合汇编代码的处罚来改变GCC默认行为 为了驱动最大性能，频繁使用的代码段通常是手写汇编。然而，当不同代码段由不同人或在不同时间编写时，使用的指令可能来自不兼容的汇编指令集，如Intel®高级矢量扩展512（Intel® AVX-512）和流SIMD扩展（SSE）。如果未适当地编译，混合代码会导致性能处罚。了解更多关于混合Intel AVX和SSE指令的信息。
你可以轻松确定是否使用了混合模式汇编代码且未编译代码并使用VZEROUPPER，从而导致性能处罚。可以通过如下命令观察sudo perf stat -e ‘assists.sse_avx_mix/event=0xc1,umask=0x10/’ <工作负载>。如果你的操作系统不支持该事件，使用cpu/event=0xc1,umask=0x10,name=assists_sse_avx_mix/。
Clang编译器默认插入VZEROUPPER，避免任何混合模式处罚。但GCC编译器仅在指定-O2或-O3编译器标志时插入VZEROUPPER。我们联系了GCC团队并解释了这一问题，他们现在默认正确处理混合模式汇编代码。
开始优化你的向量数据库 向量数据库在GenAI中扮演着重要角色，并且它们正在变得越来越大，以生成更高质量的响应。在优化方面，AI应用与其他软件应用没有什么不同，当使用标准性能分析工具以及基准测试框架和压力输入时，它们会揭示它们的秘密。
使用这些工具，我们发现了与不必要的内存分配、未能预取指令和使用不正确的编译器选项有关的性能陷阱。基于我们的发现，我们向Milvus、Knowhere、Redis和GCC编译器上游提交了增强，使AI更高效和可持续。向量数据库是一类值得你优化努力的重要应用程序。我们希望这篇文章能帮助你开始。