在sentence-transformers中训练新增token的嵌入向量

背景介绍

在自然语言处理领域，预训练语言模型如BERT及其变体已经成为主流。sentence-transformers是基于这些模型构建的专门用于生成句子嵌入的框架。在实际应用中，我们经常需要扩展模型的词汇表，添加一些特殊token来表示特定概念或关系。

问题场景

假设我们使用BAAI/bge-base-en-v1.5这样的预训练模型，需要添加一些表示ConceptNet关系的特殊token，如<|AtLocation|>、<|IsA|>等。然后，我们希望只训练这些新增token的嵌入向量，而保持原始模型的其他参数不变。

技术实现方案

1. 扩展tokenizer

首先需要扩展原始模型的tokenizer，添加新的特殊token。这可以通过以下步骤实现：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-base-en-v1.5")
new_tokens = ["<|AtLocation|>", "<|IsA|>", "<|AntonymOf|>"]
tokenizer.add_tokens(new_tokens)

2. 调整模型嵌入层

添加新token后，需要调整模型的嵌入层以容纳这些新token：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-base-en-v1.5")
model.resize_token_embeddings(len(tokenizer))

3. 冻结原始参数

为了确保只训练新增token的嵌入，需要冻结模型的其他参数：

for param in model.parameters():
    param.requires_grad = False

# 仅解冻新增token对应的嵌入
embedding_layer = model._first_module().auto_model.get_input_embeddings()
embedding_layer.weight.requires_grad[-len(new_tokens):] = True

4. 设计训练目标

使用ConceptNet的三元组数据(head, relation, tail)，可以设计如下训练目标：

def custom_loss(head_emb, relation_emb, tail_emb):
    return torch.norm(head_emb + relation_emb - tail_emb, p=2)

训练策略

1. 数据准备：从ConceptNet中提取大量(head, relation, tail)三元组
2. 批量处理：将head和tail转换为句子嵌入，relation转换为特殊token嵌入
3. 优化目标：最小化norm(embed(head) + embed(relation_token) - embed(tail))
4. 参数更新：仅更新新增token对应的嵌入向量

应用价值

这种方法可以：

• 保持预训练模型的核心能力不变
• 仅针对特定任务添加必要的知识表示
• 避免全模型微调带来的灾难性遗忘问题
• 实现更高效的知识注入

注意事项

1. 新增token的数量不宜过多，以免影响模型性能
2. 训练数据需要充分覆盖新增token的使用场景
3. 学习率设置要适当，通常比全模型微调时更小
4. 需要验证新增token的嵌入是否真正捕获了预期语义

通过这种针对性训练方法，我们可以有效地扩展预训练模型的能力，同时保持其原有优势。