基于bge-m3 onnx模型获取稀疏向量

前言

bge-m3 是由 BAAI 在 2024 年推出的一个比较经典的 Embedding 模型，是一个支持多语言的 Embedding 模型，除此之外，它还支持输出稀疏向量和 ColBERT 向量，因此用途比较广泛。按照目前 AI 模型的发展速度（竞争很激烈），这个推出了一年多的模型理应算是个老古董了，实际上到目前为止（2025-9-20），该模型还能在 MTEB Leaderboard 榜单上保持在第 23 位的位置，且在 hugging face 上的月下载量保持在五六百万之多：

不过话又说回来，这大概率是因为 Embedding 模型还没有隔壁 LLM 那么卷，2025 截至目前也就阿里推出的 Qwen3-Embedding 模型很有竞争力；

什么是 Embedding 模型？

简单地来说，就是把一个输入（比如文本，图像等）转为一个高维向量的模型，这个向量就代表了该文本的特征向量；具体到文本的 Embedding，这些特征大多与语义和上下文相关：

图源 ^[1]

而在一篇文章中 ^[2]，里面提及到了 Embedding 模型的两个性质，我觉得挺合适：

性质一：每一个词具有唯一量化值，不同词需要具有不同的量化值
性质二：词义相近词需要有“相近”的量化值；词义不相近的词量化值需要尽量“远离”

其实我一直挺好奇为啥要叫 Embedding（中文叫嵌入，其实挺迷惑的，一听压根不知道指的是啥，就跟鲁棒性一样 😅），然后在维基上看到一个数学用语，我觉得大概就是借鉴这个吧：

看了这个数学中关于嵌入的定义后，就能明白其实 Embedding 模型干的就是这个映射的工作，即把输入映射到一个多维向量空间中。

关于 Embedding 的具体原理其实我也有很多不懂的地方，所以这里就不再过多赘述了，不过下面这篇文章我推荐可以阅读一下：

没有思考过 Embedding，不足以谈 AI

和大部分人一样，我对自然语言处理和语言模型的了解从ChatGPT开始。 和大部分人一样，第一次接触就被ChatGPT的能力所震惊 —— 计算机确实做到了理解人类的语言。 我也产生了几乎人人都会有的疑问：怎么做到的？ …

知乎专栏

为什么要使用 onnx 模型？

ONNX（Open Neural Network Exchange，开放神经网络交换格式）是一种跨平台、开源的神经网络模型标准格式，由微软、亚马逊、Facebook 等企业联合发起。它定义了统一的计算图结构和张量数据格式，能将不同深度学习框架（如 PyTorch、TensorFlow、PaddlePaddle 等）训练的模型 “翻译” 成统一格式，实现模型在不同框架、硬件设备间的无缝迁移与部署。

很简单，由于硬件受限，我需要支持在某些没有 GPU 资源（即纯 CPU 环境）的机器中运行基于 bge-m3 模型的向量化服务；虽然像 PyTorch/TensorFlow 这类模型文件也可以在 CPU 环境进行推理运行，但是相比 onnx 模型（onnxruntime）专门对推理进行了通用的优化，只要把这些模型转为 onnx 模型就可以获得开箱即用的推理加速（相比原生模型，根据具体硬件架构的不同，可以提升 10 倍甚至上百倍），而恰好 bge-m3 官方就在模型仓库中提供了非量化版本的 onnx 模型文件：

推理性能对比

我让 AI 写了一段测试代码，用于测试基于原始模型和 onnx 模型对同样 Embedding 任务的推理速度，结果很出乎我的意料，没想到性能提升百倍以上😀：

• 原始模型：基于官方的 FlagEmbedding 库进行推理
• onnx 模型：基于 onnxruntime 库进行推理
• CPU：Intel® Xeon® E-2224G CPU @ 3.50GHz （4 核）
• python 版本：3.11

原始模型的推理速度

开始性能测试 - 嵌入类型: hybrid, 每个测试运行 5 次
==================================================
文本长度: 13 字符, 约 19 tokens
平均处理时间: 7.0167 秒
处理速度: 2.71 tokens/秒
--------------------------------------------------
文本长度: 43 字符, 约 61 tokens
平均处理时间: 20.2228 秒
处理速度: 3.02 tokens/秒
--------------------------------------------------
文本长度: 100 字符, 约 136 tokens
平均处理时间: 33.7305 秒
处理速度: 4.03 tokens/秒
--------------------------------------------------
文本长度: 390 字符, 约 541 tokens
平均处理时间: 163.0791 秒
处理速度: 3.32 tokens/秒
--------------------------------------------------

总体性能:
总处理token数: 757
总处理时间: 224.0491 秒
平均处理速度: 3.27 tokens/秒

onnx 模型的推理速度

开始性能测试 - 嵌入类型: hybrid, 每个测试运行 5 次
==================================================
文本长度: 13 字符, 约 19 tokens
平均处理时间: 0.0583 秒
处理速度: 325.99 tokens/秒
--------------------------------------------------
文本长度: 43 字符, 约 61 tokens
平均处理时间: 0.0981 秒
处理速度: 621.76 tokens/秒
--------------------------------------------------
文本长度: 100 字符, 约 136 tokens
平均处理时间: 0.1095 秒
处理速度: 1242.13 tokens/秒
--------------------------------------------------
文本长度: 390 字符, 约 541 tokens
平均处理时间: 0.3337 秒
处理速度: 1621.42 tokens/秒
--------------------------------------------------

总体性能:
总处理token数: 757
总处理时间: 0.5995 秒
平均处理速度: 952.83 tokens/秒

总结

不难看出在使用原始模型在纯 CPU 环境下进行推理的速度实在是很慢（这也是为啥要进行优化的原因，这种推理速度实在无法在实际业务场景中进行使用），而仅仅只是把模型切换为 onnx 就能得到上百倍的性能提升，只能说 onnx 模型自带的推理优化很惊人。

不过实际上在另一种 CPU 规格（海光的某款 8 核 CPU）上进行同样的测试时，发现性能提升并没有上面的那么显著，大概是缺少了某些底层矩阵加速库和指令集的支持。

bge-m3 稀疏向量的计算

使用 bge-m3 的 onnx 模型唯一的问题就是，无法直接获取到其稀疏向量和 colbert 向量的输出了，因为 bge 官方提供的推理库——FlagEmbedding 可以直接通过 return_dense、return_sparse 和 return_colbert_vecs 这三个参数来控制输出的向量类型，而 onnxruntime 得到的结果只是稠密向量。因此使用 onnx 模型时想要得到 bge-m3 的稀疏向量，就得自行按照该模型在论文中提出的稀疏向量计算公式进行实现 ^[3]：

在一开始我并不知道 bge-m3 有自己的稀疏向量计算公式，于是让 AI 随便写了一个稀疏向量计算的算法实现，然后输出的结果看起来也想那么回事，结果后面一对比之前的输出发现权重结果完全对不上，只能说有点太艺高人胆大了（一个敢写，一个敢信，这就是为啥碰到不懂的领域时会出现被 AI 蒙住的情况😂）

代码实现

def _compute_lexical_weights_v2(
        self,
        hidden_state: np.ndarray,
        input_ids: list[int],
        return_embedding: bool = True,
    ) -> Dict[str, float]:
        """
        source: https://zhuanlan.zhihu.com/p/1889076509584638806
        基于BGE-M3原始算法计算词汇权重
        实现公式: w_qt = ReLU(W_lex * H_q[i])
      
        Args:
            hidden_state: token embeddings, shape (seq_len, hidden_size)
            input_ids: token ID列表
            return_embedding: 是否返回完整的sparse embedding（当前总是返回字典）
          
        Returns:
            词汇权重字典，键为字符串形式的token ID，值为权重
        """
        # 将numpy数组转换为torch tensor
        hidden_state_tensor = torch.from_numpy(hidden_state).float()
        input_ids_tensor = torch.tensor(input_ids, dtype=torch.long)
      
        # 计算token权重：ReLU(sparse_linear(hidden_state))
        token_weights = torch.relu(self.sparse_linear(hidden_state_tensor))
        token_weights = token_weights.squeeze(-1)  # 去除最后一个维度 (seq_len, 1) -> (seq_len,)
      
        # 创建稀疏嵌入矩阵
        sparse_embedding = torch.zeros(
            self.vocab_size,
            dtype=token_weights.dtype,
            device=token_weights.device
        )
      
        # 使用scatter操作将token权重分配到对应的token ID位置
        sparse_embedding = sparse_embedding.scatter(0, input_ids_tensor, token_weights)
      
        # 转换为字典格式，只保留非零权重
        result = defaultdict(float)
        for token_id in input_ids:
            weight = sparse_embedding[token_id].item()
            if weight > 0 and token_id not in self.special_tokens:
                result[str(token_id)] = max(result[str(token_id)], weight)
      
        return result

其中 hidden_state 就是 token 级别的稠密向量：

# 分词
inputs = self.tokenizer(
	texts,
	padding="longest",
	truncation=True,
	max_length=max_length,
	return_tensors="np"
)

# 转换为ONNX输入格式
inputs_onnx = {k: ort.OrtValue.ortvalue_from_numpy(v) for k, v in inputs.items()}

# 推理
outputs = self.ort_session.run(None, inputs_onnx)

# 获取token级别的embeddings和句子级别的embedding
token_embeddings = outputs[0]  # shape: (batch_size, seq_len, hidden_size)
sentence_embeddings = outputs[1]  # shape: (batch_size, hidden_size)

sparse_linear 的作用

实际上对应的就是公式中的 $W_{lex}$，即一个特殊的线性层 ^[4]：

线性变换：在编码器栈的输出上添加了一个额外的线性层。通过此层计算每个 Token 的重要性权重，BGE-M3 得到了一组权重（$W_{lex}$）。

实际上这个特殊的线性层就是把每个 token 的 hidden state 转为一个一维的权重值，线性层是预训练得到到每个 token 在上下文中的重要性，因此基于线性层的转换可以得到该 token 在上下文中的权重。如果把这个线性层替换成自定义的权重，则可以得到自定义的稀疏向量（即可以自定义具体分词的重要性）。

该权重文件在 bge-m3 仓库的顶层目录中：

可以使用 pytorch 直接读取：

self.sparse_linear = torch.nn.Linear(
	in_features=self.hidden_size,  # 使用正确的hidden_size
	out_features=1
)
sparse_linear_path = os.path.join(model_path, "sparse_linear.pt")
if os.path.exists(sparse_linear_path):
	self.sparse_linear.load_state_dict(torch.load(sparse_linear_path, map_location='cpu', weights_only=True))
	logger.info(f"成功加载sparse_linear权重: {sparse_linear_path}")
else:
	logger.info(f"sparse_linear.pt not found, using default weights")

从打印结果来看，sparse_linear 权重文件的维度跟稠密向量输出的维度一致，即 1024：

embedding_models/bge-m3/sparse_linear.pt, torch.Size([1, 1024])

特殊字符的处理

在基于上述公式计算得到输入文本的稀疏向量时，我发现有一个特定维度的权重值异常的高，且每个输入文本该维度都存在，因此会导致基于稀疏向量的搜索相似度过高，从而匹配结果不准确的问题。后面经过排查，发现原来是有一个特殊字符并没有包含到 special_tokens_map.json 所导致的，这个特殊字符就是 SentencePiece（词汇分界标记），在 bge-m3 模型中的 token id 为 6：

// bge-m3模型的special_tokens_map.json文件
{
  "bos_token": {
    "content": "<s>",
    "lstrip": false,
    "normalized": false,
    "rstrip": false,
    "single_word": false
  },
  "cls_token": {
    "content": "<s>",
    "lstrip": false,
    "normalized": false,
    "rstrip": false,
    "single_word": false
  },
  "eos_token": {
    "content": "</s>",
    "lstrip": false,
    "normalized": false,
    "rstrip": false,
    "single_word": false
  },
  "mask_token": {
    "content": "<mask>",
    "lstrip": true,
    "normalized": false,
    "rstrip": false,
    "single_word": false
  },
  "pad_token": {
    "content": "<pad>",
    "lstrip": false,
    "normalized": false,
    "rstrip": false,
    "single_word": false
  },
  "sep_token": {
    "content": "</s>",
    "lstrip": false,
    "normalized": false,
    "rstrip": false,
    "single_word": false
  },
  "unk_token": {
    "content": "<unk>",
    "lstrip": false,
    "normalized": false,
    "rstrip": false,
    "single_word": false
  }
}

因此需要在加载模型的阶段，手动把 SentencePiece 字符加入到特殊字符列表中，然后在计算稀疏向量时就可以跳过该 token 的权重：

# 缓存特殊token ID，参考bge-m3的special_tokens_map.json
self.special_tokens = set([
	self.tokenizer.cls_token_id,
	self.tokenizer.sep_token_id,
	self.tokenizer.pad_token_id,
	self.tokenizer.mask_token_id,
	self.tokenizer.unk_token_id,
	self.tokenizer.bos_token_id,
	self.tokenizer.eos_token_id,
])
logger.info(f"特殊token集合: {self.special_tokens}")

# 添加SentencePiece边界标记token
# 检查常见的SentencePiece边界标记
boundary_tokens = [' ', '_', '】', '[UNK]']
for token in boundary_tokens:
	try:
		token_id = self.tokenizer.convert_tokens_to_ids(token)
		if token_id is not None and token_id != self.tokenizer.unk_token_id:
			self.special_tokens.add(token_id)
			logger.info(f"添加SentencePiece边界token '{token}' (ID: {token_id}) 到特殊token集合")
	except Exception as e:
		logger.error(f"添加SentencePiece边界token '{token}' 失败: {e}")

# 特别处理：如果token ID 6对应 字符，确保将其添加为特殊token
try:
	token_6_str = self.tokenizer.convert_ids_to_tokens([6])[0]
	if token_6_str in ['▁', '_']:
		self.special_tokens.add(6)
		logger.info(f"确认添加边界token ID 6 ('{token_6_str}') 为特殊token")
except:
	pass

扩展阅读：Learned Sparse Embedding

由于 bge-m3 模型中 sparse_linear 的存在，因此其计算得到的稀疏向量已经不是传统意义上基于统计方法（如 BM25）得到的分词权重了，这类基于机器学习方法计算得到的稀疏 Embedding 被称为 Learned Sparse Embedding；详细内容推荐阅读下面的文章：

BGE-M3 and Splade: Two Popular Sparse Embedding Models - Zilliz Learn

In this blog, we’ve journeyed through the intricate world of vector embeddings and explored how BGE-M3 and Splade generate learned sparse embeddings.

zilliz.com

Enhancing Information Retrieval with Sparse Embeddings | Zilliz Learn - Zilliz Learn

Explore the inner workings, advantages, and practical applications of learned sparse embeddings with the Milvus vector database

zilliz.com

spade

这是另一种不同于 bge-m3 模型计算稀疏向量的 Learned Sparse Embedding 方法，该方法不仅可以得到当前输入的 token 的权重值，还可以得到相同语义（同一上下文中）的没有出现的 token 的权重值。

该方法最后通过统计每个 token 在其位置上所有token 出现的概率得到稀疏向量权重。

---

1. 深入理解Embedding Models（嵌入模型）：从原理到实战（下） ↩︎

2. 没有思考过 Embedding，不足以谈 AI ↩︎

3. https://bge-model.com/bge/bge_m3.html ↩︎

4. https://zilliz.com/learn/bge-m3-and-splade-two-machine-learning-models-for-generating-sparse-embeddings ↩︎