dew.diffusion.discrete

Masked (absorbing-state) discrete diffusion, shaped like the Gaussian one.

The forward process replaces each token by a mask id independently, with a probability that grows along t in [0, 1]; MaskingSchedule.alpha(t) is the fraction of tokens still visible. Training is the continuous-time negative ELBO of MDLM (Sahoo et al. 2024, “Simple and Effective Masked Diffusion Language Models”): the cross entropy of the model’s prediction at the masked positions, weighted by -alpha’(t) / (1 - alpha(t)). Sampling reverses the process one interval at a time: a masked token is revealed with probability (alpha(s) - alpha(t)) / (1 - alpha(t)) and, when revealed, drawn from the model’s categorical, which is MDLM’s _ddpm_update.

DiscreteProcess has the surface dew.sampling.sample walks: a time grid, an initial state, and a denoiser whose two outputs are the model’s argmax fill of the masked positions and the log-probabilities the solver draws from, in the slots a Gaussian denoiser puts x_0 and epsilon.

Name	Summary
MDLM_STEPS	Reverse steps generate takes by default, the count MDLM samples with.
MaskingSchedule	Says how fast tokens are masked along t.
LogLinear	MDLM’s log-linear schedule, alpha(t) = 1 - (1 - eps) t.
DiscreteProcess	Masks tokens of a vocabulary whose mask token is mask_id.
DiscreteDenoiser	(x_t, t) -> (argmax fill, log-probabilities) for model under params.
Unmask	Integrates a DiscreteProcess with MDLM’s reverse step from t to s < t.
MDLM	Builds the process of Sahoo et al.

MDLM_STEPS

attribute source

MDLM_STEPS = 64

Reverse steps generate takes by default, the count MDLM samples with.

MaskingSchedule

class source

class MaskingSchedule(ABC)

Says how fast tokens are masked along t.

alpha(t) in (0, 1] is the fraction of tokens left unmasked at t, with alpha(0) = 1.

MaskingSchedule.alpha

def alpha(t) -> jax.Array

MaskingSchedule.alpha_prime

def alpha_prime(t) -> jax.Array

d alpha / dt, negative.

LogLinear

dataclass source

class LogLinear(eps: float = 0.001)

MDLM’s log-linear schedule, alpha(t) = 1 - (1 - eps) t.

The masking rate -log alpha is then linear in log space and the NELBO weight is 1 / t.

LogLinear.alpha

def alpha(t)

LogLinear.alpha_prime

def alpha_prime(t)

DiscreteProcess

dataclass source

class DiscreteProcess(schedule: MaskingSchedule, mask_id: int)

Masks tokens of a vocabulary whose mask token is mask_id.

T

The fully masked end of the time domain, as a Gaussian process names it.

DiscreteProcess.sample_t

def sample_t(key, n: int) -> jax.Array

n times stratified over [0, 1), MDLM’s antithetic draw.

One uniform offset is shared by the batch, so the weights 1 / t of one batch cover the trajectory.

DiscreteProcess.corrupt

def corrupt(key, tokens, t) -> tuple[jax.Array, jax.Array]

(masked tokens, is_masked) at t, one t per row.

DiscreteProcess.weight

def weight(t) -> jax.Array

The NELBO weight -alpha’(t) / (1 - alpha(t)) on the masked cross entropy.

It is exactly zero at t = 0, where nothing is masked, no token contributes, and the quotient itself is undefined.

DiscreteProcess.times

def times(steps: int) -> jax.Array

DiscreteProcess.noise

def noise(key, shape) -> jax.Array

x_T, with every position masked.

key goes unread: the fully masked state is one point, not a draw.

DiscreteProcess.denoiser

def denoiser(
    model: nn.Module,
    params: Variables,
    conditions: Mapping[str, Conditioning] | None = None,
    unconditional: Mapping[str, Conditioning] | None = None,
    *,
    inputs: ModelInputs | None = None,
    mutable_mask: jax.Array | None = None,
) -> DiscreteDenoiser

DiscreteProcess.generate

def generate(
    model: nn.Module,
    variables: Variables,
    inputs: ModelInputs | jax.typing.ArrayLike,
    max_new_tokens: int,
    *,
    key: jax.Array | None = None,
    seed: int | None = None,
    n: int = 1,
    steps: int = MDLM_STEPS,
    sampler: Unmask | None = None,
    eos_token_ids: tuple[int, ...] = (),
    pad_token_id: int = 0,
) -> CanvasGeneration

Runs native MDLM over one full response span.

Prompt tokens are immutable, including literal mask ids. EOS trims the completed response and does not stop bidirectional refinement early.

DiscreteDenoiser

dataclass source

class DiscreteDenoiser(
    process: DiscreteProcess,
    model: nn.Module,
    params: Variables,
    inputs: ModelInputs | None = None,
    mutable_mask: jax.Array | None = None,
)

(x_t, t) -> (argmax fill, log-probabilities) for model under params.

t goes unread here: the masked model is conditioned on the corruption it sees rather than on the time, and Unmask.step reads the time from the process instead. It stays in the signature because sample calls every denoiser as (x_t, t).

The model’s own logits at an unmasked position are irrelevant, since the position keeps its token (MDLM’s carry-over parameterization). The mask token itself carries no mass: it marks corruption, so the categorical a reveal draws from never offers it, however the model scores it.

DiscreteDenoiser.masked

def masked(tokens)

Unmask

dataclass source

class Unmask()

Integrates a DiscreteProcess with MDLM’s reverse step from t to s < t.

Each masked position is revealed with probability (alpha(s) - alpha(t)) / (1 - alpha(t)), taking a token drawn from the model’s categorical. The rest stay masked.

Unmask.init

def init(x, times, process, *, key) -> tuple

Unmask.step

def step(x, t, t_next, denoised, log_probs, state, key, process, denoise)

MDLM

dataclass source

class MDLM(mask_id: int, eps: float = 0.001)

Builds the process of Sahoo et al. 2024, on the log-linear schedule.