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 层做了第三种选择。
方法

双向注意力。在对比学习阶段直接把 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 训练后):
| Pooling | Bidirect | Causal |
|---|---|---|
| EOS | 62.68 | 60.06 |
| Mean | 64.00 | 62.32 |
| Latent-attention | 64.18 | 63.39 |
| Self-attention | 63.27 | 63.11 |
两点观察:
- 双向注意力对每种池化都更好。EOS pooling 受益最明显(+2.62),latent attention 受益最小(+0.79),说明 latent attention 一定程度上能缓解 causal mask 的限制
- 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 训练集对齐。