DiffusionGemma-26B-A4B-it-ORL

DiffusionGemma + ORL: DiffusionGemma-26B-A4B post-trained with On-Policy Residual Learning (ORL); it decodes 3-6 more tokens per forward than DiffusionGemma at comparable quality

Highlights

• Higher tokens per forward (TPF) on every benchmark: 30.0 vs 26.0 on GSM8K, 27.9 vs 24.5 on HumanEval, 28.7 vs 23.1 on MBPP
• Comparable accuracy: equal or higher on GSM8K, MATH-500 and HumanEval, 0.3 points lower on MBPP
• 1,456 tokens/s for a single request on GSM8K with dlmrt on one H100; with NVFP4, 4,214 on one B200, 7,708 on eight B200s and 1,540 on a consumer RTX 5090

Results

Accuracy (%) / TPF; each model uses its selected decoder.

Speed

Tokens per second for a single request, measured with dlmrt unless noted, with the paper's prompts and decoders (the paper reports the base model on GSM8K and HumanEval):

The 8× B200 and RTX 5090 rows use generation seed 1234.

Usage

Plain transformers generation does not run the paper's decoding. Two engines do: dlmrt, our inference engine, and the ORL repository's patched vLLM (evaluation/inference/vllm/patches) with its recorded sampler, which the paper's evaluation uses.

dlmrt

dlmrt, our inference engine, decodes one request at a time with the paper's decoder (role="orl") and measures the paper's throughput; see its README for installation, tensor parallelism and the commands of every paper cell.

from dlmrt import LLM

llm = LLM("yifanyu/DiffusionGemma-26B-A4B-it-ORL", role="orl")
out = llm.generate("Write a Python function that reverses a linked list.", max_new_tokens=2048)
print(out.text)
print(f"{out.request_tps:.0f} tokens/s")

python -m dlmrt serve --model yifanyu/DiffusionGemma-26B-A4B-it-ORL --role orl --port 18080
curl -N localhost:18080/generate -d '{"messages": [{"role": "user", "content": "Hi!"}], "max_new_tokens": 512}'

The paper's evaluation pipeline

Evaluate it with the paper's protocol (GSM8K, MATH-500, HumanEval, MBPP) using the repository's evaluation pipeline; DiffusionGemma generation is not bit-deterministic, so results vary slightly between runs:

git clone https://github.com/Introspective-Diffusion/ORL orl && cd orl/evaluation
export ORL_EVAL_WORK=/path/to/eval-work
bash scripts/build_images.sh && bash scripts/prepare.sh
hf download yifanyu/DiffusionGemma-26B-A4B-it-ORL --local-dir $ORL_EVAL_WORK/hf/DiffusionGemma-26B-A4B-it-ORL
bash scripts/main_table.sh dg-orl $ORL_EVAL_WORK/hf/DiffusionGemma-26B-A4B-it-ORL

Serve it as the evaluation does (one 80 GB GPU; an OpenAI-compatible server on 127.0.0.1:8000, model name dg, with the GSM8K context length and sampler):

apptainer exec --nv -B "$PWD,$ORL_EVAL_WORK" --env PYTHONPATH=$PWD $ORL_EVAL_WORK/images/vllm-dg-orl.sif \
  python3 -m runtime.vllm.serve --model-root $ORL_EVAL_WORK/hf/DiffusionGemma-26B-A4B-it-ORL --role orl \
    --benchmark gsm8k --out $ORL_EVAL_WORK/dg-server --execute

Training

ORL trains on fresh, verified trajectories of the current model, replays the exact residual inputs each denoising step consumed, and adds adjacent entropy contraction (AEC) so that refinement gets faster, not only more accurate. This checkpoint is step 900 of the paper run on a mixture of 819 GSM8K, 819 MATH and 819 Python problems, on two nodes of 8 × H100 (actor and rollout). training/scripts/reproduce_diffusion_gemma.sh in the code repository reruns the recipe; MoE kernels accumulate in a nondeterministic order, so reruns differ slightly.

Checkpoint identity

Paper role dg-orl, training step 900. checkpoint_manifest.json records the hashes of the weights, tokenizer and configuration, and paper_decoding.json the engine and benchmark-specific decoding settings, which downloading the model does not apply.

Related Models

Citation

@misc{yu2026orl,
  title         = {On-Policy Residual Learning for Diffusion Language Models},
  author        = {Yu, Yifan and Jian, Yuqing and Li, Jisen and Hu, Yuezhou and Zhang, Tianyi and Wu, Xiaoxia and
                   Zhu, Rui-Jie and Wang, Junxiong and Zhou, Zhongzhu and Lai, Fan and Zou, James and Xu, Chenfeng},
  year          = {2026},
  eprint        = {2610.20687},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2610.20687}
}