Blame for lecture_4/mdlm.py

"""Masked diffusion on DNA, the whole model in thirty lines.

Corruption replaces each base by the mask with probability 1 - alpha_t, and a network reads the
masked sequence and predicts the clean token at every position. To generate we reveal masked
positions with the probability the schedule allows. The corruption, the loss and the sampler below
are the three listings of Section "The whole model for DNA in thirty lines", expanded with
comments.

Chapter 4, "Masked diffusion on one token" through "The whole model for DNA in thirty lines". Time
runs from t = 0 clean to t = 1 all masks, the survival probability is alpha_t = 1 - t, and the rates
R_t(x, y) act on probability rows, p_dot = p R. Trained on the synthetic DNA of cis6270.data. Prints
the chapter's worked reverse step and loss term, then the training loss and the sample statistics.

    python lecture_4/mdlm.py                  # train and generate, about a minute on a CPU
    python lecture_4/mdlm.py --steps 20       # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math  # The worked numbers of the chapter are logarithms.

import torch  # Tensors, gradients and the categorical draws.
import torch.nn.functional as F  # Cross-entropy at every position.
from torch.distributions import Categorical  # One draw per position from the denoiser.

from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import SequenceDenoiser, count_parameters  # The denoiser d_theta(. | z, t).
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

MASK = DNA_SIZE  # A, C, G, T are 0 to 3, and token 4 is the mask m.
VOCAB = DNA_SIZE + 1  # The masked alphabet V_m has K + 1 = 5 symbols.
TIME_FLOOR = 1e-3  # The weight 1/t is unbounded at t = 0, so draws start just above it.


# %% 02. The corruption
#
# Fact 2 of the construction: the forward marginal q_t(z^i | x_0^i) = alpha_t delta(z^i, x_0^i) +
# (1 - alpha_t) delta(z^i, m) is closed form, so one draw reaches any noise level.
def corrupt(x0: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Mask each position independently with probability 1 - alpha_t. x0: [B, L], t: [B, 1]."""
    alpha = 1.0 - t  # The linear schedule, alpha_0 = 1 and alpha_1 = 0.
    keep = torch.rand(x0.shape, device=x0.device) < alpha  # [B, L] one Bernoulli survival draw.
    return torch.where(keep, x0, MASK)  # [B, L] the clean token where it survived, m where it did not.


# %% 03. The weighted cross-entropy at masked positions
#
# The per-step KL divergence collapses to -r log d_theta(a), and once we sum the steps into an
# integral, r becomes the reveal rate w(t) = -alpha_dot / (1 - alpha) = 1/t on the linear schedule.
def loss(model: torch.nn.Module, x0: torch.Tensor) -> torch.Tensor:
    """One Monte Carlo sample of L_mask, with one time and one corruption per sequence."""
    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
    zt = corrupt(x0, t[:, None])  # [B, L] the time spreads across the positions of its own sequence.
    logits = model(zt, t)  # [B, L, 5] the denoiser reads the whole noisy sequence.
    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L] -log d_theta(x_0^i).
    w = 1.0 / t  # [B] the reveal rate, or -alpha_dot / (1 - alpha) on this schedule.
    return (w[:, None] * ce * (zt == MASK)).sum(-1).mean()  # Visible positions contribute zero.


# %% 04. The sampler
#
# We reveal a masked position with probability r_{s,t} = (alpha_s - alpha_t)/(1 - alpha_t), or
# (t - s)/t here, and the letter that appears is one draw from the denoiser.
@torch.no_grad()  # Generation needs no gradient tracking.
def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 32) -> torch.Tensor:
    """Generate from all masks down to clean DNA. Returns [B, L] token ids over ACGT."""
    device = next(model.parameters()).device  # Sample on the device the weights sit on.
    z = torch.full((batch, length), MASK, device=device)  # [B, L] the all-mask start, pi = delta_m.
    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid, t down to 0.
    for t, s in zip(ts[:-1], ts[1:]):  # One interval [s, t] at a time.
        logits = model(z, t.expand(batch))  # [B, L, 5] the clean-token prediction at this level.
        logits[..., MASK] = -1e9  # The zero-mask rule again, at sampling time.
        x0 = Categorical(logits=logits).sample()  # [B, L] one candidate clean token per position.
        unmask = torch.rand(z.shape, device=device) < (t - s) / t  # [B, L] the reveal draw.
        z = torch.where((z == MASK) & unmask, x0, z)  # Carry-over unmasking: a revealed base stays.
    logits = model(z, ts[-1].expand(batch))  # The last step runs down to s = 0, where r = 1.
    logits[..., MASK] = -1e9  # The same rule once more.
    return torch.where(z == MASK, Categorical(logits=logits).sample(), z)  # [B, L] clean DNA.


# %% 05. A complete reverse step, worked
#
# With alpha_s = 0.7 and alpha_t = 0.4, reached by the linear schedule at s = 0.3 and t = 0.6,
# the reveal probability is half and the denoiser splits that half between A and C.
def reveal_probability(alpha_s: float, alpha_t: float) -> float:
    """r_{s,t} = (alpha_s - alpha_t)/(1 - alpha_t), the lost survival mass over the mask mass."""
    return (alpha_s - alpha_t) / (1.0 - alpha_t)  # A function of the schedule alone.


def reverse_step_distribution(alpha_s: float, alpha_t: float, denoiser: list[float]) -> list[float]:
    """The reverse law over (A, C, G, T, m): r times the denoiser, with 1 - r left on the mask."""
    r = reveal_probability(alpha_s, alpha_t)  # The probability of acting at all.
    return [r * d for d in denoiser] + [1.0 - r]  # [5] the five outcomes of one step.


# %% 06. One term of the loss, and the rate error it comes from
#
# At t = 0.5 the weight is w(0.5) = 2, and the clean sequence ACGT corrupted to AmGm contributes the
# negative log probabilities of C and T. The chapter reaches the same weighted cross-entropy a
# second time from the rate error, on the reveal figure's denoiser. If we sum l(A, B) over every
# destination at one masked position we are left with w(t) times the negative log probability of
# the clean token there, once the staying terms restore the normalization of the denoiser.
def reveal_rate(t: float) -> float:
    """w(t) = -alpha_dot / (1 - alpha_t), or 1/t for the linear schedule."""
    return 1.0 / t  # The rate at which a masked position reveals something, whatever the denoiser.


def rate_error(target: float, model_rate: float) -> float:
    """l(A, B) = A log(A/B) - A + B, the Bregman divergence that the path KL divergence produces."""
    if target == 0.0:
        return model_rate  # l(0, B) = B, the staying term on its own.
    return target * math.log(target / model_rate) - target + model_rate  # Zero exactly at B = A.


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    # The motif needs a longer budget than the GC fraction, so the default trains for 1000 steps
    # and measures on 64 sequences. The whole run stays inside about a minute on one core.
    parser.set_defaults(steps=1000, sample_steps=64, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)  # The same seed gives the same training batches and the same samples.
    device = pick_device(args.device)
    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.

    # Section 05, the worked reverse step at alpha_s = 0.7 and alpha_t = 0.4.
    r = reveal_probability(0.7, 0.4)  # 0.3 of recovered mass over 0.6 of mask mass.
    step_law = reverse_step_distribution(0.7, 0.4, [0.6, 0.4, 0.0, 0.0])  # [5] over (A, C, G, T, m).
    talk(f"reverse step        r = {r:.1f}  p(A, C, G, T, m) = {[round(v, 2) for v in step_law]}  "
         f"sum {sum(step_law):.1f}")

    # Section 06, one term of the loss at t = 0.5 on ACGT corrupted to AmGm.
    weight = reveal_rate(0.5)  # w(0.5) = 2.
    masked_ce = -math.log(0.7) - math.log(0.5)  # The clean tokens C and T at positions 2 and 4.
    worked_term = weight * masked_ce  # One sample of the objective.
    position_two = [0.2, 0.6, 0.1, 0.1]  # The reveal figure's denoiser at the second position.
    reveal_rates = [weight * d for d in position_two]  # [4] rates that sum to w(t) = 2.
    collapse = rate_error(weight, reveal_rates[1]) + sum(  # The clean token is C here.
        rate_error(0.0, rate) for i, rate in enumerate(reveal_rates) if i != 1)
    talk(f"loss term at t=0.5  w = {weight:.0f}  masked CE = {masked_ce:.5f}  "
         f"w CE = {worked_term:.5f}")
    talk(f"reveal rates        {[round(v, 1) for v in reveal_rates]}  sum {sum(reveal_rates):.0f}  "
         f"rate error {collapse:.4f} = w(-log 0.6) = {weight * -math.log(0.6):.4f}")

    # Train the denoiser on the shared synthetic DNA.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the bound below.
    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
    talk(f"device={device} parameters={count_parameters(model):,} "
         f"train={len(train_tokens)} length={args.length}")

    def step_loss(step: int) -> torch.Tensor:
        """One optimizer step draws its own batch of clean sequences."""
        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
        return loss(model, train_tokens[rows])  # [B, L] clean sequences into one scalar.

    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)

    # The validation bound, averaged over several corruptions of the held-out sequences.
    with torch.no_grad():
        bound = sum(float(loss(model, valid_tokens)) for _ in range(4)) / 4.0  # Nats per sequence.

    # Generate, then measure the structure the data was built with.
    torch.manual_seed(args.seed + 1)  # Fix the reveal draws, so the samples repeat.
    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
    report = {
        "reveal probability r": r,
        "reverse step on the mask": step_law[-1],
        "weight w(0.5)": weight,
        "worked loss term": worked_term,
        "rate error at position 2": collapse,
        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "validation bound, nats/seq": bound,
        "bound, nats per token": bound / args.length,
        "GC fraction": gc_fraction(samples),
        "GC fraction in the data": gc_fraction(train_tokens),
        "motif fraction": motif_fraction(samples),
        "motif fraction in the data": motif_fraction(train_tokens),
        "distinct samples": unique_fraction(samples),
        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
    }
    print_report("mdlm: masked diffusion on DNA", report, decode_dna(samples))
    save_run(args.out, model=model, config=vars(args), report=report,
             losses=losses, samples=decode_dna(samples))


# %% 08. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work the numbers, train the denoiser, and generate.