DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models

Referee: [Abstract and §4] Abstract and §4 (Experiments): the headline claim that DiffRetriever on Dream is the strongest BEIR-7 retriever after SFT rests on comparisons whose causal attribution to the parallel K-token diffusion mechanism is not yet load-bearing. The manuscript does not report identical parameter counts, pre-training corpora, or exact SFT recipes (data, epochs, learning-rate schedule) for the RepLLaMA contrastive baseline versus the diffusion SFT runs; without these controls the performance delta cannot be isolated from backbone or training differences.

Authors: We agree that cross-family comparisons to RepLLaMA cannot fully isolate the contribution of the parallel diffusion mechanism from differences in pre-training data and contrastive versus supervised fine-tuning recipes. Our primary evidence for the value of parallel multi-token decoding is therefore the within-family comparisons on the same Dream (and other diffusion) backbones: DiffRetriever consistently outperforms both the single-token diffusion baseline and the DiffEmbed encoder baseline under identical SFT conditions. In the revised manuscript we will add an explicit subsection in §4 detailing the SFT data, epochs, learning-rate schedule, and batch size used for all diffusion runs, and we will qualify the headline claim to emphasize that the strongest result is obtained by applying the parallel mechanism to a diffusion backbone rather than claiming strict superiority over every possible autoregressive training recipe. revision: partial

Referee: [§4.2 and Table 2] §4.2 (Ablations) and Table 2: while multi-token diffusion is reported to improve over single-token on every backbone, the paper does not present an ablation that holds total compute or total representation dimensionality fixed when increasing K (e.g., by comparing K=4 at hidden size d versus K=1 at hidden size 4d). This leaves open whether the observed gains are due to parallelism per se or simply to increased representational capacity.

Authors: The ablations in §4.2 hold model architecture (including hidden dimension d) fixed while varying only K; this isolates the effect of parallel decoding at constant per-token capacity and constant forward-pass compute. The suggested capacity-matched ablation (K=1 with 4d hidden size) would require retraining models with altered architecture and is outside the scope of the present study. The practical contribution of DiffRetriever is precisely that K can be increased without any increase in inference latency or model size, a property that cannot be replicated by simply widening a single-token model. We will add a short discussion paragraph in the revised §4.2 that explicitly contrasts the two forms of capacity increase and reiterates that all reported gains occur at fixed hidden dimension. revision: partial

Referee: [§3.2] §3.2 (Aggregation): the method for collapsing the K parallel tokens into a single retrieval score (or set of scores) is described only at a high level. If the aggregation involves learned parameters or additional fine-tuning, this must be stated explicitly so that readers can assess whether the reported gains are still “parameter-free” relative to the single-token baseline.

Authors: The aggregation step in §3.2 consists of mean-pooling the K decoded token embeddings to obtain the final dense representation; the same mean-pooling is applied to the single-token case (trivially). No learned parameters, projection layers, or additional fine-tuning are introduced by the aggregation. We will revise the text of §3.2 to state this procedure explicitly, including the mathematical definition of the pooled vector, thereby confirming that the multi-token gains remain parameter-free relative to the single-token baseline. revision: yes