Delete lecture_4

lecture_4/README.md

@@ -1,108 +0,0 @@
-# Lecture 4: Discrete Diffusion
-
-Chapter 4 keeps every state inside a finite vocabulary and replaces the continuous displacement of
-Chapter 3 by a decision about which tokens to erase, reveal or replace. The chapter builds masked
-diffusion by hand from a survival schedule, then rebuilds it as one member of a family: a rate
-matrix says what a token can do in the next instant, the master equation follows the marginals, the
-matrix exponential gives the corruption kernel at any time, and time reversal turns the forward
-rates into the rates that generate. Seven published models come out of that family by changing the
-rate matrix, the time variable, what the network predicts, or the state space itself, and a
-further group keeps the rates and changes the order in which positions are filled, the weighting
-of the loss, or the distribution to sample from. Every file here runs on the shared synthetic DNA
-of `cis6270.data`, with `t = 0` clean and `t = 1` all masks, survival probability `alpha_t`, and
-rates `R_t(x, y)` in the row convention `p_dot = p R`.
-
-## The files
-
-| File | Command | What it trains and generates |
-| --- | --- | --- |
-| `ctmc.py` | `python lecture_4/ctmc.py` | Nothing is trained. Works every number of the CTMC sections: the two rules, one row for DNA, waiting times, the master equation, the matrix exponential, the `R_pi` family, eigenvalues, detailed balance and time reversal. |
-| `mdlm.py` | `python lecture_4/mdlm.py` | Trains the denoiser on the weighted cross-entropy at masked positions, generates DNA from all masks. The chapter's thirty-line model. |
-| `uniform.py` | `python lecture_4/uniform.py` | Trains against the rate loss under uniform replacement, generates by tau-leaping from uniformly random tokens. |
-| `d3pm.py` | `python lecture_4/d3pm.py` | Trains against the discrete-time variational bound plus the auxiliary cross-entropy, generates by `T` ancestral steps. |
-| `sedd.py` | `python lecture_4/sedd.py` | Trains a network to output the concrete score through denoising score entropy, generates from the predicted ratios. |
-| `block.py` | `python lecture_4/block.py` | Trains masked diffusion inside blocks with a clean prefix, generates block by block. `--block` moves the dial between MDLM and autoregression. |
-| `simplex.py` | `python lecture_4/simplex.py` | Trains a denoiser that reads a Dirichlet belief at every position, generates by thinning the belief, drawing an innovation and mixing the two. |
-| `guidance.py` | `python lecture_4/guidance.py` | Trains a label-dropout denoiser and a classifier on noisy sequences, then samples unguided, with exact classifier guidance, with its gradient approximation, and classifier free. |
-| `planner.py` | `python lecture_4/planner.py` | Trains with the planner weights inside the bound, samples the same weights in a uniform order and in the planner's order with a remasking corrector. |
-| `search.py` | `python lecture_4/search.py` | Trains masked diffusion, then searches the tree of partial sequences under the h-transform and keeps a Pareto archive over two objectives. |
-
-Two constructions of the chapter have no file here, and the chapter develops both in full. Per-token
-schedules, MD4 and GenMD4, give every clean token its own rate `beta_a(t)` and change nothing but
-the weight, which becomes `w_a(t) = -alpha_dot_t(a)/(1 - alpha_t(a))` inside the sum over tokens.
-A2D2
-lets the length change by adding an insertion jump beside the unmasking jump, which enlarges the
-state space and keeps the master equation, the reversal and the rate loss exactly as they are.
-
-Every file takes the shared flags of `cis6270.runner.common_parser`, so `--steps`, `--samples`,
-`--sample-steps`, `--seed` and `--quiet` mean the same thing everywhere. `ctmc.py` trains nothing,
-so its `--steps` is the number of Euler steps in the matrix-exponential comparison and its
-`--samples` the number of simulated paths. `d3pm.py` takes
-`T` from `--timesteps`, since the same number fixes the corruption and the ancestral sampler.
-`search.py` reads `--sample-steps` as the length of one rollout, which is short because every leaf
-of the tree is scored by its own rollouts.
-
-## What a default run produces
-
-All numbers below come from `--seed 0` on a CPU. The worked numbers are exact and repeat on any
-machine; the trained numbers move a little with the hardware's floating point, and the sample
-statistics move with the seed. The training data have GC fraction 0.577 and motif fraction 0.777,
-and a sampler drawing its bases uniformly would land near 0.05 on the motif, since ACGT fits in a
-length-16 sequence at thirteen places and each one has probability `1/256`.
-
-| File | Worked numbers it prints | Trained numbers |
-| --- | --- | --- |
-| `ctmc.py` | exit rate 4 out of A, `P(A, A) = 0.6` at `h = 0.1`, Euler step `(0.16, 0.22, 0.28, 0.34)`, `exp(tau R_pi)` at `0.625` and `0.125`, Euler products `0.480`, `0.570`, `0.616`, reverse rates `2` and `0.5` out of C | none |
-| `mdlm.py` | `r = 0.5`, reverse law `(0.3, 0.2, 0, 0, 0.5)`, `w(0.5) = 2`, loss term `2.0996`, rate error `1.0217` | bound 1.25 nats per token, GC 0.56, motif 0.69 |
-| `uniform.py` | kernel `0.625` and `0.125`, clean factor 5, posterior `(0.375, 0.625)`, conditional rates 5 and 0.2, posterior average **2**, model rates `2.76` and `1.16`, rate error `0.754` | rate loss 19.1 nats per sequence, GC 0.56, motif 0.50 |
-| `d3pm.py` | `alpha_bar_2 = 0.56`, `Q_bar_2` at `0.67` and `0.11`, the step-1 posterior `(0.5795, 0.3523, 0.0341, 0.0341)`, Gaussian row `0.4088` and `0.0912`, `exp(0.3 R)` first row `(0.6747, 0.1964, 0.1122, 0.0167)` | bound 18.2 nats per sequence, GC 0.57, motif 0.72 |
-| `sedd.py` | target ratio 2, score entropy `0.0107`, `0.3863`, `0.1891` | score entropy 18.8 nats per sequence, mean predicted ratio 1.26 toward a conditional target of 5, GC 0.55, motif 0.33 |
-| `block.py` | worked bound `2.7567` nats on ACGTAC, `integral of w(t)(1 - alpha_t) = 1` | bound 1.21 nats per token at `L' = 4` and 1.29 at `L' = 1`, GC 0.57, motif 0.67 |
-| `simplex.py` | `c_t = 8` and `c_s = 20`, parameters `(5, 1, 1, 1)` and `(17, 1, 1, 1)`, standard deviation `0.161` on A, increment `(12, 0, 0, 0)`, `W ~ Beta(8, 12)` with mean `0.4`, simulated standard deviation `0.0779` against the exact `0.0779` | cross-entropy 0.42 nats per token, GC 0.57, motif 0.34 |
-| `guidance.py` | `h = (2, 1, 0.5, 0.5)`, guided rates `(1, 1.5, 0.5)`, at `gamma = 2` rates `(0.5, 4.5, 0.5)` with exit rate `5.5` and share `0.8182` on G | GC 0.58 unguided, 0.68 with the exact classifier, 0.63 with the gradient approximation, 0.74 classifier free |
-| `planner.py` | reveals positions 3 and 6, remasks position 5, both corrector fluxes 1, PAPL weights `(0.4, 2.0, 0.6, 1.0)`, inner sums `2.9` and `2.12` | motif 0.61 in a uniform order against 0.86 in the planner's order, GC 0.58 and 0.51 |
-| `search.py` | `h(z) = 0.35`, `h(z') = 0.75`, factor `2.143`, a rate of `1.2` becomes `2.571`, the archive keeps P1 and P2 | 144 rollouts scored, an archive of 3 sequences, best GC 0.88 and best motif match 1.0 |
-
-A default run of `ctmc.py` takes about two seconds, and every other file takes between thirty and
-forty seconds of processor time on one core. `--steps 20` runs each file in a few seconds.
-
-`mdlm.py --steps 3000` lowers the bound from 1.25 to 1.22 nats per token in about three times the
-time, and leaves the motif fraction inside the spread it already has across seeds, which at seeds
-0, 1 and 2 is 0.69, 0.45 and 0.72. The one-position bound of `block.py --block 1` is the exact
-autoregressive log-likelihood with no slack, and at this budget it still measures above the
-`L' = 4` bound, because the estimator draws one block of sixteen and multiplies by sixteen, which
-is sixteen times noisier at the same number of draws, and because a denoiser restricted to a left
-prefix has a harder prediction problem than one that reads both sides.
-`sedd.py` sits below the other samplers on the motif because its sampler stops at the time floor,
-as the list below says. `simplex.py --sample-steps 256` raises its motif fraction from 0.34 to
-0.38, so the reverse grid is not what holds that number down.
-
-## What is simplified
-
-- The time weight `w(t) = 1/t` is unbounded at `t = 0`, so every file draws its times above a small
-  floor. The chapter's integrand is correct as written; the floor keeps one draw from dominating a
-  batch.
-- The samplers that use tau-leaping rescale a step whose total jump probability would exceed one.
-  `uniform.py` then finishes with the exact posterior step down to `t = 0` that the chapter derives
-  for the `pi` family. `sedd.py` predicts ratios, with no posterior over clean tokens to step
-  through, so it stops at the time floor, where `alpha_t = 0.92` leaves about eight percent of the
-  positions redrawn from `pi`.
-- `uniform.py` and `sedd.py` use the constant schedule `beta(t) = 4`, which is the schedule of the
-  chapter's worked example and reaches `alpha_t = 1/2` at `t = (log 2)/4`. Its survival probability
-  at `t = 1` is `0.018` and not zero, and that residue is the `KL(q_1 || p_1)` term of the bound.
-- `d3pm.py` samples one step index per sequence and multiplies by `T`, which estimates the sum over
-  steps without evaluating all of it.
-- `block.py` samples one block per sequence and multiplies by the number of blocks, for the same
-  reason. Training a block of length `L'` hides the suffix behind the mask in place of a
-  block-causal attention mask.
-- `simplex.py` holds the concentration constant at `gamma = 1`, realizes the sum over clean tokens
-  by one categorical draw as the chapter does, and reads the final token as the corner the belief
-  has reached. Its loss is the chapter's cross-entropy, which bounds no likelihood, so its nats
-  per token do not compare with the bounds above.
-- `guidance.py` uses a mean-pooling classifier on one-hot inputs, because the gradient
-  approximation needs a derivative with respect to the sequence and the label of this data is a
-  thresholded mean over positions.
-- `planner.py` reveals the number of positions the schedule calls for at each step, which is the
-  deterministic version of the chapter's independent reveal draws, and it remasks a fixed count.
-- `search.py` scores `--sample-steps` rollout steps per completion and expands a fixed number of
-  children, so the tree stays small enough to build in a few seconds.

lecture_4/block.py

@@ -1,180 +0,0 @@
-"""Block diffusion on DNA, with the block size as the dial between diffusion and autoregression.
-
-The chain rule over blocks is exact, and each factor is a distribution over one block of tokens
-conditioned on a clean prefix, which is the kind of distribution masked diffusion models. At one
-block the bound is the MDLM loss, at blocks of a single position it is the autoregressive negative
-log-likelihood with no slack, and every size in between interpolates.
-
-Chapter 4, "Blocks between diffusion and autoregression" and "The block size as a dial between two
-models". 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 worked bound of 2.7567 nats, the one-position identity,
-and the sample statistics.
-
-    python lecture_4/block.py                 # train and generate with blocks of four
-    python lecture_4/block.py --block 1       # the autoregressive end of the dial
-    python lecture_4/block.py --steps 20      # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math  # The worked bound is a sum of logarithms.
-
-import torch  # Tensors, gradients and the categorical draws.
-import torch.nn.functional as F  # Cross-entropy at every position.
-
-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 reveal rate w(t) = 1/t is unbounded at t = 0.
-
-
-# %% 02. One block, with a clean prefix and a masked suffix
-#
-# The rate matrix acts only on the current block. Everything before it is clean and conditions the
-# prediction, and everything after it is hidden, so the denoiser never reads the future.
-def block_input(x0: torch.Tensor, block: torch.Tensor, size: int, t: torch.Tensor) -> torch.Tensor:
-    """The network input for one block: clean prefix, corrupted block, masked suffix. [B, L]."""
-    position = torch.arange(x0.shape[1], device=x0.device)[None, :]  # [1, L] the column index.
-    start = (block * size)[:, None]  # [B, 1] where the current block begins.
-    inside = (position >= start) & (position < start + size)  # [B, L] the current block.
-    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)  # [B, L] the survival draw.
-    corrupted = torch.where(keep, x0, MASK)  # [B, L] masking applied inside the block.
-    return torch.where(position < start, x0, torch.where(inside, corrupted, MASK))  # [B, L].
-
-
-def block_loss(model: torch.nn.Module, x0: torch.Tensor, size: int) -> torch.Tensor:
-    """One Monte Carlo sample of the block bound, with one block and one time per sequence."""
-    blocks = math.ceil(x0.shape[1] / size)  # B, the number of blocks the sequence splits into.
-    block = torch.randint(0, blocks, (x0.shape[0],), device=x0.device)  # [B] which factor to sample.
-    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
-    z = block_input(x0, block, size, t[:, None])  # [B, L] the denoiser's input.
-    position = torch.arange(x0.shape[1], device=x0.device)[None, :]  # [1, L].
-    start = (block * size)[:, None]  # [B, 1].
-    inside = (position >= start) & (position < start + size)  # [B, L] the current block.
-    logits = model(z, t)  # [B, L, 5] the clean-token prediction at every position.
-    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].
-    scored = ce * ((z == MASK) & inside)  # [B, L] only the masked positions of the current block.
-    # Drawing one block uniformly estimates a sum over blocks, so the term is scaled by that count.
-    return (blocks * (1.0 / t)[:, None] * scored).sum(-1).mean()  # Nats per sequence.
-
-
-# %% 03. The block size as a dial
-#
-# At a block of one position the time integral collapses. The position is masked with probability
-# 1 - alpha_t and weighted by w(t) = -alpha_dot/(1 - alpha_t), so the integrand is -alpha_dot times
-# the cross-entropy, and the remaining integral is alpha_0 - alpha_1 = 1.
-def one_position_integral(grid: int = 2048) -> float:
-    """The integral of w(t)(1 - alpha_t) over [0, 1], which the collapse makes exactly one."""
-    t = (torch.arange(grid).double() + 0.5) / grid  # [grid] midpoints of a uniform partition.
-    return float(((1.0 / t) * t).mean())  # w(t)(1 - alpha_t) = 1 at every t on the linear schedule.
-
-
-def one_position_estimate(probability: float, draws: int, generator) -> float:
-    """The same integral by Monte Carlo, which is how the training loss evaluates it."""
-    t = torch.rand(draws, generator=generator).clamp(min=TIME_FLOOR)  # [N] uniform times.
-    masked = torch.rand(draws, generator=generator) < t  # [N] the single position is masked or not.
-    return float(((1.0 / t) * masked * -math.log(probability)).mean())  # Should land on -log d.
-
-
-# %% 04. Generating block by block
-@torch.no_grad()  # Generation needs no gradient tracking.
-def sample(model: torch.nn.Module, batch: int, length: int, size: int,
-           steps: int = 32) -> torch.Tensor:
-    """Fill one block at a time by masked diffusion, with the finished blocks as a clean prefix."""
-    device = next(model.parameters()).device
-    z = torch.full((batch, length), MASK, device=device)  # [B, L] every position starts masked.
-    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid per block.
-    position = torch.arange(length, device=device)[None, :]  # [1, L] the column index.
-    for start in range(0, length, size):  # Blocks are completed in order, left to right.
-        inside = (position >= start) & (position < start + size)  # [1, L] the current block.
-        for t, s in zip(ts[:-1], ts[1:]):
-            logits = model(z, t.expand(batch))  # [B, L, 5] the prefix conditions every prediction.
-            logits[..., MASK] = -1e9  # The zero-mask rule, at sampling time.
-            x0 = torch.distributions.Categorical(logits=logits).sample()  # [B, L] candidate tokens.
-            unmask = torch.rand(z.shape, device=device) < (t - s) / t  # [B, L] the reveal draw.
-            z = torch.where((z == MASK) & unmask & inside, x0, z)  # Reveal inside this block only.
-        logits = model(z, ts[-1].expand(batch))  # The last step of the block runs down to s = 0.
-        logits[..., MASK] = -1e9
-        final = torch.distributions.Categorical(logits=logits).sample()  # [B, L].
-        z = torch.where((z == MASK) & inside, final, z)  # The block is clean before the next starts.
-    return z  # [B, L] clean DNA.
-
-
-# %% 05. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--block", type=int, default=4, help="L', the block length")
-    # Generation runs the whole reverse grid once per block, so --sample-steps stays at 16 while
-    # training rises to 1000 steps and the report measures 64 sequences, inside about a minute.
-    parser.set_defaults(steps=1000, sample_steps=16, 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, one draw of the bound on ACGTAC with blocks of two at t = 0.5.
-    weight = 1.0 / 0.5  # w(0.5) = 2 on the linear schedule.
-    contributions = [  # Block 1 masked twice, block 2 once, block 3 once.
-        weight * (-math.log(0.5) - math.log(0.8)),
-        weight * -math.log(0.7),
-        weight * -math.log(0.9),
-    ]
-    talk(f"worked block bound  {[round(v, 4) for v in contributions]}  total "
-         f"{sum(contributions):.4f} nats on ACGTAC")
-
-    # Section 03, the one-position identity that makes blocks of length one autoregressive.
-    generator = torch.Generator().manual_seed(args.seed)  # A fixed stream for the estimate.
-    exact = one_position_integral()  # The integral of w(t)(1 - alpha_t), which is one.
-    estimate = one_position_estimate(0.7, 20000, generator)  # The same number by Monte Carlo.
-    talk(f"one-position dial   integral {exact:.4f}  Monte Carlo bound {estimate:.4f}  against "
-         f"-log 0.7 = {-math.log(0.7):.4f}")
-
-    # Train the denoiser on the shared synthetic DNA at the chosen block size.
-    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)
-    blocks = math.ceil(args.length / args.block)  # How many factors the chain rule has.
-    talk(f"device={device} parameters={count_parameters(model):,} "
-         f"blocks={blocks} of length {args.block}")
-
-    def step_loss(step: int) -> torch.Tensor:
-        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
-        return block_loss(model, train_tokens[rows], args.block)  # One clean batch into one scalar.
-
-    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
-    with torch.no_grad():
-        bound = sum(float(block_loss(model, valid_tokens, args.block)) for _ in range(8)) / 8.0
-
-    torch.manual_seed(args.seed + 1)  # Fix the reveal draws, so the samples repeat.
-    samples = sample(model, args.samples, args.length, args.block, args.sample_steps)
-    report = {
-        "worked bound on ACGTAC": sum(contributions),
-        "first block of the worked bound": contributions[0],
-        "integral of w(t)(1 - alpha_t)": exact,
-        "one-position Monte Carlo bound": estimate,
-        "blocks": blocks,
-        "block length": args.block,
-        "training bound, 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(f"block: diffusion in blocks of {args.block}", report, decode_dna(samples))
-    save_run(args.out, model=model, config=vars(args), report=report,
-             losses=losses, samples=decode_dna(samples))
-
-
-# %% 06. Run the script
-if __name__ == "__main__":  # Execute when launched directly.
-    main()  # Work the block bound, train at one block size, and generate block by block.

lecture_4/ctmc.py

@@ -1,278 +0,0 @@
-"""Continuous-time Markov chains on one DNA base, the mathematics with no learning in it.
-
-A rate matrix says what a base can do in the next instant, and the master equation, the matrix
-exponential, the stationary distribution and the time reversal all follow from it. Nothing here is
-trained: every number this file prints is a number the chapter works by hand.
-
-Chapter 4, "The two rules a rate matrix obeys" through "Running the jumps backward". Time runs from
-t = 0 clean to t = 1 all masks, the survival probability is alpha_t, and the rates R_t(x, y) act on
-probability rows, p_dot = p R. Prints every worked quantity of those sections.
-
-    python lecture_4/ctmc.py                   # the worked numbers, 10 Euler steps
-    python lecture_4/ctmc.py --steps 20        # a finer Euler grid in Section 09
-"""
-
-# %% 01. Imports and configuration
-import math  # Logarithms and exponentials of the worked numbers.
-
-import torch  # Matrices, eigenvalues and the matrix exponential.
-
-from cis6270.data import DNA, DNA_SIZE  # The alphabet ACGT and K = 4.
-from cis6270.runner import common_parser, print_report, save_run, set_seed  # Flags and the report.
-
-MASK_STATES = DNA_SIZE + 1  # A, C, G, T and m, the five states of the masked alphabet.
-
-
-# %% 02. The two rules a rate matrix obeys
-def rate_matrix(off_diagonal: torch.Tensor) -> torch.Tensor:
-    """Fill the diagonal of a nonnegative array so that R_t(x, x) = -sum_{y != x} R_t(x, y)."""
-    R = off_diagonal.clone().double()  # [K, K] the chosen jump rates.
-    R.fill_diagonal_(0.0)  # The diagonal adds no information of its own.
-    return R - torch.diag(R.sum(dim=1))  # [K, K] every row now sums to zero.
-
-
-def obeys_the_two_rules(R: torch.Tensor, tol: float = 1e-9) -> bool:
-    """Nonnegative off the diagonal, and zero row sums, which is the definition."""
-    off = R - torch.diag(torch.diag(R))  # [K, K] the diagonal removed.
-    return bool((off >= -tol).all()) and bool(R.sum(dim=1).abs().max() < tol)
-
-
-def exit_rates(R: torch.Tensor) -> torch.Tensor:
-    """lambda_t(x) = -R_t(x, x), the total rate of leaving each state."""
-    return -torch.diag(R)  # [K] one exit rate per state.
-
-
-# %% 03. One row of a rate matrix for DNA
-#
-# A jumps to C at rate 2 and to G and to T at rate 1 each. The chapter gives that one row; the
-# other three repeat the pattern with the next base in ACGT order taking the rate 2.
-DNA_RATES = rate_matrix(torch.tensor([
-    [0.0, 2.0, 1.0, 1.0],  # A: rate 2 toward C, rate 1 toward G and T.
-    [1.0, 0.0, 2.0, 1.0],  # C: rate 2 toward G.
-    [1.0, 1.0, 0.0, 2.0],  # G: rate 2 toward T.
-    [2.0, 1.0, 1.0, 0.0],  # T: rate 2 toward A.
-]))
-
-
-def short_interval_row(R: torch.Tensor, state: int, h: float) -> torch.Tensor:
-    """One row of P_{t,t+h} = I + h R_t + o(h), the first-order transition probabilities."""
-    identity_row = torch.zeros(R.shape[0], dtype=R.dtype)  # [K] the row of I at this state.
-    identity_row[state] = 1.0  # Staying put is the only certainty at h = 0.
-    return identity_row + h * R[state]  # [K] probabilities, valid while the diagonal stays positive.
-
-
-def largest_valid_step(R: torch.Tensor) -> float:
-    """The bound h * sum_{y != x} R_t(x, y) <= 1, read at the state with the largest exit rate."""
-    return float(1.0 / exit_rates(R).max())  # One bad state is enough to give a negative entry.
-
-
-# %% 04. Waiting times and destinations
-#
-# A state x holds for an exponential time of rate lambda(x), because S(s + h) = S(s)(1 - lambda h)
-# gives dS/ds = -lambda S. The destination is then the row of rates normalized to sum to one.
-def sample_path(R: torch.Tensor, start: int, horizon: float, generator) -> list[tuple[float, int]]:
-    """Gillespie simulation: an exponential wait, then a destination drawn from the normalized row."""
-    time_now, state, path = 0.0, start, [(0.0, start)]  # The path begins at the starting base.
-    while True:
-        rate = float(exit_rates(R)[state])  # lambda(x) for the base we are sitting on.
-        u = float(torch.rand((), generator=generator))  # One uniform variate for the wait.
-        time_now += -math.log(u) / rate  # Inverting S(s) = exp(-lambda s) at u.
-        if time_now > horizon:
-            return path  # The last jump falls outside the interval we are watching.
-        row = R[state].clamp(min=0.0)  # [K] the nonnegative rates out of this base.
-        state = int(torch.multinomial(row / row.sum(), 1, generator=generator))  # The destination.
-        path.append((time_now, state))  # Record when the jump happened and where it went.
-
-
-# %% 05. The master equation
-def master_equation(p: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
-    """p_dot = p R, the Kolmogorov forward equation written on probability rows."""
-    return p @ R  # [K] the y-th entry collects the y-th column of R, every rate pointing into y.
-
-
-def arrivals_and_departures(p: torch.Tensor, R: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """The same derivative split into the flux j_t(x -> y) arriving at y and the flux leaving y."""
-    flux = p[:, None] * R  # [K, K] j_t(x -> y) = p_t(x) R_t(x, y), off the diagonal.
-    arriving = flux.sum(dim=0) - torch.diag(flux)  # [K] sum over x != y of p_t(x) R_t(x, y).
-    leaving = -torch.diag(flux)  # [K] p_t(y) times the exit rate of y.
-    return arriving, leaving  # Their difference is the master equation.
-
-
-# %% 06. A worked evolution of four DNA probabilities
-#
-# Every base jumps to every other at rate 1, so each diagonal entry is -3 and each coordinate obeys
-# p_dot(y) = 1 - 4 p(y), whose solution is 1/4 + (p_0(y) - 1/4) exp(-4t).
-UNIT_RATES = rate_matrix(torch.ones(DNA_SIZE, DNA_SIZE))  # [4, 4] all pairwise rates equal to 1.
-START = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.float64)  # p_t over (A, C, G, T).
-
-
-def exact_evolution(p0: torch.Tensor, t: float) -> torch.Tensor:
-    """The closed-form solution of p_dot(y) = 1 - 4 p(y) for the all-ones rate matrix."""
-    return 0.25 + (p0 - 0.25) * math.exp(-4.0 * t)  # [K] each coordinate decays toward 1/4.
-
-
-# %% 07. The matrix exponential and its composition over two intervals
-def transition_matrix(R: torch.Tensor, elapsed: float) -> torch.Tensor:
-    """P_{0,t} = exp(t R), the solution of the master equation for a fixed rate matrix."""
-    return torch.matrix_exp(elapsed * R)  # [K, K] rows sum to one and no entry is negative.
-
-
-def euler_product(R: torch.Tensor, elapsed: float, n: int) -> torch.Tensor:
-    """(I + (t/n) R)^n, the product of n short-interval steps, which approaches exp(t R)."""
-    step = torch.eye(R.shape[0], dtype=R.dtype) + (elapsed / n) * R  # [K, K] one first-order step.
-    return torch.linalg.matrix_power(step, n)  # [K, K] the steps compose by matrix multiplication.
-
-
-# %% 08. The R_pi family and its closed-form kernel
-#
-# At rate one the current token is discarded and a fresh one drawn from pi, so every row of the rate
-# matrix is pi with the current state subtracted. Because M = 1 pi^T is idempotent, the exponential
-# collapses to a mixture of staying put and drawing afresh.
-def r_pi(pi: torch.Tensor) -> torch.Tensor:
-    """R_pi = 1 pi^T - I, the one-free-distribution family."""
-    return torch.ones(len(pi), 1, dtype=pi.dtype) @ pi[None, :] - torch.eye(len(pi), dtype=pi.dtype)
-
-
-def pi_kernel(pi: torch.Tensor, alpha: float) -> torch.Tensor:
-    """Q(alpha) = alpha I + (1 - alpha) 1 pi^T, the kernel at survival probability alpha."""
-    M = torch.ones(len(pi), 1, dtype=pi.dtype) @ pi[None, :]  # [K, K] every row equals pi.
-    return alpha * torch.eye(len(pi), dtype=pi.dtype) + (1.0 - alpha) * M  # [K, K] a convex mixture.
-
-
-# %% 09. Eigenvalues, stationarity and detailed balance
-def detailed_balance_gap(R: torch.Tensor, pi: torch.Tensor) -> float:
-    """The largest mismatch in pi(x) R(x, y) = pi(y) R(y, x), zero when the pair is reversible."""
-    flux = pi[:, None] * R  # [K, K] the stationary flux along each edge.
-    return float((flux - flux.T).abs().max())  # Zero exactly when every edge balances.
-
-
-CYCLE = rate_matrix(torch.tensor([  # A -> C -> G -> A at rate 1, with no jumps the other way.
-    [0.0, 1.0, 0.0],
-    [0.0, 0.0, 1.0],
-    [1.0, 0.0, 0.0],
-]))
-
-
-# %% 10. Time reversal
-def reverse_rates(R: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
-    """R_hat_t(y, x) = R_t(x, y) p_t(x) / p_t(y), the rates of the process watched backward."""
-    reversed_rates = (R * p[:, None] / p[None, :]).T  # [K, K] transpose moves (x, y) to (y, x).
-    reversed_rates.fill_diagonal_(0.0)  # The diagonal is fixed by the row-sum rule, not by the ratio.
-    return reversed_rates - torch.diag(reversed_rates.sum(dim=1))  # [K, K] a rate matrix again.
-
-
-# %% 11. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    # Nothing trains here, so --steps is the Euler grid of Section 09 and --samples the number of
-    # Gillespie paths of Section 04. The flags that describe a network go unread.
-    parser.set_defaults(steps=10, samples=512)
-    args = parser.parse_args()
-    set_seed(args.seed)  # The same seed gives the same sampled path and the same histogram.
-    generator = torch.Generator().manual_seed(args.seed)  # Local stream for the simulations.
-    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.
-
-    # Sections 02 and 03, the row for A and its first-order transition probabilities.
-    row_probabilities = short_interval_row(DNA_RATES, 0, 0.1)  # [4] at h = 0.1.
-    exact_stay = math.exp(-4.0 * 0.1)  # The probability of never leaving A over that interval.
-    talk(f"rate row for A      {[round(float(v), 2) for v in DNA_RATES[0]]}  rules hold "
-         f"{obeys_the_two_rules(DNA_RATES)}  exit rate {float(exit_rates(DNA_RATES)[0]):.0f}")
-    talk(f"P(A, .) at h = 0.1  {[round(float(v), 2) for v in row_probabilities]}  "
-         f"exact stay {exact_stay:.4f}  valid while h <= {largest_valid_step(DNA_RATES):.2f}, "
-         f"since at h = 0.3 the staying entry is "
-         f"{float(short_interval_row(DNA_RATES, 0, 0.3)[0]):.1f}")
-
-    # Section 04, the waiting time and the destination, first from the two uniforms of the chapter.
-    wait = -math.log(0.42) / 4.0  # u = 0.42 inverted through S(s) = exp(-4s).
-    destination = DNA[int(torch.cumsum(DNA_RATES[0].clamp(min=0) / 4.0, 0).lt(0.31).sum())]
-    path = sample_path(DNA_RATES, 0, 1.0, generator)  # One sampled trajectory over the unit interval.
-    ends = torch.tensor([sample_path(DNA_RATES, 0, 1.0, generator)[-1][1]
-                         for _ in range(args.samples)])  # [N] the base each path finishes on.
-    histogram = torch.bincount(ends, minlength=DNA_SIZE).double() / args.samples  # [4] proportions.
-    exact_end = transition_matrix(DNA_RATES, 1.0)[0]  # [4] the exact law of the base at t = 1.
-    talk(f"wait from u = 0.42  {wait:.4f}  destination from v = 0.31  {destination}  "
-         f"mean {1 / 4:.2f}  median {math.log(2) / 4:.4f}")
-    talk(f"sampled path        {' -> '.join(f'{DNA[s]}@{time:.3f}' for time, s in path)}")
-    talk(f"{args.samples} paths at t = 1   {[round(float(v), 3) for v in histogram]}  against the "
-         f"exact {[round(float(v), 3) for v in exact_end]}")
-
-    # Sections 05 and 06, the master equation and the worked evolution.
-    derivative = master_equation(START, UNIT_RATES)  # [4] p_dot = (0.6, 0.2, -0.2, -0.6).
-    arriving, leaving = arrivals_and_departures(START, UNIT_RATES)  # The two halves of that number.
-    euler_step = START + 0.1 * derivative  # One step of size h = 0.1, inside the bound h <= 1/3.
-    talk(f"p_dot at A          {float(arriving[0]):.1f} in minus {float(leaving[0]):.1f} out "
-         f"= {float(derivative[0]):.1f}   p_dot {[round(float(v), 1) for v in derivative]}")
-    talk(f"Euler step h = 0.1  {[round(float(v), 2) for v in euler_step]}  exact "
-         f"{[round(float(v), 4) for v in exact_evolution(START, 0.1)]}  "
-         f"total change {float(derivative.sum()):.0e}")
-    talk(f"exact at t = 1      {[round(float(v), 4) for v in exact_evolution(START, 1.0)]}")
-
-    # Section 07, the matrix exponential and the Chapman-Kolmogorov composition.
-    whole = transition_matrix(UNIT_RATES, 0.7)  # P_{s,u} over an interval of length 0.7.
-    first, second = transition_matrix(UNIT_RATES, 0.3), transition_matrix(UNIT_RATES, 0.4)
-    halves = first @ second  # P_{s,t} P_{t,u}, the two intervals composed.
-    talk(f"exp(tR) rows        sum to {float(whole.sum(1).min()):.6f}, smallest entry "
-         f"{float(whole.min()):.6f}, composition gap {float((whole - halves).abs().max()):.2e}")
-
-    # Section 08, the R_pi family, its closed form and the multiplication of survival probabilities.
-    uniform = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE, dtype=torch.float64)  # pi uniform over ACGT.
-    tau = math.log(2.0)  # The clock the linear schedule reaches at t = 0.5, where alpha = 1/2.
-    closed_form = pi_kernel(uniform, math.exp(-tau))  # alpha I + (1 - alpha) 1 pi^T.
-    series = transition_matrix(r_pi(uniform), tau)  # The same matrix from the series.
-    composed = pi_kernel(uniform, 0.8) @ pi_kernel(uniform, 0.5 / 0.8)  # Q(0.8) Q(0.625).
-    talk(f"closed form vs exp  gap {float((closed_form - series).abs().max()):.2e}   "
-         f"Q(0.8) Q(0.625) at (A, A) {float(composed[0, 0]):.5f} = Q(0.5) entry "
-         f"{float(pi_kernel(uniform, 0.5)[0, 0]):.5f}")
-
-    # Section 09, the matrix exponential worked for DNA, and the masking kernel beside it.
-    euler_entries = {n: float(euler_product(r_pi(uniform), tau, n)[0, 0]) for n in (1, 2, args.steps)}
-    mask_pi = torch.zeros(MASK_STATES, dtype=torch.float64)  # pi = delta_m on the five states.
-    mask_pi[DNA_SIZE] = 1.0  # All of the target probability sits on the mask.
-    mask_kernel = pi_kernel(mask_pi, 0.5)  # [5, 5] the masking kernel at alpha = 1/2.
-    talk(f"exp(tau R_pi)       diagonal {float(closed_form[0, 0]):.3f}, off-diagonal "
-         f"{float(closed_form[0, 1]):.3f}, Euler " +
-         ", ".join(f"n={n}: {v:.3f}" for n, v in euler_entries.items()) +
-         f";  masking A row {[round(float(v), 2) for v in mask_kernel[0]]}, m row "
-         f"{[round(float(v), 2) for v in mask_kernel[DNA_SIZE]]}")
-
-    # Section 10, the eigenvalues and the stationary distribution, then detailed balance.
-    eigenvalues = sorted(float(v.real) for v in torch.linalg.eigvals(r_pi(uniform)))  # 0 and -1.
-    cycle_pi = torch.full((3,), 1.0 / 3.0, dtype=torch.float64)  # Uniform over the three bases.
-    talk(f"eigenvalues of R_pi {[round(v, 4) for v in eigenvalues]}, pi R_pi "
-         f"{float((uniform @ r_pi(uniform)).abs().max()):.2e}, uniform flux "
-         f"{float((uniform[:, None] * r_pi(uniform))[0, 1]):.4f} each way with gap "
-         f"{detailed_balance_gap(r_pi(uniform), uniform):.2e}")
-    talk(f"cycle of three      stationary gap {float((cycle_pi @ CYCLE).abs().max()):.2e}, "
-         f"edge gap {detailed_balance_gap(CYCLE, cycle_pi):.4f}")
-
-    # Section 11, the reverse rates and the marginals they retrace.
-    marginal = torch.tensor([0.5, 0.25, 0.125, 0.125], dtype=torch.float64)  # p_t at alpha_t = 1/2.
-    unit_forward = rate_matrix(torch.ones(DNA_SIZE, DNA_SIZE))  # Forward rate 1 between any two bases.
-    backward = reverse_rates(unit_forward, marginal)  # [4, 4] R_hat_t(y, x).
-    talk(f"reverse row for C   {[round(float(v), 3) for v in backward[1]]}  exit rate "
-         f"{float(exit_rates(backward)[1]):.1f}  mean holding time "
-         f"{1 / float(exit_rates(backward)[1]):.4f}")
-
-    report = {
-        "exit rate out of A": float(exit_rates(DNA_RATES)[0]),
-        "P(A, A) at h = 0.1": float(row_probabilities[0]),
-        "largest valid step": largest_valid_step(DNA_RATES),
-        "wait from u = 0.42": wait,
-        "p_dot at A": float(derivative[0]),
-        "Euler step at A": float(euler_step[0]),
-        "composition gap": float((whole - halves).abs().max()),
-        "exp(tau R_pi) at (A, A)": float(closed_form[0, 0]),
-        "exp(tau R_pi) off diagonal": float(closed_form[0, 1]),
-        f"Euler product, n = {args.steps}": euler_entries[args.steps],
-        "detailed balance gap, cycle": detailed_balance_gap(CYCLE, cycle_pi),
-        "reverse rate C to A": float(backward[1, 0]),
-        "reverse rate C to G": float(backward[1, 2]),
-        "p R_hat plus p R": float((marginal @ backward + marginal @ unit_forward).abs().max()),
-    }
-    print_report("ctmc: rates, marginals and their reversal", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-# %% 12. Run the script
-if __name__ == "__main__":  # Execute when launched directly.
-    main()  # Work every number of the sections above.

lecture_4/d3pm.py

@@ -1,234 +0,0 @@
-"""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. Reversing one
-finite step is Bayes' rule on 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
-#
-# Nothing in the five facts requires that every pair of tokens be treated alike. 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.
-    # Drawing one step index uniformly estimates a sum of T terms, so the divergence is scaled by 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, which is 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.

lecture_4/guidance.py

@@ -1,308 +0,0 @@
-"""Guiding masked diffusion on DNA with Doob's h-transform.
-
-Conditioning a jump process on an event observed at the end of its path multiplies every rate by
-the ratio of the value function at the destination to the value at the current state. The reverse
-process is itself such a transform, with h_t = p_t/pi, and a second transform by a classifier aims
-generation at a property. Three estimates of that classifier appear here: the exact one, its
-first-order expansion in the one-hot input, and the conditional denoiser of classifier-free
-guidance.
-
-Chapter 4, "Conditioning a jump process: Doob's h-transform", "The reverse process as an
-h-transform" and "Guiding generation toward a property". 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, whose label is a GC fraction of at
-least 0.6. Prints the worked guided jump and the GC fraction under each kind of guidance.
-
-    python lecture_4/guidance.py              # train, then sample four ways
-    python lecture_4/guidance.py --gamma 3    # a sharper transform
-    python lecture_4/guidance.py --steps 20   # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import torch  # Tensors, gradients and the categorical draws.
-import torch.nn as nn  # The classifier's layers.
-import torch.nn.functional as F  # Cross-entropy and the one-hot encoding.
-
-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 TimeEmbedding, SequenceDenoiser, count_parameters
-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.
-NULL = 2  # The null label the conditional denoiser drops to, one past the two real labels.
-DROP = 0.1  # How often training replaces the label by the null one.
-TIME_FLOOR = 1e-3  # The reveal rate w(t) = 1/t is unbounded at t = 0.
-
-
-# %% 02. Doob's h-transform
-#
-# R^h_t(x, y) = R_t(x, y) h_t(y)/h_t(x). A jump that raises the chance of the event is faster by
-# the factor by which it raises it, the graph of allowed jumps is unchanged, and a constant h
-# changes nothing, which is the statement that conditioning on a certain event does nothing.
-def h_transform(R: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
-    """Reweight every off-diagonal rate by h(y)/h(x), then refill the diagonal. Returns [K, K]."""
-    transformed = R * (h[None, :] / h[:, None])  # [K, K] one ratio per edge.
-    transformed.fill_diagonal_(0.0)  # The diagonal is fixed by the row-sum rule.
-    return transformed - torch.diag(transformed.sum(dim=1))  # [K, K] a rate matrix again.
-
-
-# %% 03. The reverse process as an h-transform
-#
-# For a forward process in detailed balance with pi, the reverse rate equals the forward rate times
-# h_t(x)/h_t(y) with h_t = p_t/pi, so generation is the noise process aimed at the states that are
-# more likely under the data than under pure noise.
-def value_from_marginal(p: torch.Tensor, pi: torch.Tensor) -> torch.Tensor:
-    """h_t = p_t / pi, defined wherever pi is positive, as uniform noise is and masking is not."""
-    return p / pi  # [K] the excess of the data marginal over the noise distribution.
-
-
-# %% 04. A classifier on noisy sequences
-#
-# The value function for a property is h_t(x) = p(c | X_t = x), estimated by a classifier trained at
-# every noise level. The input is a one-hot sequence, so a gradient with respect to it exists and
-# Section 06 can take one.
-class NoisyClassifier(nn.Module):
-    """A small network from a one-hot noisy sequence and a time to the logit of p(c | z, t)."""
-
-    def __init__(self, length: int, width: int = 128):
-        super().__init__()
-        self.embed = nn.Linear(VOCAB, width, bias=False)  # One-hot in, so the input has a gradient.
-        self.position = nn.Embedding(length, width)  # Where each token sits.
-        self.time = TimeEmbedding(width)  # How far along the corruption we are.
-        self.body = nn.Sequential(nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width),
-                                  nn.SiLU())
-        self.head = nn.Linear(width, 1)  # One logit per sequence.
-
-    def forward(self, one_hot: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-        """one_hot: [B, L, V], t: [B] -> [B] the logit of the property."""
-        positions = torch.arange(one_hot.shape[1], device=one_hot.device)  # [L].
-        h = (self.embed(one_hot) + self.position(positions)[None]
-             + self.time(t)[:, None, :])  # [B, L, width].
-        return self.head(self.body(h).mean(dim=1))[:, 0]  # [B] pooled over the positions.
-
-
-def as_one_hot(z: torch.Tensor) -> torch.Tensor:
-    """The sequence over the masked alphabet as floats, so a gradient in the input exists."""
-    return F.one_hot(z, VOCAB).to(torch.get_default_dtype())  # [B, L, V] one row per position.
-
-
-def classifier_loss(classifier: nn.Module, x0: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
-    """Binary cross-entropy on sequences corrupted to a random noise level."""
-    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one level each.
-    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)[:, None]  # [B, L] the survival draw.
-    z = torch.where(keep, x0, MASK)  # [B, L] a partially masked sequence.
-    logit = classifier(as_one_hot(z), t)  # [B] the predicted log odds of the property.
-    return F.binary_cross_entropy_with_logits(logit, labels.to(torch.get_default_dtype()))
-
-
-# %% 05. Exact classifier guidance
-#
-# The guided rate has p(c | z^{i <- a}) in it at every position and every token, so the exact
-# version costs one classifier evaluation per candidate. A length-16 DNA sequence has 64 of them.
-def exact_ratio(classifier: nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-    """p(c | z^{i <- a}) / p(c | z) at every position and clean token. Returns [B, L, K]."""
-    batch, length = z.shape
-    one_hot = as_one_hot(z)  # [B, L, V] the current state.
-    here = torch.sigmoid(classifier(one_hot, t))[:, None, None]  # [B, 1, 1] the value at z.
-    candidates = one_hot[:, None, None].repeat(1, length, DNA_SIZE, 1, 1)  # [B, L, K, L, V].
-    index = torch.arange(length, device=z.device)  # [L] the position being changed.
-    candidates[:, index, :, index] = torch.eye(VOCAB, device=z.device)[:DNA_SIZE]  # Set position i.
-    flat = candidates.reshape(batch * length * DNA_SIZE, length, VOCAB)  # One batch of candidates.
-    time = t[:, None, None].expand(batch, length, DNA_SIZE).reshape(-1)  # [B L K] the same time.
-    value = torch.sigmoid(classifier(flat, time)).view(batch, length, DNA_SIZE)  # [B, L, K].
-    return value / here.clamp(min=1e-8)  # [B, L, K] the h-transform factor on each jump.
-
-
-# %% 06. The gradient approximation
-#
-# One backward pass through the classifier gives every ratio at once, because a first-order
-# expansion of log p(c | .) in the one-hot input differs between two candidates only in the two
-# coordinates that change.
-def gradient_ratio(classifier: nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-    """exp(g_{i,a} - g_{i,z^i}) with g the gradient of log p(c | z) in the one-hot. [B, L, K]."""
-    one_hot = as_one_hot(z).requires_grad_(True)  # [B, L, V] the point of the expansion.
-    with torch.enable_grad():  # The sampler runs under no_grad, and this one pass needs gradients.
-        log_value = F.logsigmoid(classifier(one_hot, t)).sum()  # One backward covers the batch.
-        gradient, = torch.autograd.grad(log_value, one_hot)  # [B, L, V] the expansion's slopes.
-    current = gradient.gather(2, z[:, :, None])  # [B, L, 1] the coordinate that is being left.
-    return (gradient[..., :DNA_SIZE] - current).exp()  # [B, L, K] the approximated ratio.
-
-
-# %% 07. The guided sampler
-#
-# The reveal rate at a masked position is w(t) d_theta(a), so the h-transform multiplies each of
-# those rates by its own ratio. Both the probability of revealing anything and the letter that
-# appears change, and a ratio of one at every token returns the unguided sampler exactly.
-def denoiser_probabilities(model: nn.Module, z: torch.Tensor, time: torch.Tensor,
-                           labels: torch.Tensor | None) -> torch.Tensor:
-    """d_theta(. | z, t) over the four bases, with the zero-mask rule applied. Returns [B, L, K]."""
-    logits = model(z, time, labels)  # [B, L, 5] the clean-token prediction.
-    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
-    return logits.softmax(-1)[..., :DNA_SIZE]  # [B, L, K], which sums to one over ACGT.
-
-
-def guided_weights(model: nn.Module, z: torch.Tensor, time: torch.Tensor,
-                   labels: torch.Tensor | None, ratio_fn, gamma: float) -> torch.Tensor:
-    """The reveal rates over w(t): the denoiser times its h-transform ratio. Returns [B, L, K]."""
-    weights = denoiser_probabilities(model, z, time, labels)  # [B, L, K] the unguided rates.
-    if ratio_fn is None:
-        return weights  # [B, L, K] summing to one, which is the unguided exit rate w(t).
-    return weights * ratio_fn(z, time).clamp(min=1e-8) ** gamma  # [B, L, K] the h-transform.
-
-
-@torch.no_grad()  # Only the gradient ratio reopens the graph, and it does so locally.
-def sample(model: nn.Module, labels: torch.Tensor | None, batch: int, length: int,
-           steps: int = 32, ratio_fn=None, gamma: float = 1.0) -> torch.Tensor:
-    """Generate from all masks, with the reveal rates reweighted by ratio_fn. Returns [B, L]."""
-    device = next(model.parameters()).device
-    z = torch.full((batch, length), MASK, device=device)  # [B, L] the all-mask start.
-    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
-    for t, s in zip(ts[:-1], ts[1:]):
-        time = t.expand(batch)  # [B] the shared time of this step.
-        weights = guided_weights(model, z, time, labels, ratio_fn, gamma)  # [B, L, K].
-        total = weights.sum(-1, keepdim=True)  # [B, L, 1] the transformed exit rate over w(t).
-        reveal = ((t - s) / t * total).clamp(max=1.0)  # [B, L, 1] the probability of acting at all.
-        drawn = torch.distributions.Categorical(probs=weights / total).sample()  # [B, L].
-        open_now = torch.rand(z.shape, device=device)[:, :, None] < reveal  # [B, L, 1].
-        z = torch.where((z == MASK) & open_now[..., 0], drawn, z)  # Carry-over unmasking.
-    # The last step runs down to s = 0, where the schedule reveals every remaining mask. The same
-    # transform picks the letter, so the guidance reaches the positions that waited to the end.
-    last = ts[-1].expand(batch)  # [B] the final time of the grid.
-    weights = guided_weights(model, z, last, labels, ratio_fn, gamma)  # [B, L, K].
-    drawn = torch.distributions.Categorical(  # [B, L] one letter per position.
-        probs=weights / weights.sum(-1, keepdim=True)).sample()
-    return torch.where(z == MASK, drawn, z)  # [B, L] clean DNA.
-
-
-# %% 08. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--gamma", type=float, default=2.0, help="exponent on the h ratio")
-    # Two networks train and four samplers run, so each network gets 1000 steps and each sampler
-    # measures 64 sequences, which is the longest default run of the lecture.
-    parser.set_defaults(steps=1000, sample_steps=32, 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.
-
-    # Sections 02 and 03, the reverse rates read off h_t = p_t / pi at alpha_t = 1/2.
-    marginal = torch.tensor([0.5, 0.25, 0.125, 0.125])  # p_t over (A, C, G, T).
-    pi = torch.full((DNA_SIZE,), 0.25)  # Uniform noise.
-    unit = torch.ones(DNA_SIZE, DNA_SIZE) - torch.diag(torch.full((DNA_SIZE,), DNA_SIZE * 1.0))
-    values = value_from_marginal(marginal, pi)  # [4] h_t = (2, 1, 0.5, 0.5).
-    reverse = h_transform(unit, values)  # [4, 4] the reverse rates, at unit forward rates.
-    conditioned = value_from_marginal(torch.tensor([0.625, 0.125, 0.125, 0.125]), pi)  # Clean A.
-    talk(f"h from the marginal {[round(float(v), 2) for v in values]}  reverse row for C "
-         f"{[round(float(v), 2) for v in reverse[1]]}")
-    talk(f"h given a clean A   {[round(float(v), 2) for v in conditioned]}  rate C to A "
-         f"{float(conditioned[0] / conditioned[1]):.0f}  given a clean C "
-         f"{float(conditioned[1] / conditioned[0]):.1f}")
-
-    # Section 05 on paper, the worked guided jump out of C with a classifier on the four bases.
-    unguided = torch.tensor([2.0, 0.5, 0.5])  # The reverse rates toward A, G and T.
-    classifier_value = torch.tensor([0.1, 0.6, 0.2]) / 0.2  # p(c | destination) over p(c | C).
-    guided = unguided * classifier_value  # [3] the guided rates.
-    sharper = unguided * classifier_value ** 2  # [3] the same transform at gamma = 2.
-    shares = [round(float(v), 4) for v in guided / guided.sum()]  # Where the jump most likely goes.
-    talk(f"guided jump         {[round(float(v), 2) for v in guided]}  exit "
-         f"{float(guided.sum()):.0f}  destinations {shares}")
-    talk(f"at gamma = 2        {[round(float(v), 2) for v in sharper]}  exit "
-         f"{float(sharper.sum()):.1f}  destinations "
-         f"{[round(float(v), 4) for v in sharper / sharper.sum()]}  holding "
-         f"{1 / float(sharper.sum()):.4f}")
-
-    # Train the conditional denoiser and the classifier on the shared synthetic DNA.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
-    labels = splits["train_labels"].to(device)  # [N] one when the GC fraction is at least 0.6.
-    model = SequenceDenoiser(VOCAB, args.length, width=args.width, num_classes=2).to(device)
-    classifier = NoisyClassifier(args.length, args.width).to(device)
-    share = float(labels.to(torch.get_default_dtype()).mean())  # How much of the data is GC-rich.
-    talk(f"device={device} denoiser={count_parameters(model):,} "
-         f"classifier={count_parameters(classifier):,} labelled {share:.2f}")
-
-    def denoiser_loss(step: int) -> torch.Tensor:
-        """The MDLM loss, with the label replaced by the null one a tenth of the time."""
-        rows = torch.randint(0, len(tokens), (args.batch_size,), device=device)
-        x0, y = tokens[rows], labels[rows].clone()
-        y[torch.rand(len(y), device=device) < DROP] = NULL  # Dropping the label trains both heads.
-        t = torch.rand(len(x0), device=device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
-        zt = torch.where(torch.rand(x0.shape, device=device) < (1.0 - t)[:, None], x0, MASK)
-        logits = model(zt, t, y)  # [B, L, 5].
-        logits[..., MASK] = -1e9  # The zero-mask rule.
-        ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
-        return ((1.0 / t)[:, None] * ce * (zt == MASK)).sum(-1).mean()
-
-    def labelled_batch(step: int) -> torch.Tensor:
-        """One batch of clean sequences and their labels, for the classifier."""
-        rows = torch.randint(0, len(tokens), (args.batch_size,), device=device)
-        return classifier_loss(classifier, tokens[rows], labels[rows])
-
-    losses = train(model, denoiser_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
-    classifier_losses = train(classifier, labelled_batch, steps=args.steps, lr=args.lr,
-                              quiet=args.quiet)
-
-    # Section 06, the exact ratios against their first-order approximation on one noisy batch.
-    probe_t = torch.full((32,), 0.5, device=device)  # [32] halfway through the corruption.
-    probe = torch.where(torch.rand(tokens[:32].shape, device=device) < 0.5, tokens[:32], MASK)
-    with torch.no_grad():
-        exact = exact_ratio(classifier, probe, probe_t)  # [32, L, K] every candidate evaluated.
-    approximate = gradient_ratio(classifier, probe, probe_t)  # [32, L, K] one backward pass.
-    gap = exact.clamp(min=1e-6).log() - approximate.clamp(min=1e-6).log()  # [32, L, K] log ratios.
-    agreement = float(gap.abs().mean())  # How far the expansion sits from the exact ratio.
-
-    # Sample four ways, from the same weights and the same starting state.
-    wanted = torch.ones(args.samples, dtype=torch.long, device=device)  # Ask for the GC-rich label.
-    null = torch.full_like(wanted, NULL)  # The label the unconditional head reads.
-
-    def free_ratio(z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-        """d_theta(a | z, c) / d_theta(a | z), which is the classifier ratio implied by Bayes."""
-        # Both denoisers go through the zero-mask rule first. The loss overwrites the mask logit,
-        # so it never trains, and a softmax that kept it would leave a stray factor in the ratio.
-        conditional = denoiser_probabilities(model, z, t, wanted)  # [B, L, K] with the label.
-        unconditional = denoiser_probabilities(model, z, t, null)  # [B, L, K] with the null label.
-        return conditional / unconditional.clamp(min=1e-8)  # [B, L, K] the h-transform factor.
-
-    runs = {}
-    for name, ratio in (
-        ("unguided", None),
-        ("exact classifier", lambda z, t: exact_ratio(classifier, z, t)),
-        ("gradient approximation", lambda z, t: gradient_ratio(classifier, z, t)),
-        ("classifier free", free_ratio),
-    ):
-        torch.manual_seed(args.seed + 1)  # The same reveal draws for every run.
-        runs[name] = sample(model, null, args.samples, args.length, args.sample_steps,
-                            ratio, args.gamma)
-
-    report = {
-        "h at A from the marginal": float(values[0]),
-        "reverse rate C to A": float(reverse[1, 0]),
-        "guided rate C to A": float(guided[0]),
-        "guided rate C to G": float(guided[1]),
-        "guided exit rate": float(guided.sum()),
-        "guided exit rate at gamma = 2": float(sharper.sum()),
-        "share on G at gamma = 2": float((sharper / sharper.sum())[1]),
-        "denoiser loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
-        "classifier loss, last 50": sum(classifier_losses[-50:]) / len(classifier_losses[-50:]),
-        "mean log-ratio gap, exact vs gradient": agreement,
-        "GC fraction in the data": gc_fraction(tokens),
-        "motif fraction in the data": motif_fraction(tokens),
-        "gamma": args.gamma,
-    }
-    for name, drawn in runs.items():
-        report[f"GC fraction, {name}"] = gc_fraction(drawn)
-        # The label is a GC threshold, so the motif is the structure guidance leaves alone.
-        report[f"motif fraction, {name}"] = motif_fraction(drawn)
-        report[f"distinct, {name}"] = unique_fraction(drawn)
-    report["letter entropy, exact classifier"] = token_entropy(runs["exact classifier"], DNA_SIZE)
-    print_report("guidance: h-transforms of masked diffusion", report,
-                 decode_dna(runs["exact classifier"]))
-    save_run(args.out, model=model, config=vars(args), report=report, losses=losses,
-             samples=decode_dna(runs["exact classifier"]))
-
-
-# %% 09. Run the script
-if __name__ == "__main__":  # Execute when launched directly.
-    main()  # Work the guided jump, train the two networks, and sample four ways.

lecture_4/mdlm.py

@@ -1,190 +0,0 @@
-"""Masked diffusion on DNA, the whole model in thirty lines.
-
-Corruption replaces each base by the mask with probability 1 - alpha_t, a network reads the masked
-sequence and predicts the clean token at every position, and generation reveals 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 any noise level is reachable in one draw.
-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 summing the steps into an integral
-# replaces r by 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, which is -alpha_dot / (1 - alpha) on this schedule.
-    return (w[:, None] * ce * (zt == MASK)).sum(-1).mean()  # Visible positions contribute zero.
-
-
-# %% 04. The sampler
-#
-# A masked position is revealed with probability r_{s,t} = (alpha_s - alpha_t)/(1 - alpha_t), which
-# is (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, which the linear schedule reaches 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 and of nothing else.
-
-
-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: summing l(A, B) over every
-# destination at one masked position leaves 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), which is 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, which is 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.

lecture_4/planner.py

@@ -1,238 +0,0 @@
-"""Choosing the order of the reveals in masked diffusion on DNA.
-
-Every masked position leaves at the same total rate w(t), so which positions a step reveals is free
-and the marginal reverse rates stay where they were. A planner spends that freedom on the positions
-the denoiser is most sure about, a corrector returns the least supported visible tokens to the
-mask, and the loss can be reweighted toward the positions the sampler reaches first.
-
-Chapter 4, "Choosing the order of the reveals", "One sampling step with a planner" and "Bringing
-the planner into training". 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 worked selection, the worked PAPL weights, and the
-sample statistics under both samplers.
-
-    python lecture_4/planner.py                     # train with planner weights, sample both ways
-    python lecture_4/planner.py --remask 0          # confidence order with no corrector
-    python lecture_4/planner.py --steps 20          # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import torch  # Tensors, gradients and the categorical draws.
-import torch.nn.functional as F  # Cross-entropy at every position.
-
-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 reveal rate w(t) = 1/t is unbounded at t = 0.
-
-
-# %% 02. The planner score and the selection it drives
-#
-# The score condenses the distribution at a position into one number. The probability of the most
-# likely clean token is high where the denoiser has committed and low where it is spread out, and
-# the negative entropy of the same distribution is the other natural choice.
-def planner_score(probabilities: torch.Tensor, kind: str = "confidence") -> torch.Tensor:
-    """One number per position from d_theta(. | z, t). probabilities: [B, L, K] -> [B, L]."""
-    if kind == "entropy":
-        return (probabilities * probabilities.clamp(min=1e-20).log()).sum(-1)  # Negative entropy.
-    return probabilities.max(-1).values  # The probability of the most likely clean token.
-
-
-def topk_where(score: torch.Tensor, eligible: torch.Tensor, count: int) -> torch.Tensor:
-    """A boolean mask selecting the ``count`` highest-scoring eligible positions in each row."""
-    picked = torch.zeros_like(eligible)  # [B, L] nothing selected yet.
-    if count <= 0:
-        return picked  # Taking none is a valid choice, and it recovers the plain sampler.
-    ranked = score.masked_fill(~eligible, float("-inf"))  # Ineligible positions never win.
-    index = ranked.topk(min(count, score.shape[1]), dim=1).indices  # [B, count] the best positions.
-    return picked.scatter_(1, index, True) & eligible  # Rows with too few eligible slots take fewer.
-
-
-# %% 03. One sampling step with a planner
-#
-# The first four lines are the sampler of the implementation section: one forward pass, the
-# zero-mask rule, a softmax, and one categorical draw per position. Nothing in them depends on the
-# planner, so a trained checkpoint can be sampled this way with no retraining.
-def planner_step(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor,
-                 n_unmask: int, n_remask: int, kind: str = "confidence") -> torch.Tensor:
-    """One reverse step that reveals the best ``n_unmask`` masks and remasks the worst tokens."""
-    logits = model(z, t)  # [B, L, 5] the clean-token prediction at this noise level.
-    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
-    p = logits.softmax(-1)  # [B, L, 5] the denoiser as probabilities.
-    x0 = torch.distributions.Categorical(probs=p).sample()  # [B, L] one candidate per position.
-
-    score = planner_score(p[..., :DNA_SIZE], kind)  # [B, L] how sure the denoiser is here.
-    masked = z == MASK  # [B, L] the positions a reveal may act on.
-    pick = topk_where(score, masked, n_unmask)  # [B, L] the best masked positions.
-    z = torch.where(pick, x0, z)  # Unmask the best.
-
-    back = topk_where(-score, ~masked, n_remask)  # [B, L] the worst visible positions.
-    return torch.where(back, torch.full_like(z, MASK), z)  # Remask the worst, which is the corrector.
-
-
-# %% 04. The corrector and its flux balance
-#
-# C_t = R_t + R_hat_t is in detailed balance with p_t, so running it at a fixed time holds the
-# marginal in place. Under masking its forward part returns a visible token to the mask at rate
-# beta(t) and its reverse part reveals a masked position at rate w(t).
-def corrector_fluxes(t: float) -> tuple[float, float]:
-    """alpha_t beta(t) and (1 - alpha_t) w(t), the two fluxes across the visible-masked boundary."""
-    alpha = 1.0 - t  # The linear schedule.
-    return alpha * (1.0 / alpha), (1.0 - alpha) * (1.0 / t)  # Both equal -alpha_dot, which is one.
-
-
-# %% 05. Bringing the planner into training
-#
-# The MDLM loss weights every masked position equally, which is right when the order of the reveals
-# is uniform. Once a planner chooses the order, the weight lambda_i = M rho_i puts the gradient on
-# the positions the sampler will reach first, and it averages to one over the masked positions.
-def planner_weights(score: torch.Tensor, masked: torch.Tensor, strength: float) -> torch.Tensor:
-    """lambda_i = 1 + strength (M rho_i - 1), with rho the planner's softmax over masked positions."""
-    # A large negative number in place of -inf, so a row with nothing masked still has a softmax.
-    ranked = score.masked_fill(~masked, -1e4)  # [B, L] only masked positions compete.
-    rho = ranked.softmax(dim=1)  # [B, L] the planner's distribution over which position to fill.
-    count = masked.sum(dim=1, keepdim=True).clamp(min=1)  # [B, 1] M, the number of masked positions.
-    return 1.0 + strength * (count * rho - 1.0)  # [B, L] a weight of one at strength zero.
-
-
-def papl_loss(model: torch.nn.Module, x0: torch.Tensor, strength: float) -> torch.Tensor:
-    """The MDLM loss with the planner weight inside the sum over masked positions."""
-    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
-    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)[:, None]  # [B, L] the survival draw.
-    zt = torch.where(keep, x0, MASK)  # [B, L] the corrupted sequence.
-    logits = model(zt, t)  # [B, L, 5] the clean-token prediction.
-    logits[..., MASK] = -1e9  # The zero-mask rule.
-    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
-    masked = zt == MASK  # [B, L] where the loss lives.
-    score = planner_score(logits.softmax(-1)[..., :DNA_SIZE]).detach()  # No gradient through these.
-    weights = planner_weights(score, masked, strength).clamp(min=0.0)  # [B, L] lambda_i.
-    return ((1.0 / t)[:, None] * weights * ce * masked).sum(-1).mean()  # Nats per sequence.
-
-
-# %% 06. Sampling with and without the planner
-#
-# The number revealed per step follows the schedule: after a step down to s, the masked count is
-# the one the survival probability calls for, which is s L on the linear schedule.
-@torch.no_grad()  # Generation needs no gradient tracking.
-def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 32,
-           n_remask: int = 0, kind: str = "confidence", planner: bool = True) -> torch.Tensor:
-    """Generate from all masks. With planner false the reveals are drawn uniformly at random."""
-    device = next(model.parameters()).device
-    z = torch.full((batch, length), MASK, device=device)  # [B, L] the all-mask start.
-    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
-    for t, s in zip(ts[:-1], ts[1:]):
-        target = int(round(float(s) * length))  # The masked count the schedule calls for at s.
-        current = int(z.eq(MASK).sum(dim=1).max())  # The largest masked count in the batch.
-        n_unmask = max(current - target, 0) + n_remask  # Enough to reach the target after remasking.
-        if planner:
-            z = planner_step(model, z, t.expand(batch), n_unmask, n_remask, kind)
-        else:  # The uniform order of the plain sampler, with the same number of reveals.
-            z = _random_step(model, z, t.expand(batch), n_unmask)
-    logits = model(z, ts[-1].expand(batch))  # The last step runs down to s = 0, where r = 1.
-    logits[..., MASK] = -1e9
-    return torch.where(z == MASK, torch.distributions.Categorical(logits=logits).sample(), z)
-
-
-def _random_step(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor,
-                 n_unmask: int) -> torch.Tensor:
-    """The same step with the masked positions ranked by a uniform draw in place of the score."""
-    logits = model(z, t)  # [B, L, 5].
-    logits[..., MASK] = -1e9  # The zero-mask rule.
-    x0 = torch.distributions.Categorical(logits=logits).sample()  # [B, L].
-    noise = torch.rand(z.shape, device=z.device)  # [B, L] an order with no preference.
-    return torch.where(topk_where(noise, z == MASK, n_unmask), x0, z)
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--remask", type=int, default=1, help="positions the corrector sends back")
-    parser.add_argument("--planner-strength", type=float, default=1.0,
-                        help="0 gives the MDLM loss, 1 gives the PAPL weights")
-    parser.add_argument("--score", choices=["confidence", "entropy"], default="confidence")
-    # Two samplers run on the same weights, so the default trains 1000 steps and measures 64
-    # sequences in each order, which keeps the run inside about a minute.
-    parser.set_defaults(steps=1000, sample_steps=32, 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 selection on a six-position sequence.
-    scores = torch.tensor([[0.42, 0.88, 0.91, 0.55, 0.31, 0.78]])  # [1, 6] one score per position.
-    masked = torch.tensor([[True, False, True, True, False, True]])  # Positions 1, 3, 4 and 6.
-    revealed = topk_where(scores, masked, 2).nonzero()[:, 1] + 1  # The two best masked positions.
-    returned = topk_where(-scores, ~masked, 1).nonzero()[:, 1] + 1  # The worst visible position.
-    talk(f"planner selection   reveals {sorted(int(i) for i in revealed)}  remasks "
-         f"{sorted(int(i) for i in returned)}")
-
-    # Section 04, the two fluxes across the boundary between visible and masked.
-    talk("corrector flux      " + "  ".join(
-        f"t={t}: beta flux {a:.1f}, w flux {b:.1f}" for t in (0.5, 0.25)
-        for a, b in [corrector_fluxes(t)]))
-
-    # Section 05, the worked PAPL weights on a state with four masked positions.
-    rho = torch.tensor([[0.10, 0.50, 0.15, 0.25]])  # [1, 4] the planner's distribution.
-    cross_entropies = torch.tensor([[1.2, 0.3, 0.9, 0.5]])  # [1, 4] nats at each masked position.
-    lam = 4.0 * rho  # [1, 4] lambda_i = M rho_i, which averages to one.
-    talk(f"PAPL weights        {[round(float(v), 1) for v in lam[0]]}  sum {float(lam.sum()):.1f}  "
-         f"MDLM inner sum {float(cross_entropies.sum()):.2f}  PAPL inner sum "
-         f"{float((lam * cross_entropies).sum()):.2f}")
-
-    # Train the denoiser with the planner weights inside the loss.
-    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 loss below.
-    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
-    talk(f"device={device} parameters={count_parameters(model):,} strength={args.planner_strength}")
-
-    def step_loss(step: int) -> torch.Tensor:
-        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
-        return papl_loss(model, train_tokens[rows], args.planner_strength)
-
-    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
-    with torch.no_grad():
-        weighted = sum(float(papl_loss(model, valid_tokens, args.planner_strength))
-                       for _ in range(4)) / 4.0
-        plain = sum(float(papl_loss(model, valid_tokens, 0.0)) for _ in range(4)) / 4.0
-
-    # Sample the same weights twice, once in a uniform order and once in the planner's order.
-    runs = {}
-    for name, planner, remask in (("uniform order", False, 0),
-                                  ("planner order", True, args.remask)):
-        torch.manual_seed(args.seed + 1)  # The same draws for both samplers.
-        runs[name] = sample(model, args.samples, args.length, args.sample_steps,
-                            remask, args.score, planner)
-
-    report = {
-        "revealed positions": sorted(int(i) for i in revealed),
-        "remasked position": sorted(int(i) for i in returned),
-        "flux at t = 0.5": corrector_fluxes(0.5)[0],
-        "flux at t = 0.25": corrector_fluxes(0.25)[1],
-        "PAPL weight at the planner's pick": float(lam[0, 1]),
-        "MDLM inner sum": float(cross_entropies.sum()),
-        "PAPL inner sum": float((lam * cross_entropies).sum()),
-        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
-        "validation PAPL loss": weighted,
-        "validation MDLM loss": plain,
-        "corrector remasks per step": args.remask,
-    }
-    for name, drawn in runs.items():
-        report[f"GC fraction, {name}"] = gc_fraction(drawn)
-        report[f"motif fraction, {name}"] = motif_fraction(drawn)
-        report[f"distinct, {name}"] = unique_fraction(drawn)
-    report["GC fraction in the data"] = gc_fraction(train_tokens)
-    report["motif fraction in the data"] = motif_fraction(train_tokens)
-    report["letter entropy, planner order"] = token_entropy(runs["planner order"], DNA_SIZE)
-    print_report("planner: the order of the reveals", report, decode_dna(runs["planner order"]))
-    save_run(args.out, model=model, config=vars(args), report=report,
-             losses=losses, samples=decode_dna(runs["planner order"]))
-
-
-# %% 08. Run the script
-if __name__ == "__main__":  # Execute when launched directly.
-    main()  # Work the selection, train with the planner weights, and sample both ways.

lecture_4/search.py

@@ -1,250 +0,0 @@
-"""Guiding masked diffusion on DNA by search over the tree of partial sequences.
-
-For a property computed by an external scorer there is no classifier at intermediate noise levels,
-so the value function is estimated by finishing the sequence and scoring the result. The partial
-sequences form a tree whose edges are single reveals, selection walks it by an upper-confidence
-rule, a rollout scores a new leaf, and the reward propagates back up the branch. Two objectives
-that no single number orders are kept in a Pareto archive.
-
-Chapter 4, "Guiding generation by search" and "Search over the tree of partial sequences". The
-fine-tuning of "Keeping the results of the search" is left to the chapter. 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 worked value estimate, the worked archive,
-and the archive the search finds.
-
-    python lecture_4/search.py                  # train, then search
-    python lecture_4/search.py --searches 64    # a larger tree
-    python lecture_4/search.py --sample-steps 8 # longer rollouts behind every value
-    python lecture_4/search.py --steps 20       # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math  # The upper-confidence rule is a square root of a logarithm.
-
-import torch  # Tensors, gradients and the categorical draws.
-import torch.nn.functional as F  # Cross-entropy at every position.
-
-from cis6270.data import DNA_SIZE, MOTIF, decode_dna, dna_splits, encode_dna
-from cis6270.metrics import gc_fraction, motif_fraction, 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 reveal rate w(t) = 1/t is unbounded at t = 0.
-EXPLORATION = 1.0  # c_uct, the weight on the visit-count term of the selection rule.
-
-
-# %% 02. The two toy objectives
-#
-# Binding, solubility and permeability do not reduce to one number without a weighting that fixes
-# the trade-off in advance. Two scores stand in for them here, a GC fraction and the best partial
-# match to the planted motif, both in [0, 1] and neither implying the other.
-def objectives(tokens: torch.Tensor) -> torch.Tensor:
-    """Returns [N, 2]: the GC fraction of each sequence and its best match to the motif."""
-    gc = ((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean(dim=1)  # [N] C and G.
-    motif = encode_dna([MOTIF])[0].to(tokens.device)  # [4] the planted motif ACGT.
-    windows = tokens.unfold(1, len(motif), 1)  # [N, L - 3, 4] every window of the motif's length.
-    hits = (windows == motif).to(torch.get_default_dtype()).mean(dim=2)  # [N, L - 3] per window.
-    match = hits.max(dim=1).values  # [N] the best window of each sequence.
-    return torch.stack([gc, match], dim=1)  # [N, 2] one row per sequence.
-
-
-# %% 03. Rollouts as a value estimate
-#
-# h_t(z) is the probability that a sequence continuing from z ends with the property, so running
-# the sampler to completion N times and averaging the reward estimates it without bias, with a
-# variance that falls like 1/N.
-@torch.no_grad()  # Rollouts need no gradient tracking.
-def complete(model: torch.nn.Module, z: torch.Tensor, t: float, steps: int = 3) -> torch.Tensor:
-    """Run the masked sampler from z at time t down to clean DNA. z: [N, L] -> [N, L]."""
-    device = z.device
-    ts = torch.linspace(t, TIME_FLOOR, steps + 1, device=device)  # The remaining reverse grid.
-    for now, nxt in zip(ts[:-1], ts[1:]):
-        logits = model(z, now.expand(len(z)))  # [N, L, 5] the clean-token prediction.
-        logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
-        x0 = torch.distributions.Categorical(logits=logits).sample()  # [N, L].
-        unmask = torch.rand(z.shape, device=device) < (now - nxt) / now  # [N, L] the reveal draw.
-        z = torch.where((z == MASK) & unmask, x0, z)  # Carry-over unmasking.
-    logits = model(z, ts[-1].expand(len(z)))  # The last step runs down to s = 0, where r = 1.
-    logits[..., MASK] = -1e9
-    return torch.where(z == MASK, torch.distributions.Categorical(logits=logits).sample(), z)
-
-
-def value_estimate(model: torch.nn.Module, state: torch.Tensor, t: float, rollouts: int,
-                   steps: int = 3) -> tuple[torch.Tensor, torch.Tensor]:
-    """Average the reward over N completions of one partial sequence. Returns [2] and [N, L]."""
-    batch = state[None, :].repeat(rollouts, 1)  # [N, L] the same state, N independent endings.
-    finished = complete(model, batch, 1.0 if t >= 1.0 else t, steps)  # [N, L] clean sequences.
-    return objectives(finished).mean(dim=0), finished  # [2] the value, and what produced it.
-
-
-# %% 04. The Pareto archive
-#
-# A candidate worse on every axis than one already kept is dominated and discarded, and the rest
-# stay, since each of them beats every other on at least one axis.
-class ParetoArchive:
-    """Every nondominated sequence the search has produced, with its objective vector."""
-
-    def __init__(self) -> None:
-        self.entries: list[tuple[tuple[int, ...], tuple[float, ...]]] = []
-
-    def add(self, sequence: tuple[int, ...], scores: tuple[float, ...]) -> None:
-        """Keep the candidate when nothing in the archive beats it on every axis."""
-        for _, kept in self.entries:
-            if all(k >= s for k, s in zip(kept, scores)) and kept != scores:
-                return  # Dominated, so it adds nothing.
-        self.entries = [(q, kept) for q, kept in self.entries  # Drop what the newcomer dominates.
-                        if not (all(s >= k for s, k in zip(scores, kept)) and kept != scores)]
-        if sequence not in {q for q, _ in self.entries}:
-            self.entries.append((sequence, scores))
-
-
-# %% 05. The tree of partial sequences
-#
-# Selection walks from the root to a leaf by the upper-confidence rule, expansion adds children
-# drawn from the denoiser perturbed by Gumbel noise, a rollout scores one of them, and the reward
-# propagates back up the branch so that every node on it accumulates visits.
-class Node:
-    """A partially unmasked sequence, with the visits and rewards the search has accumulated."""
-
-    def __init__(self, state: torch.Tensor, depth: int, parent=None):
-        self.state, self.depth, self.parent = state, depth, parent
-        self.children: list[Node] = []
-        self.visits, self.total = 0, 0.0  # The visit count and the summed scalar reward.
-
-    @property
-    def mean(self) -> float:
-        """The mean reward at this node, which the selection rule weighs against its visit count."""
-        return self.total / self.visits if self.visits else 0.0
-
-    def best_child(self):
-        """argmax of r_bar_j + c sqrt(log n / n_j), the usual upper-confidence rule."""
-        return max(self.children, key=lambda child: child.mean + EXPLORATION * math.sqrt(
-            math.log(max(self.visits, 1)) / child.visits) if child.visits else float("inf"))
-
-
-@torch.no_grad()  # The expansion needs no gradient tracking.
-def expand(model: torch.nn.Module, node: Node, t: float, branch: int) -> None:
-    """Add children by revealing one masked position, drawn from the denoiser plus Gumbel noise."""
-    masked = (node.state == MASK).nonzero().flatten()  # [m] the positions a reveal may act on.
-    if len(masked) == 0:
-        return  # A complete sequence has no children.
-    logits = model(node.state[None], torch.tensor([t], device=node.state.device))[0]  # [L, 5].
-    logits[..., MASK] = -1e9  # The zero-mask rule.
-    gumbel = -torch.log(-torch.log(torch.rand_like(logits).clamp(min=1e-20)))  # [L, 5].
-    perturbed = (logits + gumbel).masked_fill(  # Taking maxima of this samples the denoiser.
-        (node.state != MASK)[:, None], float("-inf"))
-    flat = perturbed[:, :DNA_SIZE].flatten()  # [L K] one score per position and clean token.
-    for index in flat.topk(min(branch, int(len(masked)) * DNA_SIZE)).indices:  # Distinct reveals.
-        child_state = node.state.clone()  # [L] the parent with one more position filled.
-        child_state[int(index) // DNA_SIZE] = int(index) % DNA_SIZE
-        node.children.append(Node(child_state, node.depth + 1, node))
-
-
-def backup(node: Node, reward: float) -> None:
-    """Add one rollout's reward to every node on the branch from here to the root."""
-    while node is not None:
-        node.visits, node.total = node.visits + 1, node.total + reward
-        node = node.parent
-
-
-# %% 06. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--searches", type=int, default=24, help="tree expansions")
-    parser.add_argument("--branch", type=int, default=3, help="children added at each expansion")
-    parser.add_argument("--rollouts", type=int, default=6, help="completions scored per leaf")
-    # Every leaf is scored by its own rollouts, so --sample-steps is the length of one rollout and
-    # stays short. Training rises to 1000 steps, which keeps the run inside about a minute.
-    parser.set_defaults(steps=1000, sample_steps=3)
-    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 03, the worked value estimate and the reweighting it produces.
-    at_parent = float(torch.tensor([0.1, 0.8, 0.3, 0.2]).mean())  # Four rollouts from z.
-    at_child = float(torch.tensor([0.7, 0.9, 0.6, 0.8]).mean())  # Four rollouts from its child.
-    factor = at_child / at_parent  # The h-transform factor on that edge.
-    sibling = 0.2 / at_parent  # A child whose rollouts averaged 0.2 is slowed by this factor.
-    talk(f"value estimates     h(z) = {at_parent:.2f}  h(z') = {at_child:.2f}  factor "
-         f"{factor:.3f}  a model rate of 1.2 becomes {1.2 * factor:.3f}, and a sibling at 0.2 is "
-         f"slowed by {sibling:.3f}")
-
-    # Section 04, the worked archive on three candidates scored on two axes.
-    worked = ParetoArchive()
-    for name, scores in (("P1", (0.8, 0.3)), ("P2", (0.5, 0.7)), ("P3", (0.4, 0.2))):
-        worked.add(tuple(ord(c) for c in name), scores)  # The name stands in for a sequence.
-    kept = [chr(q[0]) + chr(q[1]) for q, _ in worked.entries]  # Which candidates survived.
-    talk(f"worked archive      keeps {kept}, so P3 is dominated by P1 on both axes")
-
-    # Train the masked diffusion model the search will guide.
-    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.
-    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
-    talk(f"device={device} parameters={count_parameters(model):,} searches={args.searches}")
-
-    def step_loss(step: int) -> torch.Tensor:
-        """The masked diffusion loss, unchanged: the search reweights rates, not weights."""
-        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
-        x0 = train_tokens[rows]
-        t = torch.rand(len(x0), device=device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
-        zt = torch.where(torch.rand(x0.shape, device=device) < (1.0 - t)[:, None], x0, MASK)
-        logits = model(zt, t)  # [B, L, 5].
-        logits[..., MASK] = -1e9  # The zero-mask rule.
-        ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
-        return ((1.0 / t)[:, None] * ce * (zt == MASK)).sum(-1).mean()
-
-    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
-
-    # Run the search, keeping every completed rollout that is not dominated.
-    torch.manual_seed(args.seed + 1)  # Fix the Gumbel draws and the rollouts.
-    root = Node(torch.full((args.length,), MASK, device=device), 0)
-    archive, scored = ParetoArchive(), []  # The archive and every rollout's objective vector.
-    for _ in range(args.searches):
-        node = root  # Selection walks from the root down the tree already built.
-        while node.children:
-            node = node.best_child()
-        time_here = max(1.0 - node.depth / args.length, TIME_FLOOR)  # The level at this depth.
-        expand(model, node, time_here, args.branch)  # Expansion adds children by one reveal each.
-        leaf = node.children[0] if node.children else node  # Score the first new child.
-        value, finished = value_estimate(model, leaf.state, time_here, args.rollouts,
-                                         args.sample_steps)
-        backup(leaf, float(value.mean()))  # One scalar reward guides the selection rule.
-        for sequence, score in zip(finished.tolist(), objectives(finished).tolist()):
-            archive.add(tuple(sequence), tuple(round(v, 4) for v in score))  # Keep the best front.
-            scored.append(score)
-
-    found = torch.tensor([s for _, s in archive.entries])  # [A, 2] the archive's objectives.
-    everything = torch.tensor(scored)  # [R, 2] every rollout the search scored.
-    best = torch.tensor([list(q) for q, _ in archive.entries], device=device)  # [A, L] sequences.
-    report = {
-        "value at the parent": at_parent,
-        "value at the child": at_child,
-        "reweighting factor": factor,
-        "guided rate from 1.2": 1.2 * factor,
-        "factor at a sibling scoring 0.2": sibling,
-        "worked archive size": len(worked.entries),
-        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
-        "rollouts scored": len(scored),
-        "archive size": len(archive.entries),
-        "best GC in the archive": float(found[:, 0].max()),
-        "best motif match in the archive": float(found[:, 1].max()),
-        "mean GC over all rollouts": float(everything[:, 0].mean()),
-        "mean motif match over all rollouts": float(everything[:, 1].mean()),
-        "rollout length, steps": args.sample_steps,
-        "GC fraction in the data": gc_fraction(train_tokens),
-        "motif fraction in the data": motif_fraction(train_tokens),
-        "motif fraction in the archive": motif_fraction(best),
-        "distinct in the archive": unique_fraction(best),
-    }
-    print_report("search: a Pareto archive over partial sequences", report, decode_dna(best))
-    save_run(args.out, model=model, config=vars(args), report=report,
-             losses=losses, samples=decode_dna(best))
-
-
-# %% 07. Run the script
-if __name__ == "__main__":  # Execute when launched directly.
-    main()  # Work the value estimate, train the model, and search under the h-transform.

lecture_4/sedd.py

@@ -1,176 +0,0 @@
-"""SEDD on DNA, a network that predicts the ratios the reverse rates are built from.
-
-The reverse rate at a position is a forward rate times the ratio p_t(z with position i set to a)
-divided by p_t(z), and the two models before this one reached that ratio through a denoiser over
-clean tokens and a sum against the forward kernel. Here the network outputs the ratio itself, and
-the rate error becomes score entropy, whose minimum over the prediction sits at the true ratio.
-
-Chapter 4, "The ratio we have to learn" and "Predicting the ratios directly", on the uniform member
-of the pi family. Time runs from t = 0 clean to t = 1 noise, the survival probability is
-alpha_t = exp(-beta 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 worked score entropy at three predictions and the sample
-statistics.
-
-    python lecture_4/sedd.py                  # train and generate
-    python lecture_4/sedd.py --steps 20       # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math  # The worked score entropies are logarithms.
-
-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  # Reused as L heads of K outputs.
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-BETA = 4.0  # The constant noise rate, so alpha_t = exp(-4t) and the forward rate is BETA/K.
-TIME_FLOOR = 2e-2  # The true ratios grow like 1/(1 - alpha_t), so sampling stops just above zero.
-MAX_RATIO = 1e4  # The predicted ratios are exponentials, kept finite for the sampler.
-
-
-# %% 02. The concrete score
-#
-# s_t(z)_{i,a} = p_t(z^{i <- a}) / p_t(z) holds L(K - 1) numbers per sequence, one per position and
-# per alternative token, which is exactly the shape of a network with L heads of K logits.
-def concrete_score(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-    """s_theta(z)_{i,a}, positive and with no normalization constraint. Returns [B, L, K]."""
-    return model(z, t).clamp(max=math.log(MAX_RATIO)).exp()  # [B, L, K] exponentiate the logits.
-
-
-def kernel_ratio(x0: torch.Tensor, z: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
-    """The conditional target r = q_t(a | x_0) / q_t(z^i | x_0), one per destination. [B, L, K]."""
-    uniform = (1.0 - alpha) / DNA_SIZE  # [B, 1] the redrawn share at this noise level.
-    numerator = uniform[..., None] + alpha[..., None] * torch.nn.functional.one_hot(
-        x0, DNA_SIZE)  # [B, L, K] q_t(a | x_0) at every destination a.
-    denominator = numerator.gather(2, z[:, :, None])  # [B, L, 1] the same kernel at the current token.
-    return numerator / denominator  # [B, L, K] the ratio the network has to reproduce.
-
-
-# %% 03. Denoising score entropy
-#
-# Writing the target rate as lambda r and the model rate as lambda s turns the rate error into
-# lambda [s - r log s + r(log r - 1)]. The first two terms hold every dependence on the prediction,
-# a linear term against a logarithmic one, so the minimum over s at fixed r sits at s = r.
-def score_entropy(target: torch.Tensor, predicted: torch.Tensor,
-                  forward_rate: float = 1.0) -> torch.Tensor:
-    """l(lambda r, lambda s), elementwise, which vanishes at s = r."""
-    constant = target * (target.clamp(min=1e-20).log() - 1.0)  # r(log r - 1), free of the prediction.
-    return forward_rate * (predicted - target * predicted.clamp(min=1e-20).log() + constant)
-
-
-def sedd_loss(model: torch.nn.Module, x0: torch.Tensor) -> torch.Tensor:
-    """One Monte Carlo sample of the score entropy, summed over destinations and positions."""
-    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
-    alpha = torch.exp(-BETA * t)  # [B, 1] the survival probability at that time.
-    keep = torch.rand(x0.shape, device=x0.device) < alpha  # [B, L] survival of the original base.
-    fresh = torch.randint(0, DNA_SIZE, x0.shape, device=x0.device)  # [B, L] a draw from pi.
-    z = torch.where(keep, x0, fresh)  # [B, L] uniform replacement, with no marker of the damage.
-    target = kernel_ratio(x0, z, alpha)  # [B, L, K] the ratio attached to this clean sequence.
-    predicted = concrete_score(model, z, t[:, 0])  # [B, L, K] the network's own ratios.
-    current = torch.nn.functional.one_hot(z, DNA_SIZE).bool()  # [B, L, K] the token already in place.
-    error = score_entropy(target, predicted, BETA / DNA_SIZE)  # [B, L, K] one term per jump.
-    return error.masked_fill(current, 0.0).sum((1, 2)).mean()  # Nats per sequence.
-
-
-# %% 04. Sampling from the predicted ratios
-#
-# The reverse rate is the forward rate of the jump that would undo it, times the predicted ratio,
-# and tau-leaping turns each rate into a jump probability over a step of length h.
-@torch.no_grad()  # Generation needs no gradient tracking.
-def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
-    """Generate from uniformly random tokens down to clean DNA. Returns [B, L] over ACGT."""
-    device = next(model.parameters()).device
-    z = torch.randint(0, DNA_SIZE, (batch, length), device=device)  # [B, L] a draw from pi.
-    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
-    for t, s in zip(ts[:-1], ts[1:]):
-        rates = (BETA / DNA_SIZE) * concrete_score(model, z, t.expand(batch))  # [B, L, K].
-        jump = (t - s) * rates  # [B, L, K] the tau-leaping probability of each jump.
-        jump.scatter_(2, z[:, :, None], 0.0)  # A jump to the current base is not a jump.
-        jump = jump / jump.sum(-1, keepdim=True).clamp(min=1.0)  # Rescale an oversized step.
-        stay = (1.0 - jump.sum(-1, keepdim=True)).clamp(min=0.0)  # [B, L, 1] the chance of staying.
-        choice = torch.distributions.Categorical(
-            probs=torch.cat([jump, stay], dim=-1)).sample()  # [B, L] destination, or K for staying.
-        z = torch.where(choice == DNA_SIZE, z, choice.clamp(max=DNA_SIZE - 1))  # [B, L].
-    return z  # [B, L] clean DNA, since uniform replacement never leaves the alphabet.
-
-
-# %% 05. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    # The ratios need a longer budget than the worked numbers do, so the default trains for 1000
-    # steps and measures on 64 sequences, which keeps the run inside about a minute.
-    parser.set_defaults(steps=1000, sample_steps=64, 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 target ratio of the running example, reached through the denoiser route.
-    posterior = torch.tensor([0.375, 0.625, 0.0, 0.0])  # p(x_0 | X_t = C) at alpha_t = 1/2.
-    per_clean = torch.tensor([5.0, 0.2, 1.0, 1.0])  # q_t(A | b) / q_t(C | b) for b over ACGT.
-    target_ratio = float(posterior @ per_clean)  # The ratio in the A direction, which comes out 2.
-
-    # Section 03, the score entropy at three predictions of that ratio.
-    errors = {s: float(score_entropy(torch.tensor(target_ratio), torch.tensor(s)))
-              for s in (1.8, 1.0, 3.0)}  # A ten percent error against two larger ones.
-    talk(f"target ratio        {target_ratio:.1f} from 0.375(5) + 0.625(0.2)")
-    at_target = float(score_entropy(torch.tensor(target_ratio), torch.tensor(target_ratio)))
-    talk("score entropy       " + "  ".join(f"s={s}: {v:.4f}" for s, v in errors.items()) +
-         f"  at s = r it is {at_target:.1e}")
-
-    # Train the network to output the ratios.
-    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 loss below.
-    model = SequenceDenoiser(DNA_SIZE, args.length, width=args.width).to(device)
-    talk(f"device={device} parameters={count_parameters(model):,} ratios per sequence "
-         f"{args.length * (DNA_SIZE - 1)}")
-
-    def step_loss(step: int) -> torch.Tensor:
-        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
-        return sedd_loss(model, train_tokens[rows])  # 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(sedd_loss(model, valid_tokens)) for _ in range(4)) / 4.0
-
-    # The ratios the trained network predicts at a half-corrupted batch, against their targets.
-    with torch.no_grad():
-        check_t = torch.full((len(valid_tokens), 1), math.log(2.0) / BETA, device=device)
-        check_alpha = torch.exp(-BETA * check_t)  # alpha_t = 1/2 at this time.
-        keep = torch.rand(valid_tokens.shape, device=device) < check_alpha
-        noisy = torch.where(keep, valid_tokens,
-                            torch.randint(0, DNA_SIZE, valid_tokens.shape, device=device))
-        predicted = concrete_score(model, noisy, check_t[:, 0])  # [M, L, K] the predicted ratios.
-        exact = kernel_ratio(valid_tokens, noisy, check_alpha)  # [M, L, K] the conditional targets.
-        toward_clean = exact > 1.0  # The jumps that point back at the clean base.
-        mean_predicted = float(predicted[toward_clean].mean())  # Should sit below the target 5.
-
-    torch.manual_seed(args.seed + 1)  # Fix the jump draws, so the samples repeat.
-    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
-    report = {
-        "target ratio at C toward A": target_ratio,
-        "score entropy at s = 1.8": errors[1.8],
-        "score entropy at s = 1.0": errors[1.0],
-        "score entropy at s = 3.0": errors[3.0],
-        "training score entropy, last 50": sum(losses[-50:]) / len(losses[-50:]),
-        "validation score entropy": held_out,
-        "mean predicted ratio toward clean": mean_predicted,
-        "conditional target toward clean": float(exact[toward_clean].mean()),
-        "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("sedd: score entropy on DNA", report, decode_dna(samples))
-    save_run(args.out, model=model, config=vars(args), report=report,
-             losses=losses, samples=decode_dna(samples))
-
-
-# %% 06. Run the script
-if __name__ == "__main__":  # Execute when launched directly.
-    main()  # Work the score entropy, train the ratios, and generate.

lecture_4/simplex.py

@@ -1,246 +0,0 @@
-"""Simplex diffusion on DNA, a belief at every position in place of a token.
-
-Every model before this one draws a token at each reverse step and discards the rest of the
-denoiser's distribution, so a position the model was unsure about looks the same as one it was
-certain about. Here the state at a position is a probability vector over the four bases, the
-corruption is a Dirichlet around the uniform-replacement kernel, and one reverse step thins the
-current belief, draws an innovation and mixes the two.
-
-Chapter 4, "From tokens to beliefs" and "Thinning, an innovation, and a mixture". Time runs from
-t = 0 clean to t = 1 noise, the survival probability is alpha_t = 1 - t, the concentration is
-c_t = gamma K / (1 - alpha_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 worked Dirichlet parameters, the worked
-thinning weight and its Monte Carlo check, then the sample statistics.
-
-    python lecture_4/simplex.py                    # train and generate
-    python lecture_4/simplex.py --mixture-draws 10000  # a coarser check of the mixture
-    python lecture_4/simplex.py --steps 20         # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import torch  # Tensors, gradients and the Dirichlet, Beta and categorical draws.
-import torch.nn as nn  # The denoiser whose input layer is the one thing that changes.
-import torch.nn.functional as F  # Cross-entropy at every position.
-
-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 transformer the belief feeds.
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-GAMMA = 1.0  # The constant of the concentration schedule c_t = gamma K / (1 - alpha_t).
-TIME_FLOOR = 1e-3  # The concentration grows like 1/(1 - alpha_t), so the grid stops above zero.
-
-
-# %% 02. The forward marginal as a Dirichlet
-#
-# Fact 2 calls for a forward marginal in closed form, and the natural requirement is that the
-# belief fluctuate around the marginal the token models already use,
-# E[P_t | x_0] = alpha_t onehot(x_0) + (1 - alpha_t) pi, with the spread set by c_t.
-def concentration(alpha: float | torch.Tensor) -> float | torch.Tensor:
-    """c_t = gamma K / (1 - alpha_t), the schedule that makes the increment nonnegative."""
-    return GAMMA * DNA_SIZE / (1.0 - alpha)  # 8 at alpha_t = 0.5 and 20 at alpha_s = 0.8.
-
-
-def belief_mean(alpha: float, x0: int) -> torch.Tensor:
-    """E[P_t | x_0] over (A, C, G, T), the uniform-replacement kernel of the token models. [K]."""
-    mean = torch.full((DNA_SIZE,), (1.0 - alpha) / DNA_SIZE)  # [K] the share from pi.
-    mean[x0] = mean[x0] + alpha  # The clean token keeps the survival mass on top of its share.
-    return mean  # [K] a distribution, since alpha + (1 - alpha) = 1.
-
-
-def dirichlet_parameters(alpha: float, x0: int) -> torch.Tensor:
-    """c_t E[P_t | x_0], the Dirichlet parameters at one position. Returns [K]."""
-    return concentration(alpha) * belief_mean(alpha, x0)  # (5, 1, 1, 1) at alpha_t = 0.5.
-
-
-def parameter_field(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
-    """The same parameters for a whole batch. x0: [B, L], alpha: [B, 1] -> [B, L, K].
-
-    Every coordinate other than the clean token has parameter gamma, constant in time, and the
-    clean coordinate has gamma (1 + K alpha_t/(1 - alpha_t)), the factor of the uniform rates.
-    """
-    off = torch.full(x0.shape + (DNA_SIZE,), GAMMA, device=x0.device)  # [B, L, K] the pi share.
-    clean = GAMMA * DNA_SIZE * alpha / (1.0 - alpha)  # [B, 1] the survival mass times c_t.
-    return off.scatter_add(2, x0[:, :, None], clean[..., None].expand_as(x0[:, :, None]))
-
-
-def corrupt(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
-    """Draw P_t ~ Dir(c_t E[P_t | x_0]) at each position. x0: [B, L], alpha: [B, 1] -> [B, L, K]."""
-    return torch.distributions.Dirichlet(parameter_field(x0, alpha)).sample()  # [B, L, K] beliefs.
-
-
-# %% 03. Thinning, an innovation, and a mixture
-#
-# A Dirichlet with parameter a + b equals W V + (1 - W) U for V ~ Dir(a), U ~ Dir(b) and
-# W ~ Beta(sum a, sum b). With this concentration schedule the increment b sits entirely on the
-# clean coordinate, so the innovation U is the one-hot vector of the clean token.
-def increment(alpha_t: float, alpha_s: float, x0: int) -> torch.Tensor:
-    """b = c_s E[P_s | x_0] - c_t E[P_t | x_0], nonnegative for alpha_s > alpha_t. Returns [K]."""
-    return dirichlet_parameters(alpha_s, x0) - dirichlet_parameters(alpha_t, x0)  # (12, 0, 0, 0).
-
-
-def thinning_law(alpha_t: float, alpha_s: float) -> tuple[float, float]:
-    """The two Beta parameters, sum a = c_t and sum b = c_s - c_t, which are 8 and 12 here."""
-    return float(concentration(alpha_t)), float(concentration(alpha_s) - concentration(alpha_t))
-
-
-def reverse_step(belief: torch.Tensor, clean: torch.Tensor,
-                 alpha_t: float, alpha_s: float) -> torch.Tensor:
-    """P_s = W P_t + (1 - W) onehot(x_0), one Beta draw per position. belief: [B, L, K]."""
-    first, second = thinning_law(alpha_t, alpha_s)  # The two parameters of the thinning weight.
-    w = torch.distributions.Beta(  # [B, L] one weight per position, with mean c_t / c_s.
-        torch.tensor(first, device=belief.device),
-        torch.tensor(second, device=belief.device)).sample(belief.shape[:2])
-    innovation = F.one_hot(clean, DNA_SIZE).to(belief.dtype)  # [B, L, K] the corner it moves to.
-    return w[..., None] * belief + (1.0 - w[..., None]) * innovation  # [B, L, K] the mixture.
-
-
-# %% 04. The denoiser that reads a belief
-#
-# The only architectural change is at the input. A token is looked up in the embedding table, and
-# a belief is the same K rows of that table combined in its own proportions, so a one-hot input
-# recovers the denoiser of every other file in this lecture exactly.
-class BeliefDenoiser(nn.Module):
-    """d_theta(. | P_t, t), the shared sequence denoiser reading K numbers at every position."""
-
-    def __init__(self, length: int, width: int = 128):
-        super().__init__()
-        self.inner = SequenceDenoiser(DNA_SIZE, length, width=width)  # Embeddings and transformer.
-
-    def forward(self, beliefs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-        """beliefs: [B, L, K], t: [B] -> logits over the clean tokens [B, L, K]."""
-        inner = self.inner
-        positions = torch.arange(beliefs.shape[1], device=beliefs.device)  # [L].
-        h = (beliefs @ inner.token_embedding.weight  # [B, L, D] the belief mixes the K embeddings.
-             + inner.position_embedding(positions)[None, :, :]  # [1, L, D] where it sits.
-             + inner.time_embedding(t)[:, None, :])  # [B, 1, D] how far along the path we are.
-        return inner.head(inner.transformer(h))  # [B, L, K] every position sees every other.
-
-
-# %% 05. The loss
-#
-# Fact 5 is the cross-entropy whose optimum is the posterior over clean tokens. The reverse step
-# carries the denoiser's answer through the mixture of Section 03, whose weight already comes from
-# the concentration schedule, so the time derivative of alpha_t multiplies nothing here.
-def loss(model: nn.Module, x0: torch.Tensor) -> torch.Tensor:
-    """One Monte Carlo sample of E[-log d_theta(x_0 | P_t, t)], summed over positions."""
-    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
-    alpha = 1.0 - t  # [B, 1] the linear schedule, alpha_0 = 1 and alpha_1 = 0.
-    beliefs = corrupt(x0, alpha)  # [B, L, K] one Dirichlet draw per position.
-    logits = model(beliefs, t[:, 0])  # [B, L, K] the denoiser reads the whole field of beliefs.
-    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
-    return ce.sum(-1).mean()  # Nats per sequence, with every position counted.
-
-
-# %% 06. The sampler
-#
-# Generation starts from Dir(gamma, ..., gamma), which at alpha_1 = 0 is the law the corruption
-# ends at. Each step draws one clean token from the denoiser, thins the current belief and mixes
-# the two, so the belief moves through the interior and reaches a corner at the last step.
-@torch.no_grad()  # Generation needs no gradient tracking.
-def sample(model: nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
-    """Generate from beliefs scattered around pi down to clean DNA. Returns [B, L] over ACGT."""
-    device = next(model.parameters()).device
-    start = torch.full((batch, length, DNA_SIZE), GAMMA, device=device)  # [B, L, K] at alpha_1 = 0.
-    belief = torch.distributions.Dirichlet(start).sample()  # [B, L, K] scattered around pi.
-    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(belief, t.expand(batch))  # [B, L, K] the clean-token prediction.
-        clean = torch.distributions.Categorical(logits=logits).sample()  # [B, L] the innovation.
-        belief = reverse_step(belief, clean, float(1.0 - t), float(1.0 - s))  # [B, L, K].
-    return belief.argmax(-1)  # [B, L] the corner each position has reached.
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--mixture-draws", type=int, default=400_000,
-                        help="draws in the check of the thinning representation")
-    # The belief carries the denoiser's uncertainty between steps, which needs enough reverse steps
-    # to show. A default run trains 1000 steps and measures 64 sequences, inside about a minute.
-    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 02, the worked Dirichlet at the two noise levels of the chapter, with a clean A.
-    at_t = dirichlet_parameters(0.5, 0)  # [K] (5, 1, 1, 1) at the later of the two times.
-    at_s = dirichlet_parameters(0.8, 0)  # [K] (17, 1, 1, 1) at the earlier one.
-    total_t, total_s = at_t.sum(), at_s.sum()  # c_t and c_s, the totals that set every variance.
-    spread_t = (at_t * (total_t - at_t) / (total_t ** 2 * (total_t + 1))).sqrt()  # [K] exact.
-    spread_s = (at_s * (total_s - at_s) / (total_s ** 2 * (total_s + 1))).sqrt()  # [K] exact.
-    talk(f"alpha_t = 0.5       c_t {concentration(0.5):.0f}  mean "
-         f"{[round(float(v), 3) for v in belief_mean(0.5, 0)]}  parameters "
-         f"{[round(float(v), 0) for v in at_t]}  std on A {float(spread_t[0]):.3f}")
-    talk(f"alpha_s = 0.8       c_s {concentration(0.8):.0f}  mean "
-         f"{[round(float(v), 2) for v in belief_mean(0.8, 0)]}  parameters "
-         f"{[round(float(v), 0) for v in at_s]}")
-
-    # Section 03, the increment, the thinning weight, and the mean the mixture lands on.
-    step = increment(0.5, 0.8, 0)  # [K] all of the increase sits on the clean coordinate.
-    first, second = thinning_law(0.5, 0.8)  # Beta(8, 12), whose mean is 0.4.
-    mean_weight = first / (first + second)  # c_t / c_s, the average share the old belief keeps.
-    mixed = mean_weight * belief_mean(0.5, 0) + (1.0 - mean_weight) * torch.eye(DNA_SIZE)[0]  # [K].
-    talk(f"increment           {[round(float(v), 0) for v in step]}  W ~ Beta({first:.0f}, "
-         f"{second:.0f})  mean {mean_weight:.1f}  mixture mean "
-         f"{[round(float(v), 2) for v in mixed]}")
-
-    # The same mixture by simulation, against the exact standard deviations of Dir(17, 1, 1, 1).
-    thinned = torch.distributions.Dirichlet(at_t).sample((args.mixture_draws,))  # [N, K] draw V.
-    weight = torch.distributions.Beta(  # [N] the Beta draw of the representation.
-        torch.tensor(first), torch.tensor(second)).sample((args.mixture_draws,))
-    drawn = weight[:, None] * thinned + (1.0 - weight[:, None]) * torch.eye(DNA_SIZE)[0]  # [N, K].
-    talk(f"mixture check       {args.mixture_draws} draws  std "
-         f"{[round(float(v), 3) for v in drawn.std(0)]}  against the exact "
-         f"{[round(float(v), 3) for v in spread_s]}")
-
-    # Train the denoiser on the shared synthetic DNA. Reseeding here keeps the training batches
-    # independent of how many draws the check above consumed.
-    set_seed(args.seed)
-    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 loss below.
-    model = BeliefDenoiser(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_index: 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)
-    with torch.no_grad():
-        held_out = sum(float(loss(model, valid_tokens)) for _ in range(4)) / 4.0
-
-    # Generate, then measure the structure the data was built with.
-    torch.manual_seed(args.seed + 1)  # Fix the Beta and categorical draws, so the samples repeat.
-    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
-    report = {
-        "c_t at alpha_t = 0.5": float(concentration(0.5)),
-        "Dirichlet parameter on A at 0.5": float(at_t[0]),
-        "standard deviation on A at 0.5": float(spread_t[0]),
-        "Dirichlet parameter on A at 0.8": float(at_s[0]),
-        "increment on A": float(step[0]),
-        "mean thinning weight": mean_weight,
-        "simulated standard deviation on A": float(drawn.std(0)[0]),
-        "exact standard deviation on A": float(spread_s[0]),
-        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
-        "validation loss, nats/seq": held_out,
-        "loss, nats per token": held_out / 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("simplex: beliefs in place of tokens", 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 Dirichlet numbers, train the belief denoiser, and generate.

lecture_4/uniform.py

@@ -1,240 +0,0 @@
-"""Uniform replacement on DNA, the member of the pi family with no mask.
-
-Corruption redraws a base from the uniform distribution over ACGT, so a noisy base may be the
-original or a replacement, and no marker of the damage appears in the sequence. Every pair of
-bases is then joined by a reverse rate, the loss keeps the staying term that masking cancelled, and
-a sampler can revise a base it has already written.
-
-Chapter 4, "A rate matrix with one free distribution", "The posterior average of conditional reverse
-rates" and "Uniform replacement and the rates it produces". Time runs from t = 0 clean to t = 1
-noise, the survival probability is alpha_t = exp(-beta 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 worked
-posterior average of 2, the worked rate error, and the sample statistics.
-
-    python lecture_4/uniform.py               # train and generate
-    python lecture_4/uniform.py --steps 20    # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math  # The worked rate errors are logarithms.
-
-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 d_theta(. | z, t).
-from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train
-
-BETA = 4.0  # The constant noise rate of the chapter's worked example, so alpha_t = exp(-4t).
-TIME_FLOOR = 2e-2  # The reverse rates grow like 1/(1 - alpha_t), so sampling stops just above zero.
-
-
-# %% 02. The uniform kernel and the rates that produce it
-#
-# Under R_pi = 1 pi^T - I with pi uniform, every pair of distinct bases is joined at rate beta/K,
-# and the matrix exponential is the kernel q_t(x | x_0) = alpha_t delta(x, x_0) + (1 - alpha_t)/K.
-def survival(t: torch.Tensor | float) -> torch.Tensor | float:
-    """alpha_t = exp(-tau(t)) for the constant schedule beta(t) = BETA, so tau(t) = BETA t."""
-    return torch.exp(-BETA * t) if torch.is_tensor(t) else math.exp(-BETA * t)
-
-
-def kernel(alpha: float | torch.Tensor) -> torch.Tensor:
-    """Q(alpha) over (A, C, G, T), whose (x_0, x) entry is q_t(x | x_0). Returns [K, K]."""
-    uniform = (1.0 - alpha) / DNA_SIZE  # The share of the redrawn mass that lands on each base.
-    return alpha * torch.eye(DNA_SIZE, dtype=torch.float64) + uniform  # [K, K] rows sum to one.
-
-
-def corrupt(x0: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
-    """Redraw each position from the uniform distribution with probability 1 - alpha_t."""
-    keep = torch.rand(x0.shape, device=x0.device) < survival(t)  # [B, L] survival of the original.
-    fresh = torch.randint(0, DNA_SIZE, x0.shape, device=x0.device)  # [B, L] a draw from pi.
-    return torch.where(keep, x0, fresh)  # [B, L] nothing marks which positions were redrawn.
-
-
-# %% 03. The conditional reverse rates and their three cases
-#
-# Dividing the conditional reverse rate by the noise rate beta/K leaves one factor, which is
-# 1 + K alpha_t/(1 - alpha_t) toward the clean token, its reciprocal away from it, and one between
-# two tokens that are both wrong.
-def clean_factor(alpha: float) -> float:
-    """1 + K alpha_t / (1 - alpha_t), the factor 5 of the worked example at alpha_t = 1/2."""
-    return 1.0 + DNA_SIZE * alpha / (1.0 - alpha)  # The ratio q_t(x_0 | x_0) / q_t(x | x_0).
-
-
-def conditional_rates(alpha: float, current: int, clean: int) -> torch.Tensor:
-    """R_hat_t(c, a | x_0) for every destination a, in units where the forward rate is one."""
-    Q = kernel(alpha)  # [K, K] the forward kernel at this survival probability.
-    return Q[clean] / Q[clean, current]  # [K] q_t(a | x_0) divided by q_t(c | x_0).
-
-
-# %% 04. A worked posterior average on one base
-#
-# With a clean base A at probability 0.75 and C at 0.25, an observed C at alpha_t = 1/2 has
-# posterior 0.375 and 0.625, and the two conditional rates 5 and 0.2 average to the marginal rate.
-def posterior_over_clean(alpha: float, observed: int, prior: torch.Tensor) -> torch.Tensor:
-    """p(x_0 = b | X_t = observed), Bayes' rule with the forward kernel as the likelihood."""
-    weights = prior * kernel(alpha)[:, observed]  # [K] prior times q_t(observed | b).
-    return weights / weights.sum()  # [K] normalized over the candidate clean bases.
-
-
-def marginal_reverse_rate(alpha: float, observed: int, destination: int,
-                          prior: torch.Tensor) -> float:
-    """The posterior average of the conditional reverse rates, which is the marginal reverse rate."""
-    posterior = posterior_over_clean(alpha, observed, prior)  # [K] the weights of the average.
-    targets = kernel(alpha)[:, destination] / kernel(alpha)[:, observed]  # [K] one target per clean b.
-    return float(posterior @ targets)  # The two together give the rate the sampler uses.
-
-
-# %% 05. The rate error and the staying term
-#
-# l(A, B) = A log(A/B) - A + B is the Bregman divergence the path KL divergence produces. The term
-# -A + B comes from the probability of staying put, and under uniform noise it does not cancel.
-def rate_error(target: torch.Tensor, model_rate: torch.Tensor) -> torch.Tensor:
-    """Elementwise l(A, B), nonnegative and zero exactly at B = A."""
-    logarithmic = torch.where(target > 0, target * (target.clamp(min=1e-12) / model_rate).log(),
-                              torch.zeros_like(target))  # A log(A/B), with l(0, B) = B.
-    return logarithmic - target + model_rate  # The staying term is the -A + B part.
-
-
-def ratio_table(current: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
-    """q_t(a | b) / q_t(c | b) for every clean base b and destination a. Returns [B, L, K, K]."""
-    uniform = (1.0 - alpha) / DNA_SIZE  # [B, 1] the redrawn share at this noise level.
-    Q = (alpha[..., None, None] * torch.eye(DNA_SIZE, device=current.device)
-         + uniform[..., None, None])  # [B, 1, K, K] the kernel q_t(a | b) at each sequence's time.
-    denominator = Q.expand(-1, current.shape[1], -1, -1).gather(
-        3, current[:, :, None, None].expand(-1, -1, DNA_SIZE, 1))  # [B, L, K, 1] q_t(c | b).
-    return Q / denominator  # [B, L, K, K] the only quantity the rates need from the kernel.
-
-
-def model_rates(denoiser: torch.Tensor, ratios: torch.Tensor) -> torch.Tensor:
-    """R_hat^theta_t(c, a) = (beta/K) sum_b d_theta(b) q_t(a | b) / q_t(c | b). Returns [B, L, K]."""
-    averaged = torch.einsum("nlb,nlba->nla", denoiser, ratios)  # [B, L, K] averaged over clean bases.
-    return (BETA / DNA_SIZE) * averaged  # [B, L, K] the forward rate times the posterior average.
-
-
-def conditional_targets(x0: torch.Tensor, ratios: torch.Tensor) -> torch.Tensor:
-    """R_hat_t(c, a | x_0), the same formula with the posterior collapsed onto the clean token."""
-    rows = x0[:, :, None, None].expand(-1, -1, 1, DNA_SIZE)  # [B, L, 1, K] select the row b = x_0.
-    return (BETA / DNA_SIZE) * ratios.gather(2, rows).squeeze(2)  # [B, L, K] one target per jump.
-
-
-def rate_loss(model: torch.nn.Module, x0: torch.Tensor) -> torch.Tensor:
-    """One Monte Carlo sample of L_cond, summed over destinations and positions."""
-    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
-    alpha = survival(t)  # [B, 1] the survival probability at that time.
-    zt = corrupt(x0, t)  # [B, L] the noisy sequence, with no marker of which positions moved.
-    ratios = ratio_table(zt, alpha)  # [B, L, K, K] the kernel ratios at this noise level.
-    denoiser = model(zt, t[:, 0]).softmax(-1)  # [B, L, K] the posterior over clean tokens.
-    predicted = model_rates(denoiser, ratios).clamp(min=1e-8)  # [B, L, K] nonnegative by construction.
-    target = conditional_targets(x0, ratios)  # [B, L, K] the rate attached to this clean sequence.
-    staying = torch.nn.functional.one_hot(zt, DNA_SIZE).bool()  # [B, L, K] the destination c itself.
-    return rate_error(target, predicted).masked_fill(staying, 0.0).sum((1, 2)).mean()
-
-
-# %% 06. Sampling by tau-leaping
-#
-# Over a step of length h a position moves to a with probability h times its reverse rate, and the
-# remaining probability keeps it where it is. The step is valid while h times the total rate is one.
-@torch.no_grad()  # Generation needs no gradient tracking.
-def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
-    """Generate from uniformly random tokens down to clean DNA. Returns [B, L] over ACGT."""
-    device = next(model.parameters()).device
-    z = torch.randint(0, DNA_SIZE, (batch, length), device=device)  # [B, L] a draw from pi.
-    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
-    for t, s in zip(ts[:-1], ts[1:]):
-        time_row = t.expand(batch, 1)  # [B, 1] the shared time of this step.
-        denoiser = model(z, time_row[:, 0]).softmax(-1)  # [B, L, K] the clean-token prediction.
-        rates = model_rates(denoiser, ratio_table(z, survival(time_row)))  # [B, L, K] per destination.
-        jump = (t - s) * rates  # [B, L, K] the tau-leaping probabilities of this step.
-        jump.scatter_(2, z[:, :, None], 0.0)  # A jump to the current base is not a jump.
-        total = jump.sum(-1, keepdim=True).clamp(min=1.0)  # Rescale when the step would exceed one.
-        jump = jump / total  # [B, L, K] now a sub-probability vector at every position.
-        stay = (1.0 - jump.sum(-1, keepdim=True)).clamp(min=0.0)  # [B, L, 1] the chance of staying.
-        choice = torch.distributions.Categorical(
-            probs=torch.cat([jump, stay], dim=-1)).sample()  # [B, L] destination, or K for staying.
-        z = torch.where(choice == DNA_SIZE, z, choice.clamp(max=DNA_SIZE - 1))  # [B, L].
-    final = model(z, ts[-1].expand(batch)).softmax(-1)  # The exact posterior step down to t = 0.
-    return torch.distributions.Categorical(probs=final).sample()  # [B, L] clean DNA.
-
-
-# %% 07. The report
-def main() -> None:
-    parser = common_parser(__doc__.split("\n")[0])
-    # The rate loss needs a longer budget than the worked numbers do, so the default trains for
-    # 1000 steps and measures on 64 sequences, which keeps the run inside about a minute.
-    parser.set_defaults(steps=1000, sample_steps=64, 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.
-
-    # Sections 03 and 04, the three cases and the worked posterior average on one base.
-    prior = torch.tensor([0.75, 0.25, 0.0, 0.0], dtype=torch.float64)  # Clean A at 0.75, C at 0.25.
-    posterior = posterior_over_clean(0.5, 1, prior)  # [4] after observing a C at alpha_t = 1/2.
-    targets = conditional_rates(0.5, current=1, clean=0)  # [4] the rates attached to a clean A.
-    average = marginal_reverse_rate(0.5, 1, 0, prior)  # The posterior average, which comes out 2.
-    marginal = prior @ kernel(0.5)  # [4] p_t over the four bases, which is (0.5, 0.25, 0.125, 0.125).
-    talk(f"kernel at alpha=0.5 clean {float(kernel(0.5)[0, 0]):.3f}  other "
-         f"{float(kernel(0.5)[0, 1]):.3f}  clean factor {clean_factor(0.5):.0f}")
-    talk(f"posterior after C   {[round(float(v), 3) for v in posterior]}  conditional rates "
-         f"{float(targets[0]):.0f} under a clean A and "
-         f"{float(conditional_rates(0.5, 1, 1)[0]):.1f} under a clean C")
-    talk(f"posterior average   {average:.1f}  marginals {[round(float(v), 3) for v in marginal]}  "
-         f"ratio {float(marginal[0] / marginal[1]):.1f}  to G and T "
-         f"{marginal_reverse_rate(0.5, 1, 2, prior):.1f}")
-
-    # Section 05, the rate error on one base, where the staying term keeps the sum from collapsing.
-    worked_denoiser = torch.tensor([[[0.5, 0.3, 0.1, 0.1]]], dtype=torch.float64)  # [1, 1, 4].
-    worked_current = torch.tensor([[1]])  # The current base is C.
-    worked_alpha = torch.tensor([[0.5]], dtype=torch.float64)  # alpha_t = 1/2.
-    worked_ratios = ratio_table(worked_current, worked_alpha)  # [1, 1, K, K] at alpha_t = 1/2.
-    predicted = model_rates(worked_denoiser, worked_ratios)[0, 0] / (BETA / DNA_SIZE)
-    target = conditional_rates(0.5, current=1, clean=0)  # The clean base is A.
-    error = float(rate_error(target[[0, 2, 3]], predicted[[0, 2, 3]]).sum())  # Skip the current base.
-    talk(f"model rates from d  {[round(float(v), 2) for v in predicted]}  targets "
-         f"{[round(float(v), 1) for v in target]}  rate error {error:.3f}")
-
-    # Train the denoiser against the conditional rate targets.
-    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 loss below.
-    model = SequenceDenoiser(DNA_SIZE, args.length, width=args.width).to(device)
-    talk(f"device={device} parameters={count_parameters(model):,} train={len(train_tokens)}")
-
-    def step_loss(step: int) -> torch.Tensor:
-        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
-        return rate_loss(model, train_tokens[rows])  # 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(rate_loss(model, valid_tokens)) for _ in range(4)) / 4.0
-
-    torch.manual_seed(args.seed + 1)  # Fix the jump draws, so the samples repeat.
-    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
-    report = {
-        "q_t(x_0 | x_0) at alpha=0.5": float(kernel(0.5)[0, 0]),
-        "clean factor at alpha=0.5": clean_factor(0.5),
-        "posterior on A after a C": float(posterior[0]),
-        "conditional rate under A": float(targets[0]),
-        "conditional rate under C": float(conditional_rates(0.5, 1, 1)[0]),
-        "posterior average rate": average,
-        "rate from the marginals": float(marginal[0] / marginal[1]),
-        "worked model rate to A": float(predicted[0]),
-        "worked rate error": error,
-        "training rate loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
-        "validation rate loss": 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("uniform: replacement 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 posterior average, train against the rate loss, and generate.