Command Palette
Search for a command to run...
DistilVDR:通过双学生蒸馏构建紧凑的端到端视觉文档检索器
DistilVDR:通过双学生蒸馏构建紧凑的端到端视觉文档检索器
Zhuchenyang Liu Ziyi Wang Yao Zhang Yu Xiao
摘要
视觉文档检索(VDR)目前由数十亿参数模型主导,这些模型在全语料库规模下索引速度慢且服务成本高昂。先前的压缩路线要么从头训练一个较小的多向量编码器,要么仅蒸馏查询端;两者均无法得到一个端到端的紧凑单向量检索器。我们提出 DistilVDR,一个 524 M 参数的端到端 VDR 系统,在逐点余弦对齐损失下,从单个 8 B 视觉-语言教师模型进行双边蒸馏。所有监督信号均来自冻结教师模型的嵌入空间,该空间本身已通过相关性监督训练,因此学生目标无需相关性标签、负采样或对比项。我们以非对称的仅编码器学生模型来匹配 VDR 的文本查询与图像文档输入不对称性,将视觉能力集中在文档端,并将查询端参数保持在 70 M。我们发布两个共享相同编码器和训练过程、仅在文档编码器的视觉图块预算上有所不同的变体:DistilVDR-HiRes 在 ViDoRe v1+v2+v3 上达到 61.74 的平均 NDCG@5(为 8 B 教师模型的 86.9%),并在对高分辨率敏感的 v3 基准上领先所有复现的 1 B 以下基线模型;而 DistilVDR-Fast 在视觉令牌预算减少 3 倍的情况下达到 59.98。两个变体均以比最强 1 B 以下多向量基线小 15.6 倍的索引存储一百万文档,并以快一个数量级的速度索引语料库。代码可在 https://github.com/Ryenhails/NanoVDR 获取。
一句话总结
阿尔托大学和一位独立研究者提出了 DistilVDR,一个 524M 参数的端到端视觉文档检索器,它通过逐点余弦对齐损失从 8B 教师模型双向蒸馏得到,采用非对称的仅编码器学生架构来匹配文本-图像的非对称性,在没有相关性标签的情况下,在 ViDoRe v1+v2+v3 上取得了 61.74 NDCG@5,索引大小比最佳 sub-1B 多向量基线小 15.6 倍。
核心贡献
- DistilVDR 通过双边余弦对齐蒸馏训练,学生嵌入回归到冻结的单个 8B 视觉-语言教师模型的嵌入,无需相关性标签、负采样或对比项。
- 一种非对称的仅编码器学生架构将视觉能力集中在文档侧,同时保持查询侧为 70M 参数,匹配视觉文档检索中文本查询/图像文档输入的非对称性。
- 对十二个检索器的统一复现与分析表明,两个发布版本 DistilVDR-HiRes 和 DistilVDR-Fast 在 ViDoRe v1+v2+v3 上最高达到 61.74 NDCG@5,超越了所有复现的 sub-1B 基线,并提供了 15.6 倍更小的索引和数量级更快的索引构建速度。
引言
视觉文档检索(VDR)将整个文档页面编码为图像,避免了 OCR 错误,并保留了布局、表格和图形,用于企业搜索和检索增强生成等任务。当前性能最佳的 VDR 系统规模庞大,参数达 2–8B,带来高昂的 GPU 内存和索引成本。先前的工作试图通过两种途径降低成本:(1) 训练从零开始的小型多向量模型,压缩编码器但使索引大小和评分延迟增加数个数量级;(2) 蒸馏,缩小查询编码器但保留大型教师模型在文档路径中,使得索引和部署成本仍然高昂。这两种方法都无法提供一个完全紧凑的单向量检索器。作者通过 DistilVDR 弥补了这一差距,这是一个 524M 参数的端到端单向量系统,在非对称的仅编码器设计下,从单个 8B 视觉-语言教师模型双向蒸馏文档和查询编码器,在实现强大检索质量的同时大幅减少了索引占用和处理时间。
方法
作者将检索任务形式化为单向量密集检索问题。给定一个文本查询 q 和一个文档图像语料库 {d1,…,dN},系统根据与 q 的相关性对语料库进行排序。查询编码器 fq 将查询映射为向量 q∈Rk,文档编码器 fd 将每个文档图像映射为同一空间中的向量 di∈Rk。两个向量都经过 L2 归一化,检索分数为点积 s(q,di)=q⊤di。文档向量在索引时一次性计算。
核心架构基于双学生蒸馏范式。一个冻结的 8B 视觉-语言教师模型为两种模态生成维度为 k=4096 的目标嵌入。两个学生独立学习复现这些目标。查询学生 fq 接收文本查询,文档学生 fd 接收文档图像。两个学生都投影到教师的输出空间并进行 L2 归一化。部署时的检索仅使用学生模型;教师在训练后被丢弃。有关该架构和两个蒸馏目标的概述,请参见框架图。
文档编码器 fd 包含 454M 参数,分四个阶段进行。首先,系统通过将每页 tile 数量限制在 Tmax 并附加一个以视觉编码器原始分辨率生成的全景缩略图来提供全局上下文,从而强制执行固定的每文档视觉 token 预算。根据页面宽高比选择匹配的网格布局。其次,每个 tile 由 InternViT-300M-448 视觉编码器编码为 1024 个维度为 768 的 patch token。所有 tile 的 patch token 被连接成一个序列。第三,视觉序列通过一个可学习的线性投影映射到 ModernBERT-base 文本主干的嵌入空间。ModernBERT 使用双向注意力重新编码投影后的视觉 token,充当视觉 token 的上下文编码器。第四,上下文化的 token 被平均池化,通过最终线性层从 768 维投影到 4096 维,并进行 L2 归一化。
查询编码器 fq 具有 70M 参数,分三个阶段操作。查询文本前缀为教师训练期间使用的相同指令字符串 π。前缀后的序列由 DistilBERT-base 文本编码器编码,并对上下文 token 表示进行平均池化。最后,池化向量通过从 768 维到 4096 维的线性投影,并进行 L2 归一化。
对于蒸馏目标,设 T 表示冻结的教师模型。对于文档图像 d,教师生成目标向量 T(d)∈R4096。对于带有指令前缀 π 的文本查询 q,它生成目标 T(π∘q)∈R4096。两个目标都经过 L2 归一化。每个学生独立地针对这些缓存的目标,在余弦对齐损失下进行训练:
Ld=1−⟨fd(d),T(d)⟩ Lq=1−⟨fq(π∘q),T(π∘q)⟩其中 ⟨⋅,⋅⟩ 表示 L2 归一化向量之间的点积。两个学生在训练期间从不共享前向传播,使得文档侧和查询侧的蒸馏完全解耦。学生目标中没有对比项、难负样本或相关性标签。
实验
论文在完整的 ViDoRe 套件上评估,涵盖英语、多语言和专业领域基准,将两个蒸馏学生检索器与广泛的基线进行比较。DistilVDR-HiRes 和 Fast 超越了所有 sub-1B 模型,HiRes 可与若干 2–3B 检索器相媲美,而效率分析表明,蒸馏的学生模型在多向量基线上取得了显著的加速和存储优势,特别是在文档编码和索引大小方面。消融实验证实,视觉 tile 对最困难的基准有益,数据规模趋于饱和,768 维输出以质量换取存储减少,并且在蒸馏后添加对比损失并不能改进纯余弦对齐。总体而言,蒸馏方法产生了强大的单向量检索,以很小的成本捕捉了教师模型的大部分质量。
DistilVDR-HiRes 和 DistilVDR-Fast 都是参数低于十亿的单向量模型,在 ViDoRe 检索基准上超越了所有 sub-1B 基线,包括多向量替代方案。HiRes 提供了最高的平均 NDCG@5,领先下一个最佳 sub-1B 检索器 8.73 分,在具有挑战性的 v3 报告数据集上具有显著优势,而 Fast 在 v1 和 v2 上以更小的视觉 token 预算实现了相似的质量。两个变体仍然与 2–3B 模型竞争,并保留了 8B 教师模型 84% 以上的性能。与顶级 sub-1B 多向量基线(colSmol-500M,平均 53.01)相比,HiRes 领先 8.73 分,Fast 领先 6.97 分。在最困难的 v3 基准(长篇专业报告)上,HiRes 达到 47.07 NDCG@5,比最佳 sub-1B 竞争对手高 13.55 分,而 Fast 得分为 43.66,领先其他任何 sub-1B 模型超过 10 分。
分解学生与教师之间的差距显示,替换文档编码器平均损失 6.03 NDCG@5,而替换查询编码器损失 4.69 分。完整的学生-学生系统比教师-教师 oracle 落后 9.31 分,这意味着两边的损失不是简单相加,残留部分由交互效应解释。文档侧替换(T×S)使平均 NDCG@5 降低 6.03 分,略高于查询侧替换(S×T)的 4.69 分。组合的学生系统(S×S)与教师 oracle 存在 9.31 分的差距,表明两侧的损失并没有简单相加。
DistilVDR 的单向量设计在多向量基线上取得了显著的效率优势:其 Fast 变体在 sub-1B 模型中实现了最高的文档吞吐量和最低的 VRAM 使用,而 Fast 和 HiRes 的查询编码速度均快于所有分析系统。与多向量替代方案相比,DistilVDR 的索引存储需求约减少 16 倍,评分速度提高两个数量级,使得端到端部署更加实用。DistilVDR-Fast 实现了最高的文档吞吐量(99 docs/s),峰值 VRAM 仅 2.1 GB,比所有多向量基线快一个数量级。DistilVDR-HiRes 恢复了部分吞吐量,同时仍比 8B 教师快 7 倍。查询编码仅需 3.4 ms,通过纯文本的 DistilBERT 路径,快于所有其他分析系统。多向量 sub-1B 模型每百万文档存储 256 GB,而 DistilVDR 仅存储 16.4 GB,减少了 16 倍。对 10,000 个文档评分,DistilVDR 仅需 9.6 ms,而多向量基线需 1.1–3.2 秒,延迟改进约 100 倍。
对 DistilVDR 设计的消融揭示了明确的权衡:将视觉 tile 预算从 2 增加到 6 个 tile,平均 NDCG@5 提高 1.76 分,在最困难的基准上提高 3.41 分,代价是视觉 token 数量增加三倍。训练数据规模呈现单调改进,完整数据集比四分之一规模的子集高出约 5 分,超过 75% 后趋于饱和。将输出维度从 4096 降低到 768 使索引缩小 5.3 倍,但损失 3.19 分,而换用轻量查询编码器可将延迟降低 5 倍,代价为 4.69 分。HiRes(6 个 tile)达到平均 NDCG@5 61.74,比 Fast(2 个 tile)提高 1.76 分,优势集中在复杂的 v3 基准(3.41 分)。将训练数据从 1.20M 图像混合的 25% 扩展到 100%,平均 NDCG@5 增加约 5 分,超过 75% 规模后质量饱和。768 维输出将索引大小减少 5.3 倍(至每百万文档 3.07 GB),但导致平均下降 3.19 分,在 v3 上更大(4.77 分)。将 524M 查询编码器替换为 105M 主干可将查询延迟降低 5 倍,同时牺牲 4.69 平均 NDCG@5。
在联合细化损失中添加对比项并不能改善纯余弦对齐的检索质量。InfoNCE 随着权重增加导致性能轻微下降,而 KL 散度保持平稳,与现有的余弦损失几乎冗余。纯余弦细化(γ=0)实现了最高的平均 NDCG@5 58.45,优于所有添加对比监督的变体。将 InfoNCE 的权重从 0.5 增加到 2.0,平均 NDCG@5 从 58.33 降至 57.95,呈持续下降趋势。KL 散度变体的 NDCG@5 得分在纯余弦基线的 ±0.14 范围内,未显示有意义增益。
在 ViDoRe 检索基准上评估,参数低于十亿的单向量 DistilVDR 模型超越了所有 sub-1B 基线,包括多向量系统,同时大幅减少了索引存储和评分延迟,并保持与更大模型的竞争力。HiRes 变体在最具挑战性的长文档报告中表现优异,Fast 变体实现了最高的吞吐量。消融实验证实,增加视觉 token 预算有利于复杂查询,数据扩展在 75% 左右饱和,简单的余弦蒸馏损失是最优的,对比目标没有带来增益。