如何使用Sentence Transformers v5微调稀疏嵌入模型以提升搜索性能
Sentence Transformers 是一个用于多种自然语言处理任务的Python库,包括检索增强生成、语义搜索、语义文本相似性、同义句挖掘等。最新发布的 v5 版本引入了多项重要的改进,特别适用于微调稀疏编码模型。这篇文章将详细介绍如何使用 Sentence Transformers 微调稀疏编码模型,并解释其应用场景。 什么是稀疏嵌入模型? 嵌入模型通常指的是将文本或其他输入转换为向量表示(嵌入)的技术,这些向量能够捕捉输入的语义信息,便于进行数学运算。常见的密集嵌入模型通常生成低维向量(如384、768或1024维度),值大多不为零。而稀疏嵌入模型则生成高维向量(如30,000+维度),大多数值为零。每个活跃维度(即非零值的维度)对应模型词汇表中的一个特定标记,增强了模型的解释性。 微调稀疏嵌入模型的原因 微调稀疏嵌入模型的主要目的是让模型更好地匹配特定领域和语言的文本。尽管预训练模型已经具备一定的词义扩展能力,但它们可能会在某些领域或语言中表现不佳。例如,“cephalalgia”应扩展为“头痛”,而不是与医学无关的词汇。通过微调,可以显著提高模型在特定任务上的性能。 训练组件 使用 Sentence Transformers 训练稀疏嵌入模型涉及以下几个组件: 模型:需要训练或微调的模型,可以是预训练的稀疏编码模型或基础模型。 数据集:用于训练和评估的数据。 损失函数:衡量模型性能并指导优化过程的函数。 训练参数(可选):影响训练性能和调试跟踪的参数。 评估器(可选):用于评估模型性能的工具。 训练器:将上述组件组合起来进行训练的类。 稀疏嵌入模型架构 Spalde 架构 Spalde 模型由两部分组成:一个 Masked Language Modeling (MLM) 变压器和一个 SpaldePooling 模块。MLM 变压器负责加载预训练模型(如 BERT、RoBERTa 或 DistilBERT),而 SpaldePooling 模块则对输出进行池化操作,生成与词汇表大小一致的稀疏嵌入。 python from sentence_transformers import models, SparseEncoder mlm_transformer = models.MLMTransformer("google-bert/bert-base-uncased") splade_pooling = models.SpaldePooling(pooling_strategy="max") model = SparseEncoder(modules=[mlm_transformer, splade_pooling]) Inference-free Spalde 架构 Inference-free Spalde 使用一个 Router 模块,分别用于查询和文档编码。文档部分采用传统的 Spalde 架构(MLM 变压器 + SpaldePooling),查询部分使用轻量级的 SparseStaticEmbedding 模块,返回查询中每个标记的预计算分数。 python from sentence_transformers import SparseEncoder, models, Router doc_encoder = models.MLMTransformer("google-bert/bert-base-uncased") router = Router.for_query_document( query_modules=[models.SparseStaticEmbedding(tokenizer=doc_encoder.tokenizer, frozen=False)], document_modules=[doc_encoder, models.SpaldePooling("max")] ) model = SparseEncoder(modules=[router], similarity_fn_name="dot") 数据集 Sentence Transformers 提供了多种方式加载数据集,可以从 Hugging Face Datasets Hub 加载公共数据集,也可以使用本地文件(如 CSV、JSON、Parquet 等格式)。 python from datasets import load_dataset train_dataset = load_dataset("sentence-transformers/natural-questions", split="train") 损失函数 选择合适的损失函数对于模型训练至关重要。主要的损失函数包括 SpladeLoss 和 CSRLoss。这些损失函数在标准损失基础上添加了稀疏正则化项。 python from sentence_transformers import SparseEncoder from sentence_transformers.sparse_encoder.losses import SpladeLoss, SparseMultipleNegativesRankingLoss model = SparseEncoder("distilbert/distilbert-base-uncased") loss = SpladeLoss( model=model, loss=SparseMultipleNegativesRankingLoss(model=model), query_regularizer_weight=5e-5, document_regularizer_weight=3e-5 ) 训练参数 使用 SparseEncoderTrainingArguments 类可以指定影响训练性能的参数。这些参数有助于实验和调试。 python from sentence_transformers import SparseEncoderTrainingArguments args = SparseEncoderTrainingArguments( output_dir="models/splade-distilbert-base-uncased-nq", num_train_epochs=1, per_device_train_batch_size=16, per_device_eval_batch_size=16, learning_rate=2e-5, learning_rate_mapping={r"SparseStaticEmbedding\.weight": 1e-3}, fp16=True, eval_strategy="steps", eval_steps=1000, save_strategy="steps", save_steps=1000 ) 评估器 评估器可以帮助监控模型在训练过程中的性能。常见的评估器有 SparseNanoBEIREvaluator、SparseEmbeddingSimilarityEvaluator 等。 python from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator dev_evaluator = SparseNanoBEIREvaluator(batch_size=16) 训练器 使用 SparseEncoderTrainer 类将所有组件整合起来进行训练。 python from sentence_transformers import SparseEncoderTrainer trainer = SparseEncoderTrainer( model=model, args=args, train_dataset=train_dataset, eval_dataset=eval_dataset, loss=loss, evaluator=dev_evaluator ) trainer.train() 评估和部署 训练完成后,可以使用多模态搜索引擎如 Qdrant 进行高效的嵌入向量存储和检索。Qdrant 提供了优秀的稀疏向量支持,可以显著提高搜索效率。 ```python from sentence_transformers import SparseEncoder from sentence_transformers.sparse_encoder.search_engines import semantic_search_qdrant sparse_model = SparseEncoder("naver/splade-cocondenser-ensembledistil") corpus = dataset["answer"][:10_000] corpus_embeddings = sparse_model.encode_document(corpus, convert_to_sparse_tensor=True) queries = dataset["query"][:2] query_embeddings = sparse_model.encode_query(queries, convert_to_sparse_tensor=True) results, search_time, corpus_index = semantic_search_qdrant(query_embeddings, corpus_index=None, corpus_embeddings=corpus_embeddings, top_k=5) ``` 结论 稀疏嵌入模型在特定领域和语言任务中表现出色,尤其是在混合搜索和重排序应用场景中。通过微调,可以显著提高模型的匹配能力和性能。业内专家认为,这种模式填补了传统词法方法(如 BM25)和密集嵌入模型之间的空白,兼具效率和效果。Sentence Transformers 的推出使得这一过程更加简便和高效,为开发者提供了强大的工具。 Sentence Transformers 由 UKP Lab 维护,该实验室在自然语言处理领域享有盛誉,致力于研发高效、准确的文本嵌入和重排序模型。该项目在 GitHub 上拥有数千星标,社区活跃,不断有新的模型和功能加入。