Blame for src/diffusion.py

"""Masked diffusion: training objectives and decoding strategies.

The central problem this file exists to study:

  MDLM training teaches each masked position its MARGINAL p(x_i | observed).
  Parallel decoding reveals k positions at once by sampling the PRODUCT of
  those marginals. The truth is the joint, and the gap is exactly the mutual
  information among the revealed positions. Quality therefore falls as the
  parallelism rises - which is the only reason to use diffusion in the first
  place.

Every published fix reorders or gates commits at INFERENCE. The model has still
never seen a partially-decoded state containing its own correlated mistakes.
`self_corrupted_loss` changes that at training time.
"""
import math
import torch
import torch.nn.functional as F


# --------------------------------------------------------------- training ---
def mdlm_loss(model, x0, tok, t_min=None, fixed_t=None):
    """Standard MDLM objective: mask each token independently with prob t,
    predict the originals, weight by 1/t."""
    B, L = x0.shape
    dev = x0.device
    t_min = t_min if t_min is not None else 1.0 / L
    t = (torch.full((B,), fixed_t, device=dev) if fixed_t
         else t_min + (1.0 - t_min) * torch.rand(B, device=dev))
    m = torch.rand(B, L, device=dev) < t[:, None]
    m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True

    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
    logits = model(xt)
    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
                         reduction="none").view(B, L)
    return ((1.0 / t) * (ce * m).sum(1) / L).mean()


def self_corrupted_loss(model, x0, tok, k_frac=0.25, t_min=None,
                        correct_weight=1.0, ramp=1.0):
    """OUR METHOD. Train on states that parallel decoding actually produces.

    1. mask as usual and take a NO-GRAD pass to get the current marginals
    2. reveal a random subset of the masked positions *in parallel*, sampling
       each independently - i.e. commit exactly the kind of jointly-inconsistent
       tokens that parallel decoding commits
    3. train on the remaining masked positions given that self-generated
       context, so the model learns p(x_i | observed, its own parallel samples)
    4. additionally train it to REPAIR the committed tokens, which turns a
       wrong commit from an absorbing error into a recoverable one

    The model is therefore optimised for the distribution it meets at inference
    rather than for ground-truth context it will never see.

    `ramp` in [0,1] scales the corruption rate. It must start near zero: a
    randomly-initialised model samples noise, so corrupting from step 0 trains
    on pure garbage context and the model never learns the task at all (it
    scored 0% even on copy, which the baseline solves trivially). This is the
    same reason scheduled sampling anneals its mixing probability.
    """
    B, L = x0.shape
    dev = x0.device
    t_min = t_min if t_min is not None else 1.0 / L
    t = t_min + (1.0 - t_min) * torch.rand(B, device=dev)
    m = torch.rand(B, L, device=dev) < t[:, None]
    m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True
    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)

    with torch.no_grad():
        probs = model(xt).softmax(-1)
        samp = torch.distributions.Categorical(probs=probs).sample()

    # commit a random subset of masked positions from the model's own samples
    commit = m & (torch.rand(B, L, device=dev) < k_frac * ramp)
    x_self = torch.where(commit, samp, xt)
    still = m & ~commit

    logits = model(x_self)
    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
                         reduction="none").view(B, L)

    # Normalise by L exactly as mdlm_loss does, so the two objectives are on the
    # same scale and any difference in results is the method, not the loss size.
    main = ((1.0 / t) * (ce * still).sum(1) / L).mean()
    rep = ((1.0 / t) * (ce * commit).sum(1) / L).mean()
    return main + correct_weight * rep


# --------------------------------------------------------------- decoding ---
@torch.no_grad()
def decode(model, shape, tok, steps, device, strategy="confidence",
           temperature=0.0, entropy_tau=None, x_init=None, keep_trace=False):
    """Reveal the canvas over `steps` passes. Fewer passes = more tokens
    committed per pass = more parallelism = more conditional-independence error.

    strategies
      confidence : reveal the most confident masked positions (the standard)
      random     : reveal a random subset (ablation: is confidence doing work?)
      entropy    : reveal only positions whose entropy is below `entropy_tau`,
                   an inference-time baseline for dependency-aware revealing
    """
    B, L = shape
    x = torch.full((B, L), tok.mask, device=device, dtype=torch.long)
    if x_init is not None:
        x = x_init.clone()
    trace = []

    for s in range(steps, 0, -1):
        masked = x == tok.mask
        if not masked.any():
            break
        logits = model(x)
        probs = logits.softmax(-1)
        if temperature > 0:
            pred = torch.distributions.Categorical(
                logits=logits / temperature).sample()
            conf = probs.gather(-1, pred[..., None]).squeeze(-1)
        else:
            conf, pred = probs.max(-1)

        if strategy == "random":
            score = torch.rand_like(conf)
        elif strategy == "entropy":
            ent = -(probs * probs.clamp_min(1e-9).log()).sum(-1)
            score = -ent
        else:
            score = conf
        score = score.masked_fill(~masked, -1e9)

        target_left = int(L * (s - 1) / steps)
        for b in range(B):
            nm = int(masked[b].sum())
            k = max(0, nm - target_left)
            if strategy == "entropy" and entropy_tau is not None and k > 0:
                ent_b = -(probs[b] * probs[b].clamp_min(1e-9).log()).sum(-1)
                allowed = int(((ent_b < entropy_tau) & masked[b]).sum())
                k = max(1, min(k, allowed)) if allowed else 1
            if k:
                idx = score[b].topk(k).indices
                x[b, idx] = pred[b, idx]
        if keep_trace:
            trace.append(x.clone())
    still = x == tok.mask
    if still.any():
        x[still] = pred[still]
    return (x, trace) if keep_trace else x


@torch.no_grad()
def nll_bpc(model, batches, tok, device, mc=8):
    """Monte-Carlo estimate of the MDLM bound, reported as bits per character -
    the standard text8 metric."""
    tot, n = 0.0, 0
    for x0 in batches:
        x0 = x0.to(device)
        B, L = x0.shape
        for _ in range(mc):
            t = (1.0 / L) + (1 - 1.0 / L) * torch.rand(B, device=device)
            m = torch.rand(B, L, device=device) < t[:, None]
            xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
            ce = F.cross_entropy(model(xt).reshape(-1, len(tok)),
                                 x0.reshape(-1), reduction="none").view(B, L)
            tot += float(((1.0 / t) * (ce * m).sum(1) / L).sum())
            n += B
    return tot / max(n, 1) / torch.log(torch.tensor(2.0)).item()


def coordination_loss(model, x0, tok, group=4, t_min=None):
    """CORE OF THE METHOD: make parallel sampling reproduce sequential sampling.

    The measured failure: on jointly-constrained-but-individually-free positions
    (our `agree` probe), each position's likelihood-optimal marginal is uniform
    over the valid values, so sampling them independently is guaranteed to be
    inconsistent - 0% accuracy at every parallelism, 100% only when decoded one
    at a time. No decoding schedule can fix that, because the marginals
    themselves carry no coordination.

    So we change the marginals. For a random subset S:

      teacher  : decode S SEQUENTIALLY with no gradient, each position
                 conditioned on the ones already chosen. Because each step sees
                 the previous choices, the result is a coherent joint sample.
      student  : the marginals S gets from ONE parallel forward pass.
      loss     : cross-entropy of the teacher's joint sample under the student's
                 parallel marginals.

    The model is its own teacher, so this needs no external model and works in
    pretraining from scratch. It deliberately trades likelihood (marginals move
    off the true conditional, toward a coordinated mode) for parallel
    decodability - a trade the standard ELBO can never make.

    COST: the teacher costs `group` sequential forward passes, so `group` is the
    knob that keeps training affordable. An early version coordinated over 32
    positions and made every step 32x more expensive, which simply never
    finished. Combined with applying the term every `coord_every` steps, the
    amortised overhead is roughly 1 + group/coord_every.
    """
    B, L = x0.shape
    dev = x0.device
    t_min = t_min if t_min is not None else 1.0 / L
    t = t_min + (1.0 - t_min) * torch.rand(B, device=dev)
    m = torch.rand(B, L, device=dev) < t[:, None]
    m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True
    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)

    # choose, per example, a subset S of masked positions to coordinate over
    score = torch.rand(B, L, device=dev).masked_fill(~m, -1.0)
    k = min(group, L)
    S = score.topk(k, dim=1).indices                      # (B, k)

    # ---- teacher: reveal S one position at a time, conditioning as we go ----
    with torch.no_grad():
        x_seq = xt.clone()
        for j in range(k):
            idx = S[:, j]
            logits = model(x_seq)
            pick = logits[torch.arange(B, device=dev), idx].argmax(-1)
            x_seq[torch.arange(B, device=dev), idx] = pick
        target = x_seq.gather(1, S)                       # the joint sample

    # ---- student: one parallel pass over the same masked state --------------
    logits = model(xt)
    sel = logits.gather(1, S[..., None].expand(-1, -1, logits.size(-1)))
    return F.cross_entropy(sel.reshape(-1, sel.size(-1)), target.reshape(-1))


def entropy_penalty(model, x0, tok, t_min=None):
    """Break the symmetry that makes free groups undecodable in parallel.

    Where a group must agree but the value is free, the likelihood-optimal
    marginal is uniform over the valid values, and independent sampling from
    uniform marginals can never coordinate. Penalising marginal entropy pushes
    each position toward a single mode; because the model is shared across the
    group, the same mode is chosen at each member and parallel sampling becomes
    self-consistent. This is the term that is *supposed* to cost likelihood -
    it is the explicit price of parallel decodability.
    """
    B, L = x0.shape
    dev = x0.device
    t_min = t_min if t_min is not None else 1.0 / L
    t = t_min + (1.0 - t_min) * torch.rand(B, device=dev)
    m = torch.rand(B, L, device=dev) < t[:, None]
    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
    p = model(xt).softmax(-1)
    ent = -(p * p.clamp_min(1e-9).log()).sum(-1)
    return (ent * m).sum() / m.sum().clamp_min(1)


def cmd_loss(model, x0, tok, alpha=0.5, beta=1.0, k_frac=0.25, ramp=1.0,
             group=4, use_coord=True, gamma=0.0, elbo_w=1.0,
             dep_weighted=True):
    """Coordinated Masked Diffusion: ELBO + self-conditioning + coordination.

    The ELBO term is always present. The two additions are auxiliaries, not
    replacements: an earlier version used self-conditioning INSTEAD of the ELBO
    and scored 35% where the plain baseline scored 100%.
    """
    total = elbo_w * mdlm_loss(model, x0, tok)
    if gamma > 0 and ramp > 0:
        total = total + gamma * ramp * (
            dependence_weighted_entropy(model, x0, tok) if dep_weighted
            else entropy_penalty(model, x0, tok))
    if alpha > 0 and ramp > 0:
        total = total + alpha * self_corrupted_loss(
            model, x0, tok, k_frac=k_frac, correct_weight=1.0, ramp=ramp)
    if beta > 0 and ramp > 0 and use_coord:
        total = total + beta * ramp * coordination_loss(model, x0, tok, group=group)
    return total


def dependence_weighted_entropy(model, x0, tok, t_min=None):
    """THE METHOD. Sharpen the marginals exactly where positions are dependent.

    Why entropy is the right quantity. Two independent draws from a marginal p
    agree with probability sum_v p(v)^2 = exp(-H2(p)), the Renyi-2 entropy. So
    minimising H2 directly maximises the probability that positions sampled
    independently - which is precisely what parallel decoding does - land on the
    same choice. For a group that must agree but whose value is free, this is
    the only way to make parallel sampling self-consistent: the likelihood
    optimum (uniform over valid values) is guaranteed to disagree.

    Why it must be WEIGHTED. Sharpening everywhere would destroy likelihood
    wherever high entropy is correct, which is most of natural text, and would
    also break `copy`, where positions are already independent and the baseline
    is perfect. So the penalty is scaled by a measured dependence:

        reveal one masked position, and see how far every other position's
        belief moves. A position whose distribution shifts a lot depends on the
        revealed one; a position that does not move is conditionally
        independent and is left alone.

    d_i = KL(p_i-after-reveal || p_i-before), normalised per example. The
    penalty is d_i * H2(p_i), so the likelihood cost is paid only where it buys
    coordination.
    """
    B, L = x0.shape
    dev = x0.device
    t_min = t_min if t_min is not None else 1.0 / L
    t = t_min + (1.0 - t_min) * torch.rand(B, device=dev)
    m = torch.rand(B, L, device=dev) < t[:, None]
    m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True
    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)

    with torch.no_grad():
        p0 = model(xt).softmax(-1)
        samp = torch.distributions.Categorical(probs=p0).sample()
        # reveal one masked position per example and measure the ripple
        pick = torch.rand(B, L, device=dev).masked_fill(~m, -1.0).argmax(1)
        x1 = xt.clone()
        ar = torch.arange(B, device=dev)
        x1[ar, pick] = samp[ar, pick]
        p1 = model(x1).softmax(-1)
        d = (p1 * (p1.clamp_min(1e-9).log() - p0.clamp_min(1e-9).log())).sum(-1)
        d = (d * m).clamp_min(0)
        d = d / d.amax(1, keepdim=True).clamp_min(1e-6)      # per-example scale

    p = model(xt).softmax(-1)
    h2 = -(p.pow(2).sum(-1).clamp_min(1e-9)).log()           # Renyi-2 entropy
    return (d * h2 * m).sum() / m.sum().clamp_min(1)


# ===========================================================================
# Conditional objectives for UNIQUE-answer tasks (prompt visible, answer
# masked). Here the target is deterministic, so unlike the free-choice probes
# there is a correct answer to aim at and coordination is learnable.
# ===========================================================================

def _cond_mask(x0, amask, tok, t):
    """Mask a fraction t of the ANSWER positions only; the prompt stays visible."""
    dev = x0.device
    r = torch.rand_like(amask, dtype=torch.float)
    m = (r < t[:, None]) & amask
    # never leave an example with nothing to predict
    empty = ~m.any(1)
    if empty.any():
        first = amask.float().argmax(1)
        m[empty, first[empty]] = True
    return m


def cond_mdlm_loss(model, x0, amask, tok, t_dist="uniform", K=None):
    """Baseline, plus the schedule-matched variant.

    t_dist="uniform" is standard MDLM: every masking ratio is trained equally.
    t_dist="matched" instead draws t from the ratios a K-pass decode actually
    visits, namely {1, (K-1)/K, ..., 1/K}. Standard training spends most of its
    capacity on nearly-complete states that high-parallelism decoding never
    sees, which is a plain train/inference mismatch.
    """
    B, L = x0.shape
    dev = x0.device
    if t_dist == "matched" and K:
        idx = torch.randint(1, K + 1, (B,), device=dev).float()
        t = idx / K
    else:
        t = torch.rand(B, device=dev).clamp_min(1.0 / L)

    m = _cond_mask(x0, amask, tok, t)
    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
    logits = model(xt)
    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
                         reduction="none").view(B, L)
    return ((ce * m).sum(1) / m.sum(1).clamp_min(1)).mean()


def sequential_distill_loss(model, x0, amask, tok, chunk=8):
    """THE METHOD: distil the sequential computation into a single pass.

    Sequential decoding is strong because each step can read the digits already
    written, letting the carry chain unroll across passes. A one-pass parallel
    prediction has to fit that whole chain inside a fixed depth, which is why
    accuracy collapses as passes are reduced.

    So we train the ONE-PASS prediction, from the fully masked answer, to match
    the answer the model produces when it is allowed to decode sequentially.
    Because the answer is unique, the teacher's target is the same every time it
    sees the same prompt - which is exactly what was missing on the free-choice
    probes, where the teacher's arbitrary pick averaged back to uniform and the
    signal cancelled.

    Ground truth is available here, so the teacher is only a convenience: the
    loss below supervises the fully-masked one-pass prediction directly against
    the true answer, which is the strongest possible version of the same idea.
    """
    B, L = x0.shape
    xt = torch.where(amask, torch.full_like(x0, tok.mask), x0)   # whole answer hidden
    logits = model(xt)
    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
                         reduction="none").view(B, L)
    return ((ce * amask).sum(1) / amask.sum(1).clamp_min(1)).mean()


def pmd_loss(model, x0, amask, tok, mode="mdlm", K=8, lam=1.0):
    """Parallelism-Matched Diffusion.

      mdlm     standard uniform-t training (baseline)
      matched  train only at the masking ratios a K-pass decode visits
      distill  uniform training + one-pass supervision (the method)
      full     matched schedule + one-pass supervision
    """
    if mode == "mdlm":
        return cond_mdlm_loss(model, x0, amask, tok, "uniform")
    if mode == "matched":
        return cond_mdlm_loss(model, x0, amask, tok, "matched", K)
    base = cond_mdlm_loss(model, x0, amask, tok,
                          "matched" if mode == "full" else "uniform", K)
    return base + lam * sequential_distill_loss(model, x0, amask, tok)


@torch.no_grad()
def cond_decode(model, x0, amask, tok, steps, device):
    """Reveal the answer over `steps` passes, prompt fixed throughout."""
    x = torch.where(amask, torch.full_like(x0, tok.mask), x0).to(device)
    am = amask.to(device)
    B, L = x.shape
    n_ans = int(am[0].sum())
    pred = None
    for s in range(steps, 0, -1):
        masked = (x == tok.mask) & am
        if not masked.any():
            break
        logits = model(x)
        conf, pred = real_probs(logits, tok).max(-1)
        conf = conf.masked_fill(~masked, -1e9)
        left = int(n_ans * (s - 1) / steps)
        for b in range(B):
            k = max(0, int(masked[b].sum()) - left)
            if k:
                idx = conf[b].topk(k).indices
                x[b, idx] = pred[b, idx]
    rem = (x == tok.mask) & am
    if rem.any() and pred is not None:
        x[rem] = pred[rem]
    return x


def pmd_progressive(model, x0, amask, tok, progress, K_max=16, lam_max=1.0,
                    anneal="log"):
    """PARALLELISM CURRICULUM - the fix for the instability at long chains.

    Supervising the fully-masked one-pass prediction from step 0 is a very hard
    objective for a long carry chain: at 8 digits every seed found it, at 12 most
    collapsed to zero. The target is right but the optimisation is not reachable
    from random initialisation in one jump.

    So we anneal along the parallelism axis instead. Early training is decoded
    over many passes, where each pass reads the digits already written and the
    chain unrolls cheaply. K then halves toward 1, each stage a small increment
    from the last, while the one-pass term is ramped in rather than imposed.
    This is progressive distillation's schedule applied to the number of
    PARALLEL commits rather than the number of noise levels, and unlike prior
    step-distillation work it runs inside pretraining from scratch.

    Addition's carry is a prefix scan, so a log-depth solution exists (carry
    lookahead) - the curriculum is about making it discoverable, not about
    adding capacity.

    progress in [0,1]; returns (loss, K_now, lam_now) so the schedule is logged.
    """
    if anneal == "log":
        # 16 -> 8 -> 4 -> 2 -> 1, equal time per halving
        K_now = max(1, int(K_max / (2 ** int(progress * (math.log2(K_max) + 1e-9)))))
    else:
        K_now = max(1, int(round(K_max * (1 - progress) + 1 * progress)))
    lam_now = lam_max * min(1.0, max(0.0, (progress - 0.25) / 0.5))

    loss = cond_mdlm_loss(model, x0, amask, tok, "matched", K_now)
    if lam_now > 0:
        loss = loss + lam_now * sequential_distill_loss(model, x0, amask, tok)
    return loss, K_now, lam_now


# ------------------------------------------------- the unified formulation ---
# Two failure modes make parallel decoding lose accuracy, and they need opposite
# fixes. Keeping them apart is the whole argument of this project.
#
#   (i)  THE MARGINALS ARE UNTRAINED AT HIGH MASK RATIO.
#        A K-pass decode only ever evaluates the model at mask ratios
#        {1, (K-1)/K, ..., 1/K}. Uniform-t training spends most of its capacity
#        on nearly-complete states that high-parallelism decoding never visits.
#        No inference-time reordering can repair a marginal that was never fit;
#        this is fixable ONLY at training time, and `cond_mdlm_loss(t_dist=
#        "matched", K=...)` is the fix. At K=1 the matched schedule degenerates
#        to t=1, which is exactly `sequential_distill_loss` - the addition
#        experiments are the K=1 special case of one principle.
#
#   (ii) THE COMMITTED POSITIONS CARRY REAL MUTUAL INFORMATION.
#        Parallel decoding samples the PRODUCT of marginals; the truth is the
#        joint. The discarded dependence is the total correlation of the
#        committed set S,
#              TC(S) = sum_i H(x_i | c)  -  H(S | c)  <=  sum_i H(x_i | c),
#        so the summed conditional entropy of what we commit is an upper bound
#        on the error we incur by committing it. If every committed conditional
#        is a point mass the bound is zero and independent sampling reproduces
#        the joint exactly. No training objective can remove this term without
#        abandoning the data distribution, so it is fixable ONLY at inference,
#        by committing less. `entropy_budget_decode` is that fix.
#
# Prior work applies inference-time fixes to both, which is why it flatlines on
# addition: there, every conditional is a point mass, (ii) is identically zero,
# and all of the loss is (i). The claim here is that the two compose - and text,
# which contains deterministic and high-entropy positions in the same sequence,
# is where that composition is tested.

@torch.no_grad()
def _entropy(logits):
    p = logits.softmax(-1)
    return -(p * p.clamp_min(1e-9).log()).sum(-1)


@torch.no_grad()
def _commit(x, take, samp):
    return torch.where(take, samp, x)


def real_probs(logits, tok, allowed=None):
    """Distribution over REAL tokens only, optionally position-typed.

    The absorbing state is an input symbol, never an output: in the SUBS
    parameterisation the denoiser places zero probability on [MASK]. Sampling
    from the raw softmax instead lets a position be "committed" as [MASK], so it
    stays masked, decoding silently stalls, and the adaptive rules stop being
    monotone in their budget. The audit caught exactly this.

    `allowed` is an optional (L, V) boolean mask of which symbols each position
    may take. Graph data needs it: a node slot must take an atom type and an
    edge slot a bond type, and that is a property of the representation rather
    than something the model should have to spend capacity rediscovering.
    """
    logits = logits.clone()
    logits[..., tok.mask] = -float("inf")
    if getattr(tok, "pad", tok.mask) != tok.mask:
        logits[..., tok.pad] = -float("inf")
    if allowed is not None:
        logits = logits.masked_fill(~allowed.to(logits.device), -float("inf"))
    return logits.softmax(-1)


@torch.no_grad()
def entropy_budget_decode(model, x, tok, device, budget=0.5, max_steps=256,
                          temperature=1.0, fillable=None, budget_cond=None,
                          allowed=None):
    """FIX (ii): commit a set whose summed conditional entropy stays under `budget`.

    At each pass, rank the still-masked positions by conditional entropy and
    commit the longest low-entropy prefix whose cumulative entropy is within
    `budget` nats. By the bound above, `budget` directly caps the dependence
    thrown away at that pass, so it is an interpretable dial rather than a tuned
    heuristic: budget -> 0 recovers near-sequential decoding, budget -> infinity
    recovers one-shot parallel decoding, and the curve between them is the
    quality/compute trade-off.

    Returns (x, nfe); nfe is data-dependent, which is the point of an adaptive
    rule - cheap sequences finish in fewer passes.
    """
    x = x.clone().to(device)
    can = (x == tok.mask) if fillable is None else fillable.to(device)
    B, L = x.shape
    ar = torch.arange(L, device=device)
    nfe = 0
    for _ in range(max_steps):
        masked = (x == tok.mask) & can
        has = masked.any(1)
        if not has.any():
            break
        logits = model(x, budget=budget_cond)
        nfe += 1
        if temperature != 1.0:
            logits = logits / temperature
        probs = real_probs(logits, tok, allowed)
        H = -(probs * probs.clamp_min(1e-9).log()).sum(-1)
        samp = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(B, L)

        # sort masked positions by entropy; unmasked go last via +inf
        Hm = H.masked_fill(~masked, float("inf"))
        order = Hm.argsort(dim=1)
        cum = Hm.gather(1, order).cumsum(1)
        k = (cum <= budget).sum(1)
        k = torch.where(has, k.clamp_min(1), torch.zeros_like(k))  # always progress
        sel = ar[None, :] < k[:, None]
        take = torch.zeros_like(masked).scatter_(1, order, sel) & masked
        x = _commit(x, take, samp)
    return x, nfe


@torch.no_grad()
def confidence_threshold_decode(model, x, tok, device, tau=0.9, max_steps=256,
                                temperature=1.0, fillable=None, budget_cond=None,
                          allowed=None):
    """The published inference-time competitor (confidence-threshold family).

    Commit every masked position whose top-1 probability exceeds `tau`. This is
    the rule the 2026 parallel-decoding accelerators use; it is included so the
    comparison is against the actual algorithm rather than a top-k heuristic. It
    is a sensible rule and on text it works - the argument here is not that it
    is wrong, but that it addresses failure mode (ii) only and therefore cannot
    recover accuracy lost to (i).
    """
    x = x.clone().to(device)
    can = (x == tok.mask) if fillable is None else fillable.to(device)
    B, L = x.shape
    nfe = 0
    for _ in range(max_steps):
        masked = (x == tok.mask) & can
        has = masked.any(1)
        if not has.any():
            break
        logits = model(x, budget=budget_cond)
        nfe += 1
        if temperature != 1.0:
            logits = logits / temperature
        probs = real_probs(logits, tok, allowed)
        conf = probs.max(-1).values
        samp = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(B, L)

        take = (conf >= tau) & masked
        # a row that clears nothing still has to advance, or decoding stalls
        stuck = has & ~take.any(1)
        if stuck.any():
            best = conf.masked_fill(~masked, -float("inf")).argmax(1)
            take[stuck, best[stuck]] = True
        x = _commit(x, take, samp)
    return x, nfe


@torch.no_grad()
def fixed_k_decode(model, x, tok, device, steps, temperature=1.0, fillable=None,
                   budget_cond=None, allowed=None):
    """Plain K-pass decoding: reveal an equal share of positions per pass,
    highest-confidence first. The non-adaptive reference point, and the setting
    in which the addition results are reported."""
    x = x.clone().to(device)
    can = (x == tok.mask) if fillable is None else fillable.to(device)
    B, L = x.shape
    total = int(can.sum(1).max())
    ar = torch.arange(L, device=device)
    nfe = 0
    for s in range(steps, 0, -1):
        masked = (x == tok.mask) & can
        has = masked.any(1)
        if not has.any():
            break
        logits = model(x, budget=budget_cond)
        nfe += 1
        if temperature != 1.0:
            logits = logits / temperature
        probs = real_probs(logits, tok, allowed)
        conf = probs.max(-1).values.masked_fill(~masked, -float("inf"))
        samp = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(B, L)

        left = int(total * (s - 1) / steps)
        k = (masked.sum(1) - left).clamp_min(0)
        k = torch.where(has, k.clamp_min(1), torch.zeros_like(k))
        order = conf.argsort(dim=1, descending=True)
        sel = ar[None, :] < k[:, None]
        take = torch.zeros_like(masked).scatter_(1, order, sel) & masked
        x = _commit(x, take, samp)
    return x, nfe


def uncond_mdlm_loss(model, x0, tok, t_dist="uniform", K=None):
    """Unconditional masked-diffusion loss with the schedule as a free choice.

    This is `cond_mdlm_loss` with every position maskable, written separately
    because text has no prompt/answer split. `t_dist="matched"` is FIX (i):
    draw t from the ratios a K-pass decode actually visits instead of from
    U(0,1). The 1/t weighting of the ELBO is kept so that `t_dist="uniform"`
    remains the standard MDLM objective and the baseline is exactly the
    published one.
    """
    B, L = x0.shape
    dev = x0.device
    if t_dist == "matched" and K:
        t = torch.randint(1, K + 1, (B,), device=dev).float() / K
    else:
        t = torch.rand(B, device=dev).clamp_min(1.0 / L)

    m = torch.rand(B, L, device=dev) < t[:, None]
    m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True
    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
    logits = model(xt)
    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
                         reduction="none").view(B, L)
    return ((1.0 / t) * (ce * m).sum(1) / L).mean()


def text_loss(model, x0, tok, mode="mdlm", K=8, lam=1.0):
    """Training objectives for text8.

      mdlm     standard uniform-t MDLM (the published baseline)
      matched  FIX (i): schedule matched to a K-pass decode
      full     matched schedule plus an explicit high-ratio term

    `full` adds one extra evaluation at t drawn from the top of the schedule,
    which is where a low-NFE decode spends its first and most damaging commits.
    """
    if mode == "mdlm":
        return uncond_mdlm_loss(model, x0, tok, "uniform")
    if mode == "matched":
        return uncond_mdlm_loss(model, x0, tok, "matched", K)
    base = uncond_mdlm_loss(model, x0, tok, "matched", K)
    return base + lam * uncond_mdlm_loss(model, x0, tok, "matched", 2)


# ------------------------------------------------ any-budget masked diffusion ---
# The limitation the fixed-K method has, stated plainly: training at a single K
# is a COMMITMENT to one decoding budget. Measured on 20-digit addition, a model
# trained for one pass scores 99.8% at one pass and 85.7% at twenty - the
# baseline's own disease with the roles swapped, because intermediate
# partially-decoded states are themselves off-distribution for it.
#
# That makes a fixed-K method unable to dominate a quality-vs-compute frontier:
# it wins at K and loses away from K. Since the frontier is precisely the axis
# the parallel-decoding literature competes on, fixing this is what turns a
# single good operating point into a curve that dominates.
#
# The fix: sample K during training and TELL the model which K it is being
# trained for. One network then serves every budget. Cost is one embedding
# lookup and no extra forward passes, so it is cheaper than the fixed-K method
# it replaces.

BUDGETS = (1, 2, 4, 8, 16, 32, 64, 128, 256)


def sample_budget(B, device, choices=BUDGETS, weights=None):
    """Draw a decoding budget per example.

    Per-example rather than per-batch so a single step covers many budgets,
    which keeps the gradient signal for every bin alive at every step.
    """
    c = torch.as_tensor(choices, device=device, dtype=torch.float)
    if weights is None:
        idx = torch.randint(0, len(c), (B,), device=device)
    else:
        w = torch.as_tensor(weights, device=device, dtype=torch.float)
        idx = torch.multinomial(w.expand(B, -1), 1).squeeze(1)
    return c[idx]


def _t_for_budget(K):
    """Draw t from the mask ratios a K-pass decode actually visits.

    A K-pass decode evaluates the model at {1, (K-1)/K, ..., 1/K}. Drawing the
    step index uniformly in [1, K] and dividing reproduces exactly that set,
    including the atom at t = 1 that continuous U(0,1) never supplies.
    """
    step = (torch.rand_like(K) * K).floor() + 1.0
    return (step / K).clamp(1e-6, 1.0)


def cond_any_budget_loss(model, x0, amask, tok, choices=BUDGETS, condition=True):
    """Any-budget objective for conditional tasks (addition, inflection)."""
    B = x0.shape[0]
    dev = x0.device
    K = sample_budget(B, dev, choices)
    t = _t_for_budget(K)

    m = _cond_mask(x0, amask, tok, t)
    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
    logits = model(xt, budget=K if condition else None)
    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
                         reduction="none").view(B, -1)
    return ((ce * m).sum(1) / m.sum(1).clamp_min(1)).mean()


def uncond_any_budget_loss(model, x0, tok, choices=BUDGETS, condition=True):
    """Any-budget objective for unconditional text.

    Keeps the 1/t ELBO weighting so that, with a uniform-t schedule and no
    conditioning, this reduces exactly to the published MDLM objective.
    """
    B, L = x0.shape
    dev = x0.device
    K = sample_budget(B, dev, choices)
    t = _t_for_budget(K)

    m = torch.rand(B, L, device=dev) < t[:, None]
    m[torch.arange(B, device=dev), torch.randint(0, L, (B,), device=dev)] = True
    xt = torch.where(m, torch.full_like(x0, tok.mask), x0)
    logits = model(xt, budget=K if condition else None)
    ce = F.cross_entropy(logits.reshape(-1, logits.size(-1)), x0.reshape(-1),
                         reduction="none").view(B, L)
    return ((1.0 / t) * (ce * m).sum(1) / L).mean()


@torch.no_grad()
def cond_decode_budget(model, x0, amask, tok, steps, device, condition=True):
    """`cond_decode`, but the model is told the budget it is being run at."""
    x = torch.where(amask, torch.full_like(x0, tok.mask), x0).to(device)
    am = amask.to(device)
    B, L = x.shape
    n_ans = int(am[0].sum())
    bud = torch.full((B,), float(steps), device=device) if condition else None
    pred = None
    for s in range(steps, 0, -1):
        masked = (x == tok.mask) & am
        if not masked.any():
            break
        conf, pred = real_probs(model(x, budget=bud), tok).max(-1)
        conf = conf.masked_fill(~masked, -1e9)
        left = int(n_ans * (s - 1) / steps)
        for b in range(B):
            k = max(0, int(masked[b].sum()) - left)
            if k:
                x[b, conf[b].topk(k).indices] = pred[b, conf[b].topk(k).indices]
    rem = (x == tok.mask) & am
    if rem.any() and pred is not None:
        x[rem] = pred[rem]
    return x