
动态超级块剪枝:加速稀疏检索的革命性技术 精华
突破性能瓶颈:动态超级块剪枝如何重塑信息检索效率
在当今数据爆炸的时代,高效的信息检索系统对于各类应用至关重要,从搜索引擎到基于检索增强的大语言模型(RAG)。随着学习型稀疏表示模型的兴起,如何在保持高检索质量的同时提升检索速度成为研究热点。本文深入探讨一项革命性技术——动态超级块剪枝(Superblock Pruning,简称SP),这一创新方法在保持高相关性的前提下,显著提升了稀疏检索的效率。
稀疏检索的挑战与机遇
稀疏检索模型如BM25和学习型稀疏表示(如SPLADE、E-SPLADE等)在仅使用CPU的服务器环境中广受欢迎,主要得益于它们能够充分利用高效的倒排索引实现。传统的稀疏检索速度优化通常采用动态秩安全索引剪枝技术,该技术能够准确跳过那些得分较低、不可能出现在最终top-k结果中的文档。
近年来,基于块的检索方法成为研究热点,这类方法将文档分配到块(或称为簇)中,并利用块级信息改进索引遍历顺序,同时剪枝低分文档组。然而,这些方法在处理大规模数据集时仍然面临效率挑战,尤其是当需要保持高相关性时。
超级块剪枝:创新的两级剪枝策略
超级块剪枝(SP)技术在现有基于块的剪枝方法基础上进行了创新性扩展。SP方法将一系列连续的文档块均匀聚合成超级块,然后以自上而下的方式进行在线索引遍历。这种设计为每个超级块分配固定数量的文档块,简化了向量化和缓存优化过程,同时提供了具有概率安全保证的两级剪枝机制。
SP的核心创新在于其两级剪枝策略:首先计算所有超级块的边界信息并进行剪枝,然后再计算块的边界并进行剪枝。具体来说,SP执行以下动态剪枝步骤:
- 对于超级块X,计算该超级块内文档的最大和平均排名分数边界
- 当超级块的最大和平均超级块边界满足特定条件时,该超级块被剪枝
- 对于文档块B,如果其边界和满足特定条件,则剪枝该块
- 对于所有未被剪枝的块,按照其边界和值的降序对相应的文档块进行排序和评分
这种方法允许SP更有效地跳过文档块,以排名安全或概率排名安全的方式加速检索。剪枝一个超级块不仅避免了计算子块的最大分数,还避免了对其子块内文档的评分,从而大幅提升检索效率。
理论保证与实现优化
SP具有与ASC类似的排名安全μ-竞争性质。可以证明,SP的平均top-k'排名分数与任何排名安全检索算法R在μ因子内相同。作为额外保障,如果我们假设文档的排名分数在每个超级块内独立同分布,SP还提供概率安全性。
在实现层面,SP采用了多项优化策略:
CPU缓存使用优化
SP使用SIMD指令计算相关公式。当顺序计算所有查询项的这些公式而不进行块跳过时,现代编译器可以轻松向量化其实现,现代CPU可以有效地预取数据。然而,由于超级块剪枝导致的不规则和非连续数据访问,编译器难以优化块级边界计算。因此,SP需要显式控制CPU缓存在计算块级边界时的重用模式。
SP采用了超级块优先的边界和计算方式,即对每个未剪枝的超级块,先对该超级块内的所有块进行完整评分,然后再处理下一个未剪枝的超级块。这种方法允许在内部循环中重用累积寄存器以获得更好的L1缓存性能。实验表明,这种方式比传统的项优先方法最高可提速1.89倍。
实验评估与性能对比
研究团队在MS MARCO段落排名数据集上进行了全面评估,该数据集包含880万个英文段落。评估采用标准指标:平均倒数排名(MRR@10)和位置1000(k=1000)或10(k=10)的召回率。所有实验在配备Intel i7-1260P、64GB RAM和AVX2指令的Linux系统上使用单线程运行。
SP与三种最先进的基于块的检索算法进行了比较:BMP、Seismic和ASC。实验结果表明,在高相关性预算要求下,SP在SPLADE和E-SPLADE上的表现显著优于这些基线方法。
SPLADE模型上的性能比较
在SPLADE上的排名安全搜索中,SP比BMP在k=10时快32%,在k=1000时快25%。与ASC相比,SP在k=10和k=1000时都快约3.3倍。对于99%的召回预算,SP比BMP快至多2.9倍,比Seismic快3.3倍,比ASC快9.1倍。
上图展示了SP和BMP在块大小b从128减小到8时的总延迟(上图)和成本细分(下图)。当b变小时,BMP能够获得更紧密的边界和,但块过滤开销增加。SP在评估小块的同时减少了块和超级块过滤的开销。
超级块剪枝的有效性
实验数据显示,即使在安全搜索(μ=1)情况下,SP也能剪枝24%的超级块(k=10)。随着μ减小,超级块级别的剪枝量显著增加,而被剪枝的块数量大致相同。这是因为块对其内部文档形成了紧密边界;SP能够避开不太可能包含相关文档的块组,从而减少计算块边界和的开销。
在100%概率安全性(η=1)下,即使在μ=0.4(k=1000)时,对Dev集、DL 19和DL 20的相关性指标影响也可忽略不计,尽管当μ=0.4时召回率开始下降。相比之下,即使在BMP中使用低估计阈值也会导致相关性大幅下降。
E-SPLADE模型上的性能
在E-SPLADE上,SP在不同高相关性召回预算下的表现也优于其他基线,比Seismic快至多16倍,比BMP快1.4倍。
超级块剪枝的优势与局限
与现有方法相比,SP具有以下显著优势:
- 更高效的块跳过:与BMP相比,SP利用其超级块结构快速跳过大量块,同时提供额外的η保障以确保概率安全性。
- 更好的缓存利用:与Anytime Ranking、ASC和Seismic相比,SP能够处理更多的块数量,并通过缓存优化的超级块剪枝克服额外开销,自然导致更紧密的边界估计。
- 高相关性保证:在保持高相关性的同时显著提升检索速度,特别适合对检索质量要求较高的应用场景。
然而,SP也存在一些局限性:
- 额外空间成本:与BMP相比,SP需要额外空间来维护每个超级块的最大和平均项权重。在MS MARCO评估中,当c=64、b=8时,额外空间约为2GB;当b=16时,额外空间约为1GB。
- 静态索引剪枝的缺失:与Seismic不同,SP没有利用静态索引剪枝,这可能在某些情况下限制其性能。
应用场景与未来展望
SP技术特别适合需要高相关性的应用场景。对于此类应用,建议将η设置接近1.0,并将μ从0.4变化到1。检索是大规模搜索系统和基于检索增强的大语言模型(如RAG)的关键组件,在低成本CPU上实现高相关性的快速检索可以产生积极影响。
未来研究方向包括:
- 探索静态索引剪枝、自定义摘要和文档邻近图等技术与SP的结合
- 研究SP与索引压缩方案的结合
- 开发针对输入复杂性的动态缩放策略
- 整合自适应推理深度控制以在推理期间平衡效率和安全性能
结论
动态超级块剪枝(SP)是一种创新的动态剪枝方案,除标准块级别外,还在超级块级别进行剪枝,并设计为利用CPU缓存局部性。实验评估表明,在SPLADE上99%或更高的召回预算下,SP比Seismic快2.3倍至3.8倍,比ASC快3.2倍至9.4倍,比BMP快至多2.9倍。对于安全搜索,SP比BMP快至多1.3倍。
随着信息检索需求的不断增长和大语言模型对高效检索系统的依赖加深,SP技术有望在未来发挥更加重要的作用,特别是在需要在保持高相关性的同时提高检索效率的场景中。
参考资料
论文:https://arxiv.org/abs/2504.17045
GitHub:https://github.com/thefxperson/hierarchical_pruning
本文转载自顿数AI,作者:可可
