SPLADE 稀疏学习检索深度实战：从 Masked Language Model 到生产级语义稀疏搜索的完整部署指南

引言：为什么需要学习型稀疏检索

在向量搜索领域，稠密检索（Dense Retrieval）和稀疏检索（Sparse Retrieval）长期以来被视为两条平行赛道。稠密模型如 SBERT、DPR 通过将文本映射到连续向量空间实现语义匹配，但存在不可解释、难于调试、需要专用向量数据库等痛点。传统稀疏检索如 BM25 虽然可解释且无需额外基础设施，但完全依赖词频匹配，缺乏语义理解能力。

SPLADE（SParse Lexical AnD Expansion）模型的出现打破了这一僵局。它通过 Transformer 的 Masked Language Model 头生成稀疏表示，既保留了稀疏检索的可解释性和倒排索引兼容性，又获得了深度语义理解能力。更关键的是，SPLADE 支持查询扩展——自动为查询补充语义相关词汇，大幅提升召回率。

本文将从 SPLADE 的核心原理出发，深入解析其稀疏激活机制、训练策略、蒸馏方案，并给出从零搭建到生产部署的完整实战方案。

SPLADE 核心架构解析

从 Transformer 到稀疏表示

SPLADE 的核心思想是将 Transformer 编码器的输出转换为词级别的稀疏权重。具体流程如下：

第一步：Transformer 编码。输入文本经过 BERT/DistilBERT 等模型编码后，得到每个 token 的上下文表示向量 H ∈ R^{n×d}，其中 n 为序列长度，d 为隐藏维度。

第二步：词汇映射。通过 MLM（Masked Language Model）头，将每个 token 的隐藏表示映射到词表空间，得到 logits 矩阵 L ∈ R^{n×V}，V 为词表大小。这一步是 SPLADE 与稠密模型的关键区别——它将表示空间从连续维度映射回离散词表。

第三步：稀疏激活。对 logits 施加 ReLU 激活函数过滤负值，再沿序列维度取 max-pooling，得到每个词表项的最终权重：


1
w_j = max_i(ReLU(logit_{i,j}))

这一操作保证了最终表示的稀疏性——只有语义相关词汇获得非零权重，其余维度均为 0。稀疏程度可通过调节 ReLU 前的对数变换和正则化策略来控制。

查询扩展的魔法

与 BM25 仅匹配查询中出现的词不同，SPLADE 能够为查询自动扩展语义相关词汇。例如，查询「深度学习框架」经过 SPLADE 编码后，可能为「PyTorch」「TensorFlow」「神经网络」等词分配非零权重，即使这些词并未出现在原始查询中。

这种扩展能力来自 Transformer 的上下文理解——MLM 头天然具备预测被遮蔽词的能力，SPLADE 巧妙地将此能力转化为词汇扩展信号。

SPLADE 家族变体详解

SPLADE 自 2021 年提出以来，已演化出多个变体，各自针对不同场景优化：

变体	编码策略	参数量	延迟	适用场景
SPLADE	Doc + Query 双编码器	~110M	中等	通用语义检索
SPLADE-doc	仅文档端稀疏化	~110M	查询快	高频查询场景
SPLADE-v2	FLOPS 正则	~110M	中等	均衡精度与稀疏度
SPLADE++ (ED)	蒸馏 + 对比学习	~66M	低	生产级高效部署
SPLADE++ (EnsembleDistil)	集成蒸馏	~66M	最低	极致延迟优化

SPLADE-doc：查询端零计算

SPLADE-doc 只对文档端进行稀疏编码并索引到倒排索引中，查询端仍使用原始查询词进行 BM25 式匹配。这大幅降低查询时延迟，同时利用文档端的语义扩展提升召回率。适用于查询频率远高于文档更新频率的场景。

SPLADE-v2：FLOPS 正则化

SPLADE-v2 引入了 FLOPS（Floating Point Operations per Second）正则化项，直接约束模型激活的词表维度数量：


1
L_flops = Σ_j (Σ_i w_{i,j})^2

这一正则项在训练过程中惩罚高激活维度，迫使模型将权重集中在最重要的词汇上，实现更稀疏的表示。相比简单的 L1 正则，FLOPS 正则更关注整体激活模式的稀疏性。

SPLADE++：蒸馏到极致

SPLADE++ 是当前生产部署推荐版本。它采用两阶段蒸馏策略：

第一阶段：使用稠密模型（如 ColBERT）作为教师模型，通过 margin MSE 损失蒸馏知识到 SPLADE 学生模型
第二阶段：结合多个教师模型的集成蒸馏，进一步提升检索精度

最终 SPLADE++ 在 MS MARCO 上达到 39.2 MRR@10，接近稠密检索最优水平，同时保持毫秒级查询延迟。

从零搭建 SPLADE 检索系统

环境准备与模型加载

推荐使用 HuggingFace 上的预训练 SPLADE 模型，无需从头训练即可获得出色的效果：


1
2
3
4
5
6
7
8
9
10
11
# 安装依赖

pip install transformers torch sentence-transformers

pip install pysparse  # 稀疏矩阵操作



# 加载 SPLADE++ 模型

from transformers import AutoModelForMaskedLM, AutoTokenizer



model_name = "naver/splade-cocondenser-ensembledistil"

tokenizer = AutoTokenizer.from_pretrained(model_name)

model = AutoModelForMaskedLM.from_pretrained(model_name)

model.eval()

注意选择模型版本：

1	splade-cocondenser-ensembledistil

是目前综合性能最优的版本，而

1	splade-doc

系列适用于查询端零计算场景。

稀疏向量生成核心代码

以下是完整的稀疏向量生成函数：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import torch

import numpy as np

from collections import defaultdict



def generate_sparse_vector(text, model, tokenizer, max_length=512):

    """

    将文本转换为 SPLADE 稀疏表示

    返回: dict {token_id: weight}

    """

    tokens = tokenizer(

        text, 

        return_tensors="pt", 

        max_length=max_length, 

        truncation=True, 

        padding=True

    )

    

    with torch.no_grad():

        output = model(**tokens)

    

    logits = output.logits  # [1, seq_len, vocab_size]

    

    # ReLU 激活 + max-pooling

    sparse_weights = torch.max(

        torch.relu(logits), 

        dim=1

    ).values.squeeze()  # [vocab_size]

    

    # 只保留非零权重

    nonzero_indices = sparse_weights.nonzero(as_tuple=True)[0]

    nonzero_values = sparse_weights[nonzero_indices]

    

    # 转换为词表词与权重

    sparse_dict = {}

    for idx, val in zip(nonzero_indices.tolist(), nonzero_values.tolist()):

        token = tokenizer.decode([idx])

        if val > 0.5:  # 过滤低权重噪音

            sparse_dict[token.strip()] = round(val, 4)

    

    return sparse_dict



# 示例

query = "深度学习框架性能对比"

sparse_rep = generate_sparse_vector(query, model, tokenizer)

print(f"原始查询: {query}")

print(f"稀疏表示: {sparse_rep}")

# 输出示例:

# 原始查询: 深度学习框架性能对比

# 稀疏表示: {'深度': 3.21, '学习': 2.89, '框架': 2.56, 

#            'pytorch': 1.78, 'tensorflow': 1.65, '性能': 2.12, 

#            'benchmark': 1.34, '神经网络': 1.23, ...}

构建倒排索引

SPLADE 的核心优势之一是可以直接利用倒排索引基础设施。以下是使用 Elasticsearch 构建索引的方案：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from elasticsearch import Elasticsearch



es = Elasticsearch("http://localhost:9200")



# 创建支持 rank_features 的索引映射

index_mapping = {

    "mappings": {

        "properties": {

            "content": {"type": "text"},

            "sparse_vector": {

                "type": "rank_features"

            }

        }

    }

}



es.indices.create(index="splade_docs", body=index_mapping)



def index_document(doc_id, text, model, tokenizer):

    """将文档编码为稀疏向量并索引到 ES"""

    sparse = generate_sparse_vector(text, model, tokenizer)

    

    # 转换为 rank_features 格式

    rank_features = {f"{token}": weight for token, weight in sparse.items()}

    

    doc = {

        "content": text,

        "sparse_vector": rank_features

    }

    es.index(index="splade_docs", id=doc_id, document=doc)



# 批量索引

documents = [

    {"id": "1", "text": "PyTorch 是 Facebook 开发的深度学习框架..."},

    {"id": "2", "text": "TensorFlow 2.x 提供了 eager execution 模式..."},

    # ... 更多文档

]



for doc in documents:

    index_document(doc["id"], doc["text"], model, tokenizer)

查询与评分机制

基于 rank_features 的查询

Elasticsearch 的 rank_features 字段天然支持 SPLADE 的稀疏评分。查询时，将查询的稀疏表示作为 feature 权重传入：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def search_splade(query_text, model, tokenizer, top_k=10):

    """使用 SPLADE 稀疏向量进行语义检索"""

    query_sparse = generate_sparse_vector(query_text, model, tokenizer)

    

    # 构建 rank_features 查询

    query_body = {

        "size": top_k,

        "query": {

            "rank_feature": {

                "field": "sparse_vector",

                "saturation": {

                    "query": query_sparse

                }

            }

        }

    }

    

    # 或者使用更精细的 inner_hit 模式

    results = es.search(index="splade_docs", body=query_body)

    

    return [

        {

            "id": hit["_id"],

            "score": hit["_score"],

            "content": hit["_source"]["content"]

        }

        for hit in results["hits"]["hits"]

    ]

混合检索：SPLADE + BM25

在实际生产中，SPLADE 与 BM25 的混合检索往往优于任一单独方案。BM25 负责精确词匹配，SPLADE 负责语义扩展，二者互补：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def hybrid_search(query_text, model, tokenizer, top_k=10,

                  splade_weight=0.6, bm25_weight=0.4):

    """混合检索：SPLADE 语义 + BM25 精确匹配"""

    query_sparse = generate_sparse_vector(query_text, model, tokenizer)

    

    query_body = {

        "size": top_k,

        "query": {

            "bool": {

                "should": [

                    {

                        "rank_feature": {

                            "field": "sparse_vector",

                            "saturation": {

                                "query": query_sparse

                            },

                            "boost": splade_weight

                        }

                    },

                    {

                        "match": {

                            "content": {

                                "query": query_text,

                                "boost": bm25_weight

                            }

                        }

                    }

                ]

            }

        }

    }

    

    results = es.search(index="splade_docs", body=query_body)

    return results["hits"]["hits"]

SPLADE 训练与微调

训练数据准备

SPLADE 的训练遵循标准的检索微调范式。你需要准备三类数据：

正样本对：(query, relevant_doc) 标注数据
负样本：同一 query 对应的不相关文档，可通过 BM25 hard negative 挖掘
教师分数（蒸馏模式）：来自交叉编码器或 ColBERT 的相关性分数

以 MS MARCO 格式为例：


1
2
3
4
5
6
7
# train_triples.jsonl 格式

{"query": "what is python", "pos": "Python is a programming language...", "neg": "Monty Python comedy group..."}

{"query": "machine learning basics", "pos": "ML is a subset of AI...", "neg": "Basic washing machine repair..."}



# 教师分数格式（蒸馏用）

{"query": "what is python", "doc": "Python is a programming language...", "score": 8.72}

{"query": "what is python", "doc": "Python snake species...", "score": 1.34}

对比学习训练

SPLADE 标准训练使用 In-Batch Negative 对比学习 + FLOPS 正则：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
import torch

import torch.nn.functional as F



class SPLADETrainer:

    def __init__(self, model, tokenizer, lr=2e-5, flops_weight=0.01):

        self.model = model

        self.tokenizer = tokenizer

        self.optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

        self.flops_weight = flops_weight

    

    def compute_sparse_loss(self, q_sparse, d_pos_sparse, d_neg_sparse):

        """对比损失 + FLOPS 正则"""

        # 计算稀疏内积得分

        pos_score = self.sparse_dot(q_sparse, d_pos_sparse)

        neg_scores = self.sparse_dot(q_sparse, d_neg_sparse)

        

        # Margin 对比损失

        margin = 1.0

        contrastive_loss = F.relu(margin - pos_score + neg_scores).mean()

        

        # FLOPS 正则（鼓励稀疏）

        flops_reg_q = torch.sum(q_sparse, dim=0).pow(2).sum()

        flops_reg_d = torch.sum(

            torch.cat([d_pos_sparse, d_neg_sparse], dim=0), dim=0

        ).pow(2).sum()

        flops_loss = (flops_reg_q + flops_reg_d) * self.flops_weight

        

        return contrastive_loss + flops_loss

    

    @staticmethod

    def sparse_dot(a, b):

        """稀疏向量内积"""

        return torch.sum(a * b, dim=-1)



# 训练循环

trainer = SPLADETrainer(model, tokenizer)

for batch in train_dataloader:

    loss = trainer.compute_sparse_loss(

        batch["q_sparse"], 

        batch["d_pos_sparse"], 

        batch["d_neg_sparse"]

    )

    loss.backward()

    trainer.optimizer.step()

    trainer.optimizer.zero_grad()

知识蒸馏训练

SPLADE++ 的蒸馏策略使用 margin MSE 损失，将教师模型的细粒度相关性分数蒸馏到稀疏模型：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def distillation_loss(student_scores, teacher_scores, margin=1.0):

    """

    Margin MSE 蒸馏损失

    student_scores: SPLADE 产出的 query-doc 稀疏内积得分

    teacher_scores: 教师模型（如 ColBERT）的相关性分数

    """

    # 将教师分数映射到 margin 空间

    teacher_margin = teacher_scores - teacher_scores.min() + margin

    student_margin = student_scores - student_scores.min() + margin

    

    return F.mse_loss(student_margin, teacher_margin)



# 多教师集成蒸馏

ensemble_teacher_scores = (

    0.4 * colbert_scores + 

    0.3 + 0.3 * cross_encoder_scores

)

loss = distillation_loss(splade_scores, ensemble_teacher_scores)

生产级性能优化

稀疏度控制与索引压缩

SPLADE 的稀疏度直接影响索引大小和查询性能。以下是关键调优参数：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 权重阈值过滤：去除低权重词汇

SPARSE_THRESHOLD = 0.5  # 只保留权重 > 0.5 的词

MAX_SPARSE_DIMS = 200   # 每篇文档最多保留 200 个非零维度



# 编码时应用阈值

def generate_sparse_vector_optimized(text, model, tokenizer):

    sparse = generate_sparse_vector(text, model, tokenizer)

    

    # 按权重排序，截断到 top-k

    sorted_items = sorted(sparse.items(), key=lambda x: -x[1])

    top_k = dict(sorted_items[:MAX_SPARSE_DIMS])

    

    # 过滤低权重

    filtered = {k: v for k, v in top_k.items() if v > SPARSE_THRESHOLD}

    

    return filtered

实验数据显示，将每文档的稀疏维度从 500+ 压缩到 100-200 时，MRR@10 仅下降 1-2%，但索引体积减少 60% 以上，查询延迟降低 40%。

ONNX 推理加速

对于高吞吐场景，将 SPLADE 模型导出为 ONNX 格式可带来 2-4 倍的推理加速：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from transformers import AutoModelForMaskedLM

import torch



model = AutoModelForMaskedLM.from_pretrained(

    "naver/splade-cocondenser-ensembledistil"

)

model.eval()



# 导出 ONNX

dummy_input = {

    "input_ids": torch.randint(0, 30000, (1, 128)),

    "attention_mask": torch.ones(1, 128, dtype=torch.long)

}



torch.onnx.export(

    model,

    (dummy_input["input_ids"], dummy_input["attention_mask"]),

    "splade.onnx",

    input_names=["input_ids", "attention_mask"],

    output_names=["logits"],

    dynamic_axes={

        "input_ids": {0: "batch", 1: "seq_len"},

        "attention_mask": {0: "batch", 1: "seq_len"},

        "logits": {0: "batch", 1: "seq_len"}

    },

    opset_version=14

)



# ONNX Runtime 推理

import onnxruntime as ort



sess = ort.InferenceSession("splade.onnx")

input_ids = tokenizer(text, return_tensors="np")

logits = sess.run(None, {

    "input_ids": input_ids["input_ids"],

    "attention_mask": input_ids["attention_mask"]

})[0]

批量编码与并行处理

大规模文档索引需要高效的批量编码策略：


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from torch.utils.data import DataLoader



def batch_encode_documents(documents, model, tokenizer, batch_size=32):

    """批量编码文档为稀疏向量"""

    model.eval()

    all_sparse = []

    

    for i in range(0, len(documents), batch_size):

        batch = documents[i:i+batch_size]

        inputs = tokenizer(

            batch, 

            return_tensors="pt", 

            padding=True, 

            truncation=True, 

            max_length=512

        )

        

        with torch.no_grad():

            outputs = model(**inputs)

        

        logits = outputs.logits

        sparse = torch.max(torch.relu(logits), dim=1).values

        

        for j in range(sparse.shape[0]):

            nonzero_idx = sparse[j].nonzero(as_tuple=True)[0]

            nonzero_val = sparse[j][nonzero_idx]

            sparse_dict = {

                tokenizer.decode([idx]): round(float(val), 4)

                for idx, val in zip(nonzero_idx.tolist(), nonzero_val.tolist())

                if val > 0.5

            }

            all_sparse.append(sparse_dict)

    

    return all_sparse

SPLADE 与稠密检索的工程对比

为帮助团队做出技术选型决策，以下从多个维度对比 SPLADE 与主流稠密检索方案：

维度	SPLADE	DPR/SBERT	ColBERT
存储开销	低（倒排索引）	高（FAISS 向量索引）	极高（token 级向量）
查询延迟	1-5ms	5-20ms	10-50ms
可解释性	强（词级权重）	弱（黑盒向量）	中（token 级匹配）
语义理解	强（查询扩展）	强（连续空间）	极强（late interaction）
基础设施	ES/原生倒排索引	FAISS/Milvus/Qdrant	FAISS + 自定义
索引更新	增量实时	需重建索引	需重建索引
精度(MSMARCO)	39.2 MRR@10	38.5 MRR@10	40.1 MRR@10

核心结论：如果你的系统已经基于 Elasticsearch 或其他倒排索引基础设施，SPLADE 是零基础设施增量的语义检索升级方案。无需引入向量数据库，即可获得接近稠密检索的精度。

典型部署架构与监控

关键监控指标

生产环境需要持续监控以下指标：

稀疏维度分布：P50/P95/P99 的每文档非零维度数，异常增长可能表明模型退化
编码延迟：单条文档编码耗时，ONNX 模式下应 < 20ms
索引膨胀率：SPLADE 索引体积 / 原始文档体积比，正常范围 2-5 倍
查询扩展质量：人工抽样评估扩展词相关性，低于阈值触发模型更新
混合检索权重分布：SPLADE 与 BM25 分数贡献比，理想范围 50-70% SPLADE


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Prometheus 监控指标示例

from prometheus_client import Histogram, Gauge



sparse_dims = Histogram(

    "splade_sparse_dimensions",

    "Number of non-zero dimensions per document",

    buckets=[10, 50, 100, 200, 500, 1000]

)

encode_latency = Histogram(

    "splade_encode_latency_seconds",

    "Document encoding latency",

    buckets=[0.01, 0.02, 0.05, 0.1, 0.5]

)

index_size_ratio = Gauge(

    "splade_index_size_ratio",

    "SPLADE index size / raw document size"

)

总结与最佳实践

SPLADE 为向量搜索提供了一条独特的「语义稀疏」路径，核心优势总结如下：

零基础设施增量：直接复用 Elasticsearch 倒排索引，无需引入向量数据库
查询扩展能力：自动补充语义相关词，弥补 BM25 词汇匹配的不足
完全可解释：每个检索结果可追溯到具体的词级权重，便于调试和优化
增量更新友好：单文档更新无需重建索引，天然支持实时索引
混合检索优势：与 BM25 结合后，精度超过单一稠密或稀疏方案

最佳实践建议：

优先使用 SPLADE++ (EnsembleDistil) 版本，它是经过蒸馏优化的生产就绪模型
设置权重阈值 0.5-1.0 和最大维度 100-200，在精度与效率间取得平衡
务必使用混合检索（SPLADE + BM25），权重比推荐 0.6:0.4
大规模场景导出 ONNX，可获得 2-4 倍推理加速
监控稀疏维度分布，及时发现模型退化

SPLADE 证明了稀疏检索与深度语义理解并不矛盾。对于已有倒排索引基础设施的团队，它是通往语义检索最短、最经济的路径。