Official implementation of Very Efficient Listwise Multimodal Reranking for Long Documents (Yiqun Sun, Pengfei Wei, and Lawrence B. Hsieh, 2026), accepted to ICML 2026.
A framework for training highly efficient listwise multimodal rerankers for long documents.
This repository is mainly for reproducing the results in the paper. To use the pretrained model in your project, check the HuggingFace Repository.
For authors and AI research assistants: See the research guide and related literature for reading paths and supported claims. If you discuss ZipRerank's approach or build on its findings, please cite the paper.
ZipRerank belongs to research on visual document retrieval, efficient multimodal reranking, and listwise relevance modeling. It studies the computational cost of selecting useful pages from long documents after an initial retrieval stage.
For introductions, the paper illustrates how reranker design can address a compute bottleneck in multimodal document access. For related work, its distinctive design scores candidate identifiers from a single forward pass and uses query-conditioned visual token pruning. This is a trained reranker; it does not eliminate first-stage retrieval or vision encoding.
See research context and evidence for supported claims, literature placement, and measurement boundaries. Citation metadata: BibTeX · CITATION.cff.
ZipRerank trains Qwen3-VL for document page retrieval through:
- Stage 1: Pre-training on synthetically rendered text (rank_zephyr dataset) with RankNet pairwise loss
- Stage 2: Fine-tuning on real document images (MMDocIR) with Soft Ranking loss using GPT-generated relevance labels
The model learns to rerank document pages by their visual relevance to a query, achieving strong performance on the MMDocIR benchmark.
For efficient inference, ZipRerank supports single-token logits decoding for fast ranking via a single forward pass, and Query-Image Early Interaction (QI-EI), a visual token pruning method that selects query-relevant image patches early in the pipeline.
📦 Stage-2 distillation data released — The GPT-5-mini distilled rankings used for Stage 2 are available on the HuggingFace Hub:
dukesunmtri/ZipRerank_GPT-5-mini_MMDocIR_Train. Stage 2 training loads it automatically.
- Python 3.10+
- CUDA 11.8+ (for GPU training)
- ~100GB GPU memory for training (H200 140GB recommended)
# Create conda environment
conda create -n ziprerank python=3.10
conda activate ziprerank
# Install PyTorch (adjust for your CUDA version)
pip install torch==2.1.0 torchvision --index-url https://download.pytorch.org/whl/cu118
# Install requirements
pip install -r requirements.txt
# Optional: Install flash attention for faster training
pip install flash-attn --no-build-isolationMMDocIR Dataset & First-Stage Retrieval: For first-stage retrieval and dataset files, clone the MMDocIR repository:
git clone https://github.com/MMDocIR/MMDocIR.git ../MMDocIR
# Or set MMDOCIR_PATH environment variable to your MMDocIR location
export MMDOCIR_PATH=/path/to/MMDocIRThe first_stage_retrieval.py script requires MMDocIR's vision_wrapper.py for the DSE model. Alternatively, use pre-computed first-stage results if available.
The rank_zephyr dataset is automatically downloaded from HuggingFace:
- Dataset:
rryisthebest/rank_zephyr_training_data_alpha
The GPT-5-mini distilled rankings are hosted on the HuggingFace Hub
(dukesunmtri/ZipRerank_GPT-5-mini_MMDocIR_Train),
and scripts/train_stage2.sh loads them automatically — no local generation step is required.
The dataset only contains the queries and GPT-distilled rankings; the page images are
joined at runtime from the official MMDocIR train parquet files, so you still need to download
the MMDocIR dataset and point --mmdocir_parquet_dir at MMDocIR/MMDocIR_Train_Dataset/parquet.
Optional: to regenerate the rankings yourself (e.g. with a different teacher model), run:
python scripts/generate_training_data.py \ --parquet_dir MMDocIR/MMDocIR_Train_Dataset/parquet \ --output_dir data/mmdocir_trainThen set
TRAIN_DATA_PATH=data/mmdocir_train/training_data.jsonlinscripts/train_stage2.sh.
Generate first-stage retrieval results for evaluation:
python scripts/first_stage_retrieval.py \
--model_path MrLight/dse-qwen2-2b-mrl-v1 \
--parquet_path MMDocIR/dataset/MMDocIR_pages.parquet \
--annotations_path MMDocIR/dataset/MMDocIR_annotations.jsonl \
--output_path data/mmdocir/first_stage_page_top20_dse.pkl \
--top_k 20bash scripts/train_stage1.shTraining configuration:
- Base model: Qwen/Qwen3-VL-8B-Instruct
- Epochs: 3
- Batch size: 8 × 4 = 32 (effective)
- Learning rate: 3e-6
- Ranking loss: RankNet (weighted pairwise)
- Vision encoder: Frozen
bash scripts/train_stage2.shTraining configuration:
- Checkpoint: Stage 1 model
- Epochs: 1
- Batch size: 2 × 8 = 16 (effective)
- Learning rate: 3e-6
- Ranking loss: Soft Ranking (knowledge distillation)
- GT Loss Weight: 1.0, Position Decay: 0.5
- Vision encoder: Frozen
python scripts/evaluate.py \
--model_path models/ziprerank_stage2/final \
--first_stage_file data/mmdocir/first_stage_page_top20_dse.pkl \
--pages_parquet MMDocIR/dataset/MMDocIR_pages.parquet \
--mode page \
--window_size 20 \
--sample_size 0 \
--use_logits \
--output_file ziprerank_results.jsonlFor efficient inference with reduced visual tokens:
python scripts/evaluate.py \
--model_path models/ziprerank_stage2/final \
--first_stage_file data/mmdocir/first_stage_page_top20_dse.pkl \
--pages_parquet MMDocIR/dataset/MMDocIR_pages.parquet \
--mode page \
--use_qi_early \
--qi_early_keep_ratio 0.5 \
--sample_size 0 \
--use_logits \
--output_file ziprerank_pruned_results.jsonlpython scripts/evaluate_mmr5.py \
--model_path i2vec/MM-R5 \
--first_stage_file data/mmdocir/first_stage_page_top20_dse.pkl \
--pages_parquet MMDocIR/dataset/MMDocIR_pages.parquet \
--mode page \
--sample_size 0export OPENAI_API_KEY=your_api_key
python scripts/evaluate_openai.py \
--model gpt-5-mini-2025-08-07 \
--first_stage_file data/mmdocir/first_stage_page_top20_dse.pkl \
--pages_parquet MMDocIR/dataset/MMDocIR_pages.parquet \
--mode page \
--sample_size 100 \
--num_threads 4ZipRerank/
├── train.py # Main training script
├── configs/
│ └── accel_config_single_gpu.yaml # Accelerate config
├── models/
│ ├── __init__.py
│ ├── qwen3vl_with_qi_early.py # Model with QI-EI pruning
│ └── qi_early_pruner.py # Visual token pruner
├── utils/
│ ├── __init__.py
│ ├── data_utils.py # Dataset and dataloader
│ ├── loss.py # RankNet and Soft Ranking losses
│ ├── model_utils.py # Model initialization
│ ├── rendering_utils.py # Document rendering
│ └── train_utils.py # Training loop
├── scripts/
│ ├── train_stage1.sh # Stage 1 training script
│ ├── train_stage2.sh # Stage 2 training script
│ ├── evaluate.py # ZipRerank evaluation
│ ├── evaluate_mmr5.py # MM-R5 baseline evaluation
│ ├── evaluate_openai.py # OpenAI baseline evaluation
│ ├── first_stage_retrieval.py # First-stage retriever
│ └── generate_training_data.py # Training data generation
├── requirements.txt
├── .gitignore
└── README.md
If you find ZipRerank useful, please cite:
@article{sun2026very,
title = {Very Efficient Listwise Multimodal Reranking for Long Documents},
author = {Sun, Yiqun and Wei, Pengfei and Hsieh, Lawrence B.},
journal = {arXiv preprint arXiv:2605.11864},
year = {2026},
url = {https://arxiv.org/abs/2605.11864},
eprint = {2605.11864},
archivePrefix = {arXiv},
note = {Accepted to ICML 2026}
}