Skip to content

README.md

Latest commit

 

History

History
175 lines (130 loc) · 6.09 KB

README.md

File metadata and controls

175 lines (130 loc) · 6.09 KB
Raw

Embedder

The Embedder is a VLM-based bi-encoder that produces multimodal embeddings for retrieval. It supports dual embedding modes via special tokens:

  • <dis_emb> — Discriminative embedding (fast, no reasoning required)
  • <gen_emb> — Generative embedding (extracted after chain-of-thought reasoning)

Architecture

The Embedder wraps a Vision-Language Model (Qwen2-VL / Qwen2.5-VL / Qwen3-VL) and extracts embeddings via configurable pooling strategies:

Pooling Mode Description
last Use the last token's hidden state
eos Use the EOS token's hidden state
special_token Extract embedding at <dis_emb> or <gen_emb> token position

Query Format

When using special token mode with emb_mode='both':

[query] <dis_emb> [CoT instruction] <think>reasoning</think> Answer: [answer] <gen_emb>
  • <dis_emb> produces a discriminative embedding without reasoning
  • <gen_emb> produces a generative embedding after CoT reasoning

Loss Function

The DualContrastiveLoss computes:

  1. Gen_emb loss: contrastive loss on <gen_emb> embeddings (CoT samples only)
  2. Dis_emb loss: contrastive loss on <dis_emb> embeddings (all samples), weighted by dis_emb_loss_weight

Hard negatives are supported via [B, B+K] logits.

Training

Quick Start

cd Embedder

export MODEL_NAME="Qwen/Qwen3-VL-2B-Instruct"
export EXP_NAME="unime_embedder_qwen3vl_2b"

bash shell/train/train_v2-qwen3vl-2B_fullv2_multinode.sh

Multi-Node Training

export NNODES=2
export NPROC_PER_NODE=8
export MASTER_ADDR="<head-node-ip>"
export MASTER_PORT=2207
export NODE_RANK=0  # increment for each worker node

# Head node
bash shell/train/train_v2-qwen3vl-2B_fullv2_multinode.sh

# Worker nodes (same command, different NODE_RANK)
export NODE_RANK=1
bash shell/train/train_v2-qwen3vl-2B_fullv2_multinode.sh

Key Training Arguments

Argument Type Default Description
--model_name str — Model path or HuggingFace ID
--model_type str — qwen2_vl, qwen2_5_vl, qwen3_vl
--pooling str eos last, eos, special_token
--use_special_emb_tokens bool False Enable <dis_emb> / <gen_emb>
--emb_mode str both both, dis_only, gen_only
--dis_emb_loss_weight float 0.0 Dis. loss weight in dual mode
--num_hardneg int 0 Hard negatives per sample
--grad_cache bool False Enable GradCache
--gc_q_chunk_size int — Query chunk size for GradCache
--gc_p_chunk_size int — Passage chunk size for GradCache
--lora bool False Enable LoRA fine-tuning
--lora_r int 16 LoRA rank
--temperature float 0.02 Contrastive temperature
--normalize bool True L2-normalize embeddings
--interleave_batch_size int — Interleaved batch size for mixed datasets

See src/arguments.py for the complete list.

Training Data Configuration

Training data is configured via a YAML file (--dataset_config). See shell/train/train_alltasks_cot_v2_meged.yaml for the format:

DatasetName:
    dataset_parser: parser_name      # e.g., mmeb, vidore, visrag
    dataset_path: /path/to/data.parquet
    image_root: /path/to/images/
    weight: 1.0                       # sampling weight
    # ... parser-specific options

Evaluation

Supported Benchmarks

Config Benchmark Tasks
image.yaml MMEB Classification, VQA, Retrieval, Visual Grounding
video.yaml UVRB Video Classification, Retrieval, Moment Retrieval, Video QA
visdoc.yaml ViDoRe Document Retrieval
image_caption_retrieval.yaml Zero-shot Flickr30K, COCO2014, ShareGPT4V, Urban1k

Running Evaluation

cd Embedder

# Set your model checkpoint
export MODEL_NAME="path/to/checkpoint"

# Image evaluation
bash shell/eval/eval_vlm2vecv1_qwen3.sh

# Video evaluation
bash shell/eval/eval_vlm2vecv1_qwen3_video_uvrb.sh

# Image caption retrieval
bash shell/eval/eval_vlm2vecv1_qwen3_image_caption.sh

Metrics

Supported ranking metrics: Hit@1, Hit@5, Hit@10, NDCG@10, Precision@10, Recall@10, F1, MAP, MRR.

Results are saved as JSON in the --encode_output_path directory.

Special Token Initialization

When --use_special_emb_tokens True, the Embedder adds two special tokens to the vocabulary:

  • <dis_emb> — initialized from <|object_ref_start|> embedding
  • <gen_emb> — initialized from <|object_ref_end|> embedding

With LoRA, these embedding layers are saved in modules_to_save.

Directory Structure

Embedder/
├── train.py                    # Training entry point
├── eval.py                     # Main evaluation script
├── eval_image_caption.py       # Image-caption retrieval evaluation
├── eval_video_uvrb.py          # Video UVRB evaluation
├── convert_model.py            # Weight format conversion
├── src/
│   ├── arguments.py            # All training/eval arguments
│   ├── trainer.py              # GradCache trainer
│   ├── loss.py                 # Contrastive loss functions
│   ├── model/
│   │   ├── model.py            # MMEBModel (main bi-encoder)
│   │   ├── processor.py        # Backbone dispatcher
│   │   ├── special_tokens.py   # Special token management
│   │   └── vlm_backbone/       # Qwen2-VL, Qwen2.5-VL, Qwen3-VL backbones
│   ├── data/
│   │   ├── collator/           # Train & eval data collators
│   │   ├── dataset/            # Training dataset parsers
│   │   ├── eval_dataset/       # Evaluation dataset parsers
│   │   └── loader/             # Mixed dataset loading
│   ├── grad_cache/             # GradCache implementation
│   ├── prompt/                 # Prompt templates
│   └── utils/                  # Evaluation metrics, vision utils, etc.
└── shell/                      # Shell scripts and YAML configs
    ├── train/
    └── eval/