基于扩散预训练的稠密与上下文嵌入

论文信息

标题: Diffusion-Pretrained Dense and Contextual Embeddings

作者: Sedigheh Eslami, Maksim Gaiduk, Markus Krimmel, et al.

发布日期: 2026-02-11

arXiv ID: 2602.11151v1

PDF 链接: 下载 PDF

3 分钟速览

研究问题：如何在保证嵌入质量的同时，提升文本嵌入模型对全局上下文的建模能力，并实现原生量化以降低大规模检索系统的存储和计算开销。
核心方法：将自回归大模型通过扩散目标继续预训练，转换为双向注意力编码器，再经过对训练、上下文训练和三元组训练的多阶段对比学习，并在训练中直接输出 INT8 甚至二值嵌入。
关键结果：PPLX-4B 在 MTEB 多语言 v2 检索任务上以 INT8 量化取得 69.66% 的平均 nDCG@10，存储效率约为同级别 Qwen3-Embedding-4B 的 4 倍；上下文模型 PPLXcontext-4B 在 ConTEB 基准上达到 81.96%，刷新纪录（见表 1、表 3）。
主要局限：二值量化对小尺寸模型（0.6B）造成的性能损失较大（约 2–5 个百分点），且扩散预训练得到的模型丧失了生成能力，仅用于嵌入；论文未详细分析不同预训练步数对下游任务的影响。
适合读者：从事信息检索、嵌入模型训练、量化部署以及 RAG 系统开发的研究与工程人员。

1 论文背景与动机

稠密文本嵌入能将查询和文档映射到统一的向量空间，是高效近似最近邻检索的基础。近年来，预训练大语言模型（LLM）在嵌入任务中表现出色，但主流的自回归解码器架构使用因果注意力，每个 token 只能看到上文，难以完整捕获长文档的全局上下文。为此，许多工作采用最后 token 池化来获取序列表示，但这仍然是一种妥协方案。

扩散语言模型（Diffusion Language Model）提供了另一种思路：它天然基于双向注意力编码器，能够在序列内捕获更全面的上下文信息。论文作者观察到，这一架构特性对检索任务尤其有利，因为理解长文档的全貌对于判断相关性至关重要。此外，实际部署中嵌入向量的存储和计算效率直接关系到成本，因此训练阶段就融入量化策略具有重要工程价值。

基于以上动机，该工作提出了 PPLX Embedding 模型家族，将扩散预训练、多阶段对比学习与原生量化感知训练结合起来，旨在提供高性能且高效的大规模检索嵌入。

2 核心技术：多阶段对比学习与扩散预训练

论文的整体训练流程（见图 1）将多个训练范式按分支执行，最后通过模型合并得到最终权重。

2.1 扩散继续预训练

作者选取 Qwen3-0.6B 和 Qwen3-4B 作为基础模型，禁用因果注意力掩码，将其转换为双向编码器。训练目标采用连续时间扩散过程与吸收态（[MASK]）：对于每个序列，独立采样时间步 $t \sim \mathcal{U}(0.001, 1)$ ，将每个 token 以概率 $t$ 替换为掩码 token。模型学习去噪重构原始 token，损失函数为加权交叉熵的期望（证据下界）。训练数据由 FineWeb-Edu 英语网页和 FineWeb2 的 29 种语言数据组成，总训练量约 2500 亿 token。

2.2 原生量化与池化策略

双向注意力架构使得均值池化成为可能，论文提出一种将均值池化与量化耦合的方法：先对 token 级嵌入做均值，再通过 $\tanh$ 函数压缩，最后线性映射并取整到 $[-127, 127]$ 的 8 位整数范围。这一量化操作在所有对比训练阶段就参与前向计算，并通过直通梯度估计器反向传播，使嵌入从一开始就适配 INT8 精度。同时，论文还支持训练后二值量化（按符号转为 $\pm 1$ ），宣称几乎无额外损失。

2.3 多阶段对比学习

对比学习分三步进行：

对训练（Pair Training）：使用大规模查询-文档对，以 InfoNCE 损失最大化正样本相似度，并通过阈值掩码抑制假阴性样本。训练分阶段逐步引入非英语和多语言数据。
上下文训练（Contextual Training）：将长文档切分为多个块（chunk），同时训练局部块级对比和全局文档级对比。局部损失组合了序列内对比和批内对比，全局损失则建模查询与整个文档的相似性，并引入重复文档掩码和假阴性掩码。全局损失权重 $\beta$ 从 0.2 余弦退火至 0.5，让模型先学局部语义再融合全局信息。该阶段的输出即为 PPLXcontext 模型。
三元组训练（Triplet Training）：在训练中加入挖掘的难负样本，进一步精细化相似和不相关文档的边界。训练后采用球形线性插值将上下文模型和三元组检查点合并，得到最终的 PPLX 模型。

3 创新点与贡献

扩散预训练+双向注意力：首次将扩散语言模型用于检索嵌入训练，利用编码器的双向注意力全面捕获段落内上下文，为均值池化和 “late chunking” 策略提供了基础。
完整的原生量化嵌入方案：训练时就引入 INT8 量化，使模型直接输出紧凑向量，省去后处理量化，且支持二值化，存储效率成倍提升（例如 4B 模型的 INT8 嵌入每 MB 可存 390 个，二值化则高达 3125 个，见表 1）。
多阶段对比与上下文嵌入统一：提出了一种分支训练、合并权重的复合流程，同时产出了标准检索模型和能感知文档全局信息的上下文模型，并刷新了 ConTEB 基准。
真实场景验证：除了公共基准，还构建了内部的大规模查询-查询和查询-文档基准（PPLXQ2Q, PPLXQ2D），在千万至亿级文档规模下评估召回率，更接近生产环境。

4 实验结果分析

4.1 公共基准性能

多语言与代码检索：在 MTEB 多语言 v2 18 个检索任务上，PPLX-4B INT8 平均 nDCG@10 为 69.66%，与 Qwen3-Embedding-4B（69.60%）持平（见表 1），但存储效率约是其 4 倍。PPLX-0.6B INT8（65.41%）也显著超过同量级对手。在 MTEB Code 基准上，PPLX-4B 达到 78.73%，优于 text-embedding-3-large 和 gemini-embedding-001。MIRACL 语种子集上，PPLX-0.6B 甚至以 68.6% 的平均 nDCG@10 超过自己的 4B 模型（66.2%），显示出色的多语言泛化能力（见表 2）。

上下文检索：在 ConTEB 的 8 个数据集上，PPLXcontext-4B INT8 平均 nDCG@10 为 81.96%，显著超越 voyage-context-3（79.45%）和 Anthropic Contextual（72.4%）（见表 3）。这验证了双向注意力结合上下文训练的有效性。

RAG 与工具检索：在 BERGEN 基准上，以 Qwen2.5-32B-Instruct 作为生成器，仅取前 5 个检索片段，PPLX-0.6B 在五个 QA 任务中有三个超过 Qwen3-Embedding-4B（见表 4），展现了小模型在 RAG 管线的潜力。ToolRet 工具检索中，PPLX-4B 以 44.45% 的平均 nDCG@10 超越 NV-Embed-v1（42.71%）等更大模型，有助于减少上下文膨胀。

4.2 内部大规模评估与量化影响

在 PPLXQ2Q（100K 查询，240 万目标查询）上，PPLX-4B INT8 的 Recall@10 达 73.46%，比 Qwen3-Embedding-4B 高出 5.56 个百分点（见表 6）。PPLXQ2D（1.5 万查询，3000 万文档）上，PPLX-4B 的 Recall@1000 在英语和多语言上分别为 88.23% 和 91.66%，均明显优于竞品（见表 7、8）。

量化带来的性能代价也经过细致分析：二值化对 4B 模型造成的损失约 0.96–1.6 个百分点，而对 0.6B 模型则高达 2–5 个百分点（见 3.3 节）。论文认为这既与参数量有关，也与 4B 模型输出维度更高（2560 vs 1024）因而能保留更多信息有关。

5 实践建议

对于希望将类似技术落地到实际检索系统的团队，该工作给出了几条清晰路径：

优先考虑双向注意力编码器：如果从零训练或继续预训练，将自回归模型转为扩散预训练编码器能够显著提升全局上下文捕获能力，对长文档检索和上下文嵌入尤其有益。若资源有限，也可直接使用已发布的 PPLX 模型。
采用原生量化感知训练：直接在对比训练时应用 INT8 量化，可以让模型在压缩状态下完成学习，避免部署后再量化的精度损失。实际部署时，根据延迟和存储需求选择 INT8 或二值版本即可，且二值化对 4B 级模型损失极小。
结合多阶段训练与上下文模块：对于需要从长文档中精确检索片段的应用，建议引入上下文训练并采用 “late chunking” 策略，即一次前向处理整个文档后再池化各块嵌入，这样既保持全局上下文又可获得细粒度表示。
评估召回率而非仅看 nDCG：在大规模生产中，嵌入模型通常作为第一级粗筛，高召回率（如 Recall@1000）比 nDCG@10 更关键。建议参照论文内部基准的设计，构造真实查询-文档集，在目标索引规模上测试召回性能。
注意小模型的量化极限：如果选择 0.6B 左右的轻量模型，二值量化可能导致召回率下降 3–5 个百分点，此时可保留 INT8 精度或在更高维度下设计模型以缓解信息损失。