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 整合，以及他對專案的總體反饋。