俄罗斯套娃嵌入

密集嵌入模型通常生成固定大小的嵌入，例如 768 或 1024。所有后续计算（聚类、分类、语义搜索、检索、重排序等）都必须在这些完整嵌入上完成。Matryoshka Representation Learning 重新审视了这个想法，并提出了一种解决方案来训练嵌入模型，这些模型的嵌入在截断为更小尺寸后仍然有用。这可以显著加快（批量）处理速度。

用例

一个特别有趣的用例是将处理过程分为两个步骤：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 数据集上进行了训练，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：此示例评估了在 matryoshka_nli.py 中使用 MatryoshkaLoss 训练的嵌入模型在 STSBenchmark 数据集的测试集上的性能，并将其与非俄罗斯套娃训练的模型进行比较。
matryoshka_sts.py：此示例将 CoSENTLoss 与 MatryoshkaLoss 结合使用，以在 STSBenchmark 数据集的训练集上训练嵌入模型。它是 STS 文档的改编。

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

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