Matryoshka Embeddings (套娃嵌入)

密集嵌入模型通常生成固定大小的嵌入，例如 768 或 1024。所有后续计算（聚类、分类、语义搜索、检索、重排等）都必须在这些完整的嵌入上进行。Matryoshka 表示学习 (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 数据集上训练的，该数据集是 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: 此示例在 STSBenchmark 数据集的测试集上评估了使用 matryoshka_nli.py 中的 MatryoshkaLoss 训练的嵌入模型，并将其与未使用 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 文档的改编版。