Blame for lecture_4/d3pm.py

"""D3PM on DNA, discrete time with one transition matrix per step.

Corruption multiplies a one-hot token by a transition matrix once per step, and the product of
those matrices stays inside the pi family, so the cumulative kernel is closed form. To reverse one
finite step we apply Bayes' rule to three factors, and the continuous-time rate matrices of the
rest of the chapter are the limit of these products.

Chapter 4, "Discrete time, as in D3PM", "Reversing one finite step" and "Rate matrices chosen for
the data". Time runs from t = 0 clean to t = T noise, the cumulative survival probability is
alpha_bar_t, and rates act on probability rows, p_dot = p R. Trained on the synthetic DNA of
cis6270.data. Prints the worked cumulative kernel, the worked reverse step, and the sample
statistics.

    python lecture_4/d3pm.py                  # train and generate
    python lecture_4/d3pm.py --steps 20       # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math  # The continuous-time limit is an exponential.

import torch  # Tensors, gradients and the categorical draws.

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 over clean tokens.
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

AUXILIARY = 0.01  # lambda_aux, the coefficient of the cross-entropy D3PM adds to the bound.
UNIFORM = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE)  # pi uniform over the four bases.


# %% 02. The transition matrices and their product
#
# Q_t = (1 - beta_t) I + beta_t 1 pi^T is the one-step version of the pi family, and two members
# multiply to a third whose survival probability is the product of the two.
def one_step(beta: float) -> torch.Tensor:
    """Q_t, with beta the probability of discarding the token and redrawing from pi. Returns [K, K]."""
    return (1.0 - beta) * torch.eye(DNA_SIZE) + beta * UNIFORM.expand(DNA_SIZE, DNA_SIZE)


def cumulative(alpha_bar: float) -> torch.Tensor:
    """Q_bar_t = alpha_bar_t I + (1 - alpha_bar_t) 1 pi^T, the product of the first t matrices."""
    return alpha_bar * torch.eye(DNA_SIZE) + (1.0 - alpha_bar) * UNIFORM.expand(DNA_SIZE, DNA_SIZE)


def schedule(timesteps: int) -> tuple[torch.Tensor, torch.Tensor]:
    """A linear cumulative survival alpha_bar_t = 1 - t/T, and the per-step beta_t it implies."""
    alpha_bar = 1.0 - torch.arange(timesteps + 1).to(torch.get_default_dtype()) / timesteps  # [T + 1] with alpha_bar_0 = 1.
    beta = 1.0 - alpha_bar[1:] / alpha_bar[:-1].clamp(min=1e-12)  # [T] the redraw probability.
    return alpha_bar, beta  # The two descriptions of the same corruption.


# %% 03. Reversing one finite step
#
# The numerator holds the two halves of the path through the candidate x_{t-1}, one from the clean
# token forward to the candidate and one from the candidate forward to the observation.
def exact_posterior(observed: torch.Tensor, clean_distribution: torch.Tensor,
                    step_matrix: torch.Tensor, previous: torch.Tensor) -> torch.Tensor:
    """q(x_{t-1} | x_t, x_0) when clean_distribution is a point mass, averaged over it otherwise.

    observed: [B, L] token ids. clean_distribution: [B, L, K] over candidate clean tokens.
    step_matrix and previous: [B, K, K], one Q_t and one Q_bar_{t-1} per sequence.
    Returns [B, L, K], normalized over the candidate token at step t - 1.
    """
    length = observed.shape[1]
    columns = step_matrix[:, None].expand(-1, length, -1, -1)  # [B, L, K, K] the step matrix.
    to_observation = columns.gather(  # [B, L, K] the column of Q_t at the observed token.
        3, observed[:, :, None, None].expand(-1, -1, DNA_SIZE, 1)).squeeze(-1)
    from_clean = torch.bmm(clean_distribution, previous)  # [B, L, K] the row of Q_bar_{t-1}.
    joint = to_observation * from_clean  # [B, L, K] the two halves of the path, multiplied.
    return joint / joint.sum(-1, keepdim=True).clamp(min=1e-20)  # [B, L, K] a distribution.


# %% 04. Rate matrices chosen for the data
#
# The five facts let us treat every pair of tokens differently. A discretized Gaussian
# keeps a pixel near its own intensity, and a nearest-neighbour graph does the same for tokens.
def gaussian_row(levels: int = 5, center: int = 2, sigma: float = 1.0) -> torch.Tensor:
    """The off-diagonal weights out of one intensity level, normalized to a total rate of one."""
    offsets = torch.arange(levels).to(torch.get_default_dtype()) - center  # [levels] the distance to each level.
    weights = torch.exp(-offsets ** 2 / (2.0 * sigma ** 2))  # [levels] nearby levels weigh more.
    weights[center] = 0.0  # The diagonal is fixed by the row-sum rule.
    return weights / weights.sum()  # [levels] a row of rates summing to one.


def graph_rates(adjacency: torch.Tensor) -> torch.Tensor:
    """R(i, j) = A_ij off the diagonal and R(i, i) = -sum_l A_il, from a token-similarity graph."""
    R = adjacency.clone()  # [K, K] the graph weights.
    R.fill_diagonal_(0.0)  # Only the off-diagonal entries are rates.
    return R - torch.diag(R.sum(dim=1))  # [K, K] every row sums to zero.


# %% 05. The discrete-time bound
#
# The loss is the variational bound of Section "A variational bound for the whole path", one KL
# divergence per step and per position, plus the auxiliary cross-entropy on the clean tokens.
def d3pm_loss(model: torch.nn.Module, x0: torch.Tensor, matrices: tuple) -> torch.Tensor:
    """One sample of L_vb + lambda_aux CE, with one step index and one corruption per sequence."""
    steps, cumulatives, alpha_bar = matrices  # [T, K, K], [T + 1, K, K] and [T + 1].
    timesteps = len(steps)
    index = torch.randint(1, timesteps + 1, (x0.shape[0],), device=x0.device)  # [B] the step t.
    probabilities = cumulatives[index][  # [B, L, K] the row of Q_bar_t at each clean token.
        torch.arange(len(x0), device=x0.device)[:, None], x0]
    xt = torch.distributions.Categorical(probs=probabilities).sample()  # [B, L] the noisy tokens.

    time = index.to(torch.get_default_dtype()) / timesteps  # [B] the step index as a time in [0, 1] for the network.
    denoiser = model(xt, time).softmax(-1)  # [B, L, K] the posterior over clean tokens.
    one_hot = torch.nn.functional.one_hot(x0, DNA_SIZE).to(torch.get_default_dtype())  # [B, L, K] the true clean token.
    step_matrix, previous = steps[index - 1], cumulatives[index - 1]  # [B, K, K] each.
    target = exact_posterior(xt, one_hot, step_matrix, previous)  # [B, L, K] the exact reverse step.
    predicted = exact_posterior(xt, denoiser, step_matrix, previous)  # [B, L, K] the model's own.
    divergence = (target * (target.clamp(min=1e-20)  # KL(q || p_theta) at every position.
                            / predicted.clamp(min=1e-20)).log()).sum((1, 2))  # [B].
    cross_entropy = -(one_hot * denoiser.clamp(min=1e-20).log()).sum((1, 2))  # [B] the auxiliary term.
    # We draw one step index uniformly to estimate a sum of T terms, so the divergence carries a T.
    return (timesteps * divergence + AUXILIARY * cross_entropy).mean()  # Nats per sequence.


# %% 06. Ancestral sampling
@torch.no_grad()  # Generation needs no gradient tracking.
def sample(model: torch.nn.Module, batch: int, length: int, matrices: tuple) -> torch.Tensor:
    """Walk the chain from x_T drawn from pi down to x_0. Returns [B, L] over ACGT."""
    steps, cumulatives, _ = matrices
    device = next(model.parameters()).device
    timesteps = len(steps)
    x = torch.distributions.Categorical(  # [B, L] the prior, pi at every position.
        probs=UNIFORM.to(device).expand(batch, length, DNA_SIZE)).sample()
    for t in range(timesteps, 0, -1):  # One finite reverse step at a time.
        time = torch.full((batch,), t / timesteps, device=device)  # [B] the step as a time.
        denoiser = model(x, time).softmax(-1)  # [B, L, K] the prediction of the clean token.
        posterior = exact_posterior(  # [B, L, K] the reverse step, averaged over the denoiser.
            x, denoiser, steps[t - 1].expand(batch, -1, -1), cumulatives[t - 1].expand(batch, -1, -1))
        x = torch.distributions.Categorical(probs=posterior).sample()  # [B, L] the earlier tokens.
    return x  # [B, L] clean DNA.


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--timesteps", type=int, default=32, help="T, the number of discrete steps")
    # T fixes both the corruption and the ancestral sampler, so this file reads it from --timesteps
    # and leaves --sample-steps alone. A default run trains 1000 steps on 64 sequences, inside
    # about a minute.
    parser.set_defaults(steps=1000, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.

    # Section 02, the worked cumulative kernel with beta_1 = 0.2 and beta_2 = 0.3.
    first, second = one_step(0.2), one_step(0.3)  # [K, K] the two transition matrices.
    product = first @ second  # [K, K] the cumulative kernel after two steps.
    alpha_bar_two = 0.8 * 0.7  # The product of the per-step survival probabilities.
    talk(f"Q_bar_2             alpha_bar {alpha_bar_two:.2f}  diagonal {float(product[0, 0]):.2f}  "
         f"off-diagonal {float(product[0, 1]):.2f}  closed-form gap "
         f"{float((product - cumulative(alpha_bar_two)).abs().max()):.2e}")

    # Section 02 again, the continuous-time limit of the same product at a constant rate.
    limits = {n: (1.0 - 1.0 / n) ** n for n in (10, 100, 1000)}  # The products climb toward exp(-1).
    talk("constant rate 1     " + "  ".join(f"T={n}: {v:.4f}" for n, v in limits.items()) +
         f"  exp(-1) = {math.exp(-1):.4f}")

    # Section 03, reversing one finite step from an observed C under a clean A.
    clean_a = torch.tensor([[[1.0, 0.0, 0.0, 0.0]]])  # [1, 1, K] the clean base A.
    observed_c = torch.tensor([[1]])  # [1, 1] the base observed at step 2.
    history = exact_posterior(observed_c, clean_a, second[None], first[None])[0, 0]  # [K] at step 1.
    talk(f"observed C, clean A column of Q_2 {[round(float(v), 3) for v in second.T[1]]}  row of "
         f"Q_bar_1 {[round(float(v), 2) for v in first[0]]}")
    talk(f"q(x_1 | x_2, x_0)   {[round(float(v), 4) for v in history]}  "
         f"evidence {float(product[0, 1]):.2f}")

    # Section 04, two rate matrices built for the data they corrupt.
    neighbours = gaussian_row()  # [5] the row out of a mid-grey pixel.
    adjacency = torch.tensor([  # A to C, C to G and G to T at weight 1, A to G at weight 0.5.
        [0.0, 1.0, 0.5, 0.0], [1.0, 0.0, 1.0, 0.0],
        [0.5, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 0.0]])
    R = graph_rates(adjacency)  # [K, K] the rate matrix of the token graph.
    graph_kernel = torch.matrix_exp(0.3 * R)  # [K, K] its transition matrix at clock 0.3.
    flux = UNIFORM[:, None] * R  # [K, K] the stationary flux along each edge.
    talk(f"gaussian row        neighbour {float(neighbours[1]):.4f}  two away "
         f"{float(neighbours[0]):.4f}  ratio {float(neighbours[1] / neighbours[0]):.1f}")
    talk(f"exp(0.3 R) first row {[round(float(v), 4) for v in graph_kernel[0]]}  detailed balance "
         f"gap {float((flux - flux.T).abs().max()):.2e}")

    # Train the denoiser against the discrete-time bound.
    alpha_bar, beta = schedule(args.timesteps)
    matrices = (torch.stack([one_step(float(b)) for b in beta]).to(device),
                torch.stack([cumulative(float(a)) for a in alpha_bar]).to(device),
                alpha_bar.to(device))
    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(DNA_SIZE, args.length, width=args.width).to(device)
    talk(f"device={device} parameters={count_parameters(model):,} T={args.timesteps}")

    def step_loss(step: int) -> torch.Tensor:
        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
        return d3pm_loss(model, train_tokens[rows], matrices)  # One clean batch into one scalar.

    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
    with torch.no_grad():
        held_out = sum(float(d3pm_loss(model, valid_tokens, matrices)) for _ in range(4)) / 4.0

    torch.manual_seed(args.seed + 1)  # Fix the ancestral draws, so the samples repeat.
    samples = sample(model, args.samples, args.length, matrices)  # [B, L] over ACGT.
    report = {
        "alpha_bar after two steps": alpha_bar_two,
        "Q_bar_2 on the diagonal": float(product[0, 0]),
        "Q_bar_2 off the diagonal": float(product[0, 1]),
        "product at T = 1000": limits[1000],
        "q(x_1 = A | C, A)": float(history[0]),
        "q(x_1 = C | C, A)": float(history[1]),
        "gaussian neighbour weight": float(neighbours[1]),
        "gaussian far weight": float(neighbours[0]),
        "exp(0.3 R) at (A, C)": float(graph_kernel[0, 1]),
        "graph detailed balance gap": float((flux - flux.T).abs().max()),
        "training bound, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "validation bound, nats/seq": held_out,
        "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("d3pm: discrete-time 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 finite-step numbers, train against the bound, and generate.