Unbalanced Optimal Transport for Efficient Visual Document Retrieval

AI summary

Overview

Research area: Computer Vision / information retrieval — specifically visual document retrieval (VDR) with multi-vector Vision-Language Model (VLM) retrievers, combined with optimal transport theory.

Technical level: Advanced. The paper assumes familiarity with ColBERT-style late interaction (MaxSim), vision transformer token representations, and optimal transport / Sinkhorn algorithms; it also contains measure-theoretic-flavored theorems.

Scope: A single paper presenting an Unbalanced Optimal Transport (UOT) formulation for compressing the visual tokens of a document page into a small set of representative anchors, with theoretical error bounds and retrieval experiments on ViDoRe V1 and V3.

What This Paper Is About

Multi-vector visual retrievers such as ColPali represent each document page as N patch-level visual tokens and score a query by matching each of its M tokens to the most similar page token (MaxSim), which costs O(NM) operations per query–document pair and makes large corpora expensive to store and search. Existing token-compression methods (attention-based pruning, clustering-based merging) use a single heuristic and can collapse many distinct patches into one vector, losing the fine-grained detail late interaction depends on. The paper recasts compression as an Unbalanced Optimal Transport problem that simultaneously accounts for semantic redundancy, token importance and per-anchor coverage capacity, and proves that minimizing this transport cost bounds the resulting MaxSim retrieval error.

Key Contributions

1. A VDR-specific compression principle that explicitly balances three objectives — token importance, semantic redundancy, and the coverage capacity of each representative token — rather than optimizing one heuristic in isolation.
2. A UOT-based token compression framework in which the transport cost encodes semantic redundancy, the source marginal encodes attention-based token importance, and the (relaxed) target marginal encodes coverage capacity, solved by an EM-style alternating optimization with a generalized Sinkhorn routine.
3. Theoretical guarantees: a bound showing the MaxSim degradation is upper-bounded by the transport cost (Theorem 4.1), a worst-case-query robustness bound tied to the geometric covering radius (Theorem 4.2), and pairwise/top-k ranking stability conditions (Theorem 4.3), plus asymptotic and complexity analyses.
4. Experiments on ViDoRe V1 (academic and practical splits) and ViDoRe V3 (English-only tasks) across Qwen2-VL-2B, PaliGemma-3B and Qwen2.5-VL-3B backbones, showing a better efficiency–effectiveness trade-off than attention pruning and clustering baselines at merging factors 5, 10, 25 and 50.

Main Findings

• Near-parity at lenient budgets: At merging factor (MF) 5, the method is close to the uncompressed retriever. On ViDoRe V1 academic tasks with Qwen2-VL-2B, it scores 79.6/71.9 nDCG@5/Recall@1 versus 80.1/72.4 for uncompressed ColQwen2, while query time drops from 0.63 ms to 0.16 ms.
• Large gains under aggressive compression: At MF=50 on ViDoRe V1 practical tasks (Qwen2-VL-2B), the method reaches 79.7/60.6 average versus 60.3/39.0 for Light-ColQwen2 and 63.6/44.4 for ColQwen2+Prune. The paper reports gains of +19.4 nDCG@5 over Light-ColQwen baselines on V1 practical (Qwen2) and +23.1 nDCG@5 on ViDoRe V3 English tasks (Qwen2.5) at MF=50.
• Robustness across backbones: On ViDoRe V3 English tasks with Qwen2.5-VL-3B at MF=50, the method averages 49.3/52.1 versus 26.2/25.6 for Light-ColQwen2.5 and 31.9/33.6 for pruning; with Qwen2-VL-2B at MF=50 it averages 46.2/48.9 versus 31.6/33.4 and 34.9/36.4 respectively. Uncompressed ColQwen2.5 is 57.6/60.4.
• Heuristic baselines collapse at extreme compression: The paper attributes baseline failures at MF=50 to unconstrained merging that forcibly combines disparate visual elements and induces representation collapse.
• Ablation supports all three objectives: On ViDoRe V1 academic tasks at MF=50, a hierarchical-agglomerative-clustering baseline using distance only reaches 57.9 nDCG@5; adding attention raises it to 62.6, adding a capacity constraint to 65.4, a balanced OT variant reaches 71.0, and the unbalanced OT formulation reaches 71.3. At MF=25 the same progression is 69.3, 71.5, 72.5, 75.1 and 75.4.
• Unbalanced relaxation helps: Relaxing the target marginal lets anchors adapt capacity to local semantic density, which the authors report as the best overall performance versus balanced OT.
• Qualitative preservation: A t-SNE visualization at MF=25 shows Light-ColQwen2 representatives deviating from dense original-token regions, while the proposed anchors stay closer; an attention histogram shows mass per representative token against the uniform share of approximately 0.033.
• Theoretical link: Minimizing the transport cost is shown to upper-bound the per-document MaxSim deviation, which in turn yields an explicit sufficient condition for preserving pairwise rankings and the exact top-k set.

Methodology in Plain English

The page is encoded by a frozen VLM into N ℓ2-normalized patch tokens. The goal is to replace them with K much smaller representatives (K ≪ N), where the merging factor is N/K. The procedure runs offline, before any query arrives.

First, K starting anchors are chosen with k-means++ and ℓ2-normalized. Then, three signals are set up:

• Cost: the cosine distance between each original token and each anchor, so similar patches are cheap to group together — this targets redundancy.
• Source weights: each original token gets a share proportional to the total incoming attention it receives, averaged over all transformer layers, so visually informative regions carry more mass — this targets importance.
• Target capacities: each anchor gets an equal share of 1/K, so no single anchor can absorb everything — this targets coverage. Because the transport is unbalanced, these capacity constraints are soft: mass can be partially discarded or redistributed instead of being forced to match exactly.

The problem is then solved by alternating two steps until convergence: an E-step that computes the soft assignment matrix (transport plan) with a generalized Sinkhorn algorithm run in the log domain for numerical stability, and an M-step that updates each anchor as the mass-weighted average of the tokens assigned to it, re-normalized. Settings reported are 20 EM iterations, 3 Sinkhorn iterations, and a KL penalty weight of 5; the entropy weight is 0.04 for Qwen2/Qwen2.5 and 0.01 for PaliGemma. At query time, standard MaxSim runs over only K anchors instead of N tokens.

The theory section argues that the transport cost term in the bound is exactly the quantity the optimization minimizes, and that attention-weighted mass concentrates on the tokens likely to be the MaxSim maximizers, tightening the bound further.

Why This Matters

Research impact. The paper reframes token compression as a principled assignment problem rather than a heuristic pipeline, and connects a retrieval-level error metric (MaxSim degradation) to an optimization objective that can actually be minimized. It also brings unbalanced optimal transport — normally used for matching or clustering — into visual document retrieval, and shows that relaxing exact marginal constraints is beneficial rather than a compromise.

Real-world applications (each grounded in the benchmark domains the paper evaluates):

• Enterprise and academic document search over dense, layout-heavy pages such as scientific figures, infographics and tables (the arXivQA, InfoVQA and TabFQuAD tasks).
• Retrieval over transactional and administrative documents such as receipts and forms requiring key-field localization (the DocVQA setting).
• Domain-specific professional corpora — computer science, finance, HR reports, industrial and pharmaceutical documents — where candidate pages are highly similar and ranking is hard (the ViDoRe V3 subsets).
• Multilingual and region-specific collections such as French environmental reports and French table queries (the Shift. and TabFQuAD subsets).

Industry relevance. The cost driver in production multi-vector search is index size and per-query scoring, both of which scale with the number of stored vectors. Reducing N to N/MF directly shrinks storage and latency: reported query times fall from 2.50 ms to 0.63 ms at MF=5 and to 0.26 ms at MF=50 on ViDoRe V3 with Qwen2-VL-2B. Because the method is training-free and operates on frozen backbones, it can be dropped onto an existing ColPali-style index without retraining.

Future Directions

• Cross-lingual evaluation: The main experiments use English queries only, with results on the remaining language queries deferred to supplementary material; how well capacity regulation holds across the 6 languages covered by ViDoRe V3 remains to be shown in the main text.
• Combining with two-stage pipelines: The related work notes filtering-plus-reranking pipelines as orthogonal; pairing UOT compression with such a first-stage filter is an obvious unexplored combination.
• Adaptive compression budgets: The current formulation fixes K globally before optimization; letting K or the target marginal vary per page based on content density is a natural extension of the capacity-regulation idea.
• Beyond page retrieval: The bound connects transport cost to MaxSim error, which could transfer to other multi-vector late-interaction settings (e.g., image or video retrieval) and to other VLM backbones not tested here. The paper does not report a limitations section in the supplied content, so additional failure modes are not stated.

Target Audience

Researchers and engineers working on multi-vector retrieval, ColBERT/ColPali-style late interaction, or efficient VLM inference; practitioners building large-scale document search systems who need to shrink index size without retraining; and readers interested in optimal-transport formulations of representation compression, who will find the theoretical bounds as relevant as the benchmark results. Comfort with MaxSim scoring and Sinkhorn iterations is recommended.

Authors’ abstract

Vision-Language Models(VLMs)have advanced Visual Document Retrieval (VDR) by enabling OCR-free, multi-vector representations. However, these models generate excessive tokens to preserve dense visual details, which severely limits retrieval scalability. Existing compression methods mitigate this overhead but rely on isolated heuristics that fail to balance token importance, semantic redundancy, and coverage capacity. We propose a visual token compression framework formulated as an Unbalanced Optimal Transport (UOT) problem. Our approach unifies these three objectives into a single alternating optimization process. Furthermore, we theoretically show that minimizing the transport cost upper-bounds the MaxSim retrieval error, ensuring ranking stability. Extensive experiments on the ViDoRe V1 and V3 benchmarks demonstrate that our method preserves fine-grained document semantics and improves the efficiency-effectiveness trade-off over state-of-the-art base-lines. Code is available at https://github.com/shhhhhyy/Unbalanced-Optimal-Transport-for-EVDR.