套娃式嵌入

密集嵌入模型通常生成固定大小的嵌入，例如 768 或 1024。所有后续计算（聚类、分类、语义搜索、检索、重排等）都必须在这些完整的嵌入上进行。套娃式表示学习重新审视了这一思想，并提出了一种解决方案，用于训练嵌入模型，其嵌入在截断到更小尺寸后仍然有用。这使得（批量）处理速度大大加快。

用例

一个特别有趣的用例是将处理分为两个步骤：1）使用小得多的向量进行预处理，然后 2）将剩余的向量作为完整尺寸进行处理（也称为“短名单和重排”）。此外，套娃式模型将允许您根据所需的存储成本、处理速度和性能来扩展您的嵌入解决方案。

结果

让我们看看套娃式嵌入模型与常规嵌入模型相比，我们可能期望的实际性能。对于这个实验，我训练了两个模型

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_eval_stsb.py 获得的结果绘制在下图中

results

在顶部图中，您可以看到套娃式模型在所有维度上都达到了比标准模型更高的 Spearman 相似度，这表明套娃式模型在此任务中表现更优。

此外，套娃式模型的性能下降速度远低于标准模型。这在第二张图中清晰显示，该图显示了嵌入维度相对于最大性能的性能。即使在嵌入尺寸的 8.3% 时，套娃式模型也保留了 98.37% 的性能，远高于标准模型的 96.46%。

这些发现表明，通过套娃式模型截断嵌入可以：1）显著加速检索等下游任务，2）显著节省存储空间，所有这些都不会对性能造成显著影响。

训练

使用套娃式表示学习（MRL）进行训练非常基础：我们不仅在完整尺寸的嵌入上应用一些损失函数，还在嵌入的截断部分上应用相同的损失函数。例如，如果模型的嵌入维度默认为 768，它现在可以在 768、512、256、128、64 和 32 上进行训练。这些损失将全部相加，可选地带有一些权重

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

参考：MatryoshkaLoss

此外，这可以与 AdaptiveLayerLoss 结合，使得最终模型不仅可以在输出维度大小上减小，还可以在层数上减小以加快推理速度。有关减少模型层数的更多信息，请参阅自适应层。在 Sentence Transformers 中，这两种损失的组合被称为 Matryoshka2dLoss，并提供了简写以简化训练。

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

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

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

参考：Matryoshka2dLoss

推理

使用套娃式损失训练模型后，您可以使用 SentenceTransformers.encode 对其进行推理。

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

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

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.",
    ]
)
assert embeddings.shape[-1] == matryoshka_dim

similarities = model.similarity(embeddings[0], embeddings[1:])
# => tensor([[0.7839, 0.4933]])

如您所见，尽管应用了非常小的套娃维度，搜索查询与正确文档之间的相似度远高于与不相关文档的相似度。请随意将此脚本复制到本地，修改 matryoshka_dim，并观察相似度的差异。

注意：尽管嵌入更小，套娃式模型的训练和推理速度并不更快，内存效率不高，模型也不更小。只有结果嵌入的处理和存储会更快、更便宜。

代码示例

请参阅以下脚本，了解如何在实践中应用 MatryoshkaLoss

matryoshka_nli.py：此示例使用 MultipleNegativesRankingLoss 和 MatryoshkaLoss 训练一个强大的嵌入模型，使用自然语言推理 (NLI) 数据。它是 NLI 文档的改编。
matryoshka_nli_reduced_dim.py：此示例使用 MultipleNegativesRankingLoss 和 MatryoshkaLoss 训练一个强大的嵌入模型，最大输出维度为 256。它使用自然语言推理 (NLI) 数据进行训练，是 NLI 文档的改编。
matryoshka_eval_stsb.py：此示例在 STSBenchmark 数据集的测试集上评估使用 MatryoshkaLoss 训练的嵌入模型（在 matryoshka_nli.py 中），并将其与未经 Matryoshka 训练的模型进行比较。
matryoshka_sts.py：此示例使用 CoSENTLoss 和 MatryoshkaLoss 在 STSBenchmark 数据集的训练集上训练嵌入模型。它是 STS 文档的改编。

以及以下脚本，了解如何应用 Matryoshka2dLoss

2d_matryoshka_nli.py：此示例使用 MultipleNegativesRankingLoss 和 Matryoshka2dLoss 训练一个强大的嵌入模型，使用自然语言推理 (NLI) 数据。它是 NLI 文档的改编。
2d_matryoshka_sts.py：此示例使用 CoSENTLoss 和 Matryoshka2dLoss 在 STSBenchmark 数据集的训练集上训练嵌入模型。它是 STS 文档的改编。