Hugging Face 的 LOGO

🪆 Matryoshka 嵌入模型简介

发布于 2024 年 2 月 23 日
在 GitHub 上更新
赞
154
Tom Aarsen's avatar
Tom Aarsen
tomaarsen
Joshua's avatar
Joshua
Xenova
Omar Sanseviero's avatar
奥马尔·桑塞维耶罗 (Omar Sanseviero)
osanseviero

在本篇博文中，我们将向您介绍 Matryoshka 嵌入的概念，并解释它们为何有用。我们将讨论这些模型在理论上是如何训练的，以及您如何使用 Sentence Transformers 来训练它们。

此外，我们将提供关于如何使用 Matryoshka 嵌入模型的实用指导，并分享 Matryoshka 嵌入模型与常规嵌入模型的比较。最后，我们邀请您体验我们的交互式演示，它展示了这些模型的强大功能。

目录

理解嵌入

嵌入是自然语言处理中功能最丰富的工具之一，它能帮助从业者解决各种各样的任务。从本质上讲，嵌入是一个更复杂对象（如文本、图像、音频等）的数值表示。

embedding model

嵌入模型将始终产生固定大小的嵌入。然后，您可以通过计算相应嵌入的相似度来计算复杂对象的相似度！

embedding similarity

这有大量的用例，并构成了推荐系统、检索、一次性或少次性学习、异常值检测、相似性搜索、释义检测、聚类、分类等的基础！

🪆 Matryoshka 嵌入

随着研究的进展，新的顶尖（文本）嵌入模型开始产生具有越来越高输出维度的嵌入，即每个输入文本都用更多的值来表示。尽管这提高了性能，但也牺牲了搜索或分类等下游任务的效率。

因此，Kusupati 等人 (2022) 受此启发，创建了可以合理缩小嵌入尺寸而性能不会大幅下降的嵌入模型。

matryoshka model

这些 Matryoshka 嵌入模型经过训练，使得这些小的截断嵌入仍然有用。简而言之，Matryoshka 嵌入模型可以生成不同维度的有用嵌入。

🪆 俄罗斯套娃

对于不熟悉的人来说，“Matryoshka 娃娃”，也被称为“俄罗斯套娃”，是一套大小递减的木制娃娃，一个套一个。类似地，Matryoshka 嵌入模型旨在将更重要的信息存储在靠前的维度中，而将不太重要的信息存储在靠后的维度中。Matryoshka 嵌入模型的这一特性使我们能够截断模型产生的原始（大）嵌入，同时仍然保留足够的信息以在下游任务中表现良好。

matryoshka models

为什么要使用 🪆 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 的示例：

  • 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]])

您可以随意尝试使用不同的 matryoshka_dim 值，并观察这如何影响相似性。您可以在本地运行此代码，也可以在云端（如 Google Colab）运行，或者查看演示。

参考文献

点击此处查看如何使用 Nomic v1.5 Matryoshka 模型

注意：Nomic 特别要求在嵌入截断之前进行 F.layer_norm。因此，以下代码片段使用手动截断到所需维度。对于所有其他模型，您可以在构造函数中使用 truncate_dim 选项，如前一个示例所示。

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 嵌入模型与常规嵌入模型相比，我们可能期望的实际性能。为了这个实验，我训练了两个模型

这两个模型都在 AllNLI 数据集上进行了训练，该数据集是 SNLI 和 MultiNLI 数据集的拼接。我已经在 STSBenchmark 测试集上使用多个不同的嵌入维度评估了这些模型。结果绘制在下图中

results

在上图中，您可以看到 Matryoshka 模型在所有维度上都达到了比标准模型更高的斯皮尔曼相似度，这表明 Matryoshka 模型在此任务中更优越。

此外，Matryoshka 模型的性能下降速度比标准模型慢得多。这在第二张图中清晰地显示出来，该图显示了嵌入维度相对于最大性能的性能。即使在嵌入大小为 8.3% 的情况下，Matryoshka 模型仍保留了 98.37% 的性能，远高于标准模型的 96.46%。

这些发现表明，通过 Matryoshka 模型截断嵌入可以：1) 显著加快检索等下游任务的速度，2) 显著节省存储空间，而性能没有明显下降。

演示

在这个演示中，您可以动态地缩小 nomic-ai/nomic-embed-text-v1.5 Matryoshka 嵌入模型的输出维度，并观察它如何影响检索性能。所有的嵌入都是在浏览器中使用 🤗 Transformers.js 计算的。

参考文献

更多博客文章

使用 Sentence Transformers v5 训练和微调稀疏嵌入模型

作者  tomaarsen 2025年7月1日 • 106

使用 Sentence Transformers v4 训练和微调 Reranker 模型

作者  tomaarsen 2025年3月26日 • 155

社区

nbroad

@tomaarsen ,

很棒的文章！关于归一化：应该在截断之前还是之后进行？

·
tomaarsen
 文章作者
•
编辑于2月20日

好问题！嵌入应该在 Matryoshka 截断之后进行（重新）归一化。如果你只在截断前进行归一化，截断后的部分将不会有确切的期望均值和标准差，而是会有一点点偏差。

Sifal

感谢分享（评论有点晚，哈哈），我对非 Matryoshka 嵌入的鲁棒性感到相当惊讶，是我理解错了吗？

·
umarbutler

我也觉得很惊讶，这真是个很酷的副作用。你会在《Matryoshka Representation Learning》的第5页找到类似的支持性结果。事实上，那页的图表表明，MRL-E/tied MRL 在低维度上的表现不如他们的 FF 模型，更不用说，他们的基线和 MRL 之间的差异并不大（在那些特定的图表中）。有趣的是，Sentence Transformers 实现的是 MRL-E/tied MRL，而不是 untied MRL。

@tomaarsen 的结果似乎有很大不同，因为 MRL-E 胜过了没有 MRL 的情况。

HalteroXHunter

你们觉得将 MRL 引入像 BAAI/bge-m3 这样的模型有意义吗？

注册 或 登录 以发表评论

赞
154