Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

CausalEmbed: Auto-Regressive Multi-Vector Generation in Latent Space for Visual Document Embedding

Jiahao Huo, Yu Huang, Yibo Yan, Ye Pan, Yi Cao, Mingdong Ou, Philip S. Yu, Xuming Hu

TL;DR

CausalEmbed tackles the storage and latency bottlenecks of visual document retrieval by reframing multi-vector embeddings as an auto-regressive latent generation process. By generating a compact sequence of latent vectors and employing a gradient-flow aware training objective that includes contrastive alignment, progressive refinement, and diversity regularization, it achieves substantial compression (up to ~30x) with competitive or superior retrieval accuracy across diverse backbones and ViDoRe benchmarks. The approach exhibits test-time scaling, where more generated tokens yield progressively better performance, and demonstrates training efficiency with rapid convergence aided by KV caching and full-context gradient propagation. Practically, CausalEmbed enables flexible deployment strategies from low-latency to high-accuracy regimes while maintaining scalability for production VDR systems. This generative embedding paradigm advances the practicality of visual document understanding in real-world applications, particularly where storage and latency constraints are critical.

Abstract

Although Multimodal Large Language Models (MLLMs) have shown remarkable potential in Visual Document Retrieval (VDR) through generating high-quality multi-vector embeddings, the substantial storage overhead caused by representing a page with thousands of visual tokens limits their practicality in real-world applications. To address this challenge, we propose an auto-regressive generation approach, CausalEmbed, for constructing multi-vector embeddings. By incorporating iterative margin loss during contrastive training, CausalEmbed encourages the embedding models to learn compact and well-structured representations. Our method enables efficient VDR tasks using only dozens of visual tokens, achieving a 30-155x reduction in token count while maintaining highly competitive performance across various backbones and benchmarks. Theoretical analysis and empirical results demonstrate the unique advantages of auto-regressive embedding generation in terms of training efficiency and scalability at test time. As a result, CausalEmbed introduces a flexible test-time scaling strategy for multi-vector VDR representations and sheds light on the generative paradigm within multimodal document retrieval.

CausalEmbed: Auto-Regressive Multi-Vector Generation in Latent Space for Visual Document Embedding

TL;DR

CausalEmbed tackles the storage and latency bottlenecks of visual document retrieval by reframing multi-vector embeddings as an auto-regressive latent generation process. By generating a compact sequence of latent vectors and employing a gradient-flow aware training objective that includes contrastive alignment, progressive refinement, and diversity regularization, it achieves substantial compression (up to ~30x) with competitive or superior retrieval accuracy across diverse backbones and ViDoRe benchmarks. The approach exhibits test-time scaling, where more generated tokens yield progressively better performance, and demonstrates training efficiency with rapid convergence aided by KV caching and full-context gradient propagation. Practically, CausalEmbed enables flexible deployment strategies from low-latency to high-accuracy regimes while maintaining scalability for production VDR systems. This generative embedding paradigm advances the practicality of visual document understanding in real-world applications, particularly where storage and latency constraints are critical.

Abstract

Although Multimodal Large Language Models (MLLMs) have shown remarkable potential in Visual Document Retrieval (VDR) through generating high-quality multi-vector embeddings, the substantial storage overhead caused by representing a page with thousands of visual tokens limits their practicality in real-world applications. To address this challenge, we propose an auto-regressive generation approach, CausalEmbed, for constructing multi-vector embeddings. By incorporating iterative margin loss during contrastive training, CausalEmbed encourages the embedding models to learn compact and well-structured representations. Our method enables efficient VDR tasks using only dozens of visual tokens, achieving a 30-155x reduction in token count while maintaining highly competitive performance across various backbones and benchmarks. Theoretical analysis and empirical results demonstrate the unique advantages of auto-regressive embedding generation in terms of training efficiency and scalability at test time. As a result, CausalEmbed introduces a flexible test-time scaling strategy for multi-vector VDR representations and sheds light on the generative paradigm within multimodal document retrieval.
Paper Structure (32 sections, 1 theorem, 42 equations, 12 figures, 4 tables, 1 algorithm)

This paper contains 32 sections, 1 theorem, 42 equations, 12 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Multimodal Embedding
  4. Visual Document Retrieval
  5. Preliminary
  6. CausalEmbed: Efficient Auto-regressive Multi-vector Generation
  7. Rethinking Multi-vector Embeddings from the Perspective of Gradient Flow
  8. Overall Training Objective
  9. Experiment
  10. Experiment Setup
  11. Main Results
  12. Test-time Scaling of CausalEmbed
  13. Framework Analysis
  14. Conclusion
  15. Ethical Considerations.
  16. ...and 17 more sections

Key Result

Theorem 2.1

Under a MaxSim-based late-interaction objective, assume the index of the maximally matched document token for each query token is uniformly distributed over its valid range. Let $L^{\mathrm{forward}}$ and $L^{\mathrm{causal}}$ denote the total preceding-token coverage under forward-based and auto-re Moreover, when $N_q \approx N_t$ and $N_d \ll N_v$, it holds that

Figures (12)

  • Figure 1: Comparison of traditional multi-vector embeddings (e.g., ColPali/ColQwen faysse2024colpali) with our auto-regressive paradigm for multi-vector generation in the VDR domain.
  • Figure 2: Overview of our overall framework.
  • Figure 3: Test-time scaling curves of CausalEmbed. We train CausalQwen with sequence lengths of 32, 64, and 128, and evaluate its test-time performance under different inference budgets.
  • Figure 4: Training and evaluation loss curves of CausalQwen and CausalPali over one epoch.
  • Figure 5: Trade-off between retrieval performance and latency on ViDoRe V2. Bubble size indicates adaptation overhead.
  • ...and 7 more figures

Theorems & Definitions (2)

  • Theorem 2.1: Preceding-Token Coverage
  • proof