Source-linked AI summary
Spatial Matryoshka Training for Multi-Granularity Visual Document Retrieval
Trishan Singha Roy, Arkadeep Acharya, Vishwajeet Kumar, Jaydeep Sen, Sachindra Joshi
TL;DR
Multi-vector late-interaction retrievers preserve visual-document evidence but incur high storage costs, while existing compression often fixes one level at indexing time. ColSNAP spatially pools a backbone’s patch grid into jointly trained nested tiers without changing the encoder, enabling index-time storage–accuracy choices. Across Nemotron ColEmbed 4B and ColPali v1.3, it retains substantial retrieval quality under large compression ratios, with most benefits available through lightweight adaptation.
Problem
Multi-vector visual retrievers preserve localized evidence but require many stored patch vectors, and existing compression methods commonly fix one compression level at indexing time.
Method
ColSNAP spatially average-pools a backbone patch grid into nested tiers and jointly trains them with contrastive and distillation objectives without changing the encoder architecture.
Results
ColSNAP retained 96.9% of full-resolution retrieval quality at 7.3× fewer stored embeddings on Nemotron ColEmbed 4B and 89.4% at 44.8× compression on ColPali v1.3.
Takeaways & Limitations
A single retriever can offer multiple vector cardinalities, allowing deployment to trade compression for accuracy at indexing time rather than retraining for each storage budget.
Takeaways & Limitations
Under aggressive compression without hard negatives, ColPali retains 88.5% and 76.2% of full-resolution performance at L2 and L3, below Nemotron’s 92.7% and 91.5%.
Abstract
from arXiv · showhide
Multi-modal late-interaction retrievers achieve strong retrieval on visually rich documents by representing each page as per patch embeddings and matching at the token level. However, this approach incurs high storage costs. Existing compression methods typically fix a single compression level at indexing time, limiting flexibility. We present ColSNAP (Spatial Nested Average Pooling)1, a training method that generates a nested hierarchy of compression levels directly from a backbone's patch grid. By spatially pooling patch embeddings into pro- gressively coarser tiers and training all tiers simultaneously, a single model learns to support retrieval at multiple compression levels without architectural changes. Crucially, a single encoding pass yields every tier, enabling the accuracy-storage trade-off to be configured at indexing time to match avail- able storage budgets, rather than being fixed during training. We demonstrate that models trained using ColSNAP maintain near full-resolution retrieval performance under substantial compression and that ColSNAP transfers effectively across multiple late-interaction backbones, and achieves most of its improvements via a lightweight adaptation stage applied to a pre-trained retriever.
1 Introduction
Visual document retrieval preserves layout-sensitive evidence through per-patch late interaction, but this creates storage costs that motivate flexible multi-level compression. ColSNAP addresses this by jointly learning nested spatially pooled tiers from one patch grid, allowing index-time selection of compression level.
- Retrieval challenge: Multi-vector retrievers preserve localized evidence by matching query tokens to relevant page regions, unlike single-vector methods that discard fine-grained document evidence.This improves handling of tables, charts, figures, and other layout-dependent content, but requires storing many patch vectors per page.
- Research questions: RQ1 asks whether retrieval performance can be maintained while substantially reducing the number of vectors stored for each document page.
- Research questions: RQ2 asks whether compression generalizes across retrievers with different model sizes and retrieval performance.
- ColSNAP: ColSNAP spatially pools patch embeddings into nested tiers, trains all tiers jointly with two distillation signals, and changes neither token count nor encoder architecture.The frozen-teacher signal anchors tiers to the pretrained model, while the cross-level signal aligns coarse tiers with the live model’s finest tier.
- ColSNAP: Up to ∼70× smaller storage on Nemotron ColEmbed 4B and ∼130× smaller storage on ColPali v1.3 are achieved with compression selected at indexing time.
2 Related Work
Late-interaction visual retrieval retains layout-sensitive evidence but stores hundreds to thousands of patch vectors per page. Prior compression methods often fix one indexing-time level, whereas ColSNAP uses the patch grid’s spatial structure to build a nested multi-vector hierarchy.
- Late-interaction retrieval: Late-interaction retrieval represents queries and documents as embedding sequences and uses MaxSim to preserve fine-grained matching that pooled vectors discard.
- Visual document retrieval: ColPali applies late interaction to rendered page images, avoiding OCR brittleness while representing each page with multi-vector patch embeddings.The approach was later carried over to stronger backbones such as Nemotron ColEmbed.
- Compression challenge: Visual late-interaction systems can require hundreds to thousands of patch vectors per page, increasing index size and retrieval latency.
- Prior compression: Existing text and visual compression methods include quantization, pruning, pooling, retention, clustering, and retriever or module fine-tuning.
- Prior compression: Most visual compression methods fix one compression level at indexing time, while the patch grid’s spatial layout remains largely unused as a hierarchy axis.
- ColSNAP: ColSNAP is presented as the first visual-document-retrieval method to build a nested multi-vector hierarchy from a backbone patch grid, orthogonally to embedding-dimensionality compression.
3 ColSNAP
ColSNAP creates nested, spatially pooled views of a page’s patch grid and jointly trains them for late-interaction retrieval. A single encoded representation can support multiple storage and retrieval operating points without changing the encoder or retraining.
- Spatial Pooling Pyramid: ColSNAP spatially average-pools a 32 × 32 patch grid into nested tiers with 1024, 128, 16, and 1 patch embeddings.Each coarser embedding summarizes a contiguous spatial block of finer embeddings, and the pooled tiers reuse the same underlying embeddings without additional parameters.
- Multi-Tier Scoring: The document representation at each tier replaces patch embeddings with pooled counterparts while preserving auxiliary tokens unchanged.The finest tier is identical to the original unpooled retriever; only the patch count varies across tiers.
- Joint Training via Spatial Matryoshka Learning: ColSNAP trains all tiers jointly with ranking, frozen-teacher, and live cross-tier objectives.The frozen-teacher term anchors tiers to the pretrained full-resolution model, while the live term aligns coarser tiers with the finest tier; score distributions are z-scored before softmax.
- Indexing and Retrieval Flexibility: A single model decouples compression from training: pages can be stored at the tier matching the storage budget, with retrieval cost determined by k + n_ℓ embeddings.A corpus encoded at L0 can be pooled to any coarser tier on demand, without retraining or separate checkpoints.
4 Experimental Setup
The experiments adapt two late-interaction visual document retrievers of different scales and evaluate them across in-distribution and out-of-distribution benchmarks. The setup also compares compression tiers and reports tier-level retrieval quality and retention.
- Backbones: ColSNAP is applied to Nemotron ColEmbed VL 4B and ColPali v1.3, which produce 2560- and 128-dimensional embeddings, respectively.The two backbones differ in scale and pretraining, enabling evaluation across retrievers with different capacities.
- Backbones: Both backbones use a 32 × 32 patch grid, while their 14 and 7 auxiliary tokens remain unchanged across compression tiers.This permits the same pooling pyramid to be applied to both models without modification.
- Training Data: Experiments use ColPali-Train with a 95:5 train:validation split and matched query–document batches of 128, with additional mined-hard-negative experiments.Hard-negative training pairs each query with one negative sampled from its top-10 mined candidates.
- Optimization and Hyperparameters: Models are trained jointly across four tiers for five epochs using LoRA adaptation, AdamW, linear warmup-decay, bfloat16, and four NVIDIA A100 GPUs.The effective batch size is 128; the tier-weight hyperparameters differ by backbone.
- Benchmarks and Metrics: Evaluation covers ViDoRe v1, v2, and public ViDoRe V3 splits, spanning in-distribution and out-of-distribution document types and domains.The reported metrics are nDCG@5 for v1/v2 and nDCG@10 for v3, plus retention relative to full-resolution L0.
- Benchmarks and Metrics: Figure 3 compares untrained pooling, a full-resolution-only control, ColSNAP, and ColSNAP with mined hard negatives across four compression tiers.The tiers correspond to 1×, 7.3×, 34.6×, and 69.2× compression on Nemotron.
5 Results and Analysis
ColSNAP preserves retrieval quality across substantial compression, transfers across late-interaction backbones, and reduces storage and scoring costs. It also adapts pretrained retrievers quickly, although extreme compression is more limited for lower-dimensional backbones.
- Nemotron ColEmbed 4B: On Nemotron, ViDoRe v3 nDCG@10 declines from 61.7 at full resolution to 59.7, 57.4, and 56.7 across L1–L3 tiers.The L1 tier uses roughly seven times fewer vectors while retaining over 96% of full-resolution performance.
- Nemotron ColEmbed 4B: Compressed Nemotron tiers remain competitive with full-resolution retrievers, with L2 scoring 57.4 versus 57.5 for jina-embeddings-v4 and 56.4 for colnomic-embed-multimodal-3b on ViDoRe v3.Even L3 scores 56.7, trailing llama-nemoretriever-colembed-3b-v1 by only 0.4 nDCG.
- Robustness under distribution shift: Across ViDoRe v1–v3, Nemotron L3 retains 97.2%, 92.5%, and 91.5% of full-resolution performance, while L1 retains at least 95.4%.Together, ColSNAP preserves over 91% of full-resolution quality while storing up to 69× fewer embeddings per page.
- Transferability to ColPali v1.3: On ColPali, L1 reduces vectors by 7.6× while retaining 100.2%, 98.0%, and 99.8% of full-resolution quality on ViDoRe v1, v2, and v3.The full-resolution control leaves coarse tiers near their untrained values, attributing compressed-tier gains to the nested objective rather than additional training alone.
- Capacity and hard negatives: ColPali degrades more than Nemotron under extreme compression, with ViDoRe v3 L2 and L3 retention of 88.5% and 76.2% versus 92.7% and 91.5% for Nemotron.The paper attributes this gap to ColPali’s 128-dimensional embeddings versus Nemotron’s 2560 dimensions; mined hard negatives raise ColPali L2 and L3 retention to 89.4% and 80.7%.
- Efficiency and convergence: ColSNAP converges rapidly, achieving most gains within the first 1–3 epochs and reducing Nemotron index size and scoring cost by 7.3× at L1 while retaining 96.9% quality.On ColPali, L1 and L2 match their final five-epoch performance after one epoch.
6 Conclusion
ColSNAP converts a pretrained late-interaction retriever’s patch grid into nested retrieval granularities, allowing compression to be selected at deployment rather than fixed in advance. It retains high retrieval quality across substantial embedding reductions and can be applied through lightweight adaptation.
- ColSNAP creates nested vector cardinalities from a pretrained retriever’s patch grid without adding tokens or changing the encoder architecture.
- Compression levels can be chosen at indexing time, enabling users to trade storage reduction against retrieval accuracy on demand.
- 96.9% of full-resolution retrieval quality remains at a 7.3× reduction in stored embeddings on Nemotron ColEmbed 4B.
- The method retains 93.1% and 92.0% quality at 34.6× and 69.2× reductions, respectively, while most benefits come from a short targeted training run.