Unbalanced Optimal Transport for Efficient Visual Document Retrieval¶

会议: ECCV 2026
论文: ECCV 2026 Poster
代码: https://github.com/shhhhhyy/Unbalanced-Optimal-Transport-for-EVDR
领域: 多模态 VLM
关键词: 视觉文档检索, 视觉Token压缩, 非平衡最优传输, Late Interaction, MaxSim

一句话总结¶

针对以 ColPali 为代表的多向量视觉文档检索中 Patch Token 爆炸与晚期交互（MaxSim）计算开销巨大的瓶颈，本文提出基于非平衡最优传输（UOT）的无训练 Token 压缩框架，将语义冗余最小化、注意力重要性保留与锚点容量约束统一为交替优化问题，并在理论上证明其上界约束了检索误差。

研究背景与动机¶

视觉文档检索（Visual Document Retrieval, VDR）旨在根据给定的文本查询从海量候选文档图像中精准定位相关页面。早期的系统重度依赖光学字符识别（OCR）配合传统的文本检索器，但 OCR 流水线不仅带来高昂的离线索引延迟，在面对复杂的表格结构、图表趋势、信息图以及非规则页面排版时更屡屡失效。为此，近期以 ColPali 和 ColQwen 为代表的基于视觉语言模型（VLM）的检索器转向完全端到端的 OCR-free 范式，通过 ViT 将文档图像编码为数百至上千个 Patch 级别的多向量表征，并借助晚期交互（Late Interaction）中的 MaxSim 操作让查询词精准匹配最相似的图像局部区域，显著提升了细粒度匹配精度。

然而，保留全量视觉 Token 导致了严重的存储爆炸与在线检索延迟。若一篇文档包含 \(N\) 个视觉 Token，长度为 \(M\) 的查询计算 MaxSim 的时间复杂度为 \(\mathcal{O}(NM)\)，在大规模语料库下的打分与多向量索引难以扩展。现有的压缩方法尝试将 \(N\) 个 Token 剪裁或聚类合并为 \(K\) 个代表锚点（如基于注意力的硬剪枝或层次聚类），但它们普遍将“重要性筛选”与“语义冗余聚合”割裂为孤立的启发式步骤，并且完全忽略了代表锚点的覆盖容量限制（Coverage Capacity）。缺乏约束的合并往往会将大量异构 Patch 盲目坍缩进少数几个中心，破坏了晚期交互所必需的局部特征多样性，导致高压缩比下召回能力急剧崩溃。

本文的核心思考在于：Token 压缩本质上是一个在特征流形上的质量重新分配过程。核心 idea：将多向量文档的视觉 Token 压缩形式化为一个非平衡最优传输（Unbalanced Optimal Transport, UOT）问题，把传输代价映射为语义冗余消除、源端边际分布映射为注意力重要性赋权、目标端边际先验映射为锚点容量正则化，并通过 EM 式交替优化联合求解，在理论与经验上严格压制 MaxSim 检索误差。

方法详解¶

整体框架¶

输入为 VLM 最后一层提取的文档图像 \(N\) 个 patch-level 视觉 Token 矩阵 \(\mathbf{X} \in \mathbb{R}^{N \times D}\)（经 \(\ell_2\) 归一化）以及各层自注意力平均矩阵 \(\mathbf{A} \in \mathbb{R}^{N \times N}\)。目标是将其压缩为 \(K\) 个具有代表性的锚点向量 \(\mathbf{Y} \in \mathbb{R}^{K \times D}\)（其中 \(K \ll N\)）。整个流水线离线进行，分为边际与锚点初始化、E-step 求解非平衡最优传输分配矩阵 \(\mathbf{P}\)、以及 M-step 更新锚点向量 \(\mathbf{Y}\) 的循环迭代。在线阶段直接使用紧凑的 \(\mathbf{Y}\) 与查询 Token 矩阵 \(\mathbf{Q} \in \mathbb{R}^{M \times D}\) 执行常规的 MaxSim 晚期交互打分。

整体流程如下图所示：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入视觉 Token X 与自注意力矩阵 A"] --> B["边际分布与初始锚点构建<br/>注意力聚合为 mu，均匀容量为 nu，k-means++ 初始化 Y"]
    B --> C["E-step：求解非平衡最优传输<br/>余弦距离代价矩阵 C，广义 Sinkhorn 求解传输计划 P"]
    C --> D["M-step：加权重心与归一化更新<br/>质量加权汇聚更新锚点 Y，L2 投影重回超球面"]
    D -->|循环迭代 T 轮| E["最终压缩代表 Token Y"]
    E --> F["在线端晚期交互检索<br/>与查询 Q 快速执行高效 MaxSim 打分"]

关键设计¶

1. 边际分布与代价设计：统合冗余、重要性与覆盖容量

传统的硬聚类只看特征几何距离，而传统的剪枝只看注意力排序。为了在一套数学框架中自然兼顾多维目标，本文将 Token 压缩严格定义为一个质量重分配的非平衡最优传输：

\[\min_{\mathbf{P} \ge 0} \langle \mathbf{P}, \mathbf{C} \rangle_F + \gamma H(\mathbf{P}) + \rho \mathrm{KL}(\mathbf{P}\mathbf{1}_K \parallel \boldsymbol{\mu}) + \rho \mathrm{KL}(\mathbf{P}^\top\mathbf{1}_N \parallel \boldsymbol{\nu})\]

其中代价矩阵 \(\mathbf{C} \in \mathbb{R}^{N \times K}\) 的元素定义为原始 Token 与锚点之间的余弦距离 \(c_{ij} = 1 - \langle \mathbf{x}_i, \mathbf{y}_j \rangle\)，最小化输送代价 \(\langle \mathbf{P}, \mathbf{C} \rangle_F\) 直接推动语义极其相似的冗余 Patch 聚合到同一锚点。源端边际分布 \(\boldsymbol{\mu} \in \mathbb{R}^N\) 依靠 VLM 全层注意力矩阵的列向求和 \(\tilde{a}_i = \sum_{n=1}^N a_{ni}\) 归一化生成，赋予高注意力显著区域（如文字、关键实体、图表核心）更多的质量预算。目标端边际分布 \(\boldsymbol{\nu} \in \mathbb{R}^K\) 则设定为均匀分布 \(\nu_j = 1/K\)，为每个代表向量划定均衡的容量预算基线。通过引入软惩罚系数 \(\rho\) 驱动的 KL 散度松弛，模型既防止了所有 Token 塌缩成少数庞大超簇，又避免了严格平衡 OT 强行给空白背景分配过多配额的死板。

2. EM 式交替优化：对数域广义 Sinkhorn 与质心更新

代表锚点 \(\mathbf{Y}\) 与传输配额 \(\mathbf{P}\) 互相耦合，因此本文采用 EM 思想进行 \(T\) 轮交替求解。在初始化阶段，利用 \(k\)-means++ 从 \(\mathbf{X}\) 中播种 \(K\) 个初始向量并做 \(\ell_2\) 归一化。在 E-step 中，固定锚点 \(\mathbf{Y}\) 并计算距离代价 \(\mathbf{C}\)，利用对数域广义 Sinkhorn 算法（Generalized Sinkhorn）在 \(L\) 步行/列更新内快速收敛得出最优软分配矩阵 \(\mathbf{P} \in \mathbb{R}^{N \times K}\)；在 M-step 中，固定分配计划 \(\mathbf{P}\)，将每个锚点更新为其所分配原始 Token 的质量加权重心并重新投影至单位超球面：

\[\mathbf{y}_j = \frac{\sum_{i=1}^N p_{ij} \mathbf{x}_i}{\left\| \sum_{i=1}^N p_{ij} \mathbf{x}_i \right\|_2}, \quad \forall j \in \{1, \dots, K\}\]

这一更新确保了代表锚点严格位于高密度流形上，既保留了显著 Token 的主导方向，又借助行和约束限制了单个代表向量的最大吞噬范围。

3. 检索稳定性理论保障：MaxSim 误差上界与几何覆盖半径

本文给出了坚实的理论证明。设查询词 \(q_m\) 在原始 Token 中的最佳匹配项为 \(i_m^\star = \arg\max_i \langle \mathbf{q}_m, \mathbf{x}_i \rangle\)，对应行质量 \(r_m = \sum_j p_{i_m^\star j} > 0\)。定理证明压缩前后的 MaxSim 误差满足严格上界：

\[s(\mathbf{Q}, \mathbf{X}) - s(\mathbf{Q}, \mathbf{Y}) \le \sqrt{2} \sum_{m=1}^M \sqrt{\frac{1}{r_m} \sum_{j=1}^K p_{i_m^\star j} c_{i_m^\star j}}\]

这表明在 UOT 目标中最小化传输代价直接紧致化了决定最终得分的核心匹配项的误差。更进一步，对于任意单位范数的极端查询，最坏情况误差 \(\Delta(\mathbf{X}, \mathbf{Y})\) 由几何覆盖半径 \(\max_i \min_j \|\mathbf{x}_i - \mathbf{y}_j\|_2\) 所界定，并通过传输代价实现整体压制，从而在文档对间排序稳定性与 Top-\(k\) 集合一致性上给出了充分保证条件。

损失函数 / 训练策略¶

本方法完全属于训练后即插即用（Training-Free / Post-Hoc）的无监督压缩框架，无需任何额外的微调或反向传播梯度。实验统一超参数设置为：EM 交替轮数 \(T = 20\)，Sinkhorn 迭代步数 \(L = 3\)，KL 边际松弛系数 \(\rho = 5.0\)。熵正则化系数 \(\gamma\) 针对 Qwen2-VL / Qwen2.5-VL 设为 \(0.04\)，针对 PaliGemma 设为 \(0.01\)。所有计算均可在 GPU 上高效并行完成。

实验关键数据¶

主实验¶

论文在 ViDoRe V1（学术与实用任务）及 ViDoRe V3（更具挑战的多级相关性领域检索基准）上，跨越 Qwen2-VL-2B、PaliGemma-3B 与 Qwen2.5-VL-3B 三大骨干网络，对比了全量 Token 模型、注意力剪枝（Prune）与层次聚类（Light-ColPali / Light-ColQwen）。合并因子（Merging Factor, MF）取值 5、10、25、50，对应的检索延迟显著降低。

下表摘录了 ViDoRe V1 学术任务上各方法在不同压缩比下的整体检索指标对比：

模型 / 方法	合并因子 (MF)	平均耗时 (ms/q)	InfoVQA (nDCG@5 / R@1)	TabFQuAD (nDCG@5 / R@1)	平均性能 (nDCG@5 / R@1)
ColQwen2-2B (上限)	1 (未压缩)	0.63	90.6 / 86.2	88.7 / 81.4	80.1 / 72.4
ColQwen2 + Prune	5	-	86.7 / 80.2	85.8 / 78.2	76.5 / 68.2
Light-ColQwen2	5	0.16	87.4 / 81.6	87.3 / 79.6	78.1 / 69.6
Ours-ColQwen2	5	0.16	89.0 / 84.2	89.5 / 83.2	79.6 / 71.9
ColQwen2 + Prune	25	-	72.3 / 62.2	76.8 / 67.9	63.9 / 54.7
Light-ColQwen2	25	0.08	72.9 / 63.0	86.5 / 78.6	69.3 / 60.3
Ours-ColQwen2	25	0.08	85.8 / 79.2	87.7 / 81.4	75.4 / 66.8
ColQwen2 + Prune	50	-	65.7 / 54.4	65.0 / 54.3	55.2 / 45.6
Light-ColQwen2	50	0.07	65.8 / 55.8	81.1 / 71.8	61.6 / 52.3
Ours-ColQwen2	50	0.07	81.2 / 73.6	86.8 / 79.6	71.3 / 62.4

在更为严苛的 ViDoRe V1 Practical 任务与 ViDoRe V3 领域基准中，极限压缩（MF=50）的优势进一步拉大： - ViDoRe V1 实用任务（Qwen2, MF=50）：Ours 平均达到 79.7 / 60.6，相比 Light-ColQwen2 的 60.3 / 39.0 获得 +19.4 nDCG@5 的巨大领先。 - ViDoRe V3 英文任务（Qwen2.5-VL-3B, MF=50）：Ours 平均达到 49.3 / 52.1，相比 Light-ColQwen2.5 的 26.2 / 25.6 暴涨 +23.1 nDCG@5。

消融实验¶

为验证三项核心设计（距离消除冗余、注意力强化重要性、目标容量约束）的独立贡献，作者基于 Qwen2 在 ViDoRe V1 学术子集上，从单变量层次聚类（HAC）逐步叠加组件，并对比平衡 OT 与非平衡 UOT：

配置	机制说明	MF=25 (nDCG@5 / R@1)	MF=50 (nDCG@5 / R@1)
(1) HAC: Dist	仅按质心欧氏距离合并（纯消除冗余）	69.3 / 60.3	61.6 / 52.3
(2) HAC: Dist+Attn	距离合并 + 平均注意力重要性引导	71.5 / 62.4	62.6 / 52.7
(3) HAC: Dist+Attn+Cap.	距离 + 注意力 + 启发式限制每簇总容量上限	72.5 / 63.5	65.4 / 56.1
(4) Balanced OT	严格质量守恒的最优传输配给	75.1 / 66.5	71.0 / 62.1
(5) Unbalanced OT (Ours)	松弛边际的非平衡最优传输（自适应容量弹性）	75.4 / 66.8	71.3 / 62.4

关键发现¶

容量正则化是防止表征崩溃的核心护城河：由消融表可见，仅靠距离聚类在 MF=50 时得分仅 61.6；硬性加入注意力虽有提升但依旧脆弱。而一旦对锚点赋予容量约束（无论是启发式还是 OT 的边际先验），极端压缩下的性能跌幅显著收窄。t-SNE 与注意力质量直方图可视化证实：无约束聚类会使少量大簇抢占全部质量，而 UOT 能将代表向量的注意力质量紧密约束在期望理论均值（\(1/K \approx 0.033\)）附近。
非平衡松弛（UOT）优于严格平衡（OT）：严格平衡 OT 假定源与目标质量绝对守恒，强迫无信息的空白背景 Patch 占用均等的锚点容量；UOT 允许通过 KL 散度适度丢弃/平移无用背景信息，使具有高语义密度的文本与图表细节获得更自然的质量自适应，在各类高压缩配置下均超越 Balanced OT。

亮点与洞察¶

将 Token 压缩严格映射为非平衡最优传输：把“降重”、“保要”与“防塌缩”这三个以往割裂调参的直觉设计，完美对应至 UOT 的代价函数、源边际分布和目标边际分布，数学形式优雅且闭环。
提供了 MaxSim 检索指标下降的严格误差界：许多压缩工作仅依赖经验指标，本文从数学上推导了单查询词、全句 MaxSim、几何覆盖半径及 Top-\(k\) 排序稳定性的上界，为检索系统的性能兜底提供了理论支撑。
即插即用且完全免训练：无需对 VLM 进行昂贵的检索对比学习或再蒸馏，离线端仅需通过数轮极其轻量的 GPU 对数域 Sinkhorn 迭代，便可使向量索引体积缩减 10 到 50 倍，显著铺平了多向量检索落地的工程路径。

局限与展望¶

离线压缩阶段的额外矩阵运算开销：虽然在线查询时计算量由 \(\mathcal{O}(NM)\) 降至 \(\mathcal{O}(KM)\) 且延迟大幅降低，但离线建库时需要提取自注意力矩阵并运行 20 轮 EM 与 Sinkhorn 迭代，对于超大规模工业级离线文档建库而言，索引吞吐量略有增加。
超参数 \(\rho\) 与 \(\gamma\) 的跨模型迁移性：实验中 Qwen 系列与 PaliGemma 采用了不同的熵正则化系数（0.04 vs 0.01），面对不同注意力分布特性的视觉骨干网络时，可能仍需要细微调优以达到最佳 Pareto 前沿。
未来方向：可进一步探索如何将 UOT 压缩思路下放至 VLM 浅层特征前向过程中，实现从视觉编码前向到多向量检索的全链路推理加速。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次将多向量 VDR 的视觉 Token 压缩提炼为非平衡最优传输问题，建立了与 MaxSim 误差理论界的形式化联系。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 ViDoRe V1（学术与实用）及 V3 领域基准，跨越三大 VLM 骨干网络，对比细致且消融深入。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严密，图文对照自洽，算法伪代码与实验分析兼备。
价值: ⭐⭐⭐⭐⭐ 对推动 OCR-free 多向量文档检索的大规模工业落地与低成本高并发部署具有直接而重要的实用价值。