深入理解UKPLab/sentence-transformers中的自适应层技术-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01195/article/details/148375703

深入理解UKPLab/sentence-transformers中的自适应层技术

sentence-transformers Multilingual Sentence & Image Embeddings with BERT 项目地址: https://gitcode.com/gh_mirrors/se/sentence-transformers

引言

在自然语言处理领域，句子嵌入模型已成为各类下游任务的基础组件。传统嵌入模型通常需要完整执行所有网络层才能获得最终嵌入表示，这在某些对延迟敏感的应用场景中可能成为瓶颈。本文将深入探讨一种创新性的训练方法——自适应层技术（Adaptive Layers），它能够让模型在不同层数下都能产生高质量的嵌入表示，从而实现推理速度与模型性能的灵活权衡。

自适应层技术原理

基本概念

自适应层技术的核心思想源于2D Matryoshka Sentence Embeddings（后更名为ESE: Espresso Sentence Embeddings）的研究工作。传统嵌入模型如all-mpnet-base-v2（12层）或all-MiniLM-L6-v2（6层）必须完整执行所有层才能获得最终嵌入表示。而自适应层技术通过特殊训练方法，使模型在仅使用部分层时仍能保持较好的性能表现。

技术实现

自适应层训练包含两个关键组件：

多层损失应用：不仅对最后一层的输出应用损失函数，同时对中间层的输出也应用相同的损失函数
KL散度约束：通过KL散度损失使非最后一层的嵌入表示尽可能接近最后一层的表示

这种设计可以看作是一种特殊的知识蒸馏形式，其中最后一层作为教师模型，而前面的层作为学生模型。

性能优势

实验数据表明，采用自适应层技术的模型展现出显著优势：

性能保持：当层数减少时，自适应层模型能保持更高比例的原始性能
推理加速：减少一半层数可获得约2倍的推理速度提升，而性能损失仅约15%
极端情况：即使仅使用1层，仍能保留约80%的原始性能

实际应用

训练配置

在Sentence Transformers框架中，实现自适应层训练非常简单。以下是使用CoSENTLoss结合AdaptiveLayerLoss的示例：

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)

二维自适应（Matryoshka2d）

更进一步的，可以将自适应层技术与Matryoshka嵌入技术结合，实现同时在层数和嵌入维度上的自适应：

from sentence_transformers.losses import Matryoshka2dLoss

loss = Matryoshka2dLoss(
    model=model,
    loss=base_loss,
    matryoshka_dims=[768, 512, 256, 128, 64]
)

这种组合允许模型在推理时根据需求同时调整使用的层数和输出维度。

推理优化

训练完成后，可以通过截断模型层数来优化推理速度。具体操作需要根据模型结构进行调整。以MPNet模型为例：

model = SentenceTransformer("tomaarsen/mpnet-base-nli-adaptive-layer")
new_num_layers = 3  # 只保留前3层
model[0].auto_model.encoder.layer = model[0].auto_model.encoder.layer[:new_num_layers]