arXiv:2605.07210cs.IRcs.CL2026-05

用扩散语言模型的掩码预测能力,实现高效精准的文本检索。

DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

论文配图:DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models
图 1 · 摘自论文原文
  • 直接利用扩散模型的掩码位置预测生成检索向量,无需平均池化。
  • 单掩码版本超越DiffEmbed,多掩码支持细粒度匹配且延迟极低。
  • 适合追求高精度与低延迟的检索系统,尤其适用于大模型场景。

本文展示了扩散语言模型(DLM)可作为高效有效的检索器。现有基于DLM的检索器(如DiffEmbed)采用BERT式编码,将查询或段落表示为单一均值池化向量,忽略了DLM通过双向注意力进行掩码位置预测的训练机制,而该机制本可提供更强的检索信号。我们提出DiffRetriever,直接利用DLM的原生掩码位置预测进行检索:对每个查询或段落添加一个或多个掩码位置,以输出作为单次前向传播的检索表示。使用一个掩码位置时,单表示DiffRetriever已在相同骨干模型上优于DiffEmbed。DiffRetriever还自然扩展至多表示检索:DLM联合处理多个掩码位置,实现类似ColBERT的细粒度匹配,额外编码延迟极小。而在自回归LLM检索器中,同一策略需串行解码,导致更高延迟。在匹配对比实验中,DiffRetriever取得最强整体效果,优于DiffEmbed、PromptReps和RepLLaMA。训练数据中选定的掩码数量可良好迁移至不同数据集,而每查询的动态调整则提示仍有优化空间。代码已公开于https://github.com/ielab/diffretriever。

原文摘要 · Abstract (English)

This paper shows how diffusion language models (DLMs) can be used as effective and efficient retrievers. Existing DLM-based retrievers (e.g., DiffEmbed) follow BERT-style encoding, representing each query or passage as a single mean-pooled vector. This ignores how DLMs are trained to generate responses through masked-position prediction under bidirectional attention, a capability that can provide stronger retrieval signals. We propose DiffRetriever, which uses the DLM's native masked-position prediction directly for retrieval. For each query or passage, DiffRetriever appends one or more masked positions, using the outputs as retrieval representations in a single forward pass. With one masked position, single-representation DiffRetriever already improves over DiffEmbed on the same backbones. DiffRetriever also naturally extends to multi-representation retrieval: DLMs process multiple masked positions jointly, enabling ColBERT-style fine-grained matching with little additional encoding latency. In autoregressive LLM retrievers, the same multi-representation strategy requires sequential decoding and therefore incurs much higher latency. DiffRetriever obtains the strongest aggregate effectiveness within our matched comparison, outperforming DiffEmbed, PromptReps, and RepLLaMA. Masked-position counts selected on training data transfer well across datasets, while per-query variation suggests headroom for adaptive allocation. Code is available at https://github.com/ielab/diffretriever.

检索扩散模型多表示

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。