Unbalanced Optimal Transport for Efficient Visual Document Retrieval

Conference: ECCV 2026
Paper: ECCV 2026 Poster
Code: https://github.com/shhhhhyy/Unbalanced-Optimal-Transport-for-EVDR
Area: Multimodal VLM
Keywords: Visual Document Retrieval, Visual Token Compression, Unbalanced Optimal Transport, Late Interaction, MaxSim

TL;DR

To resolve the token explosion and excessive late-interaction (MaxSim) overhead in multi-vector visual document retrieval models like ColPali, this paper presents a training-free token compression framework via Unbalanced Optimal Transport (UOT) that unifies redundancy reduction, attention importance, and coverage capacity into an alternating optimization process while theoretically upper-bounding retrieval error.

Background & Motivation

Visual Document Retrieval (VDR) aims to retrieve relevant document images from massive candidate collections given a natural language text query. Early VDR systems relied extensively on optical character recognition (OCR) pipelines coupled with conventional lexical or dense text retrievers. However, OCR engines impose heavy offline indexing latency and consistently fail when dealing with complex page layouts, dense numerical tables, visual trend charts, and infographics. Consequently, recent vision-language model (VLM) retrievers, such as ColPali and ColQwen, have transitioned to an OCR-free paradigm. By encoding entire page images directly into hundreds or thousands of patch-level multi-vector representations, they employ late interaction via the MaxSim operator to match query tokens against fine-grained local document regions, achieving remarkable retrieval accuracy.

Nevertheless, retaining all visual tokens introduces a severe computational and storage bottleneck. With \(N\) visual tokens per page and \(M\) tokens per query, evaluating the MaxSim score incurs an \(\mathcal{O}(NM)\) computational complexity per document, rendering multi-vector indexing and real-time ranking prohibitive at enterprise scale. Existing token reduction heuristics, such as attention-based pruning and hierarchical clustering, attempt to compress \(N\) tokens into \(K\) representative vectors (\(K \ll N\)). However, these approaches isolate token importance from redundancy elimination and completely disregard the coverage capacity of individual representatives. Unconstrained clustering inevitably collapses large sets of disparate patches into a handful of oversized centroids, destroying the localized feature diversity indispensable for late-interaction matching and triggering catastrophic recall drops under aggressive compression.

The foundational insight of this work is that visual token compression can be naturally understood as a mass redistribution problem over the feature manifold. Core idea: formulate visual token compression as an Unbalanced Optimal Transport (UOT) problem, where transport cost minimizes semantic redundancy, source marginals prioritize attention-based token importance, and target marginals enforce uniform coverage capacity to prevent representation collapse, thereby strictly upper-bounding the MaxSim retrieval error degradation.

Method

Overall Architecture

The input consists of \(N\) patch-level visual token embeddings \(\mathbf{X} \in \mathbb{R}^{N \times D}\) extracted from the final layer of a frozen VLM (each \(\ell_2\)-normalized) alongside the multi-layer average self-attention matrix \(\mathbf{A} \in \mathbb{R}^{N \times N}\). The goal is to compress \(\mathbf{X}\) into \(K\) representative anchors \(\mathbf{Y} \in \mathbb{R}^{K \times D}\) (\(K \ll N\)). The offline pipeline initializes marginal distributions and anchor seeds, and then executes an Expectation-Maximization (EM)-style alternating optimization: the E-step solves for the optimal soft assignment plan \(\mathbf{P}\) via the generalized Sinkhorn algorithm, and the M-step updates the anchor vectors \(\mathbf{Y}\) as mass-weighted barycenters. At query time, the compact anchors \(\mathbf{Y}\) interact directly with the query token matrix \(\mathbf{Q} \in \mathbb{R}^{M \times D}\) via standard MaxSim operations.

The overall workflow is illustrated in the diagram below:

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Visual Tokens X and Attention Matrix A"] --> B["Marginal Distributions and Anchor Initialization<br/>Attention aggregation for mu, uniform capacity for nu, k-means++ seeding for Y"]
    B --> C["E-step: Unbalanced Optimal Transport<br/>Cosine distance cost C, Generalized Sinkhorn yields transport plan P"]
    C --> D["M-step: Barycentric Anchor Update<br/>Mass-weighted aggregation updates Y, followed by L2-normalization"]
    D -->|Iterate T loops| E["Final Compressed Representative Tokens Y"]
    E --> F["Online Late-Interaction Retrieval<br/>Fast MaxSim scoring against query tokens Q"]

Key Designs

1. Marginal and Cost Formulation: Unifying Redundancy, Saliency, and Capacity

Unlike heuristic baselines that decouple semantic clustering from pruning, the authors cast visual token compression into a unified Unbalanced Optimal Transport objective:

\[\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})\]

Here, the cost matrix \(\mathbf{C} \in \mathbb{R}^{N \times K}\) measures the cosine distance \(c_{ij} = 1 - \langle \mathbf{x}_i, \mathbf{y}_j \rangle\) between original tokens and anchors; minimizing \(\langle \mathbf{P}, \mathbf{C} \rangle_F\) encourages semantically redundant patches to converge into the same anchor. The source marginal \(\boldsymbol{\mu} \in \mathbb{R}^N\) is derived from the column-wise sum of the VLM self-attention matrix \(\tilde{a}_i = \sum_{n=1}^N a_{ni}\), assigning greater mass to visually salient regions such as text lines and diagram contours. The target marginal \(\boldsymbol{\nu} \in \mathbb{R}^K\) is defined as a uniform distribution \(\nu_j = 1/K\), establishing an equitable baseline capacity for each representative anchor. Relaxing these marginal constraints via soft KL penalties (\(\rho > 0\)) prevents the formation of dominant super-clusters while avoiding the rigid distortion of classical balanced OT that over-allocates budget to blank background tokens.

2. Alternating Optimization: Log-Domain Sinkhorn and Barycentric Projection

Because the transport plan \(\mathbf{P}\) and the anchor set \(\mathbf{Y}\) are interdependent, the framework optimizes them iteratively over \(T\) EM iterations. After initializing \(\mathbf{Y}\) via \(k\)-means++ and \(\ell_2\)-normalization, the E-step fixes \(\mathbf{Y}\) and solves the UOT objective using the generalized Sinkhorn algorithm in the log domain over \(L\) scaling iterations. In the subsequent M-step, the transport plan \(\mathbf{P}\) is held fixed, and each representative anchor is updated to the mass-weighted barycenter of its assigned tokens, projected back onto the unit hypersphere:

\[\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\}\]

This continuous barycentric adjustment keeps the anchors strictly on the data manifold while restricting unconstrained aggregation, ensuring that high-attention tokens guide the representations without overwhelming localized feature granularity.

3. Theoretical Guarantees: MaxSim Error Bound and Geometric Robustness

The authors establish rigorous mathematical bounds linking the transport plan directly to downstream retrieval accuracy. For any query token \(q_m\), letting \(i_m^\star = \arg\max_i \langle \mathbf{q}_m, \mathbf{x}_i \rangle\) be its optimal match among uncompressed tokens with positive row mass \(r_m = \sum_j p_{i_m^\star j} > 0\), the MaxSim degradation is provably bounded:

\[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}}\]

Because minimizing \(\langle \mathbf{P}, \mathbf{C} \rangle\) directly tightens this bound over dominant tokens, informative features are guaranteed to stay close to their assigned anchors. Furthermore, the worst-case retrieval gap across arbitrary unit-norm queries is bounded by the geometric covering radius \(\max_i \min_j \|\mathbf{x}_i - \mathbf{y}_j\|_2\), which in turn is bounded by the transport cost. These token-level bounds translate into explicit sufficient conditions that preserve document pairwise ranking orders and top-\(k\) set stability.

Loss & Training

The proposed method is strictly training-free and post-hoc, requiring no downstream fine-tuning, contrastive re-training, or back-propagation. The default hyperparameters are fixed across models: \(T = 20\) EM iterations, \(L = 3\) Sinkhorn scaling steps, and KL penalty weight \(\rho = 5.0\). The entropic regularization parameter \(\gamma\) is set to \(0.04\) for Qwen2-VL / Qwen2.5-VL and \(0.01\) for PaliGemma. The entire procedure runs entirely on a GPU in tens of milliseconds per document during offline index building.

Key Experimental Results

Main Results

The framework is comprehensively evaluated on the ViDoRe V1 benchmark (academic and practical tasks) and the ViDoRe V3 benchmark across three backbones: Qwen2-VL-2B, PaliGemma-3B, and Qwen2.5-VL-3B. Comparisons are made against uncompressed multi-vector models, attention-based pruning (Prune), and hierarchical clustering (Light-ColPali / Light-ColQwen2) across Merging Factors (MF) \(\in \{5, 10, 25, 50\}\).

The table below summarizes retrieval effectiveness and query latency across different compression ratios on ViDoRe V1 Academic Tasks:

Model / Method	Merging Factor (MF)	Query Time (ms/q)	InfoVQA (nDCG@5 / R@1)	TabFQuAD (nDCG@5 / R@1)	Overall Avg. (nDCG@5 / R@1)
ColQwen2-2B (Upper Bound)	1 (Uncompressed)	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

On more challenging domain-specific benchmarks under extreme compression (MF=50), the advantages of UOT are even more pronounced: - ViDoRe V1 Practical Tasks (Qwen2, MF=50): Ours achieves 79.7 / 60.6 (nDCG@5 / Recall@1), outperforming Light-ColQwen2 (60.3 / 39.0) by +19.4 nDCG@5. - ViDoRe V3 English Tasks (Qwen2.5-VL-3B, MF=50): Ours reaches 49.3 / 52.1, surpassing Light-ColQwen2.5 (26.2 / 25.6) by a remarkable margin of +23.1 nDCG@5.

Ablation Study

To isolate the contributions of the three core objectives (redundancy reduction, attention saliency, and coverage regulation), the authors incrementally incorporated these constraints into a Hierarchical Agglomerative Clustering (HAC) baseline and compared against Balanced OT and Unbalanced OT on ViDoRe V1 Academic Tasks:

Configuration	Mechanism	MF=25 (nDCG@5 / R@1)	MF=50 (nDCG@5 / R@1)
(1) HAC: Dist	Distance-only clustering (redundancy reduction)	69.3 / 60.3	61.6 / 52.3
(2) HAC: Dist+Attn	Distance clustering + attention saliency weighting	71.5 / 62.4	62.6 / 52.7
(3) HAC: Dist+Attn+Cap.	Distance + attention + heuristic cluster capacity bound	72.5 / 63.5	65.4 / 56.1
(4) Balanced OT	Strict mass conservation optimal transport	75.1 / 66.5	71.0 / 62.1
(5) Unbalanced OT (Ours)	Relaxed marginal optimal transport (adaptive capacity)	75.4 / 66.8	71.3 / 62.4

Key Findings

• Capacity regulation prevents representation collapse: As shown in the ablation table, relying solely on distance-based clustering leads to severe degradation at MF=50 (61.6 nDCG@5). Introducing capacity constraints (either heuristically or via OT target marginals) prevents disparate tokens from being absorbed into a single dominant centroid. Saliency mass histograms verify that Light-ColQwen2 results in a heavily skewed distribution, whereas UOT tightly centers anchor mass around the theoretical uniform expectation (\(1/K \approx 0.033\)).
• Unbalanced relaxation outperforms balanced formulation: Balanced OT forces strict conservation, causing non-informative background patches to consume equal anchor capacity. UOT relaxes this constraint via KL penalties, dynamically allocating capacity in proportion to local semantic density, which consistently achieves superior retrieval accuracy under high compression ratios.

Highlights & Insights

• Principled formulation of token compression as UOT: Instead of stacking independent heuristics, the paper maps redundancy, saliency, and capacity directly to the transport cost, source marginal, and target marginal of an unbalanced optimal transport problem.
• Direct theoretical connection to MaxSim degradation: Establishes rigorous mathematical bounds showing that minimizing the transport cost directly bounds the worst-case degradation of the late-interaction MaxSim score and guarantees top-\(k\) ranking stability.
• Seamless deployment with dramatic speedups: Operates in a strictly training-free manner with negligible offline overhead, slashing query-time latency by up to \(9\times\) and vector memory by up to \(50\times\) while retaining over 90% of original retrieval accuracy.

Limitations & Future Work

• Offline indexing computation: Computing the self-attention matrix and running 20 EM iterations introduces additional offline indexing latency compared to one-shot heuristic pooling, which may add overhead for billion-scale document ingestion pipelines.
• Cross-architecture parameter tuning: The entropy regularization parameter \(\gamma\) required different values between Qwen2 (\(0.04\)) and PaliGemma (\(0.01\)) to achieve optimal performance, indicating slight sensitivity to the attention entropy of the underlying backbone.
• Future directions: Integrating UOT compression directly into early-layer VLM vision encoders to accelerate both forward vision encoding and late-interaction retrieval concurrently.

Related Work & Insights

• vs ColPali [15]: ColPali established the SOTA for OCR-free multi-vector retrieval but suffers from massive token storage. This paper serves as an ideal compression companion, preserving over 98% of retrieval accuracy at MF=10 while drastically reducing query latency.
• vs Light-ColPali / Light-ColQwen2 [27]: Light-ColPali relies on unconstrained agglomerative clustering and experiences representation collapse at extreme compression (MF=50). The proposed UOT framework overcomes this failure mode via capacity regularization, leading by over 19~23 points in nDCG@5.
• vs DocPruner / FitPrune [45, 48]: Pruning-only methods simply discard lower-attention tokens, losing dense text tokens embedded in small table cells. This work blends salient retention with redundant merging, guaranteeing broad geometric coverage without blind spots.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Formulates visual token compression as an Unbalanced Optimal Transport problem with formal MaxSim error bounds.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Comprehensive evaluation across ViDoRe V1 and V3 benchmarks using three distinct VLM backbones under multiple compression ratios.
• Writing Quality: ⭐⭐⭐⭐⭐ Rigorous mathematical derivations, well-structured visual flowcharts, and clear qualitative analyses.
• Value: ⭐⭐⭐⭐⭐ Highly practical, training-free approach that bridges the gap between academic multi-vector retrieval models and large-scale industrial deployment.