利用 Optimum Intel 和 fastRAG 在 CPU 上优化文本嵌入

嵌入模型的应用场景非常广泛，检索、重排、聚类、分类，几乎都能看到它们的身影。近几年，这个领域的技术进展相当迅猛，不少基于语义的应用因此受益匪浅。像 BGE、GTE、E5 这些模型，长期霸占着各类评测基准的榜单头部，有些场景下甚至比一些付费的私有嵌入服务还要能打。Hugging Face 的模型仓库里，嵌入模型的尺寸选择也很丰富，从轻量级的 100-350M 参数小模型，到 7B 参数的大模型，一应俱全。很多做语义搜索的应用，倾向于选用基于编码器架构的轻量级模型，这时候，CPU 就成了一个非常有吸引力的运行选择——典型的场景就是，完全可以用 CPU 把这类模型的推理高效地跑起来。

嵌入模型与信息检索

嵌入模型的核心能力，是把文本转换成稠密向量，通过这些向量，文本的语义和上下文信息被高度浓缩。这种上下文感知的文本表征，比传统的关键词匹配更精准。在实际检索中，我们通常用向量之间的余弦相似度来衡量语义上的接近程度。

但问题来了：信息检索是不是只用稠密向量就万事大吉了？这里其实需要一些权衡：

• 稀疏检索
  ，像 BM25，思路是把文本集拆成 n-元组、短语或者元数据的集合，然后在这个集合上做高效率、大规模的搜索。好处是速度快、可解释性强，但缺点是，如果查询和文档在用词上对不上，很容易漏掉一些语义相关但字面不匹配的内容。

• 语义检索
  ，则是走另一条路，把文本编码成稠密向量。相比“词袋”模型，它能更好地理解上下文和同义词。即使查询和文档里没有共同的关键词，只要意思相近，就能被检索出来。不过，代价也不小——计算量更大，延迟更高，底层依赖复杂的编码模型。

嵌入模型在 RAG 中扮演的角色

在 RAG（检索增强生成）应用中，嵌入模型可以说是贯穿始终：

• 离线索引
  ：当构建或更新文档数据库时，需要用嵌入模型把文档批量编码成稠密向量。

• 查询编码
  ：每当用户发起查询，嵌入模型需要实时将查询转换成向量，才能去向量库里搜。

• 重排
  ：首轮检索得到的候选文档列表，通常还需要再“过一遍”嵌入模型，把它们的向量与查询向量重新比较，以完成重排，把最相关的结果提到最前面。

所以，优化 RAG 流程中的嵌入模型环节，直接关系到整个系统的效率：

• 文档索引与更新
  ：追求高吞吐，这样大型文档集的编码和索引才能快速完成，大幅缩短建库和更新的时间。

• 查询编码
  ：低延迟是保证检索实时性的关键。高吞吐则能支撑更高的查询并发，让系统有更好的扩展性。

• 重排
  ：首轮检索后，嵌入模型需要快速对候选文档进行编码。低延迟意味着重排更快，能更好地满足时效性敏感的应用；高吞吐则允许对更大的候选集进行并行重排，让最终的排序结果更全面。

用 Optimum Intel 和 IPEX 给嵌入模型提速

Optimum Intel 是一个开源库，专门针对英特尔硬件优化基于 Hugging Face 库构建的端到端流程。它集成了多种模型加速技术，比如低比特量化、权重剪枝、知识蒸馏以及运行时优化。在优化过程中，它能充分利用英特尔 CPU 中的 A VX-512、VNNI 以及 AMX（高级矩阵扩展）指令集来加速模型。每个 CPU 核心内部都内置了 BF16 和 INT8 的 GEMM（通用矩阵乘法）翻跟斗，这对深度学习推理任务来说是实打实的硬件“外设”。此外，PyTorch 2.0 及其扩展库 IPEX，也在软件层面充分利用了 AMX 指令，进一步提升了推理性能。

用 Optimum Intel 来优化预训练模型的推理任务，流程很清晰。官方的文档和示例代码也提供了不少现成的参考。

实操案例：优化 BGE 嵌入模型

这篇文章我们主要聚焦近期频繁亮相的 BGE 系列嵌入模型，它们在知名的 MTEB 排行榜上表现非常亮眼。

BGE 模型的技术要点

BGE 系列属于双编码器模型，采用 Transformer 编码器架构。它的训练目标很简单：让语义相似的两个文本，其嵌入向量之间的距离（通常是余弦相似度）最大化。举个常见的例子，可以用 BERT 作为基础模型，通过微调使其成为一个嵌入模型，用来为文档生成向量。从模型的输出中提取文本嵌入的方法有很多种，比如直接取 [CLS] 词元的嵌入，或者对所有输入词元的嵌入取平均值。

双编码器设计的优势在于结构相对简单，它只对单个文档的上下文进行编码，无法像交叉编码器那样对“查询-文档”或“文档-文档”进行联合建模。但有意思的是，即便是这样看似“偷懒”的架构，最先进的双编码器模型也能在性能上极具竞争力。再加上架构简单带来的速度优势，这类模型自然成了许多场景下的“当红炸子鸡”。

下文主要关注 3 个 BGE 模型：BGE-small、BGE-base 和 BGE-large，它们的参数量分别为 45M、110M 和 355M，对应的嵌入维度分别是 384、768 和 1024。顺便说一句，下面展示的优化过程是通用的，完全可以套用到其他嵌入模型上（无论是双编码器还是交叉编码器）。

模型量化：分步指南

下面我们演示如何提升嵌入模型在 CPU 上的性能，主要关注的指标是降低延迟（batch size 为 1）和提高吞吐（每秒处理的文档数）。我们会用到 Optimum-Intel 和 Intel Extension for Transformers 来量化模型，并用 IPEX 优化在英特尔硬件上的运行时间。

第一步：安装环境

运行以下命令安装必要的库：

pip install -U optimum[neural-compressor] intel-extension-for-transformers

第二步：训练后静态量化

训练后静态量化需要一个校准数据集，来确定模型权重和激活值的动态范围。在校准过程中，模型会在一些有代表性的样本上跑一遍，收集统计信息，然后根据这些信息对模型进行量化，目标是让精度损失降到最低。

下面是量化的代码片段：

def quantize(model_name: str, output_path: str, calibration_set: "datasets.Dataset"):
    model = AutoModel.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    def preprocess_function(examples):
        return tokenizer(examples["text"], padding="max_length", max_length=512, truncation=True)

    vectorized_ds = calibration_set.map(preprocess_function, num_proc=10)
    vectorized_ds = vectorized_ds.remove_columns(["text"])

    quantizer = INCQuantizer.from_pretrained(model)
    quantization_config = PostTrainingQuantConfig(approach="static", backend="ipex", domain="nlp")
    quantizer.quantize(
        quantization_config=quantization_config,
        calibration_dataset=vectorized_ds,
        sa ve_directory=output_path,
        batch_size=1,
    )
    tokenizer.sa ve_pretrained(output_path)

在本例中，我们使用了数据集的一个子集作为校准集。

第二步（续）：加载与推理

量化完成后，加载模型非常简单：

from optimum.intel import IPEXModel

model = IPEXModel.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")

然后用 `transformers` 的 API 把句子编码成向量：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Intel/bge-small-en-v1.5-rag-int8-static")
inputs = tokenizer(sentences, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    # 取 [CLS] token 作为句子嵌入
    embeddings = outputs[0][:, 0]

关于如何正确配置 CPU 以获得最佳性能，我们在后面的评估部分会详细说明。

用 MTEB 评估模型精度

把模型权重从 fp32 量化到 int8，必然会损失一部分精度。为了量化这个损失，我们在两项任务上对量化模型和原始模型做了对比：

• 检索
  ：对语料库进行编码并建立索引，然后对给定的查询执行搜索，找出并排序最相似的文本。

• 重排
  ：对初步检索的结果进行重新排序，提升与查询相关性排名的质量。

下表总结了量化模型（int8）与原始模型（fp32）在多个数据集上的平均准确度差异（重排任务用 MAP 评估，检索任务用 NDCG@10 评估）。

可以看到，量化模型的重排任务精度损失都在 0.3% 以内，检索任务的损失也低于 1.55%。这种精度的损失，换来的通常是数倍的性能提升，在不少实际场景中是完全可接受的。

速度与延迟实测

我们将量化模型的推理性能，与另外两种常见的模型推理方案做了对比：

1. 用 PyTorch 和 Hugging Face transformers 库，以 bf16 精度运行原始模型（基线）。
2. 用 IPEX 以 bf16 精度运行模型，并配合 torchscript 对模型进行图化。

测试环境如下：

• 硬件
  ：双路第四代英特尔至强 8480+ CPU，每路 56 个核心。

• PyTorch 基线
  ：使用单路 CPU 的 56 个核心。

• IPEX/Optimum 测例
  ：通过 ipexrun 使用单路 CPU，核心数在 22 到 56 之间变化。

• 内存管理
  ：所有测例都开启了 TCMalloc 并配置了相关环境变量。

延迟性能

延迟测试主要关注模型的响应速度，这直接关系到 RAG 流程中查询编码的快慢。我们固定 batch size 为 1，测量了不同文档长度下的延迟。

总的来说，量化模型的延迟是所有方案中最低的。BGE-small 和 BGE-base 的延迟低于 10 毫秒，BGE-large 也控制在 20 毫秒以内。相比原始 bf16 模型，量化模型的延迟最高提升了 4.5 倍。

图 1: 各尺寸 BGE 模型的延迟对比

吞吐性能

吞吐量测试的目标是找到模型在编码任务上的峰值性能——即每秒能处理的文档数。我们固定文本长度为 256 个 token（代表 RAG 流程中的平均文档长度），并在不同的 batch size（4、8、16、32、64、128、256）下进行测试。

结果很清晰：量化模型的吞吐量全面领先。在 batch size 达到 128 时，几乎所有模型都达到了吞吐峰值。总体上，对于三种尺寸的模型，量化模型在各个 batch size 下的吞吐量，都比基线 bf16 模型高出大约 4 倍。

图 2: BGE small 模型的吞吐量

图 3: BGE base 模型的吞吐量

图 4: BGE large 模型的吞吐量

在 fastRAG 中使用量化嵌入模型

说了这么多，最后还是得上手试试。我们通过一个例子来演示如何把优化后的检索/重排模型集成到 fastRAG 框架中——当然，这套方法也同样适用于 LangChain 或 LlamaIndex 等其他框架。

fastRAG 是一个高效且专门针对检索增强生成进行研究探索的框架，它集成了最前沿的 LLM 和信息检索算法，并且与 Haystack 生态完全兼容。fastRAG 实现了多种新型的高效 RAG 模块，非常适合在英特尔硬件上部署。

我们需要将上面优化好的双编码器嵌入模型，用于 fastRAG 中的两个核心模块：

1. 稠密检索器
   ：用于创建稠密向量索引，以及从索引中检索文档。

2. 重排器
   ：在文档重排流程中调用，用于对检索结果进行优化排序。

用优化检索器实现快速索引

首先，创建一个内存中的文档存储：

from haystack.document_store import InMemoryDocumentStore

document_store = InMemoryDocumentStore(use_gpu=False, use_bm25=False, embedding_dim=384, return_embedding=True)

然后加入几个示例文档：

from haystack.schema import Document

# 待索引的示例文档
examples = [
   "There is a blue house on Oxford Street.",
   "Paris is the capital of France.",
   "The first commit in fastRAG was in 2022"
]

documents = []
for i, d in enumerate(examples):
    documents.append(Document(content=d, id=i))
document_store.write_documents(documents)

接着，用量化后的双编码器模型初始化检索器，并对文档库中的所有文档进行编码：

from fastrag.retrievers import QuantizedBiEncoderRetriever

model_id = "Intel/bge-small-en-v1.5-rag-int8-static"
retriever = QuantizedBiEncoderRetriever(document_store=document_store, embedding_model=model_id)
document_store.update_embeddings(retriever=retriever)

用优化排名器实现高效重排

下面的代码展示了如何把量化模型加载为排序器，它会为检索器返回的每个文档重新编码和排序：

from haystack import Pipeline
from fastrag.rankers import QuantizedBiEncoderRanker

ranker = QuantizedBiEncoderRanker("Intel/bge-large-en-v1.5-rag-int8-static")

p = Pipeline()
p.add_node(component=retriever, name="retriever", inputs=["Query"])
p.add_node(component=ranker, name="ranker", inputs=["retriever"])
results = p.run(query="What is the capital of France?")

# 打印检索到的文档
print(results)

到这里，一个完整的 RAG 流水线就跑通了。它从文档库中检索文档，并用另一个（也是量化后的）嵌入模型对结果进行重排，整个过程完全跑在 CPU 上，而且性能相当可观。更多完整的例子可以在 fastRAG 的官方示例中找到。