🪆 Matryoshka 嵌入模型簡介

Tom Aarsen

tomaarsen

關注

Joshua

Xenova

關注

奧馬爾·桑塞維耶羅 (Omar Sanseviero)

osanseviero

關注

此外，我們將提供關於如何使用 Matryoshka 嵌入模型的實用指導，並分享 Matryoshka 嵌入模型與常規嵌入模型的比較。最後，我們邀請您體驗我們的互動式演示，它展示了這些模型的強大功能。

目錄

• 理解嵌入
• 🪆 Matryoshka 嵌入
• 🪆 俄羅斯套娃
• 為什麼要使用 🪆 Matryoshka 嵌入模型？
• 🪆 Matryoshka 嵌入模型是如何訓練的？
  • 理論上
  • 在 Sentence Transformers 中
• 如何使用 🪆 Matryoshka 嵌入模型？
  • 理論上
  • 在 Sentence Transformers 中
• 結果
• 演示
• 參考文獻

理解嵌入

嵌入是自然語言處理中功能最豐富的工具之一，它能幫助從業者解決各種各樣的任務。從本質上講，嵌入是一個更復雜物件（如文字、影像、音訊等）的數值表示。

嵌入模型將始終產生固定大小的嵌入。然後，您可以透過計算相應嵌入的相似度來計算複雜物件的相似度！

這有大量的用例，並構成了推薦系統、檢索、一次性或少次性學習、異常值檢測、相似性搜尋、釋義檢測、聚類、分類等的基礎！

🪆 Matryoshka 嵌入

隨著研究的進展，新的頂尖（文字）嵌入模型開始產生具有越來越高輸出維度的嵌入，即每個輸入文字都用更多的值來表示。儘管這提高了效能，但也犧牲了搜尋或分類等下游任務的效率。

因此，Kusupati 等人 (2022) 受此啟發，建立了可以合理縮小嵌入尺寸而效能不會大幅下降的嵌入模型。

這些 Matryoshka 嵌入模型經過訓練，使得這些小的截斷嵌入仍然有用。簡而言之，Matryoshka 嵌入模型可以生成不同維度的有用嵌入。

🪆 俄羅斯套娃

對於不熟悉的人來說，“Matryoshka 娃娃”，也被稱為“俄羅斯套娃”，是一套大小遞減的木製娃娃，一個套一個。類似地，Matryoshka 嵌入模型旨在將更重要的資訊儲存在靠前的維度中，而將不太重要的資訊儲存在靠後的維度中。Matryoshka 嵌入模型的這一特性使我們能夠截斷模型產生的原始（大）嵌入，同時仍然保留足夠的資訊以在下游任務中表現良好。

為什麼要使用 🪆 Matryoshka 嵌入模型？

這種可變尺寸的嵌入模型對從業者來說非常有價值，例如

1. 篩選和重排序：您不必在完整的嵌入上執行下游任務（例如，最近鄰搜尋），而是可以將嵌入縮小到較小的尺寸，並非常高效地“篩選”您的嵌入。之後，您可以使用剩餘嵌入的完整維度對其進行處理。
2. 權衡：Matryoshka 模型允許您根據所需的儲存成本、處理速度和效能來擴充套件您的嵌入解決方案。

🪆 Matryoshka 嵌入模型是如何訓練的？

理論上

Matryoshka 表示學習（MRL）方法幾乎可以應用於所有嵌入模型訓練框架。通常，嵌入模型的訓練步驟包括為您的訓練批次（例如，文字）生成嵌入，然後使用某個損失函式來建立一個代表所生成嵌入質量的損失值。最佳化器將在整個訓練過程中調整模型權重以降低損失值。

對於 Matryoshka 嵌入模型，訓練步驟也涉及為您的訓練批次生成嵌入，但隨後您使用某個損失函式來確定不僅是全尺寸嵌入的質量，還有不同維度下嵌入的質量。例如，輸出維度可以是 768、512、256、128 和 64。每個維度的損失值會相加，得到最終的損失值。然後，最佳化器將嘗試調整模型權重以降低這個損失值。

在實踐中，這激勵模型將最重要的資訊預載入到嵌入的起始部分，以便在嵌入被截斷時仍能保留這些資訊。

在 Sentence Transformers 中

Sentence Transformers 是一個常用於訓練嵌入模型的框架，它最近實現了對 Matryoshka 模型的支援。使用 Sentence Transformers 訓練 Matryoshka 嵌入模型非常簡單：我們不僅對全尺寸嵌入應用某個損失函式，還對嵌入的截斷部分應用相同的損失函式。

例如，如果一個模型的原始嵌入維度為 768，現在可以在 768、512、256、128 和 64 維度上進行訓練。這些損失中的每一個都會被加在一起，可以選擇性地加上一些權重。

from sentence_transformers import SentenceTransformer
from sentence_transformers.losses import CoSENTLoss, MatryoshkaLoss

model = SentenceTransformer("microsoft/mpnet-base")

base_loss = CoSENTLoss(model=model)
loss = MatryoshkaLoss(
    model=model,
    loss=base_loss,
    matryoshka_dims=[768, 512, 256, 128, 64],
    matryoshka_weight=[1, 1, 1, 1, 1],
)

model.fit(
    train_objectives=[(train_dataset, loss)],
    ...,
)

使用 MatryoshkaLoss 訓練並不會在訓練時間上產生顯著的開銷。

參考文獻

• MatryoshkaLoss
• CoSENTLoss
• SentenceTransformer
• SentenceTransformer.fit
• Matryoshka 嵌入 - 訓練

請參閱以下完整指令碼作為如何在實踐中應用 MatryoshkaLoss 的示例：

• matryoshka_nli.py：此示例使用 MultipleNegativesRankingLoss 和 MatryoshkaLoss，利用自然語言推斷（NLI）資料來訓練一個強大的嵌入模型。它是 NLI 文件的改編版。
• matryoshka_nli_reduced_dim.py：此示例使用 MultipleNegativesRankingLoss 和 MatryoshkaLoss 來訓練一個最大輸出維度為 256 的小型強大嵌入模型。它使用自然語言推斷（NLI）資料進行訓練，是 NLI 文件的改編版。
• matryoshka_sts.py：此示例使用 CoSENTLoss 和 MatryoshkaLoss 在 STSBenchmark 資料集的訓練集上訓練一個嵌入模型。它是 STS 文件的改編版。

如何使用 🪆 Matryoshka 嵌入模型？

理論上

在實踐中，從 Matryoshka 嵌入模型獲取嵌入的方式與常規嵌入模型相同。唯一的區別是，在接收到嵌入後，我們可以選擇將它們截斷到更小的維度。請注意，如果嵌入已經歸一化，那麼截斷後它們將不再是歸一化的，因此您可能需要重新歸一化。

截斷後，您可以直接將其應用於您的用例，或者儲存它們以便以後使用。畢竟，在您的向量資料庫中使用較小的嵌入應該會帶來相當大的速度提升！

請記住，儘管處理較小的嵌入進行下游任務（檢索、聚類等）會更快，但從模型中獲取較小嵌入的速度與獲取較大嵌入的速度一樣快。

在 Sentence Transformers 中

在 Sentence Transformers 中，您可以像載入任何其他模型一樣載入 Matryoshka 嵌入模型，但您可以使用 truncate_dim 引數指定所需的嵌入大小。之後，您可以使用 SentenceTransformers.encode 函式執行推理，嵌入將自動截斷到指定的大小。

讓我們嘗試使用我用 matryoshka_nli.py 和 microsoft/mpnet-base 訓練的一個模型。

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

matryoshka_dim = 64
model = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka", truncate_dim=matryoshka_dim)

embeddings = model.encode(
    [
        "The weather is so nice!",
        "It's so sunny outside!",
        "He drove to the stadium.",
    ]
)
print(embeddings.shape)
# => (3, 64)

# Similarity of the first sentence to the other two:
similarities = cos_sim(embeddings[0], embeddings[1:])
print(similarities)
# => tensor([[0.8910, 0.1337]])

• 模型連結: tomaarsen/mpnet-base-nli-matryoshka

您可以隨意嘗試使用不同的 matryoshka_dim 值，並觀察這如何影響相似性。您可以在本地執行此程式碼，也可以在雲端（如 Google Colab）執行，或者檢視演示。

參考文獻

• SentenceTransformer
• SentenceTransformer.encode
• util.cos_sim
• Matryoshka 嵌入 - 推理

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
import torch.nn.functional as F

model = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True)

matryoshka_dim = 64
embeddings = model.encode(
    [
        "search_query: What is TSNE?",
        "search_document: t-distributed stochastic neighbor embedding (t-SNE) is a statistical method for visualizing high-dimensional data by giving each datapoint a location in a two or three-dimensional map.",
        "search_document: Amelia Mary Earhart was an American aviation pioneer and writer.",
    ],
    convert_to_tensor=True,
)
# The Nomic team uses a custom architecture, making them recommend Layer Normalization before truncation
embeddings = F.layer_norm(embeddings, normalized_shape=(embeddings.shape[1],))
embeddings[..., :matryoshka_dim]  # Shrink the embedding dimensions

similarities = cos_sim(embeddings[0], embeddings[1:])
# => tensor([[0.7154, 0.4468]])

結果

既然已經介紹了 Matryoshka 模型，讓我們來看看 Matryoshka 嵌入模型與常規嵌入模型相比，我們可能期望的實際效能。為了這個實驗，我訓練了兩個模型

• tomaarsen/mpnet-base-nli-matryoshka：透過執行 matryoshka_nli.py 和 microsoft/mpnet-base 訓練。
• tomaarsen/mpnet-base-nli：透過執行 matryoshka_nli.py 的修改版本訓練，其中訓練損失僅為 MultipleNegativesRankingLoss，而不是在 MultipleNegativesRankingLoss 之上的 MatryoshkaLoss。我也使用 microsoft/mpnet-base 作為基礎模型。

這兩個模型都在 AllNLI 資料集上進行了訓練，該資料集是 SNLI 和 MultiNLI 資料集的拼接。我已經在 STSBenchmark 測試集上使用多個不同的嵌入維度評估了這些模型。結果繪製在下圖中

在上圖中，您可以看到 Matryoshka 模型在所有維度上都達到了比標準模型更高的斯皮爾曼相似度，這表明 Matryoshka 模型在此任務中更優越。

此外，Matryoshka 模型的效能下降速度比標準模型慢得多。這在第二張圖中清晰地顯示出來，該圖顯示了嵌入維度相對於最大效能的效能。即使在嵌入大小為 8.3% 的情況下，Matryoshka 模型仍保留了 98.37% 的效能，遠高於標準模型的 96.46%。

這些發現表明，透過 Matryoshka 模型截斷嵌入可以：1) 顯著加快檢索等下游任務的速度，2) 顯著節省儲存空間，而效能沒有明顯下降。

演示

在這個演示中，您可以動態地縮小 nomic-ai/nomic-embed-text-v1.5 Matryoshka 嵌入模型的輸出維度，並觀察它如何影響檢索效能。所有的嵌入都是在瀏覽器中使用 🤗 Transformers.js 計算的。

參考文獻

• Kusupati, A., Bhatt, G., Rege, A., Wallingford, M., Sinha, A., Ramanujan, V., ... & Farhadi, A. (2022). Matryoshka representation learning. Advances in Neural Information Processing Systems, 35, 30233-30249. https://arxiv.org/abs/2205.13147
• Matryoshka Embeddings — Sentence-Transformers documentation. (n.d.). https://sbert.net/examples/training/matryoshka/README.html
• UKPLab. (n.d.). GitHub. https://github.com/UKPLab/sentence-transformers
• Unboxing Nomic Embed v1.5: Resizable Production Embeddings with Matryoshka Representation Learning. (n.d.). https://blog.nomic.ai/posts/nomic-embed-matryoshka