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.

python

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

bibtex

@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:

bibtex

@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}
}