Matryoshka 嵌入模型簡介

Matryoshka 嵌入模型使得能夠創建可在不顯著降低效能的情況下截斷至更小維度的嵌入向量。此方法讓使用者能根據儲存成本、處理速度與準確度需求，動態調整其嵌入解決方案的規模。

Matryoshka 嵌入的工作原理

Matryoshka 嵌入基於 Matryoshka 表徵學習 (MRL)，這項技術鼓勵模型將最關鍵的資訊儲存在嵌入向量的前幾個維度。就像俄羅斯套娃一樣，模型被訓練以確保較小、被截斷的嵌入仍保有足夠資訊，足以應用於後續任務。

應用與優勢

可變尺寸的嵌入為 AI 使用者提供兩大主要優勢：

  1. 篩選與重新排序：使用者可以利用小尺寸、已截斷的嵌入執行初步且高效的「篩選」搜尋。當候選集合確定後，可使用完整維度的嵌入對結果進行重新排序，以提升精確度。
  2. 資源最佳化：Matryoshka 模型允許根據部署環境的具體限制，自訂儲存成本、處理速度與效能之間的取捨。

訓練 Matryoshka 嵌入模型

理論框架

在標準的嵌入訓練中，損失函數會套用於完整尺寸的嵌入。相較之下，Matryoshka 訓練會同時計算完整尺寸嵌入以及不同截斷維度（例如 768、512、256、128、64）的嵌入損失值。這些損失會相加形成最終損失，優化器據此調整模型權重，從而「前置」最重要的資訊。

在 Sentence Transformers 中的實作

Sentence Transformers 已透過 MatryoshkaLoss 實作對 Matryoshka 模型的支援。此包裝損失函數會將基礎損失（例如 CoSENTLoss）套用於嵌入的多個截斷部分。使用 MatryoshkaLoss 進行訓練不會產生顯著的訓練時間開銷。

使用 Matryoshka 嵌入模型

一般推論

推論的運作方式與標準嵌入模型相似。唯一的差異在於，產生嵌入後可選擇將其截斷至較小的維度。若嵌入原本已正規化，截斷後應重新正規化以保持一致性。

在 Sentence Transformers 中的實作

使用 SentenceTransformer 類別時，使用者可在模型初始化時指定 truncate_dim 參數。encode 函式會自動將輸出嵌入截斷至指定大小。

對於特定模型（如 nomic-ai/nomic-embed-text-v1.5），自訂架構需要在截斷嵌入之前套用 F.layer_norm。

效能結果

在 STSBenchmark 測試集上比較 Matryoshka 模型（tomaarsen/mpnet-base-nli-matryoshka）與標準模型（tomaarsen/mpnet-base-nli）的實驗結果顯示如下：

  • 更高的準確度：Matryoshka 模型在所有測試維度上皆取得更高的 Spearman 相似度。
  • 效能保持：Matryoshka 模型的效能衰減速度遠低於標準模型。即使截斷至原始大小的 8.3%，Matryoshka 模型仍保留了 98.37% 的最高效能，而標準模型僅為 96.46%。

這些結果顯示，Matryoshka 嵌入能在不顯著影響效能的前提下，大幅降低儲存需求並提升檢索速度。

Sources