HyperAIHyperAI

Command Palette

Search for a command to run...

5 天前
模型训练

仅用分数蒸馏实现紧凑型稠密检索

Kirill Dubovikov Martin Takáč Salem Lahlou

摘要

大型嵌入模型能提升检索质量,但在线部署大型编码器成本高昂。我们研究紧凑型检索器能否在不访问教师模型隐藏状态的情况下,仅从分数向量中学习教师排序行为。学生模型基于真值正例和由我们数据生成流程产生的负例候选集进行训练;我们还将学生-教师难负例挖掘作为扩展单独评估。我们采用以行为中心的分数向量目标,这是均匀全对PairMSE损失的一种内存高效实现。在固定的八任务评估面板上,我们的蒸馏方案可恢复基线与教师模型之间差距的50%。蒸馏后的0.6B学生模型在查询编码上比顺序在线教师融合快4.7倍,在文档编码上快9.7倍。蒸馏后的外部迁移性能表现不一,因此我们的证据支持在匹配检索协议下压缩教师排序。

一句话总结

阿联酋穆罕默德·本·扎耶德人工智能大学的研究人员提出了一种仅使用分数蒸馏的方法,通过基于行的中心化分数向量目标(均匀全对PairMSE损失)训练紧凑的稠密检索器,无需教师隐藏状态,最多恢复基模型与教师模型之间 50%50\%50% 的差距,得到的 0.6B0.6\text{B}0.6B 学生模型在查询编码上快 4.7×4.7\times4.7×,在文档编码上快 9.7×9.7\times9.7×,并压缩了教师排序。

核心贡献

  • 本文定义了一种黑盒分数向量蒸馏协议,仅使用教师分数训练紧凑的检索学生模型,无需访问教师隐藏状态或表示。
  • 该工作提出了一种基于行的均匀全对PairMSE实现,这是一种内存高效的目标,并在相同的行、学生模型和评估设置下,将其与仅使用标签和正负MarginMSE的对照方法进行比较。
  • 结果刻画了该协议的适用范围:分数向量蒸馏在匹配的检索协议下改进了紧凑学生模型,本文还提供了对教师融合、主动学习和难负样本挖掘等扩展的实证分析。

引言

稠密检索流程越来越依赖大型嵌入模型,但模型规模推高了服务成本、索引刷新成本和操作复杂性,这促使将这些模型压缩为紧凑的检索器。先前的检索蒸馏方法通常需要访问教师嵌入、隐藏状态或logits,或假设教师为交叉编码器,联合编码查询-文档对。当教师为黑盒嵌入模型,其唯一输出是查询-文档分数向量时,这些方法不再适用。作者提出了一种分数向量蒸馏协议,仅从教师提供的候选行分数来训练一个紧凑的双编码器学生模型。他们引入了一种基于行的均匀全对PairMSE损失重构,避免了显式的成对margin张量,并证明该协议在两种架构和八个检索任务面板上改进了紧凑学生模型。

方法

作者将分数向量蒸馏形式化为一个逐行训练过程。对于每个查询 qqq,一个候选行包含一组文档 d0,,dk1d_0, \dots, d_{k-1}d0,,dk1,其中 d0d_0d0 是标注的正例,其余 k1k-1k1 个文档为负例。学生检索器为这些对分配分数 si=fθ(q,di)s_i = f_\theta(q, d_i)si=fθ(q,di),而教师模型为相同的查询-文档对提供目标分数 tit_iti。教师目标是归一化的余弦分数,学生和教师不需要共享同一个嵌入空间。

当有多个教师可用时,作者采用平均融合:将教师分数向量逐坐标平均,得到一个单一的目标向量 ttt。这种融合仅改变目标分数,不影响候选行、学生架构和下游评估。

核心训练目标是基于行的分数向量蒸馏损失,旨在对齐学生分数与教师分数,同时丢弃行级别的加性偏移。设 r=stRkr = s - t \in \mathbb{R}^kr=stRk 为残差向量。损失定义为

Lcenter(s,t)=1krmean(r)1k22.\mathcal{L}_{\text{center}}(s, t) = \frac{1}{k} \| r - \operatorname{mean}(r) \mathbf{1}_k \|_2^2.Lcenter(s,t)=k1rmean(r)1k22.

该损失对行内所有分数添加任何常数偏移保持不变,从而使学习聚焦于决定文档排序的相对分数差异。

利用方差恒等式 i,j(rirj)2=2ki(rimean(r))2\sum_{i,j} (r_i - r_j)^2 = 2k \sum_i (r_i - \operatorname{mean}(r))^2i,j(rirj)2=2ki(rimean(r))2 并注意到 (sisj)(titj)=rirj(s_i - s_j) - (t_i - t_j) = r_i - r_j(sisj)(titj)=rirj,中心化损失等价于对有序文档对施加均匀的全对PairMSE:

Lallpairs(s,t)=1k2i=0k1j=0k1[(sisj)(titj)]2=2Lcenter(s,t).\mathcal{L}_{\text{allpairs}}(s, t) = \frac{1}{k^2} \sum_{i=0}^{k-1} \sum_{j=0}^{k-1} \left[ (s_i - s_j) - (t_i - t_j) \right]^2 = 2 \mathcal{L}_{\text{center}}(s, t).Lallpairs(s,t)=k21i=0k1j=0k1[(sisj)(titj)]2=2Lcenter(s,t).

因此,优化中心化目标隐式地施加了与全对PairMSE相同的教师偏好,但避免了构造显式的对张量。对于固定的行宽,两种损失仅相差一个常数因子。

作者将这种形式化与两个对照目标进行比较。第一个是正负MarginMSE,仅匹配标注正例与每个负例之间的margin:

Lmargin=1k1j=1k1[(s0sj)(t0tj)]2.\mathcal{L}_{\text{margin}} = \frac{1}{k-1} \sum_{j=1}^{k-1} \left[ (s_0 - s_j) - (t_0 - t_j) \right]^2.Lmargin=k11j=1k1[(s0sj)(t0tj)]2.

第二个是仅使用标签的DPR式对比目标,忽略所有教师分数。它使用固定的logit尺度 τ=0.05\tau = 0.05τ=0.05,优化标注正例与行中其他文档的对比:

LCE=logexp(s0/τ)i=0k1exp(si/τ).\mathcal{L}_{\mathrm{CE}} = -\log \frac{\exp(s_0 / \tau)}{\sum_{i=0}^{k-1} \exp(s_i / \tau)}.LCE=logi=0k1exp(si/τ)exp(s0/τ).

该对照测试了在没有教师分数监督的情况下,候选行和标注正例是否足够。整个蒸馏流程在固定的候选行上运行,损失函数驱动学生重现教师的相对评分模式。

实验

实验使用BEIR快照,训练行来自三个数据集,并在行源数据集和保留数据集上进行全语料检索,将使用教师分数向量训练的学生模型与仅使用标签和基于margin的基线进行比较。分数向量蒸馏一致地改进了紧凑检索器,超越了仅使用标签的训练,匹配行内所有候选对比基于margin的损失捕获了更多的教师信号。然而,增益是领域特定的:学生模型恢复了约四分之一到一半的教师差距,且向不相关语料的迁移效果较差,同时提供了显著的推理速度和内存改进。测试的扩展包括难负样本挖掘和主动学习,并未产生可靠的额外收益。

将教师分数向量蒸馏到紧凑的学生检索器中,相较于冻结的基模型,提升了全语料检索质量,其中E5学生模型的提升大于Qwen学生模型。在相同训练行上的仅标签对照表现比基模型差,证实了教师提供的分数信息至关重要。对Qwen而言,教师分数的平均融合并未一致地优于单教师目标,但对E5则取得了最佳结果。分数向量蒸馏将两个学生模型都提升到了各自基模型之上;E5学生模型提升了大约0.07 NDCG@10,Qwen学生模型提升了约0.02,而仅标签对照则低于冻结基模型。对于Qwen,教师分数的平均融合得到的NDCG@10为0.472,其置信区间与单教师估计重叠,因此没有明显优势;对于E5,平均融合被报告为最强设置。

消融实验表明,匹配完整的中心化教师分数向量(中心化MSE)比仅使用正负margin损失、单位归一化方向匹配或排序加权配对产生更强的检索质量。完全移除教师分数、仅使用标签交叉熵会使性能急剧下降。选择教师和学生不一致的负样本进行难负样本挖掘,在单个种子上有轻微提升,但其多种子区间与平均融合基线重叠,因此未证实具有稳健优势。中心化MSE优于正负MarginMSE、单位中心化MSE和排序加权PairMSE,表明行内所有负样本之间的相对排序承载了有用的教师信号。与使用教师分数向量相比,无教师分数的仅标签交叉熵显著降低了全语料检索质量。使用紧凑的6候选行进行难负样本挖掘取得了最佳单种子分数,但其三种子区间与平均融合的32候选基线重叠,因此未出现稳定的排序。

将教师分数向量蒸馏到紧凑学生检索器中,相较于冻结基模型,在全语料检索上有所提升,其中E5模型的增益大于Qwen,而仅标签对照表现不如基模型,证实了教师提供的负样本相对排序至关重要。消融实验显示,匹配完整的中心化教师分数向量(中心化MSE)比基于margin、基于方向或排序加权的损失产生更强的检索质量,完全丢弃教师分数会导致性能急剧下降。多个教师分数的平均融合对E5取得了最佳结果,但对Qwen没有一致的优势,难负样本挖掘也未提供比基线稳健的改进。


用 AI 构建 AI

从创意到上线——通过免费 AI 协同编码、开箱即用的环境和最优惠的 GPU 价格,加速您的 AI 开发。

AI 协同编码
开箱即用的 GPU
最优定价

HyperAI Newsletters

订阅我们的最新资讯
我们会在北京时间 每周一的上午九点 向您的邮箱投递本周内的最新更新
邮件发送服务由 MailChimp 提供