跳转至

Cluster with Auctions for Vector Search

会议: NeurIPS 2026
arXiv: 2607.13728
OpenReview: https://openreview.net/forum?id=M4uXVIjfK7
代码: https://github.com/facebookresearch/cwa
领域: 信息检索/RAG
关键词: 向量检索, 倒排索引, 聚类分配, 拍卖算法, 域外分布检索

一句话总结

将倒排索引中的数据库分区与查询探测函数严格解耦,利用拍卖算法求解带容量约束的线性指派以确保簇严格均衡,并结合查询分布端到端交替优化神经网络探测器,在分布偏移场景下吞吐量超越基线数倍。

研究背景与动机

大规模高维向量的近似最近邻搜索(ANNS)广泛应用于大模型检索增强生成(RAG)、以图搜图以及大规模数据去重等场景。基于分区的倒排索引(IVF)是工业界与学术界最常用的方案之一:将底库向量划分为若干聚类,查询到达时通过探测函数只扫描最具可能包含近邻的少数簇。然而,以经典 K-Means 为代表的传统方法,以及现有的端到端神经检索方法(如 Neural LSH、USP、BLISS),普遍将“底库向量分区”与“查询探测函数”混为一谈,强行使用同一模型或相同的几何中心距离来同时扮演底库分配与查询探测两个异质角色。

这种混淆导致了两个致命瓶颈。其一,在实际检索应用中,查询向量与底库向量往往存在严重的分布偏移(例如多模态文本到图像检索、跨模态向量空间),基于底库训练的几何聚类中心无法贴合查询分布,导致探测器在查询落入的边缘区域覆盖极差,主流神经方法性能急剧恶化甚至大幅输给经典算法;其二,由于单模型强行绑定两种任务,交替训练时更新权重会导致底库聚类与探测目标同时“漂移”,且自然聚类产生的严重簇大小不均衡会带来长尾延迟波动,显著推高最坏情况下的查询响应时间。

为了彻底突破这一架构瓶颈,Meta FAIR 团队提出显式解耦底库分区与查询探测。本文的核心 idea 是将底库向量指派与查询探测网络解耦,以真实查询分布下的近邻覆盖交叉熵构建统一目标,通过 GPU 高度并行的拍卖算法求解带配额约束的线性指派实现天然严格均衡的簇划分,与神经网络探测器的反向传播交替优化。

方法详解

整体框架

CwA(Cluster with Auctions)将整个索引系统划分为两套相互配合的组件:面向查询特征坐标的参数化神经探测网络 \(f_\theta(q) \in [0,1]^{|C|}\),以及面向固定底库向量 ID 的离散指派映射 \(h: \{1,\dots,n_{\text{index}}\} \to C\)。在离线训练阶段,系统以查询近邻在簇上的经验分布作为软标签,通过交替迭代分别推进探测器参数梯度下降与底库全局拍卖;在在线查询阶段,探测器直接输出簇推荐概率,系统挑出 Top-\(m\) 簇送入底库倒排列表进行精排比对。针对十万级别超大簇规模,算法进一步扩展为笛卡尔积架构 CwA-Prod 与图索引加速变体 CwA-HNSW。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入查询分布与真实近邻<br/>𝒟_train = (q_i, 𝒩_k'(q_i))"] --> B["查询探测网络前向传播<br/>输出候选簇概率分布 p^_θ,q_i"]
    B --> C["快速得分矩阵构建:消去Log/Softmax<br/>直接聚合原始未归一化Logits"]
    C --> D["带配额约束的并行拍卖算法<br/>底库向量竞价并严格均衡划分簇容量"]
    D --> E["反向传播更新探测网络参数<br/>基于最新划分边界最小化交叉熵损失"]
    E -->|交替迭代 A 轮直至收敛| B
    E --> F["构建完成:探测网络 f_θ + 严格均衡倒排索引 h"]
    F --> G["在线检索:Top-m 簇探测与底库精排扫描"]

关键设计

1. 目标函数解耦与查询驱动的软监督:直接对齐近邻检索召回而非重构误差

传统 K-Means 优化的是底库向量与聚类中心之间的欧氏重构误差,与最终衡量检索效果的 Recall 存在天然间隙。CwA 提出以查询分布上的真实检索效果为直接导向。设底库向量集为 \(X = (x_\ell)_{\ell=1,\dots,n_{\text{index}}}\),当前底库分区函数为 \(h\)。对于来自查询分布的训练查询 \(q_i\),其真实 \(k'\) 近邻集合记为 \(\mathcal{N}_{k'}(q_i)\)。当 \(h\) 固定时,这 \(k'\) 个近邻落在各簇的真实比例构成了经验目标分布: $\(p_{h,q_i}(c) = \frac{1}{k'} \sum_{\ell \in \mathcal{N}_{k'}(q_i)} \mathbf{1}[h(\ell) = c]\)$ 探测网络 \(f_\theta(q_i)\) 预测 query 属于各簇的概率 \(\hat{p}_{\theta,q_i}(c)\)。训练目标定义为加权交叉熵: $\(\min_{\theta, h} \sum_{i=1}^{n_{\text{train}}} \mathrm{CE}(p_{h,q_i}, f_\theta(q_i)), \quad \text{s.t.} \quad |\{\ell \mid h(\ell) = c\}| \leq Q, \quad \forall c \in C\)$ 其中 \(Q \geq n_{\text{index}}/|C|\) 为强制的单簇最大容量配额(实验中设为 \(Q = 1.05 \frac{n_{\text{index}}}{|C|}\))。在此设计下,目标函数严格惩罚将近邻分散到低概率簇的行为,且监督信号直接来自目标查询分布,即便底库与查询存在严重跨模态分布漂移,探测网络也能直接学习查询空间的几何覆盖。

2. 拍卖算法求解容量约束线性指派:GPU 并行实现全局严格均衡划分

当固定网络参数 \(\theta\) 时,最小化总交叉熵可重写为关于底库向量指派 \(h\) 的组合优化问题。展开交叉熵后,总损失严格转化为对每个底库向量与所选簇得分的求和: $\(\min_h \sum_{i=1}^{n_{\text{train}}} \mathrm{CE}(p_{h,q_i}, \hat{p}_{\theta,q_i}) \iff \max_h \sum_{\ell=1}^{n_{\text{index}}} S_{\ell, h(\ell)}, \quad \text{s.t.} \quad |\{\ell \mid h(\ell)=c\}| \leq Q\)$ 其中底库向量 \(\ell\) 放入簇 \(c\) 的基础价值得分为 \(S_{\ell,c} = -\frac{1}{k'} \sum_{i=1}^{n_{\text{train}}} \log(\hat{p}_{\theta,q_i}(c)) \mathbf{1}[\ell \in \mathcal{N}_{k'}(q_i)]\)。这是一个经典的带容量约束线性指派问题(Capacitated Linear Assignment Problem)。CwA 采用高度并行的 Bertsekas 类似对象拍卖算法(Auction Algorithm for Similar Objects):未指派向量向当前净收益(得分减去簇槽位当前价格 \(p_{c,s}\))最高的簇发起竞价,出价增量根据第一优与第二优的收益差距确定(\(\delta = v - w + \epsilon\))。高出价者获得槽位并提升该槽位价格,被挤出的向量重新参与下一轮竞拍。该机制在 GPU 上极易大规模并行展开,在数分钟内收敛,使底库各簇大小完全被配额 \(Q\) 封顶,从根本上消除了离群大簇导致的巨量扫描开销。

3. 原始 Logits 快速得分与稀疏化技巧:消除 Softmax/Log 并压降显存占用

在每次底库重新指派前,常规做法需对所有训练查询评估所有簇的 Softmax 概率并取对数,计算与显存开销极大。CwA 证明了 Softmax 的配分函数在底库更新时完全退化为常数项: $\(S_{\ell,c} = -\frac{1}{k'} \sum_{i=1}^{n_{\text{train}}} z_{i,c} \mathbf{1}[\ell \in \mathcal{N}_{k'}(q_i)] + \frac{1}{k'} \sum_{i=1}^{n_{\text{train}}} \log\left(\sum_{c'} \exp(z_{i,c'})\right) \mathbf{1}[\ell \in \mathcal{N}_{k'}(q_i)]\)$ 由于右半部分与簇选择 \(c\) 完全无关,可以直接舍弃!系统无需计算任何 Softmax 和 Log,直接将神经网络最后一层的原始 Logits \(z_{i,c}\) 沿近邻指示因子累加,即可得到保序等价的得分矩阵 \(S'_{\ell,c}\)。进一步地,针对全稠密矩阵 \(n_{\text{index}} \times |C|\) 难以放入显存的问题,CwA 对底库向量预先进行初步 K-Means 聚类,强制每个底库向量仅允许竞拍与其几何距离最近的 \(\kappa\) 个候选簇(\(\kappa \approx 5\% \sim 10\% |C|\)),其余项置为 \(-\infty\)。这使得分矩阵存储和计算减少 5 到 10 倍,实现了工业级规模的显存友好性。

4. 笛卡尔积架构 CwA-Prod 与等价欧氏重参数化 CwA-HNSW:向超大规模平滑扩展

当簇数量扩展到数万至数十万时,全量线性分类层参数爆炸且前向探测变慢。CwA 提出了两套互补的扩展架构: 其一是 CwA-Prod,受倒排多索引(IMI)启发,将输出分类层拆解为两个小规模码本的笛卡尔积。若总簇数为 \(|C| = 262,144\)(即 \(512^2\)),网络骨干输出 \(q'\) 后分别经过两组线性层映射得到 \(r, s \in \mathbb{R}^{\sqrt{|C|}}\),簇 \((i,j)\) 的得分为 \(g_\theta(r,s)_{i,j} = \gamma_{i,j} \times (r_i + s_j)\),其中 \(\gamma\) 为可学习交互矩阵。这使参数量从 \(\mathcal{O}(d|C|)\) 急剧下降为 \(\mathcal{O}(d\sqrt{|C|} + |C|)\)。在拍卖阶段,大矩阵 \(S'\) 同样严格因式分解为 \(S_{\ell, i \times \sqrt{|C|}+j}' = \gamma_{i,j}(S^1_{\ell,i} + S^2_{\ell,j})\),显存占用直降 \(\sqrt{|C|}/2\) 倍(超过 250 倍压缩)。 其二是 CwA-HNSW,针对 \(|C| \ge 10^4\) 时探测网络自身输出 Top-\(m\) 成为耗时瓶颈的问题。常规内积最大化搜索(MIPS)在图索引上效率低,作者巧妙地将最后一层输出重参数化为等价负欧氏距离:\(t_i = -e^\lambda \|q' - m_i\|^2\),其中 \(\lambda\) 为全局可学习缩放温度。这使探测层的 Top-\(m\) 挑选等价转化为几何点集 \(\{m_i\}\) 上的最近邻搜索,无缝挂载 HNSW 图索引,在对数复杂度内完成候选簇剪枝。

损失函数 / 训练策略

训练采用交替迭代框架,通常设置 \(A = 20\) 轮交互:每轮包含 1600 步基于 Adam 优化器(初始学习率 \(3 \times 10^{-3}\),配余弦退火)对探测网络参数 \(\theta\) 的反向传播训练,紧接着执行一次 10,000 步 GPU 并行竞价的拍卖指派以刷新 \(h\)。在全部交替轮次结束后,冻结最终指派 \(h\),再对探测网络做 80,000 步微调以充分收敛。 为了防止过拟合,监督近邻数 \(k'\) 普遍设置为大于检索目标 \(k\)(如 \(k=10\) 时设 \(k'=50\) 或 \(100\))。论文通过严格的实证网格搜索给出了监督邻域的尺度律公式: $\(k' = 0.78 \cdot n_{\text{index}}^{0.535} \cdot n_{\text{train}}^{-0.220}\)$ 并且推导出达到最优性能所需训练查询量的次线性规模定律: $\(n_{\text{train}} = 41331 \cdot n_{\text{index}}^{0.282}\)$ 这意味着即便将底库规模扩大 100 倍,所需查询仅增加约 4 倍,展现出极高的数据利用效率。

实验关键数据

主实验

论文在四项具有代表性的大规模基准数据集上进行评测:分布内(ID)数据集为 SIFT1M/100M (\(d=128\), \(\ell_2\)) 和 Deep1M/10M/100M (\(d=96\), \(\ell_2\));分布外(OOD)数据集为 Text-to-Image 1M/10M/100M (\(d=200\), 内积) 与 LAION 1M/10M/100M (\(d=512\), 余弦)。对比基线包含经典 K-Means/IVF、无监督空间分区 USP、图分区先验方法 Neural LSH、迭代重分区 BLISS,以及多码本方案 IMI 与 RQ。

在 1M 向量规模、固定 Recall@10 = 0.8 时,各类方法的查询吞吐量(QPS,越高越好)对比如下:

| 数据集 | 簇数量 \(|C|\) | K-Means | Neural LSH | USP | BLISS | CwA (0 FFN, 线性) | CwA (2 FFN, 深度) | |---|---|---|---|---|---|---|---| | Deep1M (ID) | 256 | 4,309 | 4,596 | 4,487 | 2,454 | 5,334 | 5,484 | | Deep1M (ID) | 1,024 | 9,161 | 8,279 | 6,058 | 3,285 | 11,664 | 11,262 | | SIFT1M (ID) | 256 | 2,495 | 2,565 | 2,552 | 340 | 2,783 | 2,831 | | SIFT1M (ID) | 1,024 | 5,274 | 4,937 | 1,637 | 886 | 6,151 | 6,118 | | TextToImage1M (OOD) | 256 | 1,338 | 910 | 453 | 229 | 2,649 | 3,240 | | TextToImage1M (OOD) | 1,024 | 3,592 | 1,766 | 764 | 550 | 6,615 | 6,881 | | LAION1M (OOD) | 256 | 276 | 55 | — | 36 | 1,225 | 1,295 | | LAION1M (OOD) | 1,024 | 811 | 87 | — | 66 | 3,210 | 2,980 |

在 100M 超大规模底库、设置 \(|C| = 512^2 = 262,144\) 超高分辨率分区的多码本对决中(CwA-Prod vs. IMI vs. RQ),不同召回率下的 QPS 如下:

数据集 目标召回 Recall@10 RQ IMI CwA-Prod (本文) 相对最优基线提升
Deep100M (ID) 0.7 1,180.7 281.3 2,729.6 +131.2% (2.31×)
Deep100M (ID) 0.9 350.8 83.0 613.7 +74.9% (1.75×)
SIFT100M (ID) 0.7 994.4 436.7 1,906.7 +91.7% (1.92×)
SIFT100M (ID) 0.9 330.0 131.5 446.5 +35.3% (1.35×)
TextToImage100M (OOD) 0.7 118.3 51.6 1,671.3 +1312.8% (14.13×)
TextToImage100M (OOD) 0.9 23.3 12.7 183.3 +686.7% (7.87×)
LAION100M (OOD) 0.7 — 11.4 1,897.7 +16546.5% (166.5×)
LAION100M (OOD) 0.9 — — 174.1 基线完全崩溃断崖

消融实验

针对网络深度架构与检索延迟构成的消融分析,以及 10M 规模下挂载 HNSW 的加速收益如下表所示:

配置 / 模块 Deep10M QPS (R@10=0.8) TextToImage10M QPS (R@10=0.8) LAION10M QPS (R@10=0.8) 机制说明
IVF (K-Means 基准) 2,514 1,667 677 标准几何质心基准
IVF-HNSW 6,271 2,662 723 仅用 HNSW 加速质心检索
CwA (纯线性层探测,无图) 2,762 2,186 1,147 分区质量更优但探测层线性遍历占瓶颈
CwA-HNSW (完整方案) 9,954 6,210 3,419 负欧氏距离重参数化 + HNSW 加速簇剪枝

关键发现

  • 分布外泛化能力产生质的跃迁:在文本到图像和多模态嵌入等真实 OOD 场景中,过往一切神经分区基线(Neural LSH、USP)由于将训练重心放在底库本身,其探测网络严重退化,QPS 甚至跌落至传统 K-Means 的数分之一;而 CwA 通过由真实查询近邻驱动的拍卖迭代,在 Text-to-Image1M 上达到 K-Means 的 2.4× 吞吐,在 LAION1M 上达到 4.7×,在 100M 尺度上更是比传统多码本基线高出一个数量级以上。
  • 纯线性层即可击败深度神经网络基线:CwA 即使完全去掉残差 FFN 模块(0 FFN 层,参数量与 K-Means 完全相同),其在所有数据集上的准确率与吞吐量依然全面超越拥有多层 MLP 的 Neural LSH 和 USP。这证明此前神经检索方法的瓶颈根本不在于模型容量不够,而在于底库分区与探测器未解耦以及缺乏严格容量约束。
  • 簇容量均衡对长尾耗时的抑制至关重要:通过拍卖算法严格将簇大小限制在 \(1.05 \times\) 平均值内,底库倒排链表长度方差趋近于零,从根本上杜绝了某些偏斜簇过大造成的查询卡顿,使固定探测簇数 \(m\) 时的总距离计算量有了确定性的上界。

亮点与洞察

  • 解耦不对称角色:彻底指出了向量检索中“底库向量静态离散聚类”与“连续高维查询动态概率打分”两者本质上的不对称性,用线性指派与参数化分类器分别建模,打破了多年来统一绑定单模型的固有思维定势。
  • 巧用拍卖算法做大规模几何配额平衡:巧妙地将 Bertsekas 经济学拍卖算法搬上现代 GPU 并行架构,成功在百万至亿级规模下毫秒级解决传统整数规划与最优传输极难求解的硬容量约束。
  • Softmax/Log 对消与 Logits 极速聚合:数学推导证明了查询侧的 Softmax 归一化分母与底库各簇选择无关,从而在底库重分区阶段完全免去昂贵耗时的逐点指数与对数运算,实现了十倍级的算法工程提速。
  • 次线性查询数据依赖律:揭示了学习型索引在工业落地中极度关键的经验规律——底库规模扩大 1000 倍仅需 7 倍查询量即可训练到位,极大增强了 CwA 部署在实际检索系统的经济可行性。

局限与展望

  • 单分区限制与多索引互补性缺失:当前工作仅评测了单套分区索引的表现;而在传统 LSH 或工业级倒排索引中,构建多套相互补充的独立分区结构是进一步换取极端召回率的常规手段,如何协同训练多套具有正交补全特性的 CwA 仍待研究。
  • 动态底库增量更新(Out-of-sample)待突破:由于指派映射 \(h\) 是直接面向底库向量 ID 的离散查表,当底库有新的动态向量插入时,无法直接推断其最优归属簇,必须通过辅助近似策略或定期触发轻量拍卖进行批次重划分。
  • 预热依赖与冷启动代价:虽然整体训练时间比 Neural LSH 短得多,但相比极速的原始 K-Means 仍慢 1 到 2 个数量级,且训练前需预先构建查询的近邻图(\(k\)-NN Graph),这在大规模底库上构成了主要的时间开销。

相关工作与启发

  • vs. K-Means / IVF (Sivic & Zisserman, 2003): K-Means 纯粹基于底库几何重构误差聚类,查询探测强行绑定质心距离,在 OOD 场景严重失真且簇大小极不平衡;CwA 彻底解耦两者,以查询近邻覆盖度为目标驱动拍卖平衡,显著提升吞吐。
  • vs. Neural LSH (Dong et al., 2020): Neural LSH 先在底库图上做图划分,再单独训练网络拟合,切断了分区与探测的双向反馈且完全忽视查询分布;CwA 建立统一交叉熵目标并通过交替迭代建立闭环优化。
  • vs. USP (Fahim et al., 2022) & BLISS (Gupta et al., 2022): USP 与 BLISS 虽采用交替迭代,但试图用同一个模型同时输出底库分区与查询打分,造成了“移动目标”的优化震荡;CwA 明确分工,网络管查询、拍卖管底库,交替收敛平稳单调。
  • vs. IMI (Babenko & Lempitsky, 2012) & RQ (Chen et al., 2010): 传统多码本乘积量化方法存在严重的簇大小长尾倾斜,导致查询耗时波动剧烈;CwA-Prod 在引入笛卡尔积大幅降维的同时依然通过因式分解保持拍卖配额均衡。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次将底库离散平衡指派与查询参数化探测深度解耦,提出 GPU 并行拍卖与神经网络交替优化的通用框架。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 1M、10M、100M 三种体量,覆盖分布内与极端跨模态分布外基准,消融实验与理论缩放律详实严密。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑主线极为清晰,数学推导直观且工程实现技巧(Logits 对消、因式分解、HNSW 重参数化)丰富落地。
  • 价值: ⭐⭐⭐⭐⭐ 为大规模向量数据库与工业检索索引设计树立了全新的技术范式,对 RAG 与多模态搜索具有极高实用价值。