💬 NLP

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

CausalEmbed is an auto-regressive approach that significantly reduces the storage overhead of visual document retrieval by generating compact multi-vector embeddings with only dozens of tokens, while maintaining competitive performance through iterative margin loss and enabling flexible test-time scaling.

Original authors: Jiahao Huo, Yu Huang, Yibo Yan, Ye Pan, Kening Zheng, Wei-Chieh Huang, Yi Cao, Mingdong Ou, Philip S. Yu, Xuming Hu

Published 2026-04-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiahao Huo, Yu Huang, Yibo Yan, Ye Pan, Kening Zheng, Wei-Chieh Huang, Yi Cao, Mingdong Ou, Philip S. Yu, Xuming Hu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Library of a Million Books"

Imagine you are building a search engine for a massive library of visual documents (like PDFs, receipts, or scientific papers).

In the past, to find a specific page, computers had to read the whole page and turn it into a single "summary note" (a single vector). But this was like trying to describe a complex painting with just one word. You'd lose all the details.

So, researchers invented a new way: Multi-Vector Embedding. Instead of one note, they broke the page into thousands of tiny puzzle pieces (patches) and wrote a detailed note for each piece.

The Good News: This is incredibly accurate. You can find exactly what you need.
The Bad News: It's a storage nightmare. If you have 10,000 documents, and each needs 1,000 notes, you now have to store 10 million notes. It's like trying to carry a library in your backpack. It's too heavy, too slow, and too expensive to run in the real world.

The Solution: CausalEmbed (The "Smart Storyteller")

The authors of this paper, CausalEmbed, asked a simple question: "What if we didn't just list all the notes at once? What if we told the story of the document, one sentence at a time?"

They created a system that acts like a smart storyteller rather than a photocopier.

1. From "The Photocopier" to "The Storyteller"

Old Way (The Photocopier): Imagine you have a document. The old system takes a photo of the whole page, cuts it into 1,000 tiny squares, and writes a description for every single square simultaneously. It's fast to take the photo, but you end up with a huge pile of paper (1,000 notes) that takes forever to sort through.
New Way (CausalEmbed): Imagine a storyteller who looks at the document and says, "Okay, the first thing I see is a big red chart. Next, I see a paragraph about profits. Then, I see a small table..."
- The storyteller generates these descriptions one by one, in order.
- Because they are telling a story, they don't need to describe every single pixel. They only generate the most important 30 sentences (tokens) that capture the whole page.
- Result: You go from 1,000 notes down to just 30 notes. That's a 30x reduction in storage!

2. The "Matryoshka Doll" Trick (Test-Time Scaling)

One of the coolest features of CausalEmbed is its flexibility, which the authors call Test-Time Scaling.

Think of the document embedding as a set of Russian Matryoshka dolls (nesting dolls).

If you need a quick, rough answer (like a fast Google search), you only open the smallest doll (the first 8 notes). It's fast and cheap.
If you need a super precise answer (like a lawyer looking for a specific clause), you open the next doll (the first 16 notes), then the next (32 notes), and so on.
Because the system generates the notes in order, you can stop whenever you want. You don't have to generate all 1,000 notes just to get a slightly better answer. You can dial the "accuracy knob" up or down instantly.

3. Why It Works Better Than "Clipping"

Other researchers tried to solve the storage problem by taking the old 1,000-note system and just clipping or clustering the notes (throwing away the "boring" ones).

The Flaw: It's like trying to make a movie by cutting out 90% of the frames. The result is choppy and loses the flow.
The CausalEmbed Advantage: Because CausalEmbed writes the notes from scratch in a sequence, it learns to be efficient naturally. It knows exactly what information is needed to tell the story. It's not throwing things away; it's distilling the essence.

The Real-World Impact

Storage: You can store 30x more documents in the same amount of space.
Speed: Searching through fewer notes is much faster.
Accuracy: Surprisingly, even with only 30 notes, it finds the right document better than the old methods that used 1,000 notes but tried to compress them clumsily.

Summary Analogy

Imagine you are trying to describe a complex city to a friend over the phone.

The Old Way: You list every single street, building, and tree in the city. (Too long, too much data).
The "Clipping" Way: You list every street, but then you cross out 90% of them and hope the remaining ones make sense. (Confusing, misses context).
CausalEmbed: You tell a story: "First, you see the giant tower. Then, you walk down the main avenue to the park. Finally, you see the museum." You use just a few key landmarks to paint a perfect picture of the whole city.

CausalEmbed is the technology that lets computers tell these "stories" about documents, making visual search faster, cheaper, and smarter.

1. Problem Statement

Visual Document Retrieval (VDR) is a critical task for retrieving relevant pages from extensive document corpora, essential for enterprise search and Retrieval-Augmented Generation (RAG). While traditional text-based retrieval relies on OCR (which loses layout information), modern approaches treat documents as visual entities.

The Bottleneck: State-of-the-art multi-vector methods (e.g., ColPali, ColQwen) generate high-quality embeddings by treating document pages as sequences of hundreds or thousands of visual patch tokens. While effective, this creates a massive storage overhead and latency bottleneck, making them impractical for large-scale production environments.
Limitations of Current Solutions: Existing compression techniques (e.g., pruning, clustering, or merging tokens) attempt to reduce the token count after generation. These methods often suffer from:
- Significant performance degradation (accuracy loss) when compressing.
- Inflexibility, as they often require fixed token budgets during training.
- Inability to adapt dynamically at inference time.

Core Question: Can the generative capabilities of Multimodal Large Language Models (MLLMs) be leveraged to synthesize compact, high-quality multi-vector embeddings directly, rather than compressing redundant ones?

2. Methodology: CausalEmbed

The authors propose CausalEmbed, a novel framework that shifts from parallel patch-based encoding to auto-regressive multi-vector generation.

A. Core Architecture

Instead of generating all document tokens in parallel (forward pass), CausalEmbed fine-tunes a pre-trained MLLM to generate a sequence of latent vectors $Z = [z_1, \dots, z_L]$ sequentially.

Process: The model conditions on the original input (visual page or text query) and the history of previously generated tokens to produce the next embedding token.
Formula: $z_k = \tau(\Psi([C, z_1, \dots, z_{k-1}]))$ , where $C$ is the initial context and $\tau$ extracts the final hidden state.
Result: This allows the model to generate a highly compressed representation (e.g., 32 tokens) that captures the document's semantics progressively.

B. Training Objective & Loss Functions

To prevent representation collapse and ensure the sequential generation is effective, the authors introduce a composite loss function:

Contrastive Margin Loss ( $L_m$ ): Standard contrastive learning to align positive query-document pairs and separate them from hard negatives using MaxSim scoring.
Progressive Refinement Loss ( $L_d$ ): Encourages the model to utilize the full sequence length. It enforces that adding more tokens to the sequence should strictly improve retrieval performance (telescoping sum of scores). This ensures early tokens capture coarse semantics while later tokens refine details.
Diversity Regularization ( $L_q$ ): Penalizes high cosine similarity between distinct tokens in the generated query sequence to prevent homogenization (repetitive patterns).

C. Theoretical Advantage (Gradient Flow)

The paper provides a theoretical analysis of gradient flow:

Forward-based (Parallel): Gradients only update the specific document tokens selected by the MaxSim operation. Most tokens receive no direct gradient signal, leading to inefficient training.
Auto-Regressive (Causal): Due to the causal dependency, gradients propagate through the entire generation chain. Every generated token influences the subsequent ones, resulting in higher training efficiency and better utilization of the model's capacity.

3. Key Contributions

Novel Paradigm: Introduces the first auto-regressive paradigm for visual document embedding, departing from the standard parallel patch-based encoding.
Extreme Compression with High Performance: Achieves a 30–155× reduction in token count (e.g., reducing ~5000 tokens to ~32) while maintaining or surpassing the performance of full-token baselines.
Test-Time Scaling: Demonstrates a unique property where retrieval precision can be dynamically adjusted by changing the number of generated tokens during inference without retraining. This allows flexible trade-offs between latency and accuracy.
Robust Generalizability: The method works effectively across diverse backbones (from strong models like Qwen2.5-VL to weaker ones like PaliGemma), often outperforming the full-resource baselines of weaker models.

4. Experimental Results

The method was evaluated on the ViDoRe benchmarks (V1, V2, V3) using backbones like ColQwen2.5 and PaliGemma.

Performance vs. Compression:
- On ViDoRe V2, CausalEmbed (using 32 tokens) achieved an nDCG@5 of 0.516, significantly outperforming the best pruning-based baselines (SemCluster: 0.455, BiQwen2: 0.464).
- On ViDoRe V3, it achieved 0.426, beating the best baseline (0.336) by a large margin.
- PaliGemma Case: When applied to the weaker PaliGemma model, CausalEmbed outperformed the full-resource ColPali baseline by nearly 5% while using 30× fewer tokens.
Test-Time Scaling: Experiments showed a clear correlation between token budget and performance. Increasing the token count from 1 to 32 steadily improved nDCG@5, approaching the upper bound of full-token models.
Latency: While auto-regressive generation adds slight query encoding overhead, the Late Interaction (MaxSim) phase is accelerated by ~11.8× due to the drastic reduction in token count, resulting in a net positive for end-to-end latency in large-scale retrieval.

5. Significance and Impact

Practicality: CausalEmbed solves the storage and latency barriers that have prevented multi-vector VDR from being widely deployed in real-world applications.
Flexibility: The "Matryoshka-like" test-time scaling allows systems to adapt to different hardware constraints (e.g., using 8 tokens for fast mobile inference vs. 128 tokens for high-accuracy server retrieval) without retraining.
Generative Paradigm: It validates that generative models can be effectively repurposed for dense retrieval tasks, offering a more efficient alternative to traditional discriminative compression methods.
Future Direction: This work bridges the gap between generative AI and retrieval systems, suggesting a future where document embeddings are synthesized on-the-fly rather than statically stored.

In conclusion, CausalEmbed represents a significant leap forward in Visual Document Retrieval, offering a scalable, efficient, and highly accurate solution that leverages the generative power of MLLMs to overcome the limitations of traditional multi-vector embeddings.