LaME: Learning to Think in Latent Space for Multimodal Embedding

📄 This repository contains the official implementation of LaME (Learning to Think in Latent Space for Multimodal Embedding via Information Bottleneck).

LaME augments Qwen-VL with learnable [REASON] tokens and a latent decoder supervision, jointly optimizing generation and embedding through an information bottleneck. Without bells and whistles, LaME achieves state-of-the-art multimodal retrieval performance on MMEB-v2 (image / video / visual-document / full aggregate) and MRMR.

If you find LaME useful, please give us a ⭐!

---

---

🚀 Updates

• [2026.06.30] 🔥 LaME code released! Training & evaluation pipelines for 2B / 3B / 7B / 8B Qwen-VL backbones are ready.
• [2026.06.25] Initial release of openlame package and training configs (Stage-1 & Stage-2).

---

Model Zoo 🏆

Supported Configurations

Backbone	Size	Params	Stage-1 Config	Stage-2 Config	Projection Dim
Qwen2-VL	2B	~2.0B	lame_s1.yaml	lame_s2.yaml	3584
Qwen2.5-VL	3B	~3.0B	lame_s1.yaml	lame_s2.yaml	2048
Qwen2-VL	7B	~7.0B	lame_s1.yaml	lame_s2.yaml	3584
Qwen3-VL	8B	~8.0B	lame_s1.yaml	lame_s2.yaml	4096

Note:

• VLM=qwen2vl|qwen3vl|qwen25vl controls the target backbone via environment variable in training scripts.
• All models are trained with 8 learnable [REASON] tokens and decoder projection.
• The LatentDecoder uses Qwen3-0.6B as the small LLM backbone.

---

Quick Start 🏃

Requirements

# Python >= 3.9, PyTorch >= 2.0 with CUDA
pip install -e .
# or with training extras
pip install -e ".[train]"

💡 Install the matching CUDA build of PyTorch and FlashAttention before multi-GPU training.

---

Usage 🛠️

Two-Stage Training

LaME adopts a two-stage training paradigm:

• Stage 1: Freeze the entire MLLM backbone, train only the projector + embedding head + learnable [REASON] tokens. The LatentDecoder decodes CoT text from the reason tokens.
• Stage 2: Unfreeze MLLM for joint contrastive + decoder-supervised training. The LatentDecoder decodes answer text.

2B / 3B Qwen-VL backbone

MODEL_ROOT=/path/to/models \
CHECKPOINT_DIR=/path/to/checkpoints \
NPROC_PER_NODE=8 \
bash startup/train_2b.sh --all       # stage1 + stage2 + eval

Select backbone via VLM (default: qwen2vl):

VLM=qwen25vl \
MODEL_ROOT=/path/to/models \
CHECKPOINT_DIR=/path/to/checkpoints \
NPROC_PER_NODE=8 \
bash startup/train_2b.sh --all

7B / 8B Qwen-VL backbone

VLM=qwen2vl \
MODEL_ROOT=/path/to/models \
CHECKPOINT_DIR=/path/to/checkpoints \
NPROC_PER_NODE=8 \
bash startup/train_7b.sh --stage1    # then run again with --stage2

Override hyperparameters via YAML configs (configs/train/lame_s1.yaml, configs/train/lame_s2.yaml) or environment variables (e.g. NUM_REASON_TOKENS, BATCH_SIZE, LEARNING_RATE).

Evaluation

startup/eval.sh auto-detects GPUs and shards encoding across ranks.

# MMEB-v2 image bench — [EMBED] pooling
bash startup/eval.sh /path/to/checkpoint image

Available benchmarks: image, video, visdoc, mmeb, mrmr.

Results are saved under eval_results/<model-name>/.

---

Citation 📚

If you find our work useful, please consider citing:

@article{wu2026lame,
  title   = {LaME: Learning to Think in Latent Space for Multimodal Embedding via Information Bottleneck},
  author  = {Wu, Peixi and Yang, Biao and Ma, Feipeng and Chai, Bosong and Lin, Bo and Yuan, Wei and Yang, Fan and Gao, Tingting and Li, Hebei and Sun, Xiaoyan},
  journal = {arXiv preprint arXiv:2606.13061},
  year    = {2026}
}

---

Acknowledgement

Built upon the awesome Qwen-VL and Transformers ecosystem. Thanks to the open-source community!

✨ Feel free to contribute and reach out if you have any questions! ✨