俄羅斯套娃 (Matryoshka) 嵌入模型概述

HuggingFace發表於2024-03-06
模型

在這篇部落格中，我們將向你介紹俄羅斯套娃嵌入的概念，並解釋為什麼它們很有用。我們將討論這些模型在理論上是如何訓練的，以及你如何使用 Sentence Transformers 來訓練它們。

除此之外，我們還會告訴你怎麼用這種像套娃一樣的俄羅斯套娃嵌入模型，並且我們會比較一下這種模型和普通嵌入模型的不同。最後，我們邀請你來玩一下我們的互動演示，看看這些模型有多厲害。

理解嵌入 (embedding)

嵌入是自然語言處理中最通用的工具之一，使從業者能夠解決大量任務。本質上，嵌入是一個更復雜數字物件的數值表示，如文字、影像、音訊等。

嵌入模型

嵌入模型總是會產生相同固定大小的嵌入。然後，你可以透過計算相應嵌入的相似性來計算複雜數字物件的相似性！

嵌入相似性

這種技術 (嵌入) 在許多領域都有應用，它是推薦系統、資訊檢索、零樣本學習或少量樣本學習、異常檢測、相似性搜尋、釋義檢測、聚類、分類等領域的基礎。

🪆 俄羅斯套娃 (Matryoshka) 嵌入

隨著研究的進展，新的最先進的 (文字) 嵌入模型開始產生具有越來越高的輸出維度，即每個輸入文字都使用更多的值來表示。儘管這提高了效能，但以下游任務 (如搜尋或分類) 的效率為代價。

因此，Kusupati 等人 (2022) 受到啟發，創造了即使嵌入尺寸合理縮小也不會在效能上遭受太大損失的嵌入模型。

俄羅斯套娃模型

這些俄羅斯套娃嵌入模型經過訓練，使得這些小的截斷嵌入仍然有用。簡而言之，俄羅斯套娃嵌入模型可以產生各種尺寸的有用嵌入。

🪆 俄羅斯套娃

對於不熟悉的人來說，“Matryoshka 娃娃”，也稱為“俄羅斯套娃”，是一組大小遞減的木製娃娃，相互巢狀。類似地，俄羅斯套娃嵌入模型旨在將更重要的資訊儲存在早期的維度中，將不太重要的資訊儲存在後面的維度中。俄羅斯套娃嵌入模型的這一特點允許我們截斷模型產生的原始 (大) 嵌入，同時仍保留足夠的資訊以在下游任務上表現良好。

俄羅斯套娃模型

為什麼使用🪆 俄羅斯套娃嵌入模型？

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

  1. 篩選和重新排序: 不必在完整嵌入上執行你的下游任務 (例如，最近鄰搜尋)，你可以縮小嵌入到更小的尺寸，並非常高效地“篩選”你的嵌入。之後，你可以使用它們的完整維度處理剩餘的嵌入。
  2. 權衡: 俄羅斯套娃模型將允許你根據所需的儲存成本、處理速度和效能來擴充套件你的嵌入解決方案。

🪆 俄羅斯套娃嵌入模型是如何訓練的？

理論上

俄羅斯套娃表示學習 (MRL) 方法幾乎可以適用於所有嵌入模型訓練框架。通常，嵌入模型的一個訓練步驟涉及為你的訓練批次 (例如文字) 產生嵌入，然後使用一些損失函式建立一個代表產生嵌入質量的損失值。最佳化器會在訓練過程中調整模型權重以減少損失值。

對於俄羅斯套娃嵌入模型，一個訓練步驟還涉及為你的訓練批次產生嵌入，但是然後你使用一些損失函式來確定不僅僅是全尺寸嵌入的質量，還有各種不同維度性下的嵌入質量。例如，輸出維度性為 768、512、256、128 和 64。每個維度性的損失值加在一起，得到最終的損失值。然後，最佳化器將嘗試調整模型權重以降低這個損失值。

實際上，這鼓勵模型在嵌入的開始部分前置最重要的資訊，這樣如果嵌入被截斷，這些資訊將得以保留。

在 Sentence Transformers 中

Sentence Tranformers 是一個常用於訓練嵌入模型的框架，它最近實現了對俄羅斯套娃模型的支援。使用 Sentence Transformers 訓練俄羅斯套娃嵌入模型非常基礎: 不是僅在全尺寸嵌入上應用一些損失函式，我們也在嵌入的截斷部分應用同樣的損失函式。

例如，如果一個模型的原始嵌入維度為 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 Embeddings - Training

請檢視以下完整指令碼，瞭解如何在實際應用中使用 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 文件的改編。

如何使用 🪆俄羅斯套娃嵌入模型？

理論上

實際上，從俄羅斯套娃嵌入模型獲取嵌入的方式與從普通嵌入模型獲取嵌入的方式相同。唯一的區別在於，在接收到嵌入後，我們可以選擇將它們截斷為更小的維度。請注意，如果嵌入已經歸一化，那麼在截斷後它們將不再歸一化，因此你可能需要重新歸一化。
截斷後，你可以直接將它們應用於你的用例，或者儲存它們以便稍後使用。畢竟，在你的向量資料庫中使用較小的嵌入應該會帶來相當大的速度提升！
請記住，儘管處理較小嵌入以進行下游任務 (檢索、聚類等) 會更快，但從模型獲取較小嵌入的速度與獲取較大嵌入的速度一樣快。

在 Sentence Transformers 中

在 Sentence Transformers 中，你可以像載入普通模型一樣載入俄羅斯套娃嵌入模型，並使用 SentenceTransformers.encode 進行推理。獲取嵌入後，我們可以將它們截斷到我們所需的尺寸，如果需要，我們還可以對它們進行歸一化。
讓我們嘗試使用我使用 matryoshka_nli.py 和 microsoft/mpnet-base 訓練的模型:

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

model = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka")

matryoshka_dim = 64
embeddings = model.encode(
    [
        "The weather is so nice!",
        "It's so sunny outside!",
        "He drove to the stadium.",
    ]
)
embeddings = embeddings[..., :matryoshka_dim] # Shrink the embedding dimensions
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 Embeddings - 推理

點選這裡檢視如何使用 Nomic v1.5 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]])
  • 模型連結: nomic-ai/nomic-embed-text-v1.5

結果

現在我們已經介紹了俄羅斯套娃模型，讓我們來看看我們可以從俄羅斯套娃嵌入模型與常規嵌入模型中實際期待的績效表現。為了這個實驗，我訓練了兩個模型:

  • 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 測試集。結果繪製在下面的圖表中:

results

在上面的圖表中，你可以看到俄羅斯套娃模型在所有維度上都達到了比標準模型更高的 Spearman 相似度，這表明俄羅斯套娃模型在此任務上是優越的。

此外，俄羅斯套娃模型的效能下降速度比標準模型要慢得多。這在第二個圖表中清晰顯示，該圖表顯示了相對於最大效能的嵌入維度的效能。 即使嵌入大小隻有 8.3%，俄羅斯套娃模型也保持了 98.37% 的效能，遠高於標準模型的 96.46%。
這些發現表明，透過俄羅斯套娃模型截斷嵌入可以:

  1. 顯著加快下游任務 (如檢索) 的速度;
  2. 顯著節省儲存空間，而且不會對效能產生顯著影響。

演示

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

https://xenova-adaptive-retrieval-web.static.hf.space

參考文獻

  • 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

英文原文: https://hf.co/blog/matryoshka

原文作者: Tom Aarsen, Joshua, Omar Sanseviero

譯者: innovation64

相關文章