Matryoshka Embeddings

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

用例

一个特别有趣的用例是将处理分为两个步骤：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 数据集上进行了训练，该数据集是 SNLI 和 MultiNLI 数据集的连接。我使用多个不同的嵌入维度在 STSBenchmark 测试集上评估了这些模型。通过运行 matryoshka_eval_stsb.py 获得的结果如下图所示

results

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

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

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

训练

使用 Matryoshka 表示学习 (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 文档的改编。