Model2Vec：从任何 Sentence Transformer 中蒸馏出小而快的模型

又名：如何让 sentence transformer 速度提升 500 倍，体积缩小 15 倍

（大）语言模型已成为特征提取的事实标准。虽然这些模型在大量任务上展现了最先进的性能，但它们也带来了沉重的资源需求：巨大的能源消耗、计算需求和更长的处理时间。尽管有很多方法可以使现有的（Sentence）Transformer 变得更快，例如量化或专用内核，但它们仍然相对较慢，尤其是在 CPU 上。如果你需要更快的速度，并且正在开发一个有时间限制的产品（例如搜索引擎），或者可用资源非常少，该怎么办？

这正是 Model2Vec 的用武之地——它提供静态嵌入，对硬件和生态友好，同时保持强大的性能。

在这篇博客中，我们将讨论 Model2Vec 是什么，它是如何工作的，如何使用它，以及它的性能。

Model2Vec 架构的可视化。

目录

• 什么是 model2vec？
• 如何使用 model2vec
• 结果
• 结论
• 致谢

什么是 Model2Vec？

Model2Vec 是一种从任何 Sentence Transformer 中蒸馏出小型、快速、高性能静态模型的技术。从宏观上看，它的工作原理是：将一个词汇表通过 sentence transformer 模型，然后使用 PCA 降低所得嵌入的维度，最后使用 zipf 权重对嵌入进行加权。这个过程不需要数据集，只需要一个模型（以及可选的词汇表）。在推理时，我们只需取句子中出现的所有词元嵌入的平均值。因此，Model2Vec 模型是完全非上下文相关的。虽然这听起来像是一个很大的缺点，但我们将展示，考虑到它的小巧和快速，它的性能仍然相当不错。

以上内容可能听起来有些复杂，让我们来详细解释一下。

Transformer 和嵌入

在 sentence transformer 的编码步骤中，一个字符串首先被切分成子词（subword）词元。这些词元的嵌入被输入到模型中，模型将它们上下文关联起来以创建高质量的句子表示。在输出端，你输入的词元数量与输出的嵌入数量相同，所以如果你的输入句子包含 10 个词元，你也会得到 10 个输出词元。然后，这些词元通过一个池化机制（pooling mechanism）转换成一个句子表示，这个机制可以是一个简单的平均值，也可以是一个特殊的池化模块。

回到 Model2Vec：这个项目最初是作为 sentence transformer 的一种缓存机制。因为 transformer 的词汇表通常只有大约 32k 个词元，像 astoundingly 这样的词会被切分成四个唯一的词元：'as', '##tou', '##nding', '##ly'，这意味着每次这个词出现时，我们都会重新计算这四个词元之间的注意力。但这个词的意义可能根本没有歧义！

然而，当我们开始实现这个想法时，我们注意到其实根本不需要缓存任何单词，只需使用单个词元的输出表示就可以获得良好的句子表示。这正是 Model2Vec 的基本工作模式：对于 sentence transformer 词汇表中的每个 32k 输入词元，我们进行一次前向传播，然后存储生成的嵌入。对于一个新句子，我们只需取我们计算出的词元嵌入的平均值。

请注意，model2vec 模型的输出词元表示是非上下文相关的。与普通的 transformer 模型不同，它无法在不同上下文中为同一个词元赋予不同的含义。虽然这看起来是一个巨大的缺点，但我们认为实际的上下文为模型提供了足够的消歧潜力。

除了这个技巧，我们还发现需要另外两个技巧才能获得最佳性能。

PCA

我们使用主成分分析（PCA）来降低生成的词元空间的维度。通常，使用 PCA 会导致性能下降，因为你会丢弃信息。然而，在我们的案例中，降低维度实际上显著提高了性能。我们认为这是因为 PCA 也对生成的空间进行了归一化，即消除了原始向量空间中的偏差，从而使从向量中学习变得更容易。

Zipf

由于我们对空间中的词元取简单的平均值，因此正确地加权向量非常重要。通常，sentence transformer 会根据上下文为我们正确地加权所有词元，但我们不再拥有这种便利。直观上，我们希望使用类似逆文档频率（IDF）的方法来降低非常频繁或不重要词语的权重。但我们无法访问可以计算文档频率的语料库。

为了克服这个问题，我们选择使用语言科学中一个众所周知的原理，即在一个按频率排序的列表中，列表中项目的频率遵循幂律分布。这被称为齐夫定律（Zipf's law）。因此，如果我们假设一个词汇表是按频率排序的，我们就可以准确地降低非常频繁项目的权重，而无需访问实际频率。由于分词器词汇表是按频率排序的，我们已经有了一个排序列表，所以这个优化可以无需任何额外工作就应用。

对嵌入应用 PCA 和 Zipf 加权效果的可视化。

用法

Model2Vec 库有两种主要的使用模式：**蒸馏** 和 **推理**。在蒸馏模式下，您可以使用任何 Sentence Transformer（以及可选的您自己的词汇表）来蒸馏自己的模型。在推理模式下，您可以使用蒸馏后的模型（或使用我们预先蒸馏的模型）以极高的速度为您的文本数据生成嵌入。

有三种方法可以蒸馏一个模型

• **输出（Output）**：行为非常像一个真正的 sentence transformer，即它使用一个子词分词器，并简单地对其词汇表中的所有词元进行编码。这种方法创建速度非常快（在 CPU 上 30 秒），体积非常小（float32 格式下 30 MB），但在某些任务上性能可能稍差。
• **词汇表（单词） (Vocab (word))**：在此模式下，您可以传入自己的词汇表来创建表示。这使您能够为您拥有的任何领域内数据创建良好的表示，并且可以作为 GloVe 或 word2vec 的直接替代品。
• **词汇表（子词） (Vocab (subword))**：在此模式下，您可以传入自己的词汇表，但它也使用子词词汇表来创建表示。这使您能够为您拥有的任何领域内数据创建良好的表示。

请注意，尽管基于词汇表的模型在 RAM 方面更大，但所有模型的速度都是一样的，因为我们的模型与词汇表大小无关。

Model2Vec 嵌入可用于各种应用，例如文本分类、聚类、构建搜索引擎或 RAG 系统。它们特别适合需要快速、轻量级嵌入且资源需求低的应用。

正如我们接下来将展示的，Model2Vec 非常易于使用。它既可以作为独立包使用，也可以直接在 Sentence Transformers 中使用。这意味着您可以轻松地将其集成到任何支持 Sentence Transformers 的流程中（例如 LangChain 和 LlamaIndex）。您还可以直接使用 Sentence Transformers 训练 model2vec 模型，保持快速的推理速度，同时直接为您的用例进行优化。

如何使用 Model2Vec

安装

Model2Vec 可以使用 pip 安装

pip install model2vec

用法

推理

开始使用 Model2Vec 最简单的方法是从我们的 HuggingFace hub 下载我们的旗舰模型之一。这些模型是预训练好的，可以直接使用。以下代码片段展示了如何加载模型并生成嵌入。

from model2vec import StaticModel

# Load a model from the HuggingFace hub (in this case the M2V_base_output model)
model_name = "minishlab/M2V_base_output"
model = StaticModel.from_pretrained(model_name)

# Make embeddings
embeddings = model.encode(["It's dangerous to go alone!", "It's a secret to everybody."])

或者蒸馏您自己的模型并直接使用它们

from model2vec import distill

# Choose a Sentence Transformer model
base_model_name = "BAAI/bge-base-en-v1.5"

# Distill an output model with the chosen dimensions
model = distill(model_name=base_model_name, pca_dims=256)

# Make embeddings
embeddings = model.encode(["supervillain Ganondorf has invaded Hyrule!"])

print(model.tokenizer.encode("supervillain Ganondorf has invaded Hyrule!", add_special_tokens=False).tokens)
# ['super', '##vill', '##ain', 'gan', '##ond', '##orf', 'has', 'invaded', 'h', '##yr', '##ule', '!']

# It looks like we split Ganondorf and Hyrule up into many subtokens
# To solve this, we can add these words to our vocabulary.
vocabulary = ["supervillain", "ganondorf", "hyrule"]

# Distill the model with the custom vocabulary.
model = distill(model_name=base_model_name, vocabulary=vocabulary, pca_dims=256)

print(model.tokenizer.encode("supervillain Ganondorf has invaded Hyrule!", add_special_tokens=False).tokens)
# ['supervillain', 'ganondorf', 'has', 'invaded', 'hyrule', '!']
# Much better.

Model2Vec 也直接在 Sentence Transformers 中得到支持。要在 Sentence Transformers 中使用 Model2Vec，您可以使用 `from_model2vec` 初始化一个 `StaticEmbedding` 类。要在 Sentence Transformers 中直接进行蒸馏，可以使用 `from_distillation` 初始化 `StaticEmbedding` 类。

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

# Initialize a StaticEmbedding module using a pre-trained model
static_embedding = StaticEmbedding.from_model2vec("minishlab/M2V_base_output")
model = SentenceTransformer(modules=[static_embedding])
embeddings = model.encode(["It's dangerous to go alone!", "It's a secret to everybody."])

# Or distill your own directly without leaving sentence-transformers
static_embedding = StaticEmbedding.from_distillation("BAAI/bge-base-en-v1.5", device="cpu", pca_dims=256)
model = SentenceTransformer(modules=[static_embedding])
embeddings = model.encode(["It's dangerous to go alone!", "It's a secret to everybody."])

结果

我们在大量任务和数据集上评估了 Model2Vec。Model2Vec 在 MTEB 以及两个额外任务上进行了评估：PEARL（一个短语表示任务）和 WordSim（一个词语相似度任务集合）。结果如下表所示。

可以看出，Model2Vec 在所有任务上都显著优于 GloVe 和 BPEmb，甚至在某些任务上超过了速度慢得多的 MiniLM 模型。

此外，我们还在一些不在 MTEB 中的分类数据集上评估了 Model2Vec。我们还用这些数据集来测试模型的速度。结果如下表所示。

再次，Model2Vec 在所有任务上都优于 GloVe 和 BPEmb，甚至表现出与 MiniLM 相似的性能。

下图显示了每秒处理的句子数与平均分类得分之间的关系。圆圈的大小对应于模型中的参数数量（越大表示参数越多）。该图表明，Model2Vec 模型比其他模型快得多，同时在分类性能方面仍能与 all-MiniLM-L6-v2 模型相媲美。

所有分类数据集的平均准确率与每秒处理句子数的对比图。圆圈大小表示模型大小。

消融研究

为了更好地理解影响 Model2Vec 性能的因素，我们进行了一系列全面的消融研究，涵盖了模型架构和预处理方法的各个方面。在这些研究中，我们检验了 PCA、Zipf 加权以及使用 Sentence Transformers 与常规 transformer 模型等关键元素的影响。我们还比较了输入嵌入与输出嵌入的性能，因为似乎这些也应该表现良好。结果如下表所示。

这些结果中有四个主要发现

1. 非 Sentence Transformer 模型效果不佳。这可以通过比较 `M2V_bert_output_nozipf_nopca`（使用BERT，一个非 Sentence Transformer）和 `M2V_base_output_nozipf_nopca`（使用BGE-base，一个 Sentence Transformer）看出。使用 Sentence Transformer 带来了约 5.2% 的性能提升。
2. PCA 对性能至关重要。这可以通过比较 `M2V_base_output_nozipf_nopca` 和 `M2V_base_output_nozipf` 看出，性能提升了约 2.8%。此外，PCA 在*所有*任务上都提高了性能。
3. Zipf 加权对性能至关重要。这可以通过比较 `M2V_base_output_nozipf_nopca` 和 `M2V_base_output_nopca` 看出，性能提升了约 3.1%。
4. 输出嵌入优于输入嵌入。这可以通过比较 `M2V_base_input` 和 `M2V_base_output` 看出，性能提升了约 6.1%。请注意，输入嵌入在某些任务上确实表现良好。我们推测这是因为输入嵌入本身是归一化的。

结论

感谢您阅读我们关于 Model2Vec 的博客文章！我们希望您觉得它信息丰富且有用。如果您有任何问题或意见，请随时与我们联系。我们仍在积极开发该项目，并且已经计划了许多新功能，敬请期待。

• 💻 代码仓库
• 🤗 HuggingFace 组织
• 🤗 HuggingFace 模型
• 👥 领英
• 📚 教程

引用

@software{minishlab2024word2vec,
  authors = {Stephan Tulkens, Thomas van Dongen},
  title = {Model2Vec: Turn any Sentence Transformer into a Small Fast Model},
  year = {2024},
  url = {https://github.com/MinishLab/model2vec},
}

致谢

我们要感谢 Tom Aarsen 将 Model2Vec 集成到 Sentence Transformers 中，并帮助我们进行 HuggingFace 集成，以及他对项目的总体反馈。