🚀 打破模态壁垒:使用多模态大语言模型进行通用嵌入学习
本项目提出了文本判别式知识蒸馏和硬负样本增强指令调优等方法,提升了多模态大语言模型(MLLM)的嵌入能力,在MMEB排行榜上取得了优异成绩。
🚀 快速开始
克隆仓库并创建环境
git clone https://github.com/deepglint/UniME.git
cd UniME
conda create -n uniME python=3.10 -y
conda activate uniME
pip install -r requirements.txt
代码使用示例
import torch
from PIL import Image
from torch.nn import functional as F
from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
base_model_path="DeepGlint-AI/UniME-LLaVA-1.6-7B"
img_prompt = "[INST] <image>\nSummary above image in one word: [/INST]"
text_prompt = "[INST] <sent>\nSummary above sentence in one word: [/INST]"
text = "A man is crossing the street with a red car parked nearby."
image_path = "figures/demo.png"
input_texts = text_prompt.replace('<sent>', text)
input_image_prompt = img_prompt
input_image = [Image.open(image_path)]
transform = LlavaNextProcessor.from_pretrained(base_model_path)
model = LlavaNextForConditionalGeneration.from_pretrained(base_model_path, device_map="cuda", torch_dtype=torch.float16, low_cpu_mem_usage=True)
transform.tokenizer.padding_side = "left"
transform.tokenizer.padding = True
inputs_text = transform(text=input_texts,
images=None,
return_tensors="pt",
padding=True)
for key in inputs_text: inputs_text[key] = inputs_text[key].to("cuda")
inputs_image = transform(text=input_image_prompt,
images=input_image,
return_tensors="pt",
padding=True).to("cuda")
with torch.no_grad():
emb_text = model(**inputs_text, output_hidden_states=True, return_dict=True).hidden_states[-1][:, -1, :]
emb_image = model(**inputs_image, output_hidden_states=True, return_dict=True).hidden_states[-1][:, -1, :]
emb_text = F.normalize(emb_text, dim=-1)
emb_image = F.normalize(emb_image, dim=-1)
Score = emb_image @ emb_text.T
print("Score: ", Score)
✨ 主要特性
文本判别式知识蒸馏
为了增强MLLM的嵌入能力,我们提出了文本判别式知识蒸馏方法。训练过程包括解耦MLLM的大语言模型(LLM)组件,使用提示“用一个词总结上述句子”处理文本,然后通过批量相似度分布上的KL散度对齐学生模型(MLLM)和教师模型(NV-Embed V2)的嵌入。值得注意的是,在此过程中仅微调LLM组件,而其他所有参数保持冻结。
硬负样本增强指令调优
之后,我们提出了硬负样本增强指令调优方法,通过提高视觉敏感度、加强跨模态对齐和提升指令遵循能力来增强多模态系统。其核心有两个关键创新:使用相似度阈值的假负样本过滤机制,以消除误导性样本;以及自动硬负样本采样策略,选择前k个相似但不匹配的示例以增加训练难度。
📚 详细文档
项目信息
属性 |
详情 |
模型类型 |
图像文本到文本 |
基础模型 |
llava-hf/llava-v1.6-mistral-7b-hf |
训练数据集 |
TIGER-Lab/MMEB-train |
评估指标 |
召回率 |
库名称 |
transformers |
许可证 |
MIT |
项目链接
🏡 项目主页 | 📄 论文 | 💻 Github
项目成果
多样化检索
MMEB排行榜
UniME在使用336×336图像分辨率进行训练时,在MMEB排行榜上名列前茅。(截图于2025年5月6日UTC+8 08:00捕获。)
引用信息
如果您发现本仓库有用,请使用以下BibTeX条目进行引用。
@misc{gu2025breakingmodalitybarrieruniversal,
title={Breaking the Modality Barrier: Universal Embedding Learning with Multimodal LLMs},
author={Tiancheng Gu and Kaicheng Yang and Ziyong Feng and Xingjun Wang and Yanzhao Zhang and Dingkun Long and Yingda Chen and Weidong Cai and Jiankang Deng},
year={2025},
eprint={2504.17432},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2504.17432},
}