Adaptive Sparsity Optimization with Learnable Soft Top-K and Per-Term Thresholding for Efficient Retrieval
Wentai Xie, Parker Carlson, Shanxiu He, Tao Yang
cs.IR
2026-10-02
UCSB 的 AdaSparse 在 Lion-SP 上加入可学习 soft top-K 与逐词阈值。MS MARCO 的 MRR@10 保持 0.408,查询与文档向量缩至复现基线的 31% 和 26.6%,Seismic 深度 1000 的延迟从 9.0 毫秒降到 2.4 毫秒,索引从 138GB 降到 33GB。
稀疏检索为每个词学一个非负权重,线上只扫描倒排索引中非零词的链,CPU 就能取回候选。SPLADE 用 BERT 把原文扩展到语义相近的词。Lion-SP 把扩展器换成 LLaMA-3,在 MS MARCO 和 BEIR 上报出领先的相关性,向量也更长。
LLaMA-3 词表有 128,256 个 token,约为 BERT 的 4.2 倍。同一英文词会拆成大小写、单复数和空格变体,多语言预训练还会点亮非英文 token。问句「how many years did william bradford serve as governor of plymouth colony?」上,SPLADE 长度 20。这条查询上 Lion-SP-1B 长度 272,「year」含 9 个英文变体和 2 个中文 token,平均权重 0.20,SPLADE 为 0.40。图 1 里,MS MARCO 上的非零分大量堆在偏低、彼此接近的区间。倒排延迟和索引体积都随非零项增长。
固定 top-K 对每条向量留同样多个词,复杂查询会被剪过头。FLOPs 正则压一个 batch 内权重的平方和,权重大时有效,小权重的梯度很快变弱,长尾留着,系数加大就伤相关性。全局阈值方法 HT 全集只学一个阈值,零样本时有剪出空向量的风险。质量比剪枝（MRP）按向量的 L1 质量留头部,比例 α 靠手调,剪多少不能学习。引言写官方 Lion-SP 的文档向量比某一版 SPLADE 长 8.9 倍,正文没有那一版的绝对长度。
AdaSparse 在 Lion-SP 的损失上加两项正则。推理不做硬 top-K,低于逐词阈值的权重直接置零,其余编码不变。
STop 是随输入变化的软上限,K = b × 原文长度。查询 b=10,文档 b=8,实际倍数可以更低,也可以略高。门控用 sigmoid,枢轴是第 K 名与第 K+1 名权重的均值。权重高于枢轴,门控接近 0,惩罚关掉;低于枢轴,门控接近 1,按权重平方下拉。温度 γ=0.04。
原文里的词用 β=2 豁免。该值高于图 1 中的常见非零权重,多半低于 1.4。查询「quantum computing impact on finance」里,quantum 只有 0.24,全局阈值 0.26 会删掉它;STop 因为它在原文中,门控同样关闭。贴着枢轴、且权重大于 0.16 的词,梯度还可以朝上,排序损失能把有用词抬回软上限内。名次若有噪声,软上限外的有用词会被错罚。
PTT 为每个 token 单学阈值,查询和文档两套。训练时用 sigmoid 把阈值以下的权重乘到接近 0,π=0.04,梯度仍能传;建索引和编码查询时改为硬置零。阈值损失系数 λT=1。某个词若在各文档上的权重很平,阈值推不上去,稀疏度会停住。
FLOPs 系数沿用 Lion-SP,λQ=0.05、λD=0.04,STop 用同一组系数。偏大的噪声主要靠 FLOPs 收缩;词落到枢轴之下后,STop 继续下推,补上小权重梯度消失的一段;排序损失再把有用的词推回去。骨干是 LLaMA-3.2 的 1B 和 8B。先在约 880 万条 MS MARCO passage 上做 masked next-token prediction 预训练,再以对比损失加知识蒸馏微调 1 个 epoch。不同损失从同一份预训练权重起步。查询平均约 8 个词,passage 平均 76 个词。
域内用 MS MARCO Dev 的 6980 条查询,报 MRR@10,看前 10 名里第一个相关文档有多靠前。TREC DL 2019 有 43 条判断查询,DL 2020 有 53 条,这两套和 BEIR 的 13 个数据集报 NDCG@10。长度和相关性的主对照,是同一设置下复现的 Lion-SP-1B。官方检查点查询 293、文档 1250、BEIR 0.535;复现版查询 208、文档 1052、BEIR 0.545。
| 方法 | MRR@10 | DL19 | DL20 | BEIR | 查询 | 文档 |
| Lion-SP-1B 复现 | 0.408 | 0.753 | 0.744 | 0.545 | 208 | 1052 |
| 加大 FLOPs | 0.400 | 0.745 | 0.743 | 0.530 | 130 | 315 |
| L1 加 FLOPs | 0.404 | 0.754 | 0.718 | 0.531 | 98 | 297 |
| 固定 top-305 | 0.372 | 0.709 | 0.690 | 0.231 | 85 | 308 |
| VDR-256 | 0.398 | 0.741 | 0.717 | 0.493 | 103 | 272 |
| MRP α=0.75 | 0.406 | 0.748 | 0.742 | 0.529 | 78 | 268 |
| HT | 0.404 | 0.754 | 0.744 | 0.536 | 189 | 391 |
| AdaSparse-1B | 0.408 | 0.755 | 0.729 | 0.541 | 65 | 280 |
| Lion-SP-8B | 0.418 | 0.760 | 0.766 | 0.552 | 476 | 1457 |
| AdaSparse-8B | 0.417 | 0.763 | 0.752 | 0.543 | 85 | 297 |
AdaSparse-1B 的 Dev MRR@10 仍是 0.408。BEIR 从 0.545 到 0.541,相对约低 0.74%。查询长度是复现版的 31%,文档是 26.6%。DL19 从 0.753 到 0.755,DL20 从 0.744 降到 0.729。固定 top-305 的 BEIR 掉到 0.231。MRP 的 BEIR 为 0.529,大约低 2.3%,长度接近。HT 查询还有 189,文档 391,比 AdaSparse 的 280 长约四成,Dev MRR 为 0.404。
FLOPs 上只加 STop,查询就从 208 收到 83,文档从 1052 收到 364,MRR 0.409,BEIR 0.540。再加 PTT,到 65 和 280,MRR 0.408,BEIR 0.541。PTT 换成全局阈值,查询能到 46,BEIR 掉到 0.530。
Seismic 在 AMD EPYC 9R45、256GB 内存上测,三次平均,只做近似检索。k=1000 时 AdaSparse-1B 为 2.390 毫秒,官方 Lion-SP-1B 为 9.048 毫秒,约 3.8 倍;k=10 时 0.292 毫秒对 1.179 毫秒,约 4 倍。索引 33GB 对 138GB,约 4.2 倍。HT 是 3.223 毫秒和 58GB。AdaSparse-8B 索引 34GB,k=1000 为 2.863 毫秒。
BEIR 十三个集平均,官方 Lion-SP-1B 查询 649、文档 1811、NDCG 0.535,AdaSparse 为 129、456、0.541。MS MARCO 上实际扩展是查询 8.1 倍、文档 3.7 倍,都低于训练用的 b。十四个集合里有十一个查询扩展不超过 10 倍,NFCorpus、SciFact、FEVER 到了 11.7 至 13.5 倍。把查询 b 调到 6、文档 b 调到 5 后,长度为 50 和 206,Dev MRR 0.409,BEIR 0.537。文档长度接近时,加大 FLOPs 只有 0.382 和 0.508,MRP 为 0.394 和 0.526。
SPLADEv3 基于 BERT 110M,查询 24、文档 170、Dev MRR 0.402、BEIR 0.517。AdaSparse-1B 查询仍约长 2.7 倍,文档约长 1.6 倍。稠密 RepLLaMA 用 LLaMA-2 7B,索引约 145GB;AdaSparse-8B 为 34GB,约小 4.3 倍,Dev 0.417 对 0.412,DL19 与 DL20 为 0.763、0.752 对 0.743、0.725,BEIR 则是 0.543 对 0.551。1B 增到 8B,Dev 从 0.408 到 0.417,BEIR 从 0.541 到 0.543,查询从 65 到 85,文档从 280 到 297。
128K 词表多出来的激活,大量是词形变体和跨语言近邻。已经在用 Lion-SP 的系统,重训这份损失并重刷文档即可,点积打分和倒排不用改。Seismic 的延迟主要跟着文档长度走,查询从 9 增到 69 时变化很小,因此文档扩展收在 3.7 倍、查询放在 8.1 倍。这个分配绑在 Seismic 上,其他倒排引擎没有测量。
相对 SPLADEv3,向量仍然更长,域内分更高。相对稠密的 RepLLaMA,索引大约小 4.3 倍,BEIR 平均低约 1.5%。预算卡在 CPU 倒排上时,这个交换成立。
STop 会受词权重名次的噪声影响。文中称为脆性悬崖,只靠较大的 b 缓解,没有单独测误伤多少有用扩展。PTT 遇到平坦的词分布就推不动:λT 从 1 加到 6,BEIR 从 0.541 降到 0.534,长度仅从查询 65、文档 280 收到 54、245。β=2、γ=0.04 来自 Lion-SP 在 MS MARCO 上的权重分布,方案也只在这一个骨干上验证。文中担心全局阈值剪出空向量,没有检查 AdaSparse 在 BEIR 上是否出现空表示。
平均分盖住了单集来回。相对复现版,DL20 从 0.744 降到 0.729;Lion-SP-8B 从 0.766 降到 0.752。相对官方 Lion-SP-1B,ArguAna 从 0.488 到 0.470,Climate-FEVER 从 0.304 到 0.281,Quora 从 0.791 升到 0.858,TREC-COVID 从 0.785 升到 0.817。
速度对照的是官方 Lion-SP,精确 MRR 0.410。更短的复现版没有 Seismic 时间,4.2 倍存储混进了检查点差异。相对复现版,长度只降到 31% 和 26.6%。HT 的查询截断是 8,AdaSparse 是 6。k=10 的近似 MRR 为 0.401,与精确值 0.408 差 0.007,超出文中写的 0.002 对齐范围;k=1000 的差为 0.002。训练也不使用 MS MARCO 标题。表 13 的 Lion 行是 8B,查询 85、文档 297;2.7 倍和 1.6 倍说的是 1B 对 SPLADEv3。