UEmbed:统一稀疏与稠密的多模态嵌入模型 UEmbed: Unified Sparse and Dense Multimodal Embeddings
在解码器模型中一次前向同时生成稀疏与稠密多模态向量
前置知识
Learned Sparse Retrieval (LSR)
学习型稀疏检索用神经网络替代 BM25 的固定词频统计,既能为已出现的词重赋权重,又能扩展出原文中没有的语义相关词。典型代表 SPLADE 对每个 token 的隐状态做词表投影,再用 max pooling 与饱和激活 $\log(1+\text{ReLU}(\cdot))$ 聚合,得到稀疏向量,可直接用倒排索引服务。
UEmbed 的核心就是把 LSR 从双向编码器搬到了因果解码器上,理解 SPLADE 的 max-pooling 与词表投影机制,才能明白为什么因果注意力会让传统池化失效。
因果注意力 (Causal Attention)
解码器模型采用单向掩码,每个位置只能看到自己和之前的 token,看不到未来。这与 BERT 的双向注意力相反。这种掩码使得每个 token 的隐状态只编码其前缀上下文,导致基于全序列 max pooling 的稀疏表示无法捕获全局信息。
本文要解决的首要难题就是因果注意力带来的信息瓶颈——单一 token 无法投影到超大词表空间,UEmbed 用 N 个特殊 token 加词表划分来绕过它。
InfoNCE 损失
对比学习常用损失:给定查询 q、正样本 d+ 和一批负样本,通过 $\exp(\text{sim}(q,d)/\tau)$ 形式最大化正样本对的相似度概率。稠密检索用余弦相似度,稀疏检索用内积,$\tau$ 为温度超参。
UEmbed 的统一目标函数同时优化稠密和稀疏两路 InfoNCE,再加 FLOPS 正则,理解 InfoNCE 才能看懂联合训练为何不会产生负迁移。
FLOPS 正则
FLOPS 正则器惩罚查询和文档侧激活词项权重的平方均值,迫使模型激活更少的词项,从而产生真正稀疏的表示,避免向量退化成稠密。系数 $\alpha_q$、$\alpha_d$ 控制稀疏强度。
没有 FLOPS 正则,SPLADE 式的词表扩展会让向量过于稠密,UEmbed 沿用了这一正则来平衡检索质量与稀疏度。
研究动机
现代搜索系统(网页检索、RAG)底层依赖稀疏词法检索与稠密语义检索。已有的 Learned Sparse Retrieval(LSR,如 SPLADE)虽然超越了 BM25 的精确匹配,但存在三大瓶颈。第一,架构受限:现有 LSR 几乎都建立在 BERT 这类双向编码器上,因为因果注意力的单向掩码使传统 max pooling 无法捕获全局上下文,要把 LLM 用于稀疏检索通常得把因果骨干改成双向、或采用专门训练流程,既增加训练成本又破坏与 vLLM 等优化推理框架的兼容性。第二,模态受限:绝大多数 LSR 只处理文本,少数多模态方法要引入额外的跨模态模块而非从单一骨干直接产出稀疏向量,难以扩展到新模态。第三,实用价值未被充分评估:现有稀疏方法只在 COCO 等传统基准上评测,对效率、服务化、Agent 搜索等实用维度的优势研究不足。
本文的目标是本文旨在构建一个能在单一因果解码器骨干内、一次前向同时产出稀疏词法向量与稠密语义向量的统一嵌入模型,并原生支持文本、图像、视频、视觉文档等多模态输入。具体目标包括:在不破坏与 vLLM 等自回归推理栈兼容性的前提下,把稀疏检索的因果化做扎实;在 MMEB-v2 多模态基准上让稀疏模式逼近稠密模式的表现(差距控制在 1 分内);在 BEIR 文本基准上保持与强稀疏/稠密基线相当的竞争力;并在混合打分、服务效率、Agent 搜索三个实用维度上验证其价值。作者还希望以 2B、4B、9B 三个尺度开源模型,确立公开数据训练模型中的新 SOTA。
与已有工作不同的是,本文的独特切入角度是直面因果解码器的信息瓶颈。以往要么把双向池化硬塞进解码器、要么借助额外跨模态模块,UEmbed 则另辟蹊径:在输入末尾追加 N 个可学习的特殊 token,每个特殊 token 在因果注意力下能看到全部前文从而充当全局摘要,再用 k-means 把词表划分成 N 个不相交子集、每个特殊 token 只负责自己那份子集的稀疏权重预测,最后拼接成完整稀疏向量。这种分区投影从根本上绕过了单 token 投影到 $|V|$ 维词表的容量瓶颈,同时让稀疏成为多模态骨干的原生产物,而非外加模块。这是首个在 MMEB-v2 上评测的稀疏模型。
核心方法
直觉上,UEmbed 把稀疏检索难题重新表述为多任务多头的预测问题:与其让一个 token 扛下整个词表,不如让一群 token 分工。技术路线基于 Qwen3.5 家族的 2B/4B/9B 解码器骨干,单次前向即同时输出稠密和稀疏向量。稠密向量取自特殊 token 之前那个 EOS token 的隐状态;仅做稠密检索时可省略特殊 token,零额外开销。稀疏向量则由 N=16 个可学习特殊 token 协同产出,每个 token 经各自的线性头预测词表子集上的权重。训练目标融合两路 InfoNCE 损失与 FLOPS 正则:$\mathcal{L} = \mathcal{L}_{InfoNCE}^{dense} + \lambda \mathcal{L}_{InfoNCE}^{sparse} + \alpha_q \mathcal{L}_{FLOPS}^q + \alpha_d \mathcal{L}_{FLOPS}^d$。训练数据来自 Echo-embedding、MLDR、MMEB 共 394 万样本,并用 Qwen3-VL-Embedding-8B 为多模态数据挖掘难负样本。
核心创新点是分区稀疏头。面对因果注意力下单 token 无法有效投影到 18 万维词表的信息瓶颈,UEmbed 追加 N 个特殊 token 并把词表按 k-means 语义聚类划分成 N 个不相交子集。对第 k 个特殊 token $\langle s_k\rangle$,其隐状态 $h_{s_k}$ 经专属线性头预测子集 $V_k$ 上每个词项的权重:$w_t^{(k)} = \log(1 + \text{ReLU}(W_t^{(k)\top} h_{s_k} + b_t^{(k)}))$。N 个子向量拼接成完整稀疏向量 $w = [w_t^{(1)}]_{t\in V_1} \oplus \cdots \oplus [w_t^{(N)}]_{t\in V_N}$。这与传统 SPLADE 的全序列 max-pooling 本质不同:它显式鼓励每个特殊 token 成为某个语义子空间的专家,而非让每个位置都参与所有词项的竞争。语义聚类划分优于随机划分与最大距离划分。
方法步骤详情
完整流程分四步。第一步词表压缩:用 NLTK 做去重音、小写化、空白折叠,把同规范形式的 token 合并(如 Hello/HELLO/héllò 合为 hello),打分时取同族最大权重,词表从 248320 压缩到 184016。第二步分区:对压缩后词嵌入做 k-means 聚类,形成 N=16 个约等大的不相交子集。第三步前向与预测:输入序列末尾追加 16 个特殊 token,每个 token 的最终隐状态经子集专属线性头计算权重,激活函数沿用 SPLADE 的 $\log(1+\text{ReLU}(\cdot))$;稠密向量取特殊 token 前 EOS 的隐状态做 last pooling。第四步联合训练:稠密路用余弦相似度的 InfoNCE(温度 $\tau$),稀疏路用内积的 InfoNCE 并用更大的稀疏温度 $\tau_s=32$(因为内积动态范围远大于 $[-1,1]$ 的余弦,共享温度会过度锐化稀疏 softmax),再加查询侧与文档侧 FLOPS 正则。超参扫描显示 $N=16$ 最优,$N=32$ 因子集过小而掉点。
技术新颖性
技术新颖性体现在三处。其一,首次在原生因果注意力下实现高质量稀疏检索,无需把骨干改成双向或设计专门课程,保留了与 vLLM、倒排索引的兼容性,解决了既有方法破坏推理框架生态的痛点。其二,分区稀疏头以可学习特殊 token 加语义聚类词表划分,将单 token 容量瓶颈转化为多头分工,且消融证明语义聚类明显优于随机/最大距离划分,说明分组语义相关 token 能让每个 token 当好软主题专家。其三,在统一骨干内首次把稀疏检索扩展到多模态(图像、视频、视觉文档),稀疏与稠密两路在同一前向内产生、几乎无性能折损(9B 模型差距仅 0.8 分),这在此前 LSR 工作中从未实现。其统一目标函数也经验性避免了多任务负迁移。
实验结果
多模态基准 MMEB-v2 上,UEmbed-9B 稠密 71.8、稀疏 71.0,三个尺度稀疏与稠密差距始终≤1分(4B 70.4/69.7,2B 66.5/65.5)。在公开数据训练模型中稠密领先,超过 RzenEmbed-V2-7B(71.1)、Ops-MM-Embed-7B(67.1),2B 稠密甚至超 7B 的 UniME(64.1);4B 稀疏(69.7)反超 7B 稠密 Ops-MM-Embed,证明稀疏多模态检索可行,视觉文档上稀疏几乎不掉点(9B 79.2 vs 79.1)。文本 BEIR(nDCG@10)上 9B 稠密 56.3、4B 56.0,超过 Qwen3-VL-Embedding-8B(55.5)与 GME-7B(53.5),9B 稀疏 55.2 追平专门稀疏模型 Echo-Mistral-SPLADE 且兼具多模态能力。消融中 UEmbed-2B 较同骨干双向 SPLADE 基线稠密 +3.2、稀疏 +2.1。Agent 搜索 BrowseComp-Plus 上稀疏平均轮数更少(9B 31.05 vs 稠密 33.68),2B 召回反超稠密。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MMEB-v2 多模态嵌入(78 数据集综合) | 平均得分 | UEmbed-9B 稠密 71.8、稀疏 71.0 | Qwen3-VL-Embedding-8B 77.8(私有大数据训练);RzenEmbed-V2-7B 71.1 | 在公开数据训练模型中稠密领先;稀疏为该基准首次报告,与自家稠密差距≤1分,4B稀疏超过7B稠密 Ops-MM-Embed |
| BEIR 文本检索(9 数据集) | nDCG@10 平均 | UEmbed-9B 稠密 56.3、稀疏 55.2 | Qwen3-VL-Embedding-8B 稠密 55.5;Echo-Mistral-SPLADE 稀疏 55.2;GME-7B 稠密 53.5 | 稠密超过所有多模态基线,稀疏追平专门稀疏模型同时兼具多模态能力 |
| MMEB-v1 图像子集(消融) | CLS/QA/RET/GRD 平均 | UEmbed-2B 稠密 64.5、稀疏 63.4 | 同骨干双向 SPLADE 基线 61.3 | 稠密 +3.2、稀疏 +2.1,IMG-QA 上稠密 +8.2、稀疏 +6.0 |
| BrowseComp-Plus Agent 搜索 | 平均搜索轮数 / 召回率 | UEmbed-9B 稀疏 31.05 轮、召回 64.33% | UEmbed-9B 稠密 33.68 轮、召回 64.72% | 稀疏模式搜索轮数更少、成本更低,2B 上召回 57.04% 反超稠密 53.47% |
局限与改进
作者承认三类局限。第一是语言与文化偏差:训练语料严重偏向英语和中文,稀疏头激活极少触及其他语言,跨语言泛化弱。第二是词表稳定性与伪影:现代 LLM 词表庞大,偶尔激活非标准 token(如 _alt、ансам),影响表示稳定性,需要进一步词表裁剪或过滤。第三是模态间性能差距:视频域稀疏与稠密差距略大(9B 视频 59.0 稠密 vs 57.5 稀疏),作者归因于视频时空信息密度高,扁平池化到稀疏向量会遇到容量瓶颈。我的观察是:UEmbed-9B 稠密在 MMEB-v2 仍显著落后于用私有大数据训练的 Qwen3-VL-Embedding-8B(71.8 vs 77.8),且仅凭公开数据难以追赶;此外特殊 token 数 N 是固定超参,对不同模态/序列长度未必最优,混合打分在图像视频上几乎无增益(Table 5),说明稀疏在自然图像上价值有限。
独立分析的弱点
弱点一:训练数据规模与多语言覆盖不足。3.94M 样本且偏英中双语,导致其他语言稀疏激活稀疏。改进方向是引入大规模多语言检索对,并对低资源语言做针对性难负样本挖掘。弱点二:固定 N=16 的分区头对短文本/长文档、不同模态可能非最优,且 N 增大时对比序列变长、子集变小而掉点。改进方向是让 N 或分区策略自适应输入,或引入层次化分区。弱点三:稀疏伪影 token 暴露词表管理问题,可能干扰下游倒排索引质量。改进方向是结合词表压缩、激活过滤或后处理校准。弱点四:混合打分在图像、视频等模态上几乎无增益(IMG +0.1、VID 0.0),说明稀疏对自然视觉信号贡献有限,视频时空信息被扁平池化。改进方向是设计时空感知的稀疏头或引入模态专属子空间。弱点五:与私有数据训练的 Qwen3-VL-Embedding 仍有 6 分差距,纯公开数据路线的天花板待突破。
未来方向
作者提出的方向包括:扩展到更广泛语言的多语言稀疏检索、改进词表稳定性与伪影过滤、针对视频等高信息密度模态设计时空感知的稀疏表示。基于本成果还可延伸若干方向。一是把分区稀疏头思想迁移到更长上下文与 Agent 工具调用场景,进一步降低检索成本(BrowseComp 已显示潜力)。二是探索稀疏向量的可解释性用于偏见审计,作者在伦理部分提到稀疏激活词能帮助审查模型偏差,可发展为可解释检索调试工具。三是与重排器、混合检索系统深度集成,验证混合打分在特定垂类(如视觉文档 OCR、视频问答)上的更大增益。四是把统一目标扩展到更多检索范式(如交互式检索、列存检索)。五是研究 N 与温度 $\tau_s$ 的自适应学习,摆脱固定超参。
复现评估
复现性整体较好。模型权重已开源,提供 2B、4B、9B 三个尺度;项目主页 https://alibaba-nlp.github.io/UEmbed 给出资源。训练数据全部来自公开数据集(Echo-embedding、MLDR、MMEB 训练集,共 3.94M),难负样本用开源的 Qwen3-VL-Embedding-8B 挖掘,方法细节如词表压缩(NLTK,248320→184016)、k-means 分区、N=16、稀疏温度 $\tau_s=32$、统一损失各项系数、FLOPS 正则、骨干选用 Qwen3.5 家族都有交代,并配有多张消融与超参扫描图。困难在于:训练 9B 级别嵌入模型算力门槛高(作者未明确披露 GPU 配置与训练时长),需要大规模对比学习训练栈与难负样本挖掘流水线;MMEB-v2 含 78 个数据集、BEIR 含 9 个,完整复现评测工程量也较大。总体方法与配方可复现,但全尺度复训练对多数团队不现实。
论文图表
跨文本、图像、视频、视觉文档四个模态可视化 top-10 激活稀疏 token。例如从天际线推断地点为新加坡、把火箭识别为 SpaceX,但也暴露伪影 token(_alt)和多语言偏向(仅英中)。
直观展示稀疏表示的跨模态语义可解释性,同时呈现作者承认的伪影与多语言局限,连接方法与 limitations。