sardi-dream-7b

SARDI (based on Dream7B)

Backbone for Self-Augmenting Retrieval for Diffusion Language Models (ICML 2026).

Paper · Code

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.

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