| """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 torch |
| import torch.nn.functional as F |
|
|
|
|
| |
| 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 = 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) |
|
|
| |
| |
| 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 |
|
|
|
|
| |
| @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) |
|
|
| |
| score = torch.rand(B, L, device=dev).masked_fill(~m, -1.0) |
| k = min(group, L) |
| S = score.topk(k, dim=1).indices |
|
|
| |
| 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) |
|
|
| |
| 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 cmd_loss(model, x0, tok, alpha=0.5, beta=1.0, k_frac=0.25, ramp=1.0, |
| group=4, use_coord=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 = mdlm_loss(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 |
|
|
