diffusion.py

import torch
import torch.nn.functional as F
from model import MASK
from data import SEQ_LEN


def corrupt(x0, t):
    mask = torch.rand_like(x0, dtype=torch.float) < t[:, None]
    return torch.where(mask, MASK, x0), mask


def loss(model, x0, eps=1e-3):
    t = torch.rand(x0.size(0), device=x0.device) * (1 - eps) + eps
    xt, mask = corrupt(x0, t)
    ce = F.cross_entropy(model(xt, t).transpose(1, 2), x0, reduction="none")
    w = (1.0 / t)[:, None]                     # MDLM linear-schedule weight
    return (ce * mask * w).sum() / mask.sum().clamp(min=1)


@torch.no_grad()
def sample(model, n, steps=64, device="cpu"):
    x = torch.full((n, SEQ_LEN), MASK, device=device)
    ts = torch.linspace(1, 0, steps + 1, device=device)
    for i in range(steps):
        t, s = ts[i], ts[i + 1]
        masked = x == MASK
        if not masked.any():
            break
        pred = torch.distributions.Categorical(logits=model(x, t.expand(n))).sample()
        unmask = masked & (torch.rand_like(x, dtype=torch.float) < (t - s) / t)
        x = torch.where(unmask, pred, x)
    return x


if __name__ == "__main__":
    from model import Denoiser
    m = Denoiser()
    x0 = torch.randint(0, 10, (32, SEQ_LEN))
    assert corrupt(x0, torch.full((32,), 0.99))[1].float().mean() > 0.8
    assert loss(m, x0).item() > 0
    assert (sample(m, 16, steps=16) != MASK).all()
    print("diffusion.py ok")