自适应层

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

注意

2DMSE 预印本后来更新并更名为《ESE: Espresso Sentence Embeddings》。Sentence Transformers 中自适应层 (Adaptive Layers) 和 Matryoshka2d (自适应层 + Matryoshka 嵌入) 的实现基于最初的预印本，我们欢迎社区贡献以实现更新后的 ESE 论文。

用例

2DMSE 论文提到，使用通过自适应层和 Matryoshka 表征学习训练的较大模型的少数几层，其性能可以超过像标准嵌入模型一样训练的较小模型。

结果

让我们看看一个自适应层嵌入模型与一个常规嵌入模型相比，可能达到的性能。在本实验中，我训练了两个模型：

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

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

adaptive_layer_results

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

最后，第三张图显示了在我的测试中 GPU 和 CPU 设备的预期加速比。如您所见，移除一半的层数大约能带来 2 倍的加速，代价是在 STSB 上的性能下降约 15%（斯皮尔曼相关系数从 ~86 降至 ~75）。当移除更多层时，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.transformers_model`
print(model.transformers_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.transformers_model.encoder.layer 下。然后我们可以对模型进行切片，只保留前几层以加速模型：

new_num_layers = 3
model.transformers_model.encoder.layer = model.transformers_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.transformers_model.encoder.layer = model.transformers_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 文档的改编版。