Col-Bandit：面向晚期交互检索的零样本查询时剪枝方法 Col-Bandit: Zero-Shot Query-Time Pruning for Late-Interaction Retrieval
用老虎机框架跳过冗余MaxSim计算，实现ColBERT级质量下最高8×加速
前置知识
晚期交互检索 (Late-Interaction Retrieval)
晚期交互是一种介于单向量稠密检索和交叉编码器之间的检索范式。以 ColBERT 为代表，查询和文档各自被编码为一组 token 级别的嵌入向量，而非单个向量。在评分阶段，对每个查询 token，找到文档中与其最相似的 token（称为 MaxSim 操作），然后将所有查询 token 的最大相似度求和得到最终得分。这种方法捕获了细粒度的语义匹配信息，比单向量方法精度更高，但计算量也更大——需要对每个文档的每个查询 token 都执行一次 MaxSim 操作。
Col-Bandit 的核心目标就是加速这个 MaxSim 计算过程，理解晚期交互的评分机制是理解本文算法的前提。
MaxSim 矩阵
给定一个查询 $Q$（含 $T$ 个 token）和一组 $N$ 个候选文档，MaxSim 矩阵 $H \in \mathbb{R}^{N \times T}$ 是一个隐式矩阵，其中每个元素 $H_{i,t}$ 表示第 $t$ 个查询 token 与第 $i$ 个文档所有 token 之间的最大余弦相似度。文档 $i$ 的最终得分就是第 $i$ 行所有元素之和 $S_i = \sum_{t=1}^{T} H_{i,t}$。传统的 Full-MaxSim 方法需要计算这个矩阵的全部 $N \times T$ 个元素。
Col-Bandit 将问题转化为「从部分观察的 MaxSim 矩阵中估计 Top-K 文档集合」，理解矩阵结构是理解算法的关键。
多臂老虎机 (Multi-Armed Bandit, MAB)
多臂老虎机是经典的序贯决策框架：一个赌徒面前有多台老虎机（臂），每次可以选择一台拉动，获得随机奖励，目标是在有限次数内最大化总奖励或以最快速度找到最好的那台。固定置信度（fixed-confidence）变体的目标是以给定的置信度 $1-\delta$ 识别出最优臂，同时尽量减少拉动总次数。UCB（Upper Confidence Bound）算法通过维护每个臂的置信区间上界来进行探索-利用权衡。
Col-Bandit 将每个文档视为一个「臂」，MaxSim 矩阵的每一行是有限总体中的确定性值，通过逐步揭示矩阵条目来淘汰低分文档，这是 MAB 框架在信息检索中的新颖应用。
Bernstein–Serfling 不等式
这是适用于无放回抽样（sampling without replacement）的集中不等式。标准的 Bernstein 不等式适用于独立同分布的随机变量，但在无放回抽样中，样本之间存在负相关性（抽了一个就少了一个），这使得方差实际上更小。Bernstein–Serfling 不等式利用了这一点，引入了一个有限总体修正项 $\rho_n$，当抽样数 $n$ 接近总体大小 $T$ 时，$\rho_n \to 0$，置信区间会比标准 bound 更紧致地收缩。Bardenet 和 Maillard（2015）给出了这个不等式的严格形式。
Col-Bandit 用这个不等式来构建每个文档得分的紧致置信区间，这是算法能在仅观察少量 token 后就安全淘汰文档的理论基础。
δ-PAC 保证
δ-PAC（Probably Approximately Correct）是一种统计可靠性保证：算法返回的 Top-K 集合与真实的穷举 Top-K 集合完全一致的概率至少为 $1-\delta$。当 $\alpha_{ef}=1$ 时，Col-Bandit 在简化的半径下具有这个保证；当 $\alpha_{ef}<1$ 时，半径被缩小，消除更激进，但失去了形式化的 PAC 保证，转而依赖经验校准。
理解这个保证的边界条件有助于判断在生产环境中应该选择什么 $\alpha_{ef}$ 值——需要严格正确性还是更看重速度。
研究动机
多向量晚期交互检索器（如 ColBERT）在质量上达到了当前最优水平，但其查询时的计算成本极高。具体来说，在标准的重排序场景中，给定一个查询和 $N$ 个候选文档，系统需要计算 $N \times T$ 个 MaxSim 操作（$T$ 通常为 32），每个操作涉及文档所有 token 的相似度比较。以 NQ-2.68M 数据集为例，这意味着对 268 万个文档执行穷举 MaxSim，在 AMD EPYC 7763 服务器上单线程需要 37.3 秒。这个计算成本成为了现代检索流水线的瓶颈，尤其在大规模语料库上，限制了晚期交互模型的实际部署。现有的加速方法——如 PLAID、MUVERA 的索引时压缩，或 token 剪枝——都必须在索引构建阶段做出决策，无法利用查询信息，因此只能保守地剪枝，保留大量冗余计算。
本文的目标是本文的目标是设计一个纯查询时（query-time）的加速算法，在不需要重新训练或修改索引的情况下，大幅减少 MaxSim 计算量，同时保持与穷举 Full-MaxSim 几乎一致的 Top-K 排序质量。具体量化目标包括：在 BEIR 和 REAL-MM-RAG 基准上，将 MaxSim FLOPs 减少约 8 倍，同时在每个语料库上保持 Overlap@5 ≥ 0.90（即 Col-Bandit 返回的 Top-5 与穷举 Top-5 的重叠度至少 90%），并将单线程墙钟时间加速 7 倍以上。
与已有工作不同的是，本文抓住了一个被已有工作忽视的关键信号：查询本身。索引时方法（PLAID、MUVERA、token 剪枝）在看到查询之前就必须决定保留或丢弃哪些信息，因此无法利用查询与文档之间的具体匹配关系。Col-Bandit 的独特视角是：只有在查询时，系统才能知道哪些文档 token 对这个特定查询是重要的，哪些 MaxSim 单元是冗余的。论文用了一个精妙的「招聘」类比来说明：一个高效的招聘经理不会让所有候选人做完所有测试，而是每轮做一小部分测试后就淘汰明显不合格的人。Col-Bandit 将晚期交互评分形式化为一个有限总体 Top-K 识别问题，逐轮揭示 MaxSim 矩阵的部分条目，利用统计置信区间安全淘汰不可能进入 Top-K 的文档，只计算真正需要的单元格。
核心方法
Col-Bandit 的核心直觉可以用一个简单的例子说明：假设要从 5 个候选人中选出得分最高的 2 个，每人有 32 项测试成绩，总分是所有成绩之和。如果我们已经看到每人 4 项成绩，就能估算每人总分的范围（置信区间）。如果某人的得分上界已经低于其他人得分下界中的第 2 大值，那他一定不是 Top-2，可以直接淘汰。Col-Bandit 将这个直觉系统化：它维护每个文档的得分下界 $LCB_i$ 和上界 $UCB_i$，每轮为所有存活文档揭示 $B=4$ 个新的 MaxSim 条目（按随机排列的顺序），更新置信区间，然后淘汰上界低于第 $K$ 大下界的文档。被淘汰的文档再也不需要计算其剩余的 MaxSim 条目，从而节省大量计算。最终，存活的 $K+M$ 个文档会用所有 $T$ 个 token 重新精确评分，保证最终结果与穷举法完全一致。
Col-Bandit 与已有方法最本质的区别在于它利用了一个「查询时才能获得的信号」来在「原子级别的 MaxSim 单元」上做自适应剪枝。索引时方法（PLAID、MUVERA）压缩的是文档表示，Ball-carving 压缩的是查询 token，但它们都不能根据查询-文档的具体交互来决定哪些 MaxSim 操作是冗余的。Col-Bandit 的创新是将 MaxSim 矩阵视为一个需要逐步揭示的隐式表格，利用经验 Bernstein–Serfling 不等式为每个文档的行和（即总分）构建紧致的有限总体置信区间。这个不等式的关键性质是：随着已观察条目数 $n_i$ 接近总条目数 $T$，置信区间会确定性地收缩到零（而非渐近收敛），这使得消除规则在有限样本下是可证明正确的。另一个关键设计是松弛参数 $\alpha_{ef} \in (0, 1]$：$\alpha_{ef}=1$ 时有形式化的 $\delta$-PAC 保证，$\alpha_{ef}=0.2$ 时（部署默认值）放弃保证换取更低的计算成本，实验表明这在所有语料库上都能保持 ≥ 0.90 的 Overlap@5。
方法步骤详情
Col-Bandit 的完整流程如下：（1）初始化：所有 $N$ 个文档进入活跃集 $A_1$，观察集 $\Omega$ 为空，为每个文档生成初始置信区间 $[aT, bT]$（基于相似度函数的已知全局支撑 $[a,b]$，如余弦相似度的 $[-1,1]$）。对 $T$ 个查询 token 生成一个均匀随机排列 $\pi$，作为所有文档共享的揭示顺序。（2）多轮淘汰循环：第 $r$ 轮，选择排列 $\pi$ 中的下一批 $B=4$ 个 token 索引 $B_r$，对所有存活文档 $i \in A_r$ 并行执行 SIMD 向量化操作：计算这 $B$ 个 MaxSim 条目 $H_{i,t}$（$t \in B_r$），加入观察集 $\Omega$，更新部分和估计 $\hat{S}_i$ 和混合置信区间 $[LCB_i, UCB_i]$。其中 $LCB_i = \max(LB_i^{\text{hard}}, \hat{S}_i - r_i^{\text{eff}})$，$UCB_i = \min(UB_i^{\text{hard}}, \hat{S}_i + r_i^{\text{eff}})$，$r_i^{\text{eff}} = \alpha_{ef} \cdot T \cdot \sqrt{2\log(cNT/\delta)} \cdot \hat{\sigma}_i / \sqrt{n_i} \cdot \sqrt{\rho_{n_i}}$。然后计算消除阈值 $\tau_r$ = 活跃集中第 $K$ 大的 $LCB_i$，淘汰所有 $UCB_i < \tau_r$ 的文档。（3）精确重评分：循环终止后（存活文档数 ≤ $K+M$ 或所有 $T$ 个 token 已揭示），对最多 $K+M=10$ 个存活文档用完整的 $T$ 个 token 执行精确 MaxSim 评分（与 Full-MaxSim 使用完全相同的融合内核，位一致），返回得分最高的 $K$ 个文档。
技术新颖性
Col-Bandit 的技术新颖性体现在多个层面。首先，在问题建模上，它是第一个将晚期交互评分过程本身（而非索引或查询压缩）建模为有限总体 Top-K 识别问题的方法。这与标准的多臂老虎机问题有一个关键差异：标准 MAB 中臂的奖励是随机变量，而 MaxSim 矩阵的每一行是 $T$ 个确定性值的有限总体，随机性仅来自揭示顺序。因此 Col-Bandit 使用 Bernstein–Serfling 不等式（适用于无放回抽样）而非标准的亚高斯界来构建置信区间，这使得区间在 $n_i \to T$ 时确定性地收缩。其次，在算法设计上，Col-Bandit 引入了混合置信区间（Eq. 8），将方差自适应的经验区间与确定性的硬支撑界 $[LB_i^{\text{hard}}, UB_i^{\text{hard}}]$ 取交集，在小样本阶段由硬界主导，大样本阶段由经验界主导。第三，在工程实现上，Col-Bandit 使用 int8 量化粗筛 + fp32 精确重评分的两阶段计算路径，配合缓存友好的打包内存布局和 SIMD 寄存器分块（$B=4$ 匹配 AVX-512/NEON 寄存器瓦片），将算法理论转化为实际的墙钟加速。
实验结果
Col-Bandit 在 BEIR（5 个文本语料库）和 REAL-MM-RAG（4 个多模态语料库）上的实验结果非常扎实。在覆盖率节省方面（Table 1），以 95% 近无损阈值恢复 Full-MaxSim 的 Overlap@5 时，Col-Bandit 在 ColBERTv2 上只需 14% 覆盖率（7.1× 节省），Jina-ColBERTv2 上 26%（3.8×），Granite Vision Embedding 上 22%（4.6×），而 Doc-Uniform 基线需要 98-100%，Ball-carving 需要 46-76%。在墙钟加速方面（Table 2），在 AMD EPYC 7763 服务器（AVX2）上，部署配置（$\alpha_{ef}=0.2, K=5$）单线程平均加速 7.0×，其中 NQ-2.68M 语料库（268 万文档）加速 8.8×（从 37.3 秒降至 4.2 秒）。在 Apple M1 Max 笔记本上（Table 3），平均加速 12.9×，其中 REAL-MM-RAG 多模态语料库加速 15.9×。在排序质量保持方面，所有 9 个语料库上 Overlap@5 ≥ 0.90（CPU-S 上除 HotpotQA-500K 外均 ≥ 0.96），Top-1 指标（Recall@1, nDCG@1, MRR@1）在 20% 覆盖率下保持 98.7-98.9% 的穷举水平。在与 GPU 的对比中（Table 21），16 线程 CPU-S 的 Col-Bandit 在 $N \geq 500K$ 的文本语料库上与 A100 GPU 的 PyTorch 密集 Full-MaxSim 延迟仅差 1.06-1.41×，而 NQ-2.68M 在 80GB GPU 上直接 OOM，CPU-S 仅需 472ms。敏感性分析表明，$K$ 从 5 增长到 1000（200×）时延迟仅增长约 2-4×（次线性），而 Full-MaxSim 延迟基本不变。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| BEIR 文本检索 (ColBERTv2, K=5) | Overlap@5 at 95% 覆盖率 | 14% 覆盖率 (Col-Bandit) | 98% (Doc-Uniform), 46% (Ball-carving) | 7.1× 计算节省 vs Doc-Uniform，3.1× vs Ball-carving |
| NQ-2.68M 全库重排序 (单线程) | 墙钟延迟 | 4,222 ms (Col-Bandit) | 37,283 ms (Full-MaxSim) | 8.8× 加速 |
| Apple M1 Max (NEON, K=5, 四语料库平均) | 单线程加速比 | 12.9× vs Full-MaxSim | 4.8× vs maxsim-cpu (Rust SIMD) | 相比已有 SIMD 基线仍有 4.8× 优势 |
| REAL-MM-RAG 多模态检索 (GVE, K=5) | Overlap@5 | 0.97 (FinSlides), 0.99 (TechSlides) | Full-MaxSim = 1.00 | 22% 覆盖率下保持 97%+ 质量 |
| Top-1 排序保真度 (20% 覆盖率) | Recall@1 相对保持率 | 98.9% (Col-Bandit) | 81.0% (Ball-carving), 55.9% (Doc-Uniform) | 相比 Ball-carving 高 17.9 个百分点 |
局限与改进
作者明确承认了三个主要局限性。（1）硬件范围受限：墙钟评估仅覆盖了两个 CPU 类别（AMD EPYC 7763 AVX2 服务器和 Apple M1 Max NEON 笔记本），尚未构建 GPU 原生的融合内核，因此 GPU 对比使用的是未融合的 PyTorch 基线，不能作为算法级对比。（2）质量上限：Col-Bandit 估计的是穷举 MaxSim 的 Top-K，它不能超越底层评分器的质量，只能以更低的成本恢复它。（3）理论与部署的差距：$\delta$-PAC 保证仅在 $\alpha_{ef}=1$ 且使用精确浮点计算时成立，而部署默认的 $\alpha_{ef}=0.2$ 使用 int8 量化消除决策，不具有形式化保证。此外，Quora 数据集的端到端指标（nDCG/Recall/MRR）因 qrels/doc-id 映射问题在所有方法上均为零，虽然 Overlap@K 不受影响，但影响了任务指标的全面评估。从独立观察来看，HotpotQA-500K 是最「难」的语料库——在 CPU-S 上 Overlap@5 仅为 0.93（其他语料库均 ≥ 0.96），在 CPU-M1 上更低至 0.90，暗示该语料库的得分分布较为平坦，Top-K 分离度不足。
独立分析的弱点
尽管 Col-Bandit 在 CPU 上表现出色，但存在几个值得关注的弱点。首先，消除决策使用 int8 量化估计而非精确浮点值，这会导致一些边界文档被错误淘汰。虽然 margin $M=5$ 吸收了大部分此类误差，但在得分分布平坦的语料库（如 HotpotQA-500K）上，Overlap@5 降至 0.90-0.93，说明量化噪声的影响在某些场景下仍然显著。改进方向：可以探索自适应量化精度——对置信区间窄的文档使用更高精度估计。其次，论文使用固定随机排列 $\pi$ 作为所有文档的共享揭示顺序，这意味着如果最重要的 token 恰好排在排列末尾，所有文档都需要等待很多轮才能揭示它。改进方向：可以设计自适应的揭示顺序，优先揭示方差贡献大的 token。第三，Col-Bandit 的加速效果在小语料库（如 ArguAna, 8.7K 文档）上不如大语料库明显（6.4× vs NQ 的 8.8×），因为消除和重评分的固定开销占比更高。改进方向：可以针对小语料库设计更轻量的消除策略。
未来方向
论文和作者提出了一些未来方向，结合本文成果可以延伸出更多研究点。（1）更紧致的置信区间：论文使用的是全局支撑界 $[a,b]$，如果引入文档级别的 token 范数或质心上界，可以进一步收紧 $LB_i^{\text{hard}}$ 和 $UB_i^{\text{hard}}$，减少冗余计算。（2）GPU 原生融合内核：当前的 CPU SIMD 内核已经非常高效，但 GPU 上的实现（使用 CUDA warp 级并行和共享内存）可能带来更大的加速，尤其是在超大语料库上。（3）与上游检索器的深度集成：Col-Bandit 目前作为独立的重排序层，但如果与 PLAID 的倒排索引或 MUVERA 的固定维编码深度集成，可以在候选生成阶段就开始剪枝。（4）自适应 $\alpha_{ef}$ 选择：当前 $\alpha_{ef}$ 是手动设置的全局参数，未来可以根据查询难度（如得分方差）自动调整每个查询的 $\alpha_{ef}$。（5）扩展到其他晚期交互模型：论文验证了 ColBERTv2、Jina-ColBERTv2 和 Granite Vision Embedding，但还可以扩展到 ColPali 等视觉文档检索模型。
复现评估
Col-Bandit 的复现条件相当友好。代码已开源发布在 GitHub（https://github.com/roipony/ColBandit），核心的 CB-NK 内核是一个基于 NUMKONG C 扩展的 SIMD 优化实现，支持 AVX-512、AVX2 和 NEON 三种后端，编译时自动选择。数据集方面，BEIR 和 REAL-MM-RAG 都是公开可用的标准基准，预训练编码器（ColBERTv2、Jina-ColBERTv2、Granite Vision Embedding）均通过 Hugging Face 发布。算力需求方面，所有 CPU 实验在消费级硬件（Apple M1 Max 笔记本）和服务器（AMD EPYC 7763）上完成，不需要 GPU，这对于资源有限的研究者非常友好。复现难度中等：算法本身不复杂（核心约 100 行），但高效的 SIMD 内核实现需要较多工程经验。部署参数默认值明确给出（$\alpha_{ef}=0.2, M=5, B=4, \delta=0.01$），无需额外调参。
论文图表
