DiffRetriever — LLaDA-8B (single-representation)

Single-representation (K=1) dense + sparse retriever fine-tuned on GSAI-ML/LLaDA-8B-Instruct, released with DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models (arXiv:2605.07210 · code).

DiffRetriever uses a diffusion language model's masked-position prediction interface directly for retrieval: it appends a single masked position (K=1) after a retrieval prompt and reads the hidden states (dense) and next-token logit vectors (sparse) from a single bidirectional forward pass (Fwd=1). With K=1 this is a fast single-vector dense + sparse retriever. The autoregressive equivalent must decode each representation sequentially.

This repo ships the LoRA adapter only (~tens of MB). The base backbone is downloaded automatically from GSAI-ML/LLaDA-8B-Instruct the first time you load the model.

Model summary

Results

Fine-tuned results. **Dense** is the recommended/headline score for this
checkpoint; sparse and hybrid are also available from the same single forward
pass when the checkpoint was trained with sparse supervision.

In-domain (MS MARCO dev, TREC DL19/DL20)

Out-of-domain — BEIR-7 (NDCG@10, dense)

See the paper for the full comparison against PromptReps, DiffEmbed, RepLLaMA, and BM25, and for latency analysis.

Usage

This repo is self-contained: the model code ships with it, so one call loads everything (the base LLaDA backbone is pulled from the Hub automatically and the LoRA adapter is attached on top).

pip install "transformers==4.54.0" peft torch    # + accelerate, safetensors

import torch
import torch.nn.functional as F
from transformers import AutoModel

# trust_remote_code runs the modeling code shipped in this repo.
model = AutoModel.from_pretrained("ielabgroup/diffretriever-llada-8b-single", trust_remote_code=True)
model.eval()

# A tiny query / passage set.
queries = ["what causes the seasons on earth?"]
passages = [
    "The tilt of Earth's axis relative to its orbital plane drives the seasons.",
    "Photosynthesis converts carbon dioxide and water into glucose using sunlight.",
]

# Encode — one forward pass per batch (tokenize() builds the prompt + masks).
def encode(texts, is_query):
    ids, mask = model.tokenize(texts, is_query=is_query)
    dev = next(model.backbone.parameters()).device
    with torch.inference_mode():
        return model.encode(ids.to(dev), mask.to(dev),
                            is_query=is_query, compute_sparse=False)

q = encode(queries,  is_query=True)
p = encode(passages, is_query=False)

# ── Scoring: single-vector dense (single_dense) ─────────────────────────────
# K=1: L2-normalize the single representation, then dot product.
qv = F.normalize(q["repr_hidden"].float(), dim=-1).mean(dim=1)   # [Q, H]
pv = F.normalize(p["repr_hidden"].float(), dim=-1).mean(dim=1)   # [P, H]
scores = qv @ pv.T                                              # [Q, P]

print(scores)   # [Q, P] — higher = more relevant

To rank a corpus, encode all passages once (offline), then encode each query and take scores.topk(k). For sharded encoding, the sparse/hybrid modes, and full BEIR/MS MARCO evaluation, see scripts/encode.py and scripts/evaluate_sweep.py in https://github.com/ielab/diffretriever.

Scoring modes

The encoder returns repr_hidden (dense, [B, K, H]) and — with compute_sparse=True — sparse_indices/sparse_values (sparse lexical weights). These support the paper's five modes: single_dense, multi_dense, sparse_max, fusion_single_sparse_max, fusion_multi_sparse_max. This checkpoint is tuned for single-vector dense (single_dense); scripts/evaluate_sweep.py runs all five in one pass.

Training details

For diffusion backbones the query/passage budgets (K_q, K_p) are selected on MS MARCO train; the paper uses (4, 16) for Dream and (4, 4) for LLaDA.

Related checkpoints

• ielabgroup/diffretriever-dream-7b-single · ielabgroup/diffretriever-dream-7b-multi-q4-p16
• ielabgroup/diffretriever-llada-8b-single · ielabgroup/diffretriever-llada-8b-multi-q4-p4

Citation

@article{wang2026diffretriever,
  title={ DiffRetriever: Parallel Representative Tokens for Retrieval with Diffusion Language Models },
  author={Wang, Shuai and Yin, Yu and Zhuang, Shengyao and Koopman, Bevan and Zuccon, Guido},
  journal={arXiv preprint arXiv:2605.07210},
  year={2026}
}

License

MIT. The base model is subject to its own license — see GSAI-ML/LLaDA-8B-Instruct.