Matryoshka Embeddings

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

用例

一个特别有趣的用例是将处理过程分为两个步骤：1) 使用小得多的向量进行预处理，然后 2) 以全尺寸处理剩余向量（也称为“候选列表和重排序”）。此外，Matryoshka 模型将允许您根据所需的存储成本、处理速度和性能来扩展您的嵌入解决方案。

结果

让我们看看我们可以从 Matryoshka 嵌入模型与常规嵌入模型中获得的实际性能。对于此实验，我训练了两个模型

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

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

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

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

训练

使用 Matryoshka Representation Learning (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

推理

在使用 Matryoshka 损失训练模型后，您可以使用 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 维度，但搜索查询与正确文档之间的相似性远高于与不相关文档的相似性。请随意在本地复制此脚本，修改 matryoshka_dim，并观察相似性的差异。

注意：尽管嵌入较小，但 Matryoshka 模型的训练和推理速度不会更快，内存效率不会更高，尺寸也不会更小。只有生成的嵌入的处理和存储速度更快且成本更低。

代码示例

请参阅以下脚本，以了解如何在实践中应用 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 训练的模型进行比较。
matryoshka_sts.py：此示例将 CoSENTLoss 与 MatryoshkaLoss 结合使用，以在 STSBenchmark 数据集的训练集上训练嵌入模型。它是 STS 文档的改编。

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

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