自适应层

嵌入模型通常是具有多层的编码器模型，例如 12 层（例如 all-mpnet-base-v2）或 6 层（例如 all-MiniLM-L6-v2）。为了获得嵌入，必须遍历所有这些层。 2D Matryoshka Sentence Embeddings (2DMSE) 预印本重新审视了这个概念，提出了一种训练嵌入模型的方法，该模型在仅使用所有层的一部分时也能表现良好。这以相对较低的性能成本实现了更快的推理速度。

注意

2DMSE 预印本后来被更新并重命名为 ESE: Espresso Sentence Embeddings。 Sentence Transformers 的自适应层和 Matryoshka2d（自适应层 + Matryoshka 嵌入）的实现基于最初的预印本，我们接受实现更新的 ESE 论文的贡献。

用例

2DMSE 论文提到，使用使用自适应层和 Matryoshka 表示学习训练的较大模型的一些层，可以优于像标准嵌入模型一样训练的较小模型。

结果

让我们看一下我们可能从自适应层嵌入模型与常规嵌入模型获得的预期性能。为了这个实验，我训练了两个模型

tomaarsen/mpnet-base-nli-adaptive-layer：通过运行 adaptive_layer_nli.py 和 microsoft/mpnet-base 训练。
tomaarsen/mpnet-base-nli：与前者几乎相同的模型，但仅使用 MultipleNegativesRankingLoss 而不是 AdaptiveLayerLoss 在 MultipleNegativesRankingLoss 之上。我也使用 microsoft/mpnet-base 作为基础模型。

这两个模型都在 AllNLI 数据集上进行了训练，AllNLI 数据集是 SNLI 和 MultiNLI 数据集的串联。我已经使用多个不同的嵌入维度在 STSBenchmark 测试集上评估了这些模型。结果绘制在下图

adaptive_layer_results

第一个图显示，当减少模型中的层数时，自适应层模型保持更高的性能。这在第二个图中也清晰地显示出来，该图显示，当层数一直减少到 1 时，性能仍然保留了 80%。

最后，第三个图显示了我的测试中 GPU 和 CPU 设备的预期加速比。正如您所看到的，移除一半的层数大致会带来 2 倍的加速，但会在 STSB 上损失约 15% 的性能（~86 -> ~75 Spearman 相关性）。当移除更多层数时，对于 CPU 来说性能提升变得更大，并且在性能损失 20% 的情况下，5 倍到 10 倍的加速是非常可行的。

训练

使用自适应层支持进行训练非常基础：与其仅在最后一层应用一些损失函数，我们还在先前层的池化嵌入上应用相同的损失函数。此外，我们采用 KL 散度损失，旨在使非最后一层的嵌入与最后一层的嵌入相匹配。这可以看作是知识蒸馏的一种引人入胜的方法，但最后一层作为教师模型，而先前的层作为学生模型。

例如，对于 12 层的 microsoft/mpnet-base，现在将对其进行训练，以便模型在 12 层中的每一层之后都产生有意义的嵌入。

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

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

base_loss = CoSENTLoss(model=model)
loss = AdaptiveLayerLoss(model=model, loss=base_loss)

参考： AdaptiveLayerLoss

请注意，使用 AdaptiveLayerLoss 进行训练并不比不使用它慢多少。

此外，这可以与 MatryoshkaLoss 结合使用，以便可以减少所得模型的层数，也可以减少输出维度的大小。另请参阅 Matryoshka 嵌入，以获取有关减少输出维度的更多信息。在 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

推理

在使用自适应层损失训练模型后，您可以将模型层截断为您所需的层数。请注意，这需要在模型本身上进行一些手术，并且每个模型的结构都略有不同，因此步骤会因模型而略有不同。

首先，我们将加载模型并访问底层的 transformers 模型，如下所示

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("tomaarsen/mpnet-base-nli-adaptive-layer")

# We can access the underlying model with `model[0].auto_model`
print(model[0].auto_model)

MPNetModel(
  (embeddings): MPNetEmbeddings(
    (word_embeddings): Embedding(30527, 768, padding_idx=1)
    (position_embeddings): Embedding(514, 768, padding_idx=1)
    (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
    (dropout): Dropout(p=0.1, inplace=False)
  )
  (encoder): MPNetEncoder(
    (layer): ModuleList(
      (0-11): 12 x MPNetLayer(
        (attention): MPNetAttention(
          (attn): MPNetSelfAttention(
            (q): Linear(in_features=768, out_features=768, bias=True)
            (k): Linear(in_features=768, out_features=768, bias=True)
            (v): Linear(in_features=768, out_features=768, bias=True)
            (o): Linear(in_features=768, out_features=768, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
          (dropout): Dropout(p=0.1, inplace=False)
        )
        (intermediate): MPNetIntermediate(
          (dense): Linear(in_features=768, out_features=3072, bias=True)
          (intermediate_act_fn): GELUActivation()
        )
        (output): MPNetOutput(
          (dense): Linear(in_features=3072, out_features=768, bias=True)
          (LayerNorm): LayerNorm((768,), eps=1e-05, elementwise_affine=True)
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (relative_attention_bias): Embedding(32, 12)
  )
  (pooler): MPNetPooler(
    (dense): Linear(in_features=768, out_features=768, bias=True)
    (activation): Tanh()
  )
)

此输出将因模型而异。我们将在编码器中寻找重复的层。对于此 MPNet 模型，它存储在 model[0].auto_model.encoder.layer 下。然后我们可以对模型进行切片，只保留前几层以加速模型

new_num_layers = 3
model[0].auto_model.encoder.layer = model[0].auto_model.encoder.layer[:new_num_layers]

然后我们可以使用 SentenceTransformers.encode 运行推理。

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("tomaarsen/mpnet-base-nli-adaptive-layer")
new_num_layers = 3
model[0].auto_model.encoder.layer = model[0].auto_model.encoder.layer[:new_num_layers]

embeddings = model.encode(
    [
        "The weather is so nice!",
        "It's so sunny outside!",
        "He drove to the stadium.",
    ]
)
# Similarity of the first sentence with the other two
similarities = model.similarity(embeddings[0], embeddings[1:])
# => tensor([[0.7761, 0.1655]])
# compared to tensor([[ 0.7547, -0.0162]]) for the full model

正如您所看到的，尽管仅使用了 3 层，但相关句子之间的相似度远高于不相关句子。随意在本地复制此脚本，修改 new_num_layers，并观察相似度的差异。

代码示例

请参阅以下脚本，作为如何在实践中应用 AdaptiveLayerLoss 的示例

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

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

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