Expand description
Masked-diffusion ancestral sampler (SUBS parameterization).
A faithful port of the noflash sample.py (Sahoo et al., NeurIPS 2024):
absorbing/masked diffusion on a linear schedule t: 1 → 0. Starting from a
fully-masked sequence (with an optional carried-over prompt prefix), each of
num_steps steps predicts the clean token at every position, forbids the
[MASK] token (SUBS), and reveals each currently-masked position to a
sampled token with probability (t - s) / t. Once revealed, a position is
carried over (never re-masked); the final step reveals everything remaining.
The sampler is backend-agnostic — it drives any
MIBackend whose forward returns [batch, seq, vocab]
logits — so it serves both the MDLM DiT and (later) decoder-style
diffusion models. Determinism is by seed: the same seed reproduces the
same unmasking schedule and tokens.
Structs§
- Diffusion
Sampling Config - Configuration for masked-diffusion ancestral sampling.
Functions§
- generate
- Generate one sequence by masked-diffusion ancestral sampling.
- generate_
trajectory - Run ancestral sampling and return the token state after every step.