Discrete Diffusion

Dive into Deep Learning · §17.11

Discrete diffusion
transition matrices · the absorbing kernel · a weighted cross-entropy · unmasking

Noise for Tokens: Transition Matrices

q(\mathbf{x}_t \mid \mathbf{x}_{t-1}) = \mathrm{Cat}(\mathbf{x}_t;\ \mathbf{x}_{t-1} Q_t),\qquad q(\mathbf{x}_t \mid \mathbf{x}_0) = \mathrm{Cat}(\mathbf{x}_t;\ \mathbf{x}_0 Q_1 \cdots Q_t)

uniform kernel: resample with probability \beta_t, stationary law uniform;
absorbing kernel: move to [\textrm{M}] with probability \beta_t, masks stay;
both collapse: \bar{Q}_t = \bar{\alpha}_t I + (1 - \bar{\alpha}_t) P with P^2 = P.

Forward Masks, Reverse Reveals

The Posterior Has Two Points

q(\mathbf{x}_{t-1} = \mathbf{x}_0 \mid \mathbf{x}_t = [\textrm{M}], \mathbf{x}_0) = \frac{\bar{\alpha}_{t-1} - \bar{\alpha}_t}{1 - \bar{\alpha}_t},\qquad q(\mathbf{x}_{t-1} = [\textrm{M}] \mid \cdots) = \frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t}

an unmasked token stays; a masked one reverts or stays masked;
the model: same two-point law, with p_{\boldsymbol{\theta}}(\mathbf{x}_0^i \mid \mathbf{x}_t) in place of \mathbf{x}_0.

The Bound Is a Weighted Cross-Entropy

L_{\textrm{VLB}} = \sum_{t=1}^{T} \frac{\bar{\alpha}_{t-1} - \bar{\alpha}_t}{1 - \bar{\alpha}_t}\; \mathbb{E}\Big[ \sum_{i \,:\, \mathbf{x}_t^i = [\textrm{M}]} -\log p_{\boldsymbol{\theta}}(\mathbf{x}_0^i \mid \mathbf{x}_t) \Big]

the KL at a masked position is -w_t \log p_{\boldsymbol{\theta}}(\mathbf{x}_0^i \mid \mathbf{x}_t);
linear schedule: w_t = 1/t, and every step carries the same expected weighted count of masks;
in continuous time, with the masking level as the network’s input, the bound is invariant to the schedule.

The Quantized Mixture Has a Known Entropy

pairs = torch.cartesian_prod(torch.arange(K), torch.arange(K))  # all 1024 pairs
patterns = torch.tensor([[True, False], [False, True], [True, True]])

@torch.no_grad()
def exact_bound(model):  # the bound averaged over the true table, all steps and patterns
    total = 0.0
    for t in range(1, T + 1):
        weight = (alpha_bar[t - 1] - alpha_bar[t]) / (1 - alpha_bar[t])
        for masked in patterns:
            masked = masked.expand(len(pairs), 2)
            z = torch.where(masked, MASK, pairs)
            logp = torch.log_softmax(model(z, torch.full((len(pairs),), t)), -1)
            logp = logp.gather(-1, pairs[..., None])[..., 0]
            n_masked = masked[0].sum()
            prob = (1 - alpha_bar[t]) ** n_masked * alpha_bar[t] ** (2 - n_masked)
            total += prob * weight * (table.flatten() * (-logp * masked).sum(1)).sum()
    return total

torch.manual_seed(1)
model = MaskedMLP()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for step in range(6000):
    if step == 4000:
        optimizer.param_groups[0]['lr'] = 2e-4
    x0 = data[torch.randint(0, len(data), (512,))]
    t = torch.randint(1, T + 1, (512,))
    loss = masked_loss(model, x0, t).mean()
    optimizer.zero_grad(), loss.backward(), optimizer.step()
    if step % 2000 == 0:
        print(f'step {step:>4}: exact bound {exact_bound(model):.3f} nats')
print(f'final exact bound {exact_bound(model):.3f} nats; entropy {entropy:.3f} nats')

step    0: exact bound 6.918 nats
step 2000: exact bound 4.874 nats
step 4000: exact bound 4.870 nats
final exact bound 4.864 nats; entropy 4.840 nats

The bound averaged over the true table cannot fall below its entropy; the gap is model error plus looseness.

Fewer Steps, More Independence

@torch.no_grad()
def unmask(model, n, steps):
    z = torch.full((n, 2), MASK)
    levels = torch.linspace(T, 0, steps + 1).round().long().tolist()
    for t, s in zip(levels[:-1], levels[1:]):
        draw = torch.distributions.Categorical(
            logits=model(z, torch.full((n,), t))).sample()
        reveal = (z == MASK) & (torch.rand(z.shape)
                                < (alpha_bar[s] - alpha_bar[t]) / (1 - alpha_bar[t]))
        z = torch.where(reveal, draw, z)
    return z

torch.manual_seed(2)
samples = {}
print(f'{"steps":>5} {"TV to table":>12} {"predicted":>10} {"TV to product":>14}')
for steps in (32, 8, 2, 1):
    samples[steps] = unmask(model, 20000, steps)
    h = histogram(samples[steps])
    print(f'{steps:>5} {tv(h, table):>12.3f} {tv(product, table) / steps:>10.3f} '
          f'{tv(h, product):>14.3f}')

steps  TV to table  predicted  TV to product
   32        0.060      0.010          0.310
    8        0.079      0.039          0.282
    2        0.175      0.157          0.163
    1        0.327      0.314          0.062

Two tokens revealed in the same step are drawn independently: with S steps the sampler mixes in the product of the marginals with weight 1/S.

One Step Draws the Product

Binary Images by Unmasking

speckles: training images 0.0037, 100-step samples 0.0077, 4-step samples 0.0262

The same U-Net, three input channels and two output channels; a hundred reveal steps, then four.

Recap

Transition matrices replace Gaussian kernels; the bound is unchanged.
Absorbing kernel: two-point posterior, weighted cross-entropy on masks.
Unmasking trades steps for an independence error, measured exactly here.
One step draws all tokens independently; one token per step is any-order autoregression.