SARDI (based on Dream7B)
Backbone for Self-Augmenting Retrieval for Diffusion Language Models (ICML 2026).
SARDI interleaves retrieval with denoising: at each step it builds a query from the partially denoised sequence, retrieves fresh evidence, and conditions the next step on it. Speculative future tokens can inform retrieval long before they are stable enough to commit to the output.
Usage
The required inference code lives in the GitHub repository.
from sardi.inference import inference, load_model
from sardi.rag.retriever import SparseBM25SRetriever
model, tokenizer = load_model("pauljngr/sardi-dream-7b")
index = "data/2wikimultihopqa/corpus/index_chunked" # any bm25s index
retriever = SparseBM25SRetriever(corpus_path=f"{index}/corpus.jsonl", index_path=index)
print(inference(model, tokenizer,
"Which city is the capital of the country where the composer "
"of The Magic Flute was born?",
retriever))
Also, see example.py for an easy runnable version.
Loading needs trust_remote_code=True and flash-attn: the remote code is based on Dream7B with a patched generation_utils.py adding the
confidence_threshold sampler. Inference needs ~18 GB of VRAM.
Training
As discussed in the paper, Dream-7B didn't reliably produce zero-shot reasoning traces out-of-the-box. We therefore provide a simple RAG-fine-tuned Dream-7B checkpoint, shipped with code for threshold-based decoding.
| Base | Dream-v0-Instruct-7B (Ye et al., 2025) |
| Data | 2WikiMultiHopQA + HotpotQA train splits, gpt-4o-mini chain-of-thought traces, gold documents in context |
| Epochs | 3 |
| LR | 2e-6 |
| Batch | 256 global, 16 per GPU |
| Max seq len | 2048 |
| Optimizer | AdamW, FSDP on 2× NVIDIA B200 |
Weights are stored in float32; inference casts to bfloat16.
License
Apache 2.0, inherited from Dream-v0-Instruct-7B.
Citation
@inproceedings{juenger2026sardi,
title = {Self-Augmenting Retrieval for Diffusion Language Models},
author = {Paul J{\"u}nger and Justin Lovelace and Linxi Zhao and Dongyoung Go and Kilian Q Weinberger},
booktitle = {International Conference on Machine Learning (ICML)},
year = {2026}
}
This checkpoint fine-tunes Dream-7B, so please cite it as well:
@article{ye2025dream,
title = {Dream 7B: Diffusion Large Language Models},
author = {Jiacheng Ye and Zhihui Xie and Lin Zheng and Jiahui Gao and
Zirui Wu and Xin Jiang and Zhenguo Li and Lingpeng Kong},
journal = {CoRR},
volume = {abs/2508.15487},
year = {2025}
}