NV-Embed:用 latent attention 池化训 decoder-only embedding 模型

NV-Embed(NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models,2024,NVIDIA,ICLR 2025)的核心思路:从 Mistral 7B 直接训起,移除 causal attention mask,在最后一层 hidden state 后接一个 latent attention 层做池化,再用两阶段对比指令微调(先纯检索数据 + in-batch negative,再混合非检索数据 + 关掉 in-batch negative)。 在 MTEB 56 个任务上 NV-Embed-v1 平均 69.32,v2 进一步用 hard-negative mining、合成数据、example-based 多类标签把分数推到 72.31,分别在 2024 年 5 月和 8 月登顶 MTEB 榜首。

问题背景

LLM2Vec GritLM 已经验证了 decoder-only LLM 改双向后做 embedding 的可行性,但每条路线都引入了额外组件:LLM2Vec 加了 MNTP 预热阶段,GritLM 加了生成损失同步训练。NV-Embed 想问的是:能不能把这些都去掉,只靠"把 causal mask 拿掉 + 换个池化层 + 分两阶段训",就拿到更好的分数。

另一条要解决的是池化方式。decoder LLM 做 embedding 的两种主流池化都有问题。最后一个 EOS token 的表示受 recency bias 影响,过度依赖末尾 token。Mean pooling 把所有 token 平均,关键短语的信息被稀释。论文用 latent attention 层做了第三种选择。

方法

NV-Embed 架构:decoder LLM 输出做 Q,可学习的 latent array 做 K/V,交叉注意力后过 MLP 再 mean pooling

双向注意力。在对比学习阶段直接把 causal mask 去掉。论文没有像 LLM2Vec 那样先用 MNTP 预热,理由是经验上"直接去掉就 work",没必要加额外训练阶段。

Latent attention 层。decoder 最后一层 hidden state 记作 $Q \in \mathbb{R}^{l \times d}$,$l$ 是序列长度,$d$ 是 hidden 维度。引入一个可学习的"字典" $K = V \in \mathbb{R}^{r \times d}$,$r$ 是字典大小。做一次交叉注意力:

$$O = \text{softmax}(QK^T)V$$

输出 $O \in \mathbb{R}^{l \times d}$ 再过一个带 GELU 的两层 MLP,最后 mean pooling 得到句向量。NV-Embed 用 $r = 512$,8 个 attention head。论文把这个设计的思想归为"字典学习",区别于 Perceiver IO(后者用 latent array 做 Q 来压缩序列长度,NV-Embed 反过来,序列做 Q,latent array 做 K/V,输出长度仍是 $l$)。

两阶段指令微调。第一阶段只用检索数据(MSMARCO、HotpotQA、NQ、SQuAD、FEVER、MIRACL 等共 16 个),开启 in-batch negative 加 mined hard negative,让模型先把检索任务学好。第二阶段混入分类、聚类、STS 三类非检索数据,关闭 in-batch negative。关掉 in-batch negative 的理由是分类任务的标签集很小(极端情况是 True/False 二分类),mini-batch 里其他样本很可能是正例而不是负例,错把正例当 hard negative 会污染对比信号。

指令模板:Instruct: {task_definition} Query: q+。论文强调训练和评估时都把 instruction token 从输出 embedding 中 mask 掉(仍参与 self-attention 影响表示,只是不进入最终池化),文档侧不加 instruction。

实验结果

MTEB。NV-Embed-v1(只用公开数据,未加合成数据和 hard-negative mining)平均 69.32,超过当时的 E5-mistral-7b-instruct 的 66.63 和 SFR-Embedding 的 67.56。NV-Embed-v2 加上 positive-aware hard-negative mining、Mixtral-8x22B 生成的 12 万条合成数据、example-based 多类标签,平均 72.31,分别在 2024 年 5 月和 8 月两次登顶 MTEB。

子任务分布上,v2 在 15 个检索任务 nDCG@10 是 62.65,11 个聚类 V-Measure 58.46,12 个分类准确率 90.37,10 个 STS 84.31,均为榜单最高。

对比 SFR-Embedding-2R(同样基于 Mistral 7B 的 70.31 分模型):SFR 是从 E5-mistral-7b-instruct 继续微调,继承了 causal mask 和 last-token pooling 的架构。NV-Embed 直接从 Mistral 7B 训起,换上双向注意力 + latent attention 池化,论文将 v2 的优势主要归为这两点架构差异加 mixed-task batch 训练策略(SFR 用 task-homogeneous batch,作者认为这会带来"zigzag"梯度更新)。

AIR-Bench。MTEB 之外的 out-of-domain 检索集,NV-Embed 在 Long Doc 部分第一,QA 部分第二,说明性能不是过拟合到 MTEB 的训练 split。

关键消融

池化方式 × 注意力 mask。论文做了 4 × 2 的完整网格(first-stage 训练后):

PoolingBidirectCausal
EOS62.6860.06
Mean64.0062.32
Latent-attention64.1863.39
Self-attention63.2763.11

两点观察:

  1. 双向注意力对每种池化都更好。EOS pooling 受益最明显(+2.62),latent attention 受益最小(+0.79),说明 latent attention 一定程度上能缓解 causal mask 的限制
  2. self-attention 池化(在 LLM 输出后加一个普通 self-attention 层)几乎不带来增益,论文的解释是 LLM 内部已经堆了很多 self-attention,再加一层冗余。latent attention 的字典学习机制带来的"外部 key-value"才是真正不一样的归纳偏置

最终配置(latent-attention + bidirect + 二阶段训练 + 全量数据增强)拿到 72.31。

两阶段 vs 单阶段。直接单阶段开 in-batch negative 训练分类样本是灾难性的(MTEB 70.83),关掉 in-batch negative 后单阶段 71.94,两阶段 72.31。两阶段相对单阶段无 in-batch 的主要增益在检索子集(BEIR 62.65 vs 61.37)。反向的"先非检索后检索"两阶段反而掉到 71.85,说明顺序不能换。论文将其归因为检索任务难度最高,需要先用 in-batch negative 充分训练,再用非检索数据微调补齐其他任务。

数据增强逐项贡献。从 S0(无增强)到 S3(全部加上):

  • S0(baseline):MTEB 70.73,BEIR 59.22
  • S1(+ hard-negative mining):BEIR 61.52(+2.3)
  • S2(+ 更多公开检索集:HoVer、SciFact、NFCorpus、MIRACL、Mr.TyDi):BEIR 62.28
  • S3(+ 12 万条 Mixtral 合成数据):MTEB 72.31,BEIR 62.65

最大的单项增益来自 hard-negative mining。合成数据的边际收益较小(约 0.24 分),论文未给出更细的解释,可能是公开数据已经足够覆盖 MTEB 的任务分布。

多类分类的正负样本构造。label-based(用类标签文本本身做正样本)vs example-based(从同类里随机抽另一条样本做正样本),example-based 在 16 个分类聚类数据集上平均 69.27,label-based 64.80。聚类任务受益最明显,比如 Reddit-Clustering 从 59.83 涨到 71.10。论文将此归因为类标签文本太短、信息密度低,用同类样本做正例能让模型学到更丰富的类内相似性。

一些细节

hard-negative mining 的阈值。用 E5-mistral-7b-instruct 当 teacher 模型给候选 negative 打分,设置 max_negative_score_threshold = pos_score × 0.95。低于这个阈值的样本才能进入 negative 池,避免把潜在的 false negative(其实是另一个正例)当 hard negative。这套技术来自同组的 NV-Retriever

Mixed batch 而非 task-homogeneous batch。SFR-Embedding-Mistral 一个 batch 只装单一任务的样本,NV-Embed 一个 batch 混装所有任务。论文将后者的优势归为避免"zigzag"梯度更新方向,但未给出更深的理论分析,主要靠 ablation 上的分数差异支撑。

Latent array 大小。$r = 512$,8 个 head。论文没有报告 $r$ 的敏感性消融,这个值的选择理由没有展开。

模型压缩。论文 appendix 给了一套从 8B 压到 3.5B 的 recipe:SparseGPT 剪枝(4:8 半结构化 70.90)+ INT8 量化 + LoRA 微调恢复 + knowledge distillation。剪枝阶段直接把 latent attention block 删掉,靠蒸馏从未压缩的 teacher 那里把精度找回来。最终压缩版在 MTEB 上达到 71.48,仍超过 Llama3.2-3B、Qwen2.5-3B、Minitron-4B 直接训练的同等规模模型。INT8 量化几乎无损(变化在 ±0.01% 以内),论文将这种鲁棒性归为 LLM-based embedding 模型的固有优势。

与 LLM2Vec/GritLM 的关系。NV-Embed 的论点是"额外训练阶段(MNTP)和混合训练目标(GRIT)都不是必需的",简单粗暴地去掉 causal mask 加上更好的池化就够。代价是 latent attention 层引入了少量额外参数(latent array $512 \times 4096$ 加 Q/K/V 投影加两层 MLP),相对 7B 主干仍可忽略。

小结

NV-Embed 把"decoder LLM 做 embedding"的路线推到了一个相对干净的形态:去 causal mask、latent attention 池化、两阶段训练、数据精细化。每一步单独看都不复杂,组合起来在 MTEB 上拿到 72.31。它的贡献不在于发明新范式,而在于把架构改动和训练流程精简到几乎只剩必要项,给后续工作提供了一个清晰的基线。值得注意的局限:v2 的 72.31 包含了从训练 split 学到的优势(MSMARCO 等数据集是 MTEB 的训练子集,AIR-Bench 上 out-of-domain 检索的相对优势比 MTEB 内部小一些),实际部署时还要看任务分布是否与 MTEB 训练集对齐。