|
|
--- |
|
|
license: apache-2.0 |
|
|
base_model: |
|
|
- google/paligemma-3b-mix-448 |
|
|
pipeline_tag: visual-document-retrieval |
|
|
library_name: transformers |
|
|
tags: |
|
|
- transformers |
|
|
- multimodal_embedding |
|
|
- embedding |
|
|
- colpali |
|
|
- multilingual-embedding |
|
|
- causalembed |
|
|
datasets: |
|
|
- vidore/colpali_train_set |
|
|
--- |
|
|
# Z1zs/CausalPali |
|
|
|
|
|
**CausalPali** is an auto-regressive multimodal embedding model based on the **PaliGemma-3B** architecture, proposed by researchers from HKUST(GZ) and Alibaba Cloud. It redefines visual document retrieval by treating embedding creation as a **sequential generation process**, mapping text queries and visual document pages into a compact and expressive latent multi-vector space. |
|
|
|
|
|
The model employs a novel **CausalEmbed** training paradigm, which fine-tunes the MLLM to synthesize latent representations token-by-token. This approach effectively distills dense visual information into a small set of highly semantic vectors, overcoming the storage bottlenecks of traditional multi-vector models while maintaining superior retrieval accuracy. On benchmarks like Vidore, **CausalPali** achieves significant performance gains over the original ColPali baseline while using a fraction of the storage. |
|
|
|
|
|
## Key Features |
|
|
|
|
|
**Model size**: 3 billion parameters. |
|
|
**Auto-regressive Generation**: Unlike parallel patch-level encoding, it generates multi-vector embeddings sequentially in a latent space, enabling better modeling of dependencies between visual elements. |
|
|
**High Efficiency**: Reduces the number of visual tokens by **30-155x** compared to standard ColPali (using only ~64 tokens vs. thousands), significantly lowering storage and memory overhead. |
|
|
**Test-time Scaling**: Supports dynamic adjustment of the number of generated embedding vectors during inference, allowing users to trade off between retrieval speed and precision without retraining. |
|
|
**Superior Performance**: Achieves a **14.6% performance uplift** on retrieval tasks compared to the full-resource ColPali baseline, demonstrating that compact, auto-regressively generated embeddings can capture richer semantics. |
|
|
**Context-aware Compression**: Incorporates iterative margin loss and progressive refinement during training to ensure each generated vector adds maximal marginal information. |
|
|
|
|
|
|
|
|
## Usage |
|
|
|
|
|
**Requirements** |
|
|
See in [CausalEmbed Repo](https://github.com/Z1zs/Causal-Embed). |
|
|
|
|
|
**Basic Usage** |
|
|
```bash |
|
|
pip install git+https://github.com/Z1zs/Causal-Embed |
|
|
``` |
|
|
|
|
|
```python |
|
|
import torch |
|
|
from PIL import Image |
|
|
from colpali_engine.models import ColPaliProcessor, CausalPali |
|
|
|
|
|
dtoken_num, qtoken_num = 64, 16 |
|
|
images = [Image.new("RGB", (32, 32), color="white"), Image.new("RGB", (16, 16), color="black")] |
|
|
queries = ["Is attention really all you need?", "What is the amount of bananas farmed in Salvador?"] |
|
|
processor = ColPaliProcessor.from_pretrained( |
|
|
"Z1zs/CausalPali", |
|
|
max_num_visual_tokens=768, |
|
|
fix_mistral_regex=True |
|
|
) |
|
|
processor.tokenizer.pad_token = processor.tokenizer.eos_token |
|
|
processor.tokenizer.add_tokens("<|latent|>") |
|
|
latent_id = processor.tokenizer.convert_tokens_to_ids("<|latent|>") |
|
|
|
|
|
model = CausalPali.from_pretrained( |
|
|
"Z1zs/CausalPali", |
|
|
torch_dtype=torch.bfloat16, |
|
|
doc_token_num=dtoken_num, # expected doc token num |
|
|
query_token_num=qtoken_num, # expected query token num |
|
|
attn_implementation="flash_attention_2", |
|
|
) |
|
|
model.latent_token_id = latent_id |
|
|
model.to(device).eval() |
|
|
|
|
|
batch_images = processor.process_images(images).to(model.device) |
|
|
batch_queries = processor.process_queries(queries).to(model.device) |
|
|
|
|
|
# Forward pass |
|
|
with torch.no_grad(): |
|
|
image_embeddings = model(**batch_images) |
|
|
query_embeddings = model(**batch_queries) |
|
|
|
|
|
scores = processor.score_multi_vector(query_embeddings, image_embeddings) |
|
|
``` |
|
|
|
|
|
## Model Performance |
|
|
|
|
|
### Vidore v1 + v2 (NDCG@5) |
|
|
|
|
|
| Model | Token Count | Vidore V1 | Vidore V2 | |
|
|
|--------------------------------------------|------|---------|---------| |
|
|
| **CausalQwen** | 64 | 81.1 | 51.6 | |
|
|
| **CausalPali** | 64 | 75.0 | 45.4 | |
|
|
|
|
|
|
|
|
### Vidore v3 (NDCG@5) |
|
|
|
|
|
| Model | Token Count | PUB AVG | |
|
|
|--------------------------------------------|------|---------| |
|
|
| **CausalQwen** | 64 | 42.6 | |
|
|
| **CausalPali** | 64 | 33.6 | |
|
|
|
|
|
|
|
|
## Citation |
|
|
|
|
|
If you use this model in your work, please cite: |
|
|
|
|
|
```bibtex |
|
|
@misc{huo2026causalembed, |
|
|
title={CausalEmbed: Auto-Regressive Multi-Vector Generation in Latent Space for Visual Document Embedding}, |
|
|
author={Jiahao Huo and Yu Huang and Yibo Yan and Ye Pan and Yi Cao and Mingdong Ou and Philip S. Yu and Xuming Hu}, |
|
|
year={2026}, |
|
|
eprint={2601.21262}, |
|
|
archivePrefix={arXiv}, |
|
|
primaryClass={cs.CL}, |
|
|
url={https://arxiv.org/abs/2601.21262}, |
|
|
} |
|
|
``` |
