File size: 10,945 Bytes
98c193b 7fa1b76 98c193b 7fa1b76 98c193b 7fa1b76 98c193b 7fa1b76 98c193b 7fa1b76 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 | """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
# --------------------------------------------------------------- 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 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
|