# DMax

A diffusion language model that decodes many tokens in parallel by revising its own earlier predictions, aiming to preserve output quality at high decoding parallelism.

DMax reformulates decoding for masked diffusion language models as a progressive refinement from mask embeddings to token embeddings, rather than the usual one-shot mask-to-token transition, so the model can revise tokens it got wrong instead of committing to them. It is trained with a method the authors call On-Policy Uniform Training, which unifies masked and uniform diffusion objectives so the model learns to recover clean tokens from both masked and self-generated erroneous input, and it adds a Soft Parallel Decoding step that interpolates between mask and token embeddings to carry confidence forward between decoding steps.

The authors report 6.0 tokens decoded per forward pass on math and reasoning tasks and 6.6 on code tasks while keeping accuracy, and release a 16B general-purpose model, DMax-16B, alongside training and evaluation code.

It is written in Python and accompanies a NeurIPS 2026 paper.

Website: https://arxiv.org/pdf/2604.08302
