如何在Sentence Transformer之上添加线性层实现小样本学习微调？

核心问题说明 #

Sentence Transformer的model.encode()是封装好的推理方法，默认不会保留计算图，因此无法触发参数更新。要实现可训练的适配网络，需要自定义模型类，直接调用其内部前向逻辑并添加线性层。

自定义可训练模型实现 #

以下是完整的代码示例，包含模型定义、训练流程：

import torch
import torch.nn as nn
from sentence_transformers import SentenceTransformer

class STAdapterModel(nn.Module):
    def __init__(self, st_model_name='all-MiniLM-L6-v2', output_dim=2):
        super().__init__()
        # 加载预训练Sentence Transformer
        self.st_backbone = SentenceTransformer(st_model_name)
        # 先冻结预训练模型参数（小样本场景下优先训练适配层）
        for param in self.st_backbone.parameters():
            param.requires_grad = False
        # 添加适配线性层，维度对应ST的输出嵌入维度
        self.adapter_layer = nn.Linear(
            self.st_backbone.get_sentence_embedding_dimension(), 
            output_dim
        )

    def forward(self, input_sentences):
        # 获取ST的句子嵌入（convert_to_tensor=True保留计算图）
        sentence_embeds = self.st_backbone.encode(
            input_sentences, 
            convert_to_tensor=True, 
            device=self.adapter_layer.weight.device
        )
        # 经过适配层得到最终输出
        outputs = self.adapter_layer(sentence_embeds)
        return outputs

# 实例化模型（根据任务调整output_dim，比如分类任务设为类别数）
model = STAdapterModel(output_dim=2)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
# 初始只优化适配层参数
optimizer = torch.optim.Adam(model.adapter_layer.parameters(), lr=1e-3)

# 训练循环示例
model.train()
for epoch in range(10):
    # 示例数据：句子列表+对应标签
    train_sentences = ["这是正面评价", "这是负面评价", "产品质量很棒", "体验很差"]
    train_labels = torch.tensor([0, 1, 0, 1])
    
    optimizer.zero_grad()
    # 前向传播
    logits = model(train_sentences)
    # 计算损失
    loss = criterion(logits, train_labels)
    # 反向传播+参数更新
    loss.backward()
    optimizer.step()
    
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

进阶：微调Sentence Transformer部分参数 #

如果训练几轮后想微调ST的部分层以提升效果，可以解冻特定层的参数并调整优化器：

# 解冻ST的最后两层（以MiniLM为例，根据模型结构调整层名）
for name, param in model.st_backbone.named_parameters():
    if 'encoder.layer.10' in name or 'encoder.layer.11' in name:
        param.requires_grad = True

# 更新优化器，同时优化适配层和解冻的ST参数
optimizer = torch.optim.Adam([
    {'params': model.adapter_layer.parameters(), 'lr': 1e-3},
    {'params': [p for n, p in model.st_backbone.named_parameters() if p.requires_grad], 'lr': 5e-5}
])

关键细节 #

• 必须在自定义模型的forward方法中调用encode()并设置convert_to_tensor=True，这样才能保留计算图，让适配层的参数参与反向传播。
• 小样本场景下优先冻结预训练模型，避免过拟合，待适配层收敛后再微调部分层效果更好。

内容的提问来源于stack exchange，提问作者Kaushal Tiwari