Unbalanced Optimal Transport for Efficient Visual Document Retrieval
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 baselines. Code is available at https://github.com/shhhhhyy/UnbalancedOptimal-Transport-for-EVDR.