Delete lecture_5

lecture_5/README.md

@@ -1,97 +0,0 @@
-# Lecture 5: Discrete Flow Matching (Chapter 5)
-
-Chapter 4 fixed a corruption process and solved for its reverse. This lecture reverses that order.
-We pick the probability path first, write down a velocity that generates it, and read the sampler
-off the velocity. On tokens the velocity is a rate matrix read from the destination's side, the
-master equation is the condition under which it generates a path, and the mixture path
-`(1 - kappa_t) delta_{x_0} + kappa_t delta_{x_1}` at each position has a conditional velocity that a
-posterior average turns into the schedule factor times a denoiser. Cross-entropy on the data token
-recovers that denoiser. Everything after that is a choice made inside the same construction: a
-schedule, a corrector, a kinetic-optimal path, edits that change the length, a relaxation of the
-token into a point of the simplex or of a sphere, a guidance tilt, or a different coupling.
-
-Two conventions hold in every file and are stated at the top of each. Time runs from the source at
-`t = 0` to the data at `t = 1`, so `kappa_t` here equals the survival probability `alpha_{1-t}` of
-Chapter 4. The probability velocity `u_t(y, x)` is the same number Chapter 4 writes as the rate
-`R_t(x, y)`, with the two arguments exchanged.
-
-## The files
-
-Everything runs on the synthetic DNA of `cis6270.data`: 512 sequences of length 16 over A, C, G, T,
-with a GC-rich background and the motif ACGT planted in three quarters of them at one of three
-positions. Those sequences have GC fraction 0.577, motif fraction 0.777 and token entropy 1.374 of
-`log 4 = 1.386`, which are the numbers a trained sampler is measured against.
-
-| File | Command | What it does |
-| --- | --- | --- |
-| `dfm.py` | `python lecture_5/dfm.py` | The coupling, the mixture path, the conditional velocity, the posterior average, the cross-entropy loss and the two-part Euler rule. Trains a denoiser and generates DNA from either source, `--source uniform` or `--source mask`. |
-| `velocity.py` | `python lecture_5/velocity.py` | The mathematics with nothing learned: a jump rate written as a velocity, the two rules, the master equation, the worked DNA velocity, survival and hazard, and the continuous and discrete velocities side by side. |
-| `corrector.py` | `python lecture_5/corrector.py` | Velocities with zero net flux, the three-state example, the corrector for the mask source, the jump counts at three corrector strengths, and the parallel-step error. |
-| `kinetic_optimal.py` | `python lecture_5/kinetic_optimal.py` | The flux and its energy, the velocity of least energy, the kinetic optimality of the mixture path, the transition distance on DNA, the worked metric path, and the energy of the linear and geodesic schedules. |
-| `edit_flows.py` | `python lecture_5/edit_flows.py` | Insertions, deletions and substitutions; alignments and the column-wise path; the posterior over two alignments; the rate loss; one sampler step; and the length statistics of a run on DNA pairs of different lengths. |
-| `dirichlet.py` | `python lecture_5/dirichlet.py` | The Dirichlet family as a conditional path, the speed from the flux across a cut, the two-letter closed form, the worked speeds on a DNA base, and the posterior average. Trains a denoiser on the simplex and integrates to `t = 8`. |
-| `fisher.py` | `python lecture_5/fisher.py` | The Fisher information as a metric, the square-root map to the sphere, geodesics, the tangent projection and the regression loss. Trains a velocity field and samples with exponential-map steps. |
-| `gumbel.py` | `python lecture_5/gumbel.py` | The Gumbel-max trick, the cooling temperature, the annealed path, the replicator equation, the inference-time velocity, and training against the marginal velocity. |
-| `guidance.py` | `python lecture_5/guidance.py` | Guidance as a tilt of the rates, the h-transform, the reachable set, the guidance strength, the classifier-free form, the Taylor expansion of the log ratio, and the guided sampler. Steers toward the GC label with `--gamma`. |
-| `multi_objective.py` | `python lecture_5/multi_objective.py` | The Pareto front, MOG-DFM's rank, direction and cone with the controller that widens it, and pCoMole's worst-objective utility with a rollout of the base flow on DNA. |
-| `rectify.py` | `python lecture_5/rectify.py` | Total correlation on two bases, ReDi's recoupling rounds run on the exact sampler, and AReUReDi's Tchebycheff tilt with its Metropolis acceptance step. |
-
-Every file takes the shared flags of `cis6270.runner`. `--steps 20 --quiet` runs each one in a few
-seconds, and `--steps` controls the training steps in the five files that train, the grid in the
-five that integrate, and the number of rollouts in `multi_objective.py`.
-
-## The numbers a default run produces
-
-Seed 0, one core. The five files that train take under half a minute each, and the other six take
-one to seven seconds. The worked numbers are exact and repeat on any machine; the trained numbers
-move a little with the hardware's floating point.
-
-| File | Headline numbers |
-| --- | --- |
-| `dfm.py` | loss 0.876 nats, GC fraction 0.578, motif fraction 0.672, token entropy 1.374 of `log 4 = 1.386` |
-| `velocity.py` | `p_dot = (-0.1, -0.2, 0.7, -0.4)`, Euler step `(0.09, 0.18, 0.37, 0.36)`, hazard 2 at `t = 0.5` and 10 at `t = 0.9`, Euler error 0.0002 against the closed form |
-| `corrector.py` | masked fraction 0.500 at `t = 0.5` for every `eta`, jumps per position 1.000, 2.991 and 10.934 at `eta = 0, 2, 10`, pair survival 0.133 measured against the predicted 0.135 at `eta = 2` |
-| `kinetic_optimal.py` | rate out of C 5.780, `p_dot = (-0.3997, -0.4774, 1.3546, -0.4774)` from the velocity and from the path, energy 27.63 and 41.45 for the linear schedule, `pi^2 = 9.8696` for the geodesic one |
-| `edit_flows.py` | averaged edit rates 0.667, 1.333 and 0.667 for a total of 2.667, example loss 3.446, untouched 0.8452 against the first-order 0.8350, DNA length 13.945 at `t = 0.5` against the predicted 13.941 |
-| `dirichlet.py` | speeds 0.232, 0.199, 0.159, 0.104 at `t = 1`, continuity residual `2.5e-08`, loss 1.027 nats, GC fraction 0.532, motif fraction 0.547 |
-| `fisher.py` | Fisher-Rao distance 1.485 against the Euclidean 0.141 and 0.013 on the two moves, arc 0.3713 from the midpoint, loss 0.278, GC fraction 0.576, motif fraction 0.203 |
-| `gumbel.py` | path `(0.0528, 0.0214, 0.8781, 0.0477)` at `t = 1`, replicator velocity `(-0.136, -0.217, 0.500, -0.147)`, loss 0.672 nats, GC fraction 0.572, motif fraction 0.516 |
-| `guidance.py` | guided destination `(0.029, 0.203, 0.478, 0.290)`, exit rate 2.50 at the exact `h`, GC fraction 0.531 unguided, 0.657 at `gamma = 2` and 0.698 at `gamma = 5` |
-| `multi_objective.py` | combined score `(-1.511, 1.541, 2.297, -2.327)`, cone angles 47.7, 37.9 and 107.7 degrees, selection `(0.371, 0.629, 0)` after the chapter's rollouts, and on DNA a base selection of `(0.417, 0.583)` that the rollout value turns into `(0.962, 0.038)` |
-| `rectify.py` | total correlation 1.386, 1.101, 0.843, 0.696 and 0.616 over rounds 0 to 4 against the limit 0.520; one-step success 25.0%, 40.3%, 52.5%, 58.0% and 60.4% against the limit 62.5% |
-
-Four thousand training steps, about two minutes a file, buy a sampler that reaches the planted
-motif more often: `dfm.py` goes from 0.672 to 0.734 against the data's 0.777, and `fisher.py` from
-0.203 to 0.453. Nothing else moves much.
-
-`tests/test_lecture_5.py` asserts the chapter's worked numbers against these files, including the
-posterior 0.625 and 0.375 at a masked position, the entropy 0.6616 that bounds the loss there, the
-velocity `(0.2, 1.2, -1.6, 0.2)` out of G, the metric-path rates 0.898 and 4.882 out of C, the
-Dirichlet speeds, the Fisher-Rao arc, the Gumbel-softmax path at four times, the guided destination
-distribution, the Pareto front, and the ReDi rounds.
-
-## What is simplified
-
-- The corrector simulation uses 40,000 positions where the chapter uses 200,000. Pass
-  `--positions 200000` to reproduce the chapter's figures exactly.
-- The parallel-step section measures the survival probability of a jointly revealed pair, which is
-  the quantity the chapter computes in closed form. At a fine grid the end-to-end error of the
-  corrector is not lower than at `eta = 0`, because the admissible step shrinks with `eta`.
-- `edit_flows.py` runs the column-wise path on one sampled alignment per pair, which the chapter's
-  marginalization theorem permits, and trains nothing. The edit intensities of the sampler section
-  are the chapter's own numbers, not a network's.
-- `dirichlet.py` evaluates the derivative of the regularized incomplete beta function by a central
-  difference on `scipy.special.betainc`, in double precision whatever the ambient dtype, since the
-  two values agree to six decimals. The continuity check differentiates the flux with a five-point
-  stencil.
-- `rectify.py` runs ReDi with the exact marginal velocity of the current coupling on a
-  two-position state space, so no network is trained. The AReUReDi section works one Metropolis
-  iteration, not a full refinement chain.
-- `guidance.py` trains the denoiser and the property model under one optimizer, and the property is
-  the GC label of `cis6270.data`, not a biological measurement.
-- `multi_objective.py` trains nothing. Its two DNA objectives are the GC fraction and the AT
-  fraction of a finished sequence, its feasible set is the sequences carrying the planted motif,
-  and the base flow a rollout continues is the mask-source sampler with the data's letter frequency
-  as its denoiser.
-- The relaxed paths replace the token lookup of `cis6270.nets.SequenceDenoiser` with a linear map on
-  the relaxed token, and leave the rest of the network untouched.

lecture_5/corrector.py

@@ -1,251 +0,0 @@
-"""Correctors: extra motion that leaves the probability path where it is.
-
-The master equation is linear in the velocity, so a second velocity with no net flux into any state
-can be added to the first without moving the marginals. Detailed balance is the easy way to arrange
-that, and the mixture path already contains such a term, the difference between the forward velocity
-and the same path run backward. Adding it with weight eta makes a position unmask more often and
-remask in between, which gives the sampler more chances to replace a token the rest of the sequence
-has ruled out.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.13 and 5.14. The denoiser of the long simulation is the letter frequency of the
-synthetic DNA of ``cis6270.data``. Prints the zero-flux check, the per-step unmasking table, the jump
-counts at three corrector strengths, the parallel-step error at 64 steps, and the chance a jointly
-revealed pair survives to t = 1 untouched.
-
-    python lecture_5/corrector.py                     # 600 Euler steps on 40,000 positions
-    python lecture_5/corrector.py --eta 10            # one corrector strength in the step table
-    python lecture_5/corrector.py --steps 20 --quiet  # a coarse grid, visibly biased jump counts
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import DNA, DNA_SIZE, dna_splits
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-EXACT = torch.double  # The worked arithmetic keeps full precision whatever the default dtype.
-MASK = DNA_SIZE  # Index 4 is the mask symbol m, outside the four-letter alphabet.
-
-
-# %% 02. Velocities that move samples without moving the path
-#
-# A velocity v_t whose net flux into every state is zero can be added to a velocity that generates
-# p_t, and the sum generates the same path. Detailed balance, equal flux in both directions along
-# every edge, is sufficient. Reversing every term turns the net flux into p_t(y) times a column of
-# rates, and those rates sum to zero.
-
-
-def net_flux(v, p):
-    """sum_x v_t(y, x) p_t(x) at every state y. v: [K, K], p: [K] -> [K]."""
-    return v @ p  # Zero at every entry when v moves no probability.
-
-
-def detailed_balance_velocity(p, chosen):
-    """Build v_t from chosen forward rates, with each reverse rate fixed by detailed balance.
-
-    ``chosen`` maps an edge (x, y) to the rate v_t(y, x). The reverse rate is then
-    v_t(x, y) = v_t(y, x) p_t(x) / p_t(y), which equalizes the flux along that edge.
-    """
-    v = torch.zeros(len(p), len(p), dtype=p.dtype)  # [K, K], destination indexing the row.
-    for (x, y), rate in chosen.items():
-        v[y, x] = rate  # The rate we pick, out of x toward y.
-        v[x, y] = rate * p[x] / p[y]  # The reverse rate detailed balance forces.
-    return v - torch.diag(v.sum(dim=0))  # Each column now sums to zero.
-
-
-# %% 03. A zero-flux velocity on three states
-#
-# On the states A, C, G with p_t = (0.2, 0.3, 0.5) we can pick the rate 3 from A to C and the rate 1
-# from A to G, and detailed balance returns 2 and 0.4. Adding this to a velocity that sends every
-# state to G at rate 1 raises the total rate of leaving A from 1 to 5 while p_dot stays put.
-
-
-def three_state_example():
-    """Section 5.13.2, returning the two velocities, the path and the two rates of change."""
-    p = torch.tensor([0.2, 0.3, 0.5], dtype=EXACT)  # [3] over A, C, G.
-    v = detailed_balance_velocity(p, {(0, 1): 3.0, (0, 2): 1.0})  # A to C at 3, A to G at 1.
-    u = torch.zeros(3, 3, dtype=EXACT)  # The velocity that does move the path.
-    for x in range(2):
-        u[2, x], u[x, x] = 1.0, -1.0  # A and C jump to G at rate 1; G never leaves.
-    return p, u, v, net_flux(u, p), net_flux(u + v, p)  # The last two agree.
-
-
-# %% 04. The corrector for the mask source
-#
-# With the linear schedule, the weights alpha_t = 1 + eta t and beta_t = eta t, and the mask source,
-# the backward velocity is the point mass on the mask and the combination collapses to two brackets.
-# A masked position unmasks at total rate (1 + eta t)/(1 - t), split over tokens by the denoiser, and
-# a visible position remasks at rate eta.
-
-
-def corrector_velocity(p1, x_i, t, eta):
-    """bar_u^i_t(., x) over the five symbols A, C, G, T, m. p1: [4] -> [5]."""
-    current = torch.zeros(DNA_SIZE + 1, dtype=EXACT)
-    current[x_i] = 1.0  # delta on the token the position holds, which may be the mask.
-    denoiser = torch.cat([p1.to(EXACT), torch.zeros(1, dtype=EXACT)])  # The data token is never m.
-    to_mask = torch.zeros(DNA_SIZE + 1, dtype=EXACT)
-    to_mask[MASK] = 1.0  # delta_m, the destination of a remasking.
-    forward = (1.0 + eta * t) / (1.0 - t) * (denoiser - current)  # The first bracket.
-    return forward + eta * (to_mask - current)  # The second bracket, the circulation back to m.
-
-
-def admissible_step(t, eta):
-    """The largest Euler step with a redraw probability at most one, h <= (1 - t)/(1 + eta t)."""
-    return (1.0 - t) / (1.0 + eta * t)  # 0.1, 0.0357 and 0.01 at t = 0.9 for eta = 0, 2, 10.
-
-
-# %% 05. Counting the jumps
-#
-# Writing M(t) for the probability that the position is masked, the two flows cancel so that
-# M_dot(t) = -1 for every eta, and M(t) = 1 - t whatever the corrector strength. Integrating the
-# unmasking rate against M and the remasking rate against 1 - M gives 1 + eta/2 unmaskings and eta/2
-# remaskings, for 1 + eta jumps per position.
-
-
-def masked_fraction_derivative(t, eta, masked=None):
-    """M_dot(t) = -M (1 + eta t)/(1 - t) + (1 - M) eta, which equals -1 at M = 1 - t."""
-    masked = 1.0 - t if masked is None else masked  # The solution we expect.
-    return -masked * (1.0 + eta * t) / (1.0 - t) + (1.0 - masked) * eta
-
-
-def simulate_jumps(p1, eta, positions, steps, generator=None):
-    """Run ``positions`` independent positions from the mask and count their jumps.
-
-    Returns the fraction still masked at t = 0.5 and the mean number of jumps per position.
-    """
-    state = torch.full((positions,), MASK, dtype=torch.long)  # Every position masked at t = 0.
-    jumps = torch.zeros(positions, dtype=torch.int32)  # One counter per position.
-    h, halfway = 1.0 / steps, None  # The increment, and the masked fraction recorded at t = 0.5.
-    for k in range(steps):
-        t = k * h
-        masked = state == MASK  # [P] which positions are still at the mask symbol.
-        if halfway is None and t >= 0.5:
-            halfway = float(masked.to(torch.get_default_dtype()).mean())  # M(0.5), which the chapter puts at 1 - t.
-        up = min(1.0, h * (1.0 + eta * t) / max(1.0 - t, 1e-12))  # Unmasking probability this step.
-        down = h * eta  # Remasking probability this step.
-        coin = torch.rand(positions, generator=generator)
-        unmask = masked & (coin < up)  # Masked positions that reveal a token.
-        remask = (~masked) & (coin < down)  # Visible positions that return to the mask.
-        revealed = int(unmask.sum())  # Draw a token only where one is needed.
-        if revealed:
-            tokens = torch.multinomial(p1, revealed, replacement=True, generator=generator)
-            state = state.masked_scatter(unmask, tokens)  # Split over tokens by the denoiser.
-        state = state.masked_fill(remask, MASK)  # The circulation back to the mask symbol.
-        jumps += (unmask | remask).to(torch.int32)  # Every reveal and remasking counts as a jump.
-    return halfway, float(jumps.double().mean())  # Near 1 - t = 0.5, and near 1 + eta.
-
-
-# %% 06. The parallel-step error
-#
-# Let the data be uniform over AA and TT and the source be all masks. With both positions masked the
-# denoiser at each is (0.5, 0.5), and a step that reveals both at once draws them independently, so
-# half of those draws land outside the data. With eta > 0 a pair revealed jointly is remasked at rate
-# 2 eta and redrawn with the other position visible, where the denoiser is a point mass.
-
-
-def joint_reveal_rate(steps):
-    """The probability that the two reveal times land in the same cell of a uniform grid, h."""
-    return 1.0 / steps  # Sum over cells of h squared, which is 1/64 at 64 steps.
-
-
-def pair_survival(eta, tau=0.5):
-    """exp(-2 eta (1 - tau)), the chance a pair revealed at tau is untouched at t = 1."""
-    return math.exp(-2.0 * eta * (1.0 - tau))  # 0.135 at eta = 2 and 4.54e-5 at eta = 10.
-
-
-def simulate_survival(eta, samples, steps, tau=0.5, generator=None):
-    """The fraction of pairs revealed at ``tau`` that neither position remasks before t = 1.
-
-    Each of the two positions returns to the mask at rate eta, so the first remasking of either
-    happens at rate 2 eta, and a pair that is remasked is redrawn with the other position visible.
-    """
-    remaining = max(1, int(steps * (1.0 - tau)))  # Grid cells left between tau and t = 1.
-    h = (1.0 - tau) / remaining  # The increment, matching the grid the sampler runs on.
-    untouched = torch.ones(samples, dtype=torch.bool)  # Pairs no remasking has reached.
-    for _ in range(remaining):
-        coin = torch.rand(samples, 2, generator=generator)  # One clock per position.
-        untouched &= ~(coin < h * eta).any(dim=1)  # Either position remasking ends the survival.
-    return float(untouched.to(torch.get_default_dtype()).mean())  # Approaches exp(-2 eta (1 - tau)).
-
-
-def simulate_pair(eta, samples, steps, generator=None):
-    """The AA-or-TT two-position problem under the corrector. Returns the fraction outside the data."""
-    state = torch.full((samples, 2), MASK, dtype=torch.long)  # Both positions masked at t = 0.
-    h = 1.0 / steps
-    for k in range(steps):
-        t, last = k * h, k == steps - 1  # The last step reveals whatever is left.
-        masked = state == MASK  # [S, 2] read before either update, so the two are simultaneous.
-        other = state.flip(1)  # What the opposite position holds, which the denoiser reads.
-        guess = torch.randint(0, 2, state.shape, generator=generator)  # The (0.5, 0.5) denoiser.
-        drawn = torch.where(other != MASK, other, guess)  # A point mass once the partner is visible.
-        up = 1.0 if last else min(1.0, h * (1.0 + eta * t) / max(1.0 - t, 1e-12))
-        down = 0.0 if last else h * eta  # No remasking on the step that ends the run.
-        reveal = masked & (torch.rand(state.shape, generator=generator) < up)
-        remask = (~masked) & (torch.rand(state.shape, generator=generator) < down)
-        state = torch.where(reveal, drawn, state)  # Unmask, in parallel at both positions.
-        state = torch.where(remask, torch.full_like(state, MASK), state)  # The circulation back to m.
-    return float((state[:, 0] != state[:, 1]).double().mean())  # Outside {AA, TT}.
-
-
-# %% 07. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--eta", type=float, default=2.0, help="corrector strength in the step table")
-    parser.add_argument("--positions", type=int, default=40_000, help="positions in the simulation")
-    parser.add_argument("--pairs", type=int, default=20_000, help="sequences in the pair simulation")
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    p, u, v, with_u, with_both = three_state_example()
-    report = {
-        "zero_flux": [round(float(f), 6) for f in net_flux(v, p)],  # Zero at A, C and G.
-        "p_dot_u": [round(float(f), 4) for f in with_u],  # (-0.2, -0.3, 0.5).
-        "p_dot_u_plus_v": [round(float(f), 4) for f in with_both],  # The same three numbers.
-        "rate_out_of_A_u": round(float(-u[0, 0]), 4),  # 1 under the velocity alone.
-        "rate_out_of_A_both": round(float(-(u + v)[0, 0]), 4),  # 5 once the corrector is added.
-    }
-
-    # Section 5.14.3: one masked position at t = 0.4 with the chapter's denoiser and h = 0.01.
-    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])  # [4] over A, C, G, T.
-    bar_u = corrector_velocity(p1, MASK, t=0.4, eta=args.eta)  # [5] rates out of the mask.
-    seen = torch.zeros(DNA_SIZE)
-    seen[1] = 1.0  # At a visible position the true denoiser is the point mass on the visible token.
-    report["unmask_rate"] = round(float(-bar_u[MASK]), 4)  # (1 + eta t)/(1 - t); 3.0 at eta = 2.
-    report["remask_rate"] = round(float(corrector_velocity(seen, 1, 0.4, args.eta)[MASK]), 4)
-    report["step_probabilities"] = [round(float(0.01 * r), 5) for r in bar_u[:DNA_SIZE]]
-    report["stay_masked"] = round(float(1.0 + 0.01 * bar_u[MASK]), 5)  # 0.97 at eta = 2.
-    report["remask_probability"] = round(0.01 * args.eta, 4)  # eta h, which is 0.02 at eta = 2.
-    report["masked_derivative"] = round(masked_fraction_derivative(0.4, args.eta), 6)  # -1.
-    report["admissible_step_at_0.9"] = round(admissible_step(0.9, args.eta), 4)  # 0.0357 at eta = 2.
-
-    # The long simulation, with the letter frequency of the synthetic DNA as the denoiser.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
-    frequency = counts / counts.sum()  # [4] the GC-rich background the lecture shares.
-    for eta in (0.0, 2.0, 10.0):
-        half, mean = simulate_jumps(frequency, eta, args.positions, args.steps)
-        report[f"masked_at_0.5_eta_{eta:g}"] = round(half, 4)  # 1 - t = 0.5 at every eta.
-        report[f"jumps_eta_{eta:g}"] = round(mean, 3)  # 1 + eta.
-
-    # Section 5.14.4: the two positions that have to agree.
-    report["joint_reveal_at_64"] = round(joint_reveal_rate(64), 4)  # 0.0156.
-    report["error_at_64_predicted"] = round(0.5 * joint_reveal_rate(64), 4)  # Half of those are wrong.
-    report["error_at_64_measured"] = f"{simulate_pair(0.0, args.pairs, 64):.5f}"  # Near 0.0078.
-    report["error_at_grid_measured"] = f"{simulate_pair(0.0, args.pairs, args.steps):.5f}"
-    for eta in (0.0, 2.0, 10.0):
-        predicted = pair_survival(eta)  # exp(-2 eta (1 - tau)) at the midpoint tau = 0.5.
-        report[f"survival_eta_{eta:g}"] = f"{predicted:.3e}"
-        if predicted * args.pairs >= 20:  # Below that, the simulation has nothing to resolve.
-            measured = simulate_survival(eta, args.pairs, args.steps)  # The same number, simulated.
-            report[f"survival_eta_{eta:g}_measured"] = f"{measured:.3e}"
-    print_report("correctors", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-if __name__ == "__main__":
-    main()

lecture_5/dfm.py

@@ -1,237 +0,0 @@
-"""Discrete flow matching, the whole construction in one file.
-
-We pick the probability path first, write down a velocity that generates it, and read the sampler off
-the velocity. Each position follows the mixture path between a source token and a data token, its
-conditional velocity is a schedule factor times the difference of two point masses, and the posterior
-average of those conditional velocities is the same factor times a denoiser. Cross-entropy on the
-data token recovers that denoiser, and one Euler step redraws a position from it with probability
-lambda_t. Chapter 4 fixed the corruption and solved for its reverse; here the corruption is a
-consequence of the path we chose.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.5 to 5.11. Trained on the synthetic DNA of ``cis6270.data``. Prints the
-chapter's worked numbers, the training loss and the measurements of a generated batch.
-
-    python lecture_5/dfm.py                     # train and generate from the uniform source
-    python lecture_5/dfm.py --source mask       # the all-mask source instead
-    python lecture_5/dfm.py --steps 20 --quiet  # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-
-import torch
-import torch.nn.functional as F
-from torch.distributions import Categorical
-
-from cis6270.data import DNA_SIZE, decode_dna, dna_splits, encode_dna
-from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
-from cis6270.nets import SequenceDenoiser, count_parameters
-from cis6270.runner import common_parser, print_report, save_run, set_seed, train
-
-
-# %% 02. The class of Section 5.11, with both sources
-#
-# The constructor holds the vocabulary size V and the schedule kappa_t with its derivative
-# kappa_dot_t, passed as functions of a time tensor. The defaults are the linear schedule
-# kappa_t = t and kappa_dot_t = 1. A mask source returns the index of the mask symbol at every
-# position, with the network's vocabulary raised by one to make room for it.
-
-
-class DiscreteFlowMatching:
-    """The source, the mixture path, the cross-entropy loss and one sampling step."""
-
-    def __init__(self, V, kappa=lambda t: t, kappa_dot=lambda t: torch.ones_like(t), mask=None):
-        self.V, self.k, self.kd = V, kappa, kappa_dot  # Alphabet size and the schedule.
-        self.mask = mask  # Mask token id under the mask source, None under the uniform source.
-
-    def source(self, shape, device="cpu"):
-        if self.mask is None:
-            return torch.randint(0, self.V, shape, device=device)  # Uniform source, [B, L].
-        return torch.full(shape, self.mask, dtype=torch.long, device=device)  # All-mask source.
-
-    def corrupt(self, x1, t):
-        """One draw from p_t(. | x_0, x_1) by the coin flip of Section 5.5.5. x1: [B, L], t: [B]."""
-        x0 = self.source(x1.shape, x1.device)  # The other end of the pair; [B, L].
-        heads = torch.rand(x1.shape, device=x1.device) < self.k(t)[:, None]  # Heads with prob kappa_t.
-        return torch.where(heads, x1, x0)  # The data token on heads and the source token on tails.
-
-
-# %% 03. The conditional velocity
-#
-# Differentiating the conditional path at one position gives kappa_dot_t times the difference of the
-# two indicators, and dividing that flux by the probability 1 - kappa_t still sitting on the source
-# token turns it into a rate out of the token the position currently holds.
-
-
-def conditional_velocity(z, x1, t, V=DNA_SIZE):
-    """u_t(a, z | x_1) = kappa_dot/(1 - kappa) [delta_{x_1}(a) - delta_z(a)]. Returns [V]."""
-    factor = 1.0 / (1.0 - t)  # Schedule factor on the linear schedule, kappa_dot_t = 1.
-    return factor * (F.one_hot(torch.tensor(x1), V) - F.one_hot(torch.tensor(z), V)).double()
-
-
-# %% 04. The posterior average, collapsed onto the data token
-#
-# The conditional velocity depends on the pair only through the data token at this position, so every
-# pair sharing that token contributes the same term and the sum over pairs collapses onto the four
-# tokens. What is left is the denoiser p_{1|t}(. | x), the only quantity a network has to learn.
-
-
-def marginal_velocity(p1, x_i, t, V=DNA_SIZE):
-    """u^i_t(., x) = kappa_dot/(1 - kappa) [p_{1|t}(. | x) - delta_{x^i}]. p1: [V] -> [V]."""
-    factor = 1.0 / (1.0 - t)  # Same schedule factor, now multiplying a distribution.
-    return factor * (p1 - F.one_hot(torch.tensor(x_i), V).double())  # Signed, sums to zero.
-
-
-def mask_posterior(data, weights, visible, kappa_t):
-    """The posterior over the data token at the masked position, computed by enumeration.
-
-    ``data`` is a list of candidate sequences and ``weights`` their probabilities under q. Position 0
-    is visible and holds ``visible``; position 1 is masked. The time factor is common to every
-    candidate and cancels, which is Proposition 6 of Discrete Flow Matching.
-    """
-    masked, seen = 1, 1  # One masked position and one visible position in this example.
-    factor = (1.0 - kappa_t) ** masked * kappa_t ** seen  # Common to every surviving candidate.
-    unnormalized = torch.zeros(DNA_SIZE, dtype=torch.double)  # One entry per base at position 1.
-    for sequence, weight in zip(data, weights):
-        if sequence[0] == visible:  # A candidate disagreeing with the visible token is excluded.
-            unnormalized[encode_dna([sequence])[0, 1]] += factor * weight
-    return unnormalized / unnormalized.sum()  # [V] the denoiser at the masked position.
-
-
-# %% 05. The cross-entropy loss and its minimizer
-#
-# The loss scores the true data token at every position, with no time weight and no rate in it, so the
-# sampler is chosen after training. Averaging over the data token given X_t splits the expected loss
-# into the entropy of the true denoiser and a KL divergence from it to the network, and only the
-# divergence depends on theta.
-
-
-def dfm_loss(logits, x1):
-    """Cross-entropy of Section 5.9, with the class dimension second. logits: [B, L, V]."""
-    return F.cross_entropy(logits.transpose(1, 2), x1)  # Mean over the batch and over positions.
-
-
-def entropy_and_gap(posterior, prediction):
-    """The minimum of the loss at one position and the excess a wrong prediction pays, in nats."""
-    entropy = -(posterior * posterior.log()).sum()  # H(p_{1|t}), the floor no network beats.
-    cross = -(posterior * prediction.log()).sum()  # The cross-entropy the prediction scores.
-    return float(entropy), float(cross - entropy)  # The gap is the KL divergence to the prediction.
-
-
-# %% 06. The Euler step as a two-part rule
-#
-# Substituting the marginal velocity into the per-position Euler step collects the two terms into a
-# mixture of the current point mass and the denoiser. With probability lambda_t the position redraws
-# from the denoiser, and with probability 1 - lambda_t it keeps the token it holds.
-
-
-def redraw_probability(t, h):
-    """lambda_t = h kappa_dot_t / (1 - kappa_t), the redraw probability of Section 5.8."""
-    return h / (1.0 - t)  # Linear schedule; above one the sampler redraws with certainty.
-
-
-@torch.no_grad()
-def dfm_step(proc, model, xt, t, h):
-    """One two-part step at every position at once. xt: [B, L], t: [B], h a scalar."""
-    p1 = model(xt, t)[..., :proc.V].softmax(-1)  # [B, L, V] denoiser, the only network call.
-    rate = proc.kd(t) / (1 - proc.k(t)).clamp(min=1e-4)  # kappa_dot/(1 - kappa), kept finite at t=1.
-    redraw = torch.rand(xt.shape, device=xt.device) < (h * rate)[:, None]  # [B, L] Bernoulli.
-    x1 = Categorical(probs=p1).sample()  # [B, L] one candidate token per position.
-    return torch.where(redraw, x1, xt)  # Keep the candidate where the coin came up true.
-
-
-# %% 07. Training and generation
-#
-# A training step draws one time per sequence, corrupts the batch at that time, and scores the true
-# tokens. Generation walks a uniform grid of ``sample_steps`` times from t = 0 to t = 1, and the grid
-# is free to change because training never chose a step size.
-
-
-@torch.no_grad()
-def generate(proc, model, batch, length, steps, device="cpu"):
-    """Run the sampler from the source to t = 1 and return the tokens. [B, L]."""
-    x = proc.source((batch, length), device)  # Draw from the source distribution p.
-    grid = torch.linspace(0.0, 1.0, steps + 1, device=device)  # The sampling times.
-    for t, s in zip(grid[:-1], grid[1:]):
-        x = dfm_step(proc, model, x, t.expand(batch), s - t)  # One forward pass per step.
-    return x
-
-
-# %% 08. The chapter's worked numbers
-def worked_numbers():
-    """The arithmetic Sections 5.5.5, 5.7.4, 5.8, 5.9.2 and 5.9.3 carry out by hand."""
-    report = {}
-
-    # Section 5.5.5: x_1 = ACGT, a uniform source draw x_0 = GTTA, kappa_t = 0.5, coins HTHT.
-    x1, x0 = encode_dna(["ACGT"])[0], encode_dna(["GTTA"])[0]  # [4] each.
-    heads = torch.tensor([True, False, True, False])  # The four coins of the example.
-    xt = torch.where(heads, x1, x0)  # Data token on heads, source token on tails.
-    report["mixture_draw"] = decode_dna(xt[None])[0]  # ATGA.
-
-    # Section 5.7.4: one position holding G, denoiser (0.1, 0.6, 0.2, 0.1), t = 0.5.
-    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1], dtype=torch.double)  # [V] over A, C, G, T.
-    u = marginal_velocity(p1, x_i=2, t=0.5)  # Factor 1/0.5 = 2 times the bracket.
-    report["velocity_at_G"] = [round(float(v), 4) for v in u]  # (0.2, 1.2, -1.6, 0.2).
-    report["velocity_sums_to"] = f"{float(u.sum()):.1e}"  # Zero, which is the second rule.
-    step = F.one_hot(torch.tensor(2), DNA_SIZE).double() + 0.1 * u  # One Euler step of size h = 0.1.
-    report["euler_step_at_G"] = [round(float(v), 4) for v in step]  # (0.02, 0.12, 0.84, 0.02).
-    report["redraw_probability"] = round(redraw_probability(0.5, 0.1), 4)  # lambda_t = 0.2.
-
-    # Section 5.9.2: the noisy input ATGA scored against the target ACGT.
-    scored = torch.tensor([0.7, 0.2, 0.6, 0.1], dtype=torch.double)  # One probability per position.
-    report["example_loss_nats"] = round(float(-scored.log().sum()), 3)  # 4.780 nats.
-
-    # Section 5.9.3: length two, the mask source, q(AC) = 0.5, q(AG) = 0.3, q(GC) = 0.2.
-    candidates, weights = ["AC", "AG", "GC"], [0.5, 0.3, 0.2]
-    half = mask_posterior(candidates, weights, visible="A", kappa_t=0.5)  # Denoiser at kappa = 0.5.
-    late = mask_posterior(candidates, weights, visible="A", kappa_t=0.8)  # The same at kappa = 0.8.
-    report["mask_posterior_C_G"] = [round(float(half[1]), 4), round(float(half[2]), 4)]
-    report["mask_posterior_unchanged"] = bool(torch.allclose(half, late))  # The time cancels.
-    live = half[half > 0]  # The two bases the posterior keeps.
-    entropy, gap = entropy_and_gap(live, torch.tensor([0.5, 0.5], dtype=torch.double))
-    report["posterior_entropy_nats"] = round(entropy, 4)  # 0.6616, the floor on the loss.
-    report["kl_gap_to_uniform"] = round(gap, 4)  # 0.0316 above that floor.
-    return report
-
-
-# %% 09. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--source", choices=["uniform", "mask"], default="uniform")
-    parser.set_defaults(steps=1000, width=64, samples=64)
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    data = splits["train_tokens"]  # [N, L] the shared synthetic DNA.
-    mask = DNA_SIZE if args.source == "mask" else None  # Index 4 is the mask symbol m.
-    vocab = DNA_SIZE + 1 if mask is not None else DNA_SIZE  # The network's vocabulary.
-    proc = DiscreteFlowMatching(DNA_SIZE, mask=mask)  # Linear schedule by default.
-    model = SequenceDenoiser(vocab, args.length, width=args.width, depth=2)
-
-    def loss_fn(step):
-        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
-        x1 = data[rows]  # [B, L] the data end of the pair.
-        t = torch.rand(args.batch_size)  # One time per sequence, uniform on [0, 1].
-        xt = proc.corrupt(x1, t)  # [B, L] a draw from the conditional path.
-        return dfm_loss(model(xt, t)[..., :DNA_SIZE], x1)  # Score the true token at every position.
-
-    if not args.quiet:
-        print(f"source={args.source} vocab={vocab} parameters={count_parameters(model):,}")
-    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
-    tokens = generate(proc, model, args.samples, args.length, args.sample_steps)
-
-    report = worked_numbers()
-    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])  # Averaged over the last steps.
-    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
-    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
-    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
-    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
-    print_report(f"discrete flow matching, {args.source} source", report, decode_dna(tokens))
-    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))
-
-
-if __name__ == "__main__":
-    main()

lecture_5/dirichlet.py

@@ -1,262 +0,0 @@
-"""Dirichlet Flow Matching: a token relaxed into a point of the simplex.
-
-A straight line toward a corner shrinks its support as t grows, and past t = 1/k at most k - 1
-corners remain possible, which makes the denoiser target discontinuous in time and in space.
-Replacing the line by the Dirichlet family Dir(1 + t e_i) keeps every corner possible at every finite
-time. The velocity of that path points at the data corner and changes only its speed, and matching
-the probability flux across the cut x_i = c solves for that speed.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y). On this path the label x_1
-still means the data sample while its subscript stops meaning time one, since time runs to infinity
-and inference stops at t = 8.
-
-Chapter 5, Sections 5.19.3 and 5.20.1 to 5.20.5. Trained on the synthetic DNA of ``cis6270.data``.
-Prints the worked speeds, the continuity check, the posterior average, and the measurements of a
-generated batch.
-
-    python lecture_5/dirichlet.py                     # train, then 100 Euler steps to t = 8
-    python lecture_5/dirichlet.py --steps 20 --quiet  # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from scipy.special import betainc, betaln
-
-from cis6270.data import DNA_SIZE, decode_dna, dna_splits
-from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
-from cis6270.nets import SequenceDenoiser, count_parameters
-from cis6270.runner import common_parser, print_report, save_run, set_seed, train
-
-T_MAX = 8.0  # The inference horizon, where the data coordinate has mean (t+1)/(t+K) = 0.75 for DNA.
-EDGE = 1e-6  # The speed divides by (1 - x_i) and by x_i, so a coordinate is held off both ends.
-
-
-# %% 02. The Dirichlet family as a conditional path
-#
-# With the concentration 1 + t in the data coordinate and 1 everywhere else, every exponent except
-# the i-th is zero and the density collapses to a single factor. At t = 0 the path is the uniform
-# distribution on the simplex, and for every finite t it is positive wherever x_i is, so no corner is
-# ever ruled out.
-
-
-def dirichlet_density(x, i, t, K=DNA_SIZE):
-    """p_t(x | x_1 = e_i) = Gamma(t + K)/Gamma(t + 1) x_i^t."""
-    prefactor = math.exp(math.lgamma(t + K) - math.lgamma(t + 1.0))  # Gamma(4) = 6 at t = 0.
-    return prefactor * float(x[i]) ** t  # Constant over the simplex at t = 0.
-
-
-def sample_dirichlet(targets, t, K=DNA_SIZE):
-    """Draw x ~ Dir(1 + t e_i) at every position. targets: [B, L], t: [B] -> [B, L, K]."""
-    alpha = torch.ones(*targets.shape, K) + t[:, None, None] * F.one_hot(targets, K)  # [B, L, K].
-    gamma = torch._standard_gamma(alpha)  # One Gamma draw per letter, with unit rate.
-    return gamma / gamma.sum(-1, keepdim=True)  # Normalizing the Gamma draws gives a Dirichlet.
-
-
-def straight_line_corners(x, t):
-    """How many corners the straight-line path still allows at x, which is the count of x_i >= t."""
-    return int((x >= t).sum())  # Proposition 5.5 bounds this by k - 1 once t > 1/k.
-
-
-# %% 03. The velocity by matching flux across a cut
-#
-# Cutting the simplex at x_i = c, the mass below the cut is the Beta(t+1, K-1) distribution function
-# at c, and its rate of change is the derivative of that function in its first parameter. Setting the
-# loss rate equal to minus the outward flux and solving for the speed leaves one scalar per
-# coordinate, with no closed form at general K.
-
-
-def d_betainc_da(x, a, b, eps=1e-6):
-    """tilde I_x(a, b), the derivative of the regularized incomplete beta function in a."""
-    return (betainc(a + eps, b, x) - betainc(a - eps, b, x)) / (2.0 * eps)  # A central difference.
-
-
-def speed(x_i, t, K=DNA_SIZE):
-    """C(x_i, t) = -tilde I_{x_i}(t+1, K-1) B(t+1, K-1) / ((1 - x_i)^{K-1} x_i^t).
-
-    The data coordinate may be a scalar or a tensor of any shape, since every factor broadcasts,
-    and a whole batch of relaxed tokens costs the two calls to the incomplete beta function that
-    one costs. The central difference subtracts two numbers that agree to about six decimals, and
-    at single precision nothing of the difference survives, so this one quantity is computed in
-    double whatever the ambient dtype and the caller casts the result back.
-    """
-    a, b = t + 1.0, float(K - 1)  # The two Beta parameters of the data coordinate.
-    beta = math.exp(betaln(a, b))  # The complete beta function, 1/12 at a = 2 and b = 3.
-    held = torch.as_tensor(x_i, dtype=torch.double).clamp(EDGE, 1.0 - EDGE)  # Off both endpoints.
-    derivative = torch.as_tensor(d_betainc_da(held.numpy(), a, b))  # Negative, so the speed is not.
-    return -derivative * beta / ((1.0 - held) ** (K - 1) * held ** t)  # Positive.
-
-
-def conditional_velocity(x, i, t):
-    """u_t(x | x_1 = e_i) = C(x_i, t)(e_i - x), tangent to the simplex. x: [K] -> [K]."""
-    direction = F.one_hot(torch.tensor(i), len(x)).to(x) - x  # Coordinates summing to zero.
-    return float(speed(float(x[i]), t, len(x))) * direction  # The speed scales the direction.
-
-
-# %% 04. Two letters, in closed form
-#
-# With K = 2 the second Beta parameter is one, the incomplete beta function is x^a, and the speed
-# comes out in elementary functions. Differentiating the density in time and the flux in space then
-# cancels term by term, which is the continuity equation at K = 2.
-
-
-def speed_two_letters(x, t):
-    """C(x_1, t) = -x_1 log x_1 / ((t + 1)(1 - x_1)), the K = 2 case of the general formula."""
-    return -x * math.log(x) / ((t + 1.0) * (1.0 - x))  # Positive on the whole open interval.
-
-
-def continuity_terms(x, i, t, h=1e-2):
-    """partial_t p and div(p u) in the chart (x_1, ..., x_{K-1}), whose sum the path drives to zero."""
-    x = x.to(torch.double)  # A stencil on top of a numerical derivative needs the extra digits.
-    K = len(x)
-
-    def density(point):
-        return dirichlet_density(point, i, t, K)  # The Dirichlet density at a simplex point.
-
-    def chart(point, coordinate):
-        return density(point) * float(conditional_velocity(point, i, t)[coordinate])  # p u_j.
-
-    time = density(x) * (_digamma(t + K) - _digamma(t + 1.0) + math.log(float(x[i])))  # partial_t p.
-    divergence = 0.0
-    for j in range(K - 1):  # The chart drops the last coordinate, which the sum constraint fixes.
-        step = torch.zeros(K, dtype=x.dtype)
-        step[j], step[K - 1] = 1.0, -1.0  # A direction inside the simplex, keeping the sum at one.
-        divergence += (  # A five-point stencil, since the speed itself is a numerical derivative.
-            -chart(x + 2 * h * step, j) + 8 * chart(x + h * step, j)
-            - 8 * chart(x - h * step, j) + chart(x - 2 * h * step, j)
-        ) / (12.0 * h)
-    return time, divergence  # Equal and opposite to the precision of the difference.
-
-
-def _digamma(z):
-    """The digamma function, used only to differentiate the Gamma prefactor in time."""
-    return float(torch.digamma(torch.tensor(z, dtype=torch.double)))
-
-
-# %% 05. A worked speed on a DNA base
-#
-# A coordinate already holding a large share has a larger speed, so the flow concentrates gradually
-# and the coordinate in the lead moves fastest toward its corner.
-
-
-def speed_vector(x, t):
-    """C(x_k, t) at every letter of one relaxed base. x: [K] -> [K]."""
-    return speed(x, t, x.shape[-1]).to(x)  # One vectorized evaluation, cast back to x's dtype.
-
-
-# %% 06. The posterior average and the loss
-#
-# The data corner is unknown at generation time, so the four conditional velocities combine with the
-# denoiser as weights. Each one is tangent to the simplex and a weighted average of tangent vectors
-# is tangent, so the learned field needs no constraint. Cross-entropy on the data token trains the
-# denoiser, and the conditional velocity appears only inside the sampler.
-
-
-def marginal_velocity(x, p_hat, t):
-    """v_t(x) = sum_k u_t(x | e_k) p_{1|t}(e_k | x). x: [K], p_hat: [K] -> [K]."""
-    columns = torch.stack([conditional_velocity(x, k, t) for k in range(len(x))], dim=1)  # [K, K].
-    return columns @ p_hat  # The posterior average of the four conditional velocities.
-
-
-def batched_velocity(x, p_hat, t):
-    """The same average at every position of a batch. x, p_hat: [B, L, K] -> [B, L, K]."""
-    speeds = speed(x, t, x.shape[-1]).to(x)  # [B, L, K] one speed per coordinate.
-    weighted = speeds * p_hat  # [B, L, K] the coefficient of each conditional velocity.
-    return weighted - weighted.sum(-1, keepdim=True) * x  # sum_k w_k (e_k - x), collected.
-
-
-def relaxed_denoiser(length, width, K=DNA_SIZE, depth=2):
-    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token."""
-    model = SequenceDenoiser(K, length, width=width, depth=depth)
-    model.token_embedding = nn.Linear(K, width, bias=False)  # A point of the simplex, not an index.
-    return model
-
-
-# %% 07. Training on DNA, and sampling by Euler steps
-#
-# Training draws t from an exponential with rate one, places a point on the conditional path, and
-# scores the true letter. Generation starts from the uniform Dirichlet and integrates the marginal
-# velocity to t = 8, projecting back onto the simplex after each step.
-
-
-@torch.no_grad()
-def generate(model, batch, length, steps, K=DNA_SIZE):
-    """Integrate the marginal velocity from the uniform simplex to t = T_MAX. Returns [B, L]."""
-    gamma = torch._standard_gamma(torch.ones(batch, length, K))  # Dir(1) is the uniform source.
-    x = gamma / gamma.sum(-1, keepdim=True)  # [B, L, K].
-    h = T_MAX / steps  # The time increment, 0.08 over 100 steps.
-    for k in range(steps):
-        t = k * h
-        p_hat = model(x, torch.full((batch,), t)).softmax(-1)  # [B, L, K] the denoiser.
-        x = x + h * batched_velocity(x, p_hat, t)  # One Euler step along the marginal velocity.
-        x = x.clamp(min=1e-8)  # A large step can push a coordinate below zero.
-        x = x / x.sum(-1, keepdim=True)  # Back onto the simplex.
-    return x.argmax(-1)  # Read the token off the largest coordinate.
-
-
-# %% 08. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=1000, width=64, samples=64, sample_steps=100)
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    report = {
-        "density_at_t_0": round(dirichlet_density(torch.full((4,), 0.25), 0, 0.0), 4),  # 6.
-        "corners_left_at_0.3": straight_line_corners(torch.tensor([0.34, 0.33, 0.33, 0.0]), 0.3),
-        "corners_left_at_0.6": straight_line_corners(torch.tensor([0.70, 0.30, 0.0, 0.0]), 0.6),
-        "speed_K2_closed": f"{speed_two_letters(0.3, 1.0):.7f}",  # 0.2579942.
-        "speed_K2_general": f"{float(speed(0.3, 1.0, K=2)):.7f}",  # The same to six decimals.
-    }
-    point = torch.tensor([0.5, 0.3, 0.2], dtype=torch.double)  # A base of three letters, at t = 0.8.
-    time_term, space_term = continuity_terms(point, 0, 0.8)
-    report["continuity_time_term"] = round(time_term, 4)  # 0.6355.
-    report["continuity_space_term"] = round(space_term, 4)  # -0.6355.
-    report["continuity_residual"] = f"{time_term + space_term:.2e}"  # Their sum, near zero.
-
-    # Section 5.20.4: K = 4, t = 1 and the relaxed base (0.40, 0.30, 0.20, 0.10).
-    x = torch.tensor([0.40, 0.30, 0.20, 0.10], dtype=torch.double)
-    report["speeds"] = [round(float(c), 3) for c in speed_vector(x, 1.0)]  # 0.232, 0.199, 0.159, 0.104.
-    toward_a = conditional_velocity(x, 0, 1.0)  # [4] the velocity toward the A corner.
-    report["velocity_toward_A"] = [round(float(v), 3) for v in toward_a]
-    report["velocity_sums_to"] = f"{float(toward_a.sum()):.1e}"  # Zero, so the step stays on it.
-
-    # Section 5.20.5: the same base with the denoiser (0.60, 0.25, 0.10, 0.05).
-    p_hat = torch.tensor([0.60, 0.25, 0.10, 0.05], dtype=torch.double)
-    v = marginal_velocity(x, p_hat, 1.0)
-    report["marginal_velocity"] = [round(float(c), 4) for c in v]  # (0.0553, -0.0133, -0.0262, -0.0158).
-    report["euler_step_h_0.1"] = [round(float(c), 4) for c in x + 0.1 * v]  # (0.4055, ...).
-    report["step_sums_to"] = round(float((x + 0.1 * v).sum()), 9)  # One, still on the simplex.
-
-    # Training on the shared DNA, then generation.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    data = splits["train_tokens"]  # [N, L].
-    model = relaxed_denoiser(args.length, args.width)
-
-    def loss_fn(step):
-        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
-        x1 = data[rows]  # [B, L] the letters the denoiser has to recover.
-        t = torch.distributions.Exponential(1.0).sample((args.batch_size,))  # t ~ Exp(1).
-        xt = sample_dirichlet(x1, t)  # [B, L, K] a point on the conditional path.
-        return F.cross_entropy(model(xt, t).transpose(1, 2), x1)  # Score the true letter.
-
-    if not args.quiet:
-        print(f"parameters={count_parameters(model):,} horizon={T_MAX}")
-    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
-    tokens = generate(model, args.samples, args.length, args.sample_steps)
-
-    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])
-    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
-    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
-    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
-    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
-    print_report("Dirichlet flow matching", report, decode_dna(tokens))
-    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))
-
-
-if __name__ == "__main__":
-    main()

lecture_5/edit_flows.py

@@ -1,295 +0,0 @@
-"""Edit Flows: paths whose sequences change length.
-
-Every path so far keeps a single length, because a substitution changes a token and never adds or
-removes one. Enlarging the set of moves to insertions, deletions and substitutions removes that
-constraint. An alignment pads the two ends with a blank until they have the same number of columns,
-the mixture construction then applies column by column, and the alignment is an auxiliary variable we
-marginalize exactly as we marginalized the data sequence. The loss scores rates in place of a
-distribution over tokens, and its minimizer is the rate at which an edit makes progress.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Section 5.17. Runs on the synthetic DNA of ``cis6270.data``, trimmed to varying lengths.
-Prints the alignment of the chapter, the posterior over two alignments, the averaged edit rates, the
-loss of one example, one sampler step, and the length statistics of a run on DNA.
-
-    python lecture_5/edit_flows.py                     # 600 sampler steps on 256 DNA pairs
-    python lecture_5/edit_flows.py --steps 20 --quiet  # a coarse grid, a visibly coarser length
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import DNA, DNA_SIZE, dna_splits, encode_dna
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-BLANK = DNA_SIZE  # Index 4 is the blank symbol epsilon that padding writes into a column.
-
-
-# %% 02. Three edits and the rate split
-#
-# At each position of the current sequence there are three edits, and each rate splits into an
-# intensity and a distribution over tokens. A deletion writes no token, so the network emits three
-# intensities and two token distributions per position, which is 3 + 2K numbers in place of the K of
-# the denoiser.
-
-
-def edit_rates(intensities, tokens=DNA_SIZE):
-    """Count the numbers a network emits per position, and the total rate of leaving the sequence."""
-    total = sum(sum(kind) for kind in intensities.values())  # -u_t(x, x), the total exit rate.
-    return 3 + 2 * tokens, total  # Three intensities and two token distributions per position.
-
-
-# %% 03. Alignments and the column-wise path
-#
-# An alignment puts the positions of the two sequences in correspondence by padding both with the
-# blank. A blank in the target column is a deletion, a blank in the source column an insertion, and
-# two different tokens a substitution. On the aligned pair each column follows the mixture path
-# independently, and removing the blanks from z_t gives the sequence x_t the model reads.
-
-
-def align(x0, x1):
-    """A minimal alignment of two token lists, as the padded pair (z_0, z_1) of equal length."""
-    n, m = len(x0), len(x1)
-    cost = [[0] * (m + 1) for _ in range(n + 1)]  # Edit distance with unit costs.
-    for i in range(n + 1):
-        cost[i][0] = i  # Deleting the first i tokens of the source.
-    for j in range(m + 1):
-        cost[0][j] = j  # Inserting the first j tokens of the target.
-    for i in range(1, n + 1):
-        for j in range(1, m + 1):
-            cost[i][j] = min(cost[i - 1][j] + 1, cost[i][j - 1] + 1,
-                             cost[i - 1][j - 1] + int(x0[i - 1] != x1[j - 1]))
-    z0, z1, i, j = [], [], n, m
-    while i or j:  # One backtrace through the table, taking the diagonal where it is minimal.
-        if i and j and cost[i][j] == cost[i - 1][j - 1] + int(x0[i - 1] != x1[j - 1]):
-            z0.append(x0[i - 1]), z1.append(x1[j - 1]), (i := i - 1), (j := j - 1)
-        elif i and cost[i][j] == cost[i - 1][j] + 1:
-            z0.append(x0[i - 1]), z1.append(BLANK), (i := i - 1)  # A deletion column.
-        else:
-            z0.append(BLANK), z1.append(x1[j - 1]), (j := j - 1)  # An insertion column.
-    return z0[::-1], z1[::-1]  # [C] and [C] over the same number of columns.
-
-
-def column_kinds(z0, z1):
-    """Name the edit each column owes: ``del``, ``ins``, ``sub``, or ``none`` when the two agree."""
-    kinds = []
-    for a, b in zip(z0, z1):
-        kinds.append("none" if a == b else "del" if b == BLANK else "ins" if a == BLANK else "sub")
-    return kinds  # [C] one label per column.
-
-
-def draw_state(z0, z1, finished):
-    """z_t from a choice of which columns already hold their target, with the blanks removed."""
-    held = [b if done else a for a, b, done in zip(z0, z1, finished)]  # [C] the aligned state.
-    return [token for token in held if token != BLANK]  # [|x_t|] the sequence the model reads.
-
-
-def expected_length(x0, x1, kappa_t):
-    """E[|x_t|] = (1 - kappa_t)|x_0| + kappa_t|x_1|, the interpolation one position up."""
-    return (1.0 - kappa_t) * len(x0) + kappa_t * len(x1)  # Linear from |x_0| to |x_1|.
-
-
-# %% 04. Averaging over the alignments we cannot see
-#
-# Many alignments produce the same pair, and x_t does not record which one the path is on. Bayes'
-# rule over the alignments weights each by the probability that its column draw produces the observed
-# x_t, and the rates average against those weights. No single alignment has to be chosen.
-
-
-def alignment_posterior(alignments, observed, kappa_t):
-    """p_t(A | x_t) over a list of aligned pairs, by enumerating the column draws of each."""
-    weights = []
-    for z0, z1 in alignments:
-        owed = [i for i, kind in enumerate(column_kinds(z0, z1)) if kind != "none"]  # Unfinished.
-        total = 0.0
-        for mask in range(1 << len(owed)):  # Every subset of the owed columns already finished.
-            finished = [True] * len(z0)
-            probability = 1.0
-            for bit, column in enumerate(owed):
-                done = bool(mask >> bit & 1)
-                finished[column] = done
-                probability *= kappa_t if done else 1.0 - kappa_t  # The column-wise mixture path.
-            if draw_state(z0, z1, finished) == observed:
-                total += probability  # This draw is consistent with what we see.
-        weights.append(total)
-    norm = sum(weights)
-    return [w / norm for w in weights]  # One posterior weight per alignment.
-
-
-def averaged_rates(alignments, observed, kappa_t, kappa_dot=1.0):
-    """The marginal rate of each edit that finishes a column, averaged over the alignments.
-
-    ``observed`` is the draw in which every column that owes an edit still holds its source symbol,
-    which is the state the worked example of Section 5.17.2 reads the rates at.
-    """
-    rate = kappa_dot / (1.0 - kappa_t)  # The weight on every unfinished column.
-    posterior, out = alignment_posterior(alignments, observed, kappa_t), {}
-    for (z0, z1), weight in zip(alignments, posterior):
-        held = [b if kind == "none" else a  # A column owing nothing is finished at every time.
-                for a, b, kind in zip(z0, z1, column_kinds(z0, z1))]
-        for i, (a, b, kind) in enumerate(zip(z0, z1, column_kinds(z0, z1))):
-            seen = sum(1 for token in held[:i] if token != BLANK)  # Position inside x_t.
-            if kind == "sub":
-                name = f"sub({seen + 1},{DNA[b]})"  # Replace the token at that position.
-            elif kind == "ins":
-                name = f"ins({seen},{DNA[b]})"  # Place a token to the right of that position.
-            elif kind == "del":
-                name = f"del({seen + 1})"  # Remove the token at that position.
-            else:
-                continue  # Nothing is owed here.
-            out[name] = out.get(name, 0.0) + weight * rate
-    return posterior, out
-
-
-# %% 05. The training loss
-#
-# The first term is the sum of the model's intensities and pushes all of them down. The second is a
-# weighted cross-entropy on the edits that move x_t toward x_1 and pushes those up. Collecting the
-# terms belonging to one edit gives lambda - w log lambda, whose stationary point is lambda = w.
-
-
-def edit_loss(total_intensity, owed, kappa_t, kappa_dot=1.0):
-    """L = sum of all intensities - sum over unfinished columns of the weight times log lambda."""
-    weight = kappa_dot / (1.0 - kappa_t)  # kappa_dot/(1 - kappa), which is 2 at t = 0.5.
-    return total_intensity - sum(weight * math.log(lam) for lam in owed)  # The model's own rates.
-
-
-def loss_minimizer(weight, grid):
-    """The intensity minimizing lambda - w log lambda on a grid, which the derivative puts at w."""
-    values = [lam - weight * math.log(lam) for lam in grid]  # The contribution of one edit.
-    return grid[min(range(len(grid)), key=values.__getitem__)]  # Near w, with w/lambda^2 > 0.
-
-
-# %% 06. One step of the sampler
-#
-# Each position inserts with probability h lambda_ins, and deletes or substitutes with probability
-# h(lambda_del + lambda_sub), the two sharing that draw in proportion to their intensities. Every
-# chosen edit is applied at once, so the step has the same second-order error as the parallel
-# step on a fixed length.
-
-
-def step_probabilities(intensities, h):
-    """The per-position insertion and delete-or-substitute probabilities of one Euler step."""
-    insert = [h * lam for lam in intensities["ins"]]  # One coin per position.
-    change = [h * (d + s) for d, s in zip(intensities["del"], intensities["sub"])]  # A second coin.
-    return insert, change
-
-
-def untouched_probability(intensities, h):
-    """The chance the step applies no edit anywhere, and the first-order estimate beside it."""
-    insert, change = step_probabilities(intensities, h)
-    exact = math.prod(1.0 - p for p in insert + change)  # Every coin coming up tails.
-    _, total = edit_rates(intensities)
-    return exact, 1.0 - h * total  # The gap is the second-order term from two edits at once.
-
-
-def expected_length_change(intensities, h):
-    """h times the insertion intensities minus the deletion intensities, in tokens."""
-    return h * (sum(intensities["ins"]) - sum(intensities["del"]))  # Nothing fixes the length.
-
-
-# %% 07. Edits on the synthetic DNA
-#
-# Each pair of sequences is aligned once, every column that owes an edit finishes at rate
-# kappa_dot/(1 - kappa), and the sequence the model reads is the aligned state with its blanks
-# removed. The realized length wanders while its mean follows the interpolation.
-
-
-def simulate_columns(pairs, steps, generator=None):
-    """Run the column-wise path on a list of aligned pairs. Returns the lengths and the end states."""
-    width = max(len(z0) for z0, _ in pairs)  # Pad every alignment to the widest one.
-    src = torch.full((len(pairs), width), BLANK)
-    tgt = torch.full((len(pairs), width), BLANK)
-    valid = torch.zeros(len(pairs), width, dtype=torch.bool)
-    for row, (z0, z1) in enumerate(pairs):
-        src[row, :len(z0)] = torch.tensor(z0)
-        tgt[row, :len(z1)] = torch.tensor(z1)
-        valid[row, :len(z0)] = True  # Columns beyond the alignment take no part.
-    done = (src == tgt) & valid  # A column whose two symbols agree owes no edit.
-    h, halfway = 1.0 / steps, None
-    for k in range(steps):
-        t = k * h
-        if halfway is None and t >= 0.5:
-            held = torch.where(done, tgt, src)  # [P, C] the aligned state at this time.
-            halfway = float(((held != BLANK) & valid).sum(dim=1).double().mean())
-        lam = 1.0 if k == steps - 1 else min(1.0, h / max(1.0 - t, 1e-12))  # lambda_t per column.
-        coin = torch.rand(src.shape, generator=generator) < lam
-        done = done | (coin & valid)  # An unfinished column takes its target symbol.
-    finished = float((done | ~valid).all(dim=1).double().mean())  # Fraction of pairs that arrived.
-    return halfway, finished
-
-
-# %% 08. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--pairs", type=int, default=256, help="DNA pairs aligned and run")
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    # Section 5.17.1: the chapter's alignment of ACGT with AGAC over five columns.
-    x0, x1 = encode_dna(["ACGT"])[0].tolist(), encode_dna(["AGAC"])[0].tolist()
-    z0 = x0 + [BLANK]  # (A, C, G, T, eps).
-    z1 = [x1[0], BLANK] + x1[1:]  # (A, eps, G, A, C).
-    kinds = column_kinds(z0, z1)
-    draw = [True, False, True, True, True]  # Column 2 behind, columns 4 and 5 already at the target.
-    report = {
-        "column_kinds": kinds,  # none, del, none, sub, ins.
-        "draw_at_kappa_0.5": "".join(DNA[token] for token in draw_state(z0, z1, draw)),  # ACGAC.
-        "draw_probability": round(0.5 ** sum(kind != "none" for kind in kinds), 4),  # 0.125.
-        "expected_length": round(expected_length(x0, x1, 0.5), 4),  # 4, both ends having length 4.
-        "numbers_per_position": edit_rates({"ins": [0.0], "del": [0.0], "sub": [0.0]})[0],  # 11.
-    }
-
-    # Section 5.17.2: the pair AC and AAC, which admits two alignments on three columns.
-    short, long = encode_dna(["AC"])[0].tolist(), encode_dna(["AAC"])[0].tolist()
-    first = ([short[0], short[1], BLANK], long)  # Substitute C with A, then insert C.
-    second = ([short[0], BLANK, short[1]], long)  # Insert A, and leave the last column alone.
-    observed = short  # The noisy sequence x_t = AC.
-    for kappa_t in (0.25, 0.5, 0.75):
-        weights = alignment_posterior([first, second], observed, kappa_t)
-        report[f"posterior_A1_kappa_{kappa_t}"] = round(weights[0], 4)  # (1-kappa)/(2-kappa).
-    posterior, rates = averaged_rates([first, second], observed, 0.5)
-    for name in sorted(rates):
-        report[f"rate_{name}"] = round(rates[name], 3)  # 0.667, 1.333 and 0.667.
-    report["total_target_rate"] = round(sum(rates.values()), 3)  # 2.667.
-
-    # Section 5.17.3: one unfinished column, a model intensity of 0.8, and 3.0 over all edits.
-    report["example_loss"] = round(edit_loss(3.0, [0.8], kappa_t=0.5), 4)  # 3.446.
-    grid = [0.01 * i for i in range(1, 801)]  # Intensities between 0.01 and 8.
-    report["loss_minimizer"] = round(loss_minimizer(2.0, grid), 2)  # 2, the weight the edit receives.
-
-    # Section 5.17.4: one whole step on ACG with the chapter's nine intensities.
-    intensities = {"ins": [0.2, 0.3, 0.4], "del": [0.1, 0.9, 0.2], "sub": [0.4, 0.3, 0.5]}
-    exact, first_order = untouched_probability(intensities, 0.05)
-    insert, change = step_probabilities(intensities, 0.05)
-    report["insert_probabilities"] = [round(p, 3) for p in insert]  # 0.010, 0.015, 0.020.
-    report["change_probabilities"] = [round(p, 3) for p in change]  # 0.025, 0.060, 0.035.
-    report["deletion_given_change"] = round(0.9 / (0.9 + 0.3), 4)  # 0.75, three times in four.
-    report["untouched"] = round(exact, 4)  # 0.8452.
-    report["untouched_first_order"] = round(first_order, 4)  # 0.835.
-    report["second_order_gap"] = round(exact - first_order, 4)  # 0.0102.
-    report["length_change"] = round(expected_length_change(intensities, 0.05), 4)  # -0.015.
-
-    # The same path on the shared DNA, trimmed so that the two ends have different lengths.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    tokens = splits["train_tokens"]
-    cuts = torch.randint(args.length - 4, args.length + 1, (2 * args.pairs,))  # Varying lengths.
-    rows = [tokens[i % len(tokens), :cuts[i]].tolist() for i in range(2 * args.pairs)]
-    pairs = [align(rows[2 * i], rows[2 * i + 1]) for i in range(args.pairs)]
-    predicted = sum(expected_length(rows[2 * i], rows[2 * i + 1], 0.5) for i in range(args.pairs))
-    measured, arrived = simulate_columns(pairs, args.steps)
-    report["dna_columns_mean"] = round(sum(len(z0) for z0, _ in pairs) / len(pairs), 3)
-    report["dna_length_at_0.5"] = round(measured, 3)  # Measured over the run.
-    report["dna_length_predicted"] = round(predicted / args.pairs, 3)  # The interpolation.
-    report["dna_pairs_arrived"] = round(arrived, 4)  # Every column finishes by t = 1.
-    report["dna_example"] = " -> ".join("".join(DNA[t] for t in row) for row in rows[:2])
-    print_report("edit flows", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-if __name__ == "__main__":
-    main()

lecture_5/fisher.py

@@ -1,257 +0,0 @@
-"""Fisher-Flow: the simplex measured by the metric statistics already uses for distributions.
-
-The straight-line distance on the simplex measures absolute differences in probability, so dividing a
-rare letter's probability by ten counts for almost nothing. The Fisher information measures how
-distinguishable two nearby distributions are, and by Chentsov's theorem it is the only metric under
-which relabeling the letters or merging two of them is an isometry. Taking the square root of every
-coordinate turns it into the ordinary metric on a sphere, where the conditional path is a
-great-circle arc and the velocity is the remaining arc divided by the remaining time.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.21.1 to 5.21.5, with the listing of Section 5.21.7. Trained on the synthetic
-DNA of ``cis6270.data``. Prints the two distances, the worked arc on a DNA base, the tangent
-projection, and the measurements of a generated batch.
-
-    python lecture_5/fisher.py                     # train, then 64 exponential-map steps
-    python lecture_5/fisher.py --steps 20 --quiet  # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-import torch.nn as nn
-
-from cis6270.data import DNA_SIZE, decode_dna, dna_splits
-from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
-from cis6270.nets import SequenceDenoiser, count_parameters
-from cis6270.runner import common_parser, print_report, save_run, set_seed, train
-
-SMOOTHING = 0.03  # Label smoothing that holds a data corner off the boundary, where the metric blows up.
-
-
-# %% 02. Small probabilities under the straight-line distance
-#
-# Moving a base from (0.5, 0.5) to (0.4, 0.6) shifts a tenth of the mass between two evenly matched
-# letters. Moving from (0.01, 0.99) to (0.001, 0.999) divides the first probability by ten, which
-# matters for a model of sequences, and the Euclidean distance calls the second move eleven times
-# smaller.
-
-
-def euclidean_distance(p, q):
-    """The straight-line distance between two probability vectors."""
-    return float((p - q).norm())  # 0.1414 and 0.0127 for the two moves above.
-
-
-def fisher_rao_distance(p, q):
-    """d_FR(p, q) = 2 arccos(<sqrt(p), sqrt(q)>), at most pi for any two distributions."""
-    inner = float((p.sqrt() * q.sqrt()).sum().clamp(max=1.0))  # The angle between square roots.
-    return 2.0 * math.acos(inner)  # 0.201 and 0.137 for the same two moves.
-
-
-# %% 03. The Fisher information as a local metric
-#
-# With the probabilities as the parameters, the derivative of the log likelihood of outcome i with
-# respect to p^j is delta_ij/p^j, and the expectation over outcomes collapses the two deltas onto the
-# diagonal. Dividing by p^i makes a change in a rare coordinate count for more.
-
-
-def fisher_matrix(p):
-    """g_jk(p) = delta_jk / p^j, the Fisher information of a categorical distribution. [K, K]."""
-    return torch.diag(1.0 / p)  # Diagonal, with the reciprocal of each probability.
-
-
-def fisher_inner(u, v, p):
-    """<u, v>_p = sum_i u^i v^i / p^i, the Fisher-Rao inner product of two tangent vectors at p."""
-    return float((u * v / p).sum())  # Tangent means the coordinates sum to zero.
-
-
-# %% 04. The square-root map to the sphere
-#
-# The simplex constraint becomes a norm constraint, so the square roots lie on the unit sphere, in
-# the octant with every coordinate nonnegative. Differentiating the square root shows the squared
-# Euclidean length there is one quarter of the Fisher-Rao length, so the map scaled by two is an
-# isometry.
-
-
-def to_sphere(p):
-    """A probability vector becomes a unit vector. The clamp holds a zero off the boundary."""
-    return p.clamp(min=1e-8).sqrt()  # simplex -> sphere.
-
-
-def to_simplex(x):
-    """A unit vector with nonnegative coordinates becomes a probability vector."""
-    return x ** 2  # sphere -> simplex.
-
-
-def smooth(one_hot, amount=SMOOTHING):
-    """Move a corner inside the simplex, which the metric needs since it diverges at the boundary."""
-    K = one_hot.shape[-1]
-    return one_hot * (1.0 - K * amount) + amount  # (0.91, 0.03, 0.03, 0.03) at amount 0.03.
-
-
-# %% 05. Geodesics on the sphere
-#
-# The shortest path between two points of a sphere is the arc of the great circle through them, and
-# the conditional path is that arc traversed at constant speed. The logarithm map has length equal to
-# the remaining arc, so dividing by the remaining time leaves the speed constant along the path.
-
-
-def log_map(x, y):
-    """The tangent vector at x reaching y in unit time, of length equal to the arc from x to y."""
-    c = (x * y).sum(-1, keepdim=True).clamp(max=1 - 1e-6)  # The cosine of the angle between them.
-    v = y - c * x  # orthogonal to x.
-    return v * torch.arccos(c) / v.norm(dim=-1, keepdim=True)  # Rescaled to the arc length.
-
-
-def exp_map(x, w):
-    """The point reached by moving along the sphere from x in the direction w for a distance ||w||."""
-    n = w.norm(dim=-1, keepdim=True).clamp(min=1e-8)  # The length of the tangent vector.
-    return torch.cos(n) * x + torch.sin(n) * (w / n)  # Unit length by cos^2 + sin^2 = 1.
-
-
-def geodesic(x0, x1, t):
-    """x_t = exp_{x_0}(t log_{x_0}(x_1)), the point a fraction t of the way along the arc."""
-    return exp_map(x0, t * log_map(x0, x1))  # the conditional path.
-
-
-def conditional_velocity(xt, x1, t):
-    """u_t(x_t | x_1) = log_{x_t}(x_1)/(1 - t), the remaining arc over the remaining time."""
-    return log_map(xt, x1) / (1 - t)  # Constant in length along the whole path.
-
-
-# %% 06. The tangent projection and the regression loss
-#
-# A linear output layer emits an arbitrary vector of R^K, and a velocity on the sphere has to be
-# orthogonal to the current point. Subtracting the component along x enforces that in one line, and
-# the minimizer of the squared loss is the average of the conditional velocities passing through the
-# point, which is the averaging every path in this chapter performs.
-
-
-def project(xt, raw):
-    """v_theta = raw - <x_t, raw> x_t, which is orthogonal to x_t since ||x_t|| = 1."""
-    return raw - (xt * raw).sum(-1, keepdim=True) * xt  # Any architecture can be used.
-
-
-def loss(model, x0, x1, t):
-    """The squared distance between the projected network output and the geodesic target."""
-    xt = geodesic(x0, x1, t[:, None, None])  # [B, L, K] a point on the arc at the sampled time.
-    u = log_map(xt, x1) / (1 - t[:, None, None])  # [B, L, K] the remaining arc over the time left.
-    raw = model(xt, t)  # [B, L, K] unconstrained.
-    v = raw - (xt * raw).sum(-1, keepdim=True) * xt  # The tangent part of that output.
-    return ((v - u) ** 2).sum(-1).mean()  # Summed over letters, averaged over positions and batch.
-
-
-def velocity_network(length, width, K=DNA_SIZE, depth=2):
-    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token.
-
-    The head emits a vector of R^K here, which ``project`` turns into a velocity on the sphere.
-    """
-    model = SequenceDenoiser(K, length, width=width, depth=depth)
-    model.token_embedding = nn.Linear(K, width, bias=False)  # A point on the sphere, not an index.
-    return model
-
-
-# %% 07. Training on DNA, and sampling along the sphere
-#
-# Source points are uniform draws from the simplex, data points are smoothed corners, and the pair is
-# joined by the arc between their square roots. Generation starts from a fresh source point and takes
-# exponential-map steps, so every state stays on the sphere and no renormalization is needed.
-
-
-def source_points(batch, length, K=DNA_SIZE):
-    """Uniform draws from the simplex, mapped to the sphere. Returns [B, L, K]."""
-    gamma = torch._standard_gamma(torch.ones(batch, length, K))  # Dir(1) over the simplex.
-    return to_sphere(gamma / gamma.sum(-1, keepdim=True))  # [B, L, K] on the unit sphere.
-
-
-@torch.no_grad()
-def generate(model, batch, length, steps, K=DNA_SIZE):
-    """Integrate the learned field from the source to t = 1 with exponential-map steps. [B, L]."""
-    x = source_points(batch, length, K)  # [B, L, K].
-    h = 1.0 / steps
-    for k in range(steps):
-        t = torch.full((batch,), k * h)
-        v = project(x, model(x, t))  # [B, L, K] the tangent part of the network output.
-        x = exp_map(x, h * v)  # Move along the sphere, which keeps the norm at one.
-    return to_simplex(x).argmax(-1)  # Read the token off the largest coordinate.
-
-
-# %% 08. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.set_defaults(steps=1000, width=64, samples=64)
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    near = torch.tensor([0.5, 0.5]), torch.tensor([0.4, 0.6])  # Two evenly matched letters.
-    rare = torch.tensor([0.01, 0.99]), torch.tensor([0.001, 0.999])  # One rare letter, divided by ten.
-    report = {
-        "euclidean_near": round(euclidean_distance(*near), 4),  # 0.1414.
-        "euclidean_rare": round(euclidean_distance(*rare), 4),  # 0.0127, eleven times smaller.
-        "fisher_rao_near": round(fisher_rao_distance(*near), 3),  # 0.201.
-        "fisher_rao_rare": round(fisher_rao_distance(*rare), 3),  # 0.137, comparable.
-        "two_letters_are_orthogonal": round(
-            fisher_rao_distance(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])), 4),  # pi.
-    }
-
-    # Section 5.21.3: a uniform base and a smoothed A, taken to the sphere.
-    p0 = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE)  # The uniform base.
-    p1 = smooth(torch.eye(DNA_SIZE)[0])  # (0.91, 0.03, 0.03, 0.03).
-    x0, x1 = to_sphere(p0), to_sphere(p1)
-    report["fisher_diagonal"] = [round(float(v), 2) for v in torch.diagonal(fisher_matrix(p0))]
-    report["sphere_x1"] = [round(float(v), 3) for v in x1]  # (0.954, 0.173, 0.173, 0.173).
-    report["inner_product"] = round(float((x0 * x1).sum()), 3)  # 0.737.
-    report["angle"] = round(math.acos(float((x0 * x1).sum())), 3)  # 0.743.
-    report["fisher_rao_distance"] = round(fisher_rao_distance(p0, p1), 3)  # 1.485, twice the angle.
-
-    # Section 5.21.4: the midpoint of the arc, and one exponential-map step from it.
-    mid = geodesic(x0, x1, torch.tensor(0.5))
-    report["midpoint_sphere"] = [round(float(v), 3) for v in mid]  # (0.780, 0.361, 0.361, 0.361).
-    report["midpoint_simplex"] = [round(float(v), 3) for v in to_simplex(mid)]  # (0.609, 0.130, ...).
-    remaining = log_map(mid, x1)
-    report["log_map_from_mid"] = [round(float(v), 4) for v in remaining]  # (0.2323, -0.1672, ...).
-    report["remaining_arc"] = round(float(remaining.norm()), 4)  # 0.3713, half the total.
-    report["constant_speed"] = round(float(conditional_velocity(mid, x1, 0.5).norm()), 3)  # 0.743.
-    stepped = exp_map(mid, 0.1 * conditional_velocity(mid, x1, 0.5))  # Delta = 0.1 from t = 0.5.
-    report["step_sphere"] = [round(float(v), 3) for v in stepped]  # (0.824, 0.327, 0.327, 0.327).
-    report["step_simplex"] = [round(float(v), 3) for v in to_simplex(stepped)]  # (0.680, 0.107, ...).
-    report["step_norm"] = round(float(stepped.norm()), 10)  # One, because the step followed the arc.
-
-    # Section 5.21.5: the tangent projection of one raw network output.
-    raw = torch.tensor([0.3, -0.2, 0.5, 0.1])
-    report["raw_inner_product"] = round(float((mid * raw).sum()), 4)  # 0.3785.
-    tangent = project(mid, raw)
-    report["projected"] = [round(float(v), 3) for v in tangent]  # (0.005, -0.337, 0.363, -0.037).
-    report["projection_inner"] = f"{float((mid * tangent).sum()):.1e}"  # Zero to machine precision.
-
-    # Training on the shared DNA, then generation.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    data = splits["train_tokens"]  # [N, L].
-    model = velocity_network(args.length, args.width)
-
-    def loss_fn(step):
-        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
-        target = to_sphere(smooth(torch.eye(DNA_SIZE)[data[rows]]))  # [B, L, K] smoothed corners.
-        start = source_points(args.batch_size, args.length)  # [B, L, K] the independent coupling.
-        t = torch.rand(args.batch_size) * 0.98  # Held off t = 1, where the target divides by zero.
-        return loss(model, start, target, t)
-
-    if not args.quiet:
-        print(f"parameters={count_parameters(model):,} smoothing={SMOOTHING}")
-    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
-    tokens = generate(model, args.samples, args.length, args.sample_steps)
-
-    report["final_loss"] = sum(losses[-20:]) / len(losses[-20:])
-    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
-    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
-    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
-    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
-    print_report("Fisher-Flow", report, decode_dna(tokens))
-    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))
-
-
-if __name__ == "__main__":
-    main()

lecture_5/guidance.py

@@ -1,234 +0,0 @@
-"""Guidance: steering a trained flow toward a property by tilting its rates.
-
-The distribution we want is the data distribution reweighted by the likelihood of a property, whose
-normalizer runs over every sequence. A rate is a ratio, so that constant cancels, and the conditioned
-process multiplies each rate by the ratio of a value function at the destination to its value at the
-current state. A first-order expansion of the logarithm of that ratio replaces one property
-evaluation per destination with one backward pass, and an exponent on the ratio gives one strength
-parameter, which the classifier-free form reaches from two trained velocities and no classifier.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.24 to 5.29, with the listing of Section 5.33. Several objectives at once are in
-``lecture_5/multi_objective.py``. Trained on the synthetic DNA of ``cis6270.data``, with the GC label
-as the property. Prints the worked guided base, the guidance strength, the Taylor tilt, and the GC
-fraction of a batch generated with and without guidance.
-
-    python lecture_5/guidance.py                     # train, then sample at gamma = 0 and gamma = 2
-    python lecture_5/guidance.py --gamma 5           # a harder tilt toward the GC-rich label
-    python lecture_5/guidance.py --steps 20 --quiet  # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from cis6270.data import DNA_SIZE, decode_dna, dna_splits
-from cis6270.metrics import gc_fraction, motif_fraction, token_entropy
-from cis6270.nets import SequenceDenoiser, TimeEmbedding, count_parameters
-from cis6270.runner import common_parser, print_report, save_run, set_seed, train
-
-V = DNA_SIZE + 1  # Four bases plus the mask symbol, which the sampler starts from.
-MASK = DNA_SIZE  # Index 4 is the mask symbol m.
-
-
-def k(t):
-    """The schedule kappa_t, linear here."""
-    return t
-
-
-def kd(t):
-    """Its derivative kappa_dot_t, which the linear schedule holds at one."""
-    return torch.ones_like(t)
-
-
-# %% 02. Guidance as a tilt of the rates
-#
-# The conditioned velocity is the unguided one multiplied by the ratio of the property probability at
-# the destination to its value at the current state. A move toward a destination the property model
-# scores higher happens sooner, and a destination with rate zero stays at rate zero, so the tilt
-# reweights the support of the model without enlarging it.
-
-
-def tilt(rates, destination_scores, current_score, gamma=1.0):
-    """u^(gamma)_t(y | x, c) = [p(c | y, t)/p(c | x, t)]^gamma u_t(y, x)."""
-    return rates * (destination_scores / current_score) ** gamma  # gamma = 1 is the exact process.
-
-
-def h_value(denoiser, destination_scores):
-    """h_t(x) = E[p(c | X_1) | X_t = x], the posterior average of the property likelihood."""
-    return float((denoiser * destination_scores).sum())  # The exact value at a masked position.
-
-
-def classifier_free(conditional, unconditional, gamma):
-    """u_c^gamma u_u^(1-gamma), the geometric mixture that reaches the same one-parameter family."""
-    return conditional ** gamma * unconditional ** (1.0 - gamma)  # Linear in the logarithms.
-
-
-def reachable(positions, letters):
-    """D(S - 1) + 1, the number of states one jump can reach, since exactly one position changes."""
-    return positions * (letters - 1) + 1  # 2001 for a 100-residue peptide over 21 symbols.
-
-
-# %% 03. A first-order expansion of the log ratio
-#
-# Embedding a sequence as one-hot vectors makes a single-position change a sparse difference, and
-# expanding the log property probability around the current sequence collapses the inner product onto
-# the difference of two gradient coordinates. One forward and one backward pass then serve every
-# destination.
-
-
-def taylor_log_ratio(gradient, current):
-    """[grad log p_phi]_{i,b} - [grad log p_phi]_{i,a} at one position. gradient: [K] -> [K]."""
-    return gradient - gradient[current]  # Zero at the letter in place, whose ratio is one.
-
-
-# %% 04. The guided sampler
-#
-# One guided step differs from the unguided one at a single line, the factor exp(gamma adv). The
-# one-hot encoding has gradients switched on, so the property network can be differentiated against
-# letters that are not currently in place, and summing over the batch before differentiating shares
-# one backward pass across it. The denoiser is read over the four bases alone, since the data token
-# is never the mask, which holds the rate back into the mask at zero and the exit rate at
-# kappa_dot/(1 - kappa) when gamma is zero.
-
-
-def guided_step(model, prop, xt, t, h, gamma=1.0):
-    """One Euler step of the tilted velocity at every position at once. xt: [B, D] -> [B, D]."""
-    oh = F.one_hot(xt, V).to(torch.get_default_dtype()).requires_grad_()
-    logp = prop(oh, t).sum()  # property on noisy x
-    g = torch.autograd.grad(logp, oh)[0]  # one backward pass
-    adv = g - (oh * g).sum(-1, keepdim=True)  # Taylor ratio
-    with torch.no_grad():
-        bases = model(xt, t)[..., :DNA_SIZE].softmax(-1)  # denoiser
-        p1 = F.pad(bases, (0, V - DNA_SIZE))  # the data token is never the mask
-        rate = kd(t) / (1 - k(t)).clamp(min=1e-4)
-        u = rate[:, None, None] * p1 * torch.exp(gamma * adv)  # every rate tilted
-        u.scatter_(-1, xt[..., None], 0.0)  # no self-transition
-        jump = torch.rand(xt.shape) < h * u.sum(-1)
-        y = torch.distributions.Categorical(probs=u).sample()
-    return torch.where(jump, y, xt)
-
-
-class PropertyModel(nn.Module):
-    """log p_phi(c | x, t) from a one-hot sequence and a time, trained on noisy intermediate states."""
-
-    def __init__(self, length, width=64):
-        super().__init__()
-        self.time_embedding = TimeEmbedding(width, scale=1.0)  # The property depends on the time.
-        self.body = nn.Sequential(
-            nn.Linear(length * V + width, width), nn.SiLU(), nn.Linear(width, 1),
-        )
-
-    def forward(self, oh, t):
-        """oh: [B, D, V] one-hot and differentiable, t: [B] -> [B] log probability of the property."""
-        h = torch.cat([oh.flatten(1), self.time_embedding(t)], dim=1)  # [B, D V + width].
-        return F.logsigmoid(self.body(h)).squeeze(-1)  # [B] log p_phi(c | x, t).
-
-
-# %% 05. Training the denoiser and the property model
-#
-# The denoiser is the one built earlier in this chapter, trained from the mask source. The property
-# model reads the same noisy sequences and predicts the GC label of the clean one, so it is a
-# function of an intermediate state and not of a clean sequence.
-
-
-def corrupt(x1, t):
-    """A draw from the conditional path with the all-mask source. x1: [B, D], t: [B] -> [B, D]."""
-    heads = torch.rand(x1.shape) < k(t)[:, None]  # Heads with probability kappa_t at each position.
-    return torch.where(heads, x1, torch.full_like(x1, MASK))  # Data token on heads, mask on tails.
-
-
-@torch.no_grad()
-def generate(model, prop, batch, length, steps, gamma):
-    """Run the guided sampler from all masks to t = 1. Returns [B, D]."""
-    x = torch.full((batch, length), MASK, dtype=torch.long)  # The all-mask source.
-    grid = torch.linspace(0.0, 1.0, steps + 1)
-    for t, s in zip(grid[:-1], grid[1:]):
-        with torch.enable_grad():  # The property gradient needs a graph inside the no-grad sampler.
-            x = guided_step(model, prop, x, t.expand(batch), float(s - t), gamma)
-    p1 = model(x, torch.ones(batch))[..., :DNA_SIZE].softmax(-1)  # The data token is never the mask.
-    drawn = torch.distributions.Categorical(probs=p1).sample()  # lambda_t reaching one on the last step.
-    return torch.where(x == MASK, drawn, x)  # Any position still masked is drawn from the denoiser.
-
-
-# %% 06. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--gamma", type=float, default=2.0, help="guidance strength")
-    parser.set_defaults(steps=1000, width=64, samples=64)
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    # Section 5.27: a masked base at t = 0.6 with the chapter's denoiser and property scores.
-    p1 = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.double)  # The denoiser over A, C, G, T.
-    rates = p1 / (1.0 - 0.6)  # The unguided unmasking rates, 2.5 times the denoiser.
-    scores = torch.tensor([0.10, 0.35, 0.55, 0.25], dtype=torch.double)  # p(c | a, t).
-    guided = tilt(rates, scores, 0.20)  # The property model scores the masked state at 0.20.
-    report = {
-        "unguided_rates": [round(float(r), 4) for r in rates],  # (0.25, 0.50, 0.75, 1.00).
-        "unguided_exit_rate": round(float(rates.sum()), 4),  # 2.50.
-        "likelihood_ratios": [round(float(r), 4) for r in scores / 0.20],  # (0.50, 1.75, 2.75, 1.25).
-        "guided_rates": [round(float(r), 4) for r in guided],  # (0.125, 0.875, 2.063, 1.250).
-        "guided_exit_rate": round(float(guided.sum()), 4),  # 4.3125.
-        "guided_destination": [round(float(r), 3) for r in guided / guided.sum()],  # (0.029, ...).
-        "bayes_products": [round(float(r), 3) for r in p1 * scores],  # (0.010, 0.070, 0.165, 0.100).
-        "exact_h": round(h_value(p1, scores), 4),  # 0.345.
-        "exit_rate_at_exact_h": round(float(tilt(rates, scores, h_value(p1, scores)).sum()), 4),  # 2.5.
-        "reachable_states": reachable(100, 21),  # 2001 destinations for a 100-residue peptide.
-    }
-
-    # Section 5.28: the same base at guidance strength two, and the classifier-free form.
-    harder = tilt(rates, scores, 0.20, gamma=2.0)
-    report["gamma_2_rates"] = [round(float(r), 4) for r in harder]  # (0.0625, 1.531, 5.672, 1.563).
-    report["gamma_2_exit_rate"] = round(float(harder.sum()), 4)  # 8.828.
-    report["gamma_2_destination"] = [round(float(r), 3) for r in harder / harder.sum()]
-    mixed = classifier_free(tilt(rates, scores, 0.20), rates, gamma=2.0)
-    report["classifier_free_agrees"] = bool(torch.allclose(mixed, harder))  # The same family.
-
-    # Section 5.29: the Taylor tilt against the exact ratio.
-    gradient = torch.tensor([-2.10, -1.15, -0.75, -1.30, -1.61], dtype=torch.double)  # With the mask.
-    approximate = taylor_log_ratio(gradient, MASK)[:DNA_SIZE]  # The mask is the letter in place.
-    report["taylor_log_ratios"] = [round(float(r), 3) for r in approximate]  # (-0.49, 0.46, ...).
-    report["exact_log_ratios"] = [round(float(r), 3) for r in (scores / 0.20).log()]  # (-0.693, ...).
-    taylor_rates = rates * approximate.exp()
-    report["taylor_destination"] = [round(float(r), 3) for r in taylor_rates / taylor_rates.sum()]
-
-    # Training the denoiser and the property model on the shared DNA, then sampling twice.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    data, labels = splits["train_tokens"], splits["train_labels"]  # [N, L] and [N] the GC label.
-    model = SequenceDenoiser(V, args.length, width=args.width, depth=2)
-    prop = PropertyModel(args.length, width=args.width)
-
-    def loss_fn(step):
-        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
-        x1 = data[rows]  # [B, L] the sequences the denoiser has to recover.
-        y = labels[rows].to(torch.get_default_dtype())  # [B] the property to steer toward.
-        t = torch.rand(args.batch_size)  # One time per sequence.
-        xt = corrupt(x1, t)  # [B, L] a draw from the conditional path.
-        denoise = F.cross_entropy(model(xt, t)[..., :DNA_SIZE].transpose(1, 2), x1)
-        logp = prop(F.one_hot(xt, V).to(torch.get_default_dtype()), t)  # [B] the property on x_t.
-        property_loss = -(y * logp + (1 - y) * torch.log1p(-logp.exp().clamp(max=1 - 1e-6))).mean()
-        return denoise + property_loss  # The two networks share one optimizer here.
-
-    if not args.quiet:
-        print(f"denoiser={count_parameters(model):,} property={count_parameters(prop):,}")
-    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet,
-                   extra_parameters=list(prop.parameters()))
-    plain = generate(model, prop, args.samples, args.length, args.sample_steps, gamma=0.0)
-    steered = generate(model, prop, args.samples, args.length, args.sample_steps, gamma=args.gamma)
-
-    report["final_loss"] = sum(losses[-20:]) / len(losses[-20:])
-    report["gc_fraction_unguided"] = gc_fraction(plain)  # The training data sit at 0.577.
-    report[f"gc_fraction_gamma_{args.gamma:g}"] = gc_fraction(steered)  # Guidance raises it.
-    report["motif_fraction_unguided"] = motif_fraction(plain)
-    report["token_entropy_unguided"] = token_entropy(plain, DNA_SIZE)
-    print_report("guidance", report, decode_dna(steered))
-    save_run(args.out, model, vars(args), report, losses, decode_dna(steered))
-
-
-if __name__ == "__main__":
-    main()

lecture_5/gumbel.py

@@ -1,226 +0,0 @@
-"""Gumbel-Softmax Flow Matching: the path comes from the way a token is sampled.
-
-Sampling a categorical is a deterministic function of noise. We add a Gumbel draw to each log
-probability and take the largest. Replacing the argmax by a softmax turns the winner into a point of
-the simplex, and cooling the temperature slides that point from the center toward a corner. With the
-noise held fixed the trajectory is a smooth curve, its time derivative is the replicator equation of
-evolutionary dynamics, and setting the noise to zero at inference collapses the velocity to a
-multiple of the direction toward the data corner.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.22.1 to 5.22.6. Trained on the synthetic DNA of ``cis6270.data``. Prints the
-Gumbel-max frequencies, the annealed path at four times, the replicator velocity, the marginal
-velocity late in a trajectory, and the measurements of a generated batch.
-
-    python lecture_5/gumbel.py                     # train, then 64 Euler steps on the simplex
-    python lecture_5/gumbel.py --steps 20 --quiet  # a few seconds, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from cis6270.data import DNA_SIZE, decode_dna, dna_splits
-from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
-from cis6270.nets import SequenceDenoiser, count_parameters
-from cis6270.runner import common_parser, print_report, save_run, set_seed, train
-
-TAU_MAX = 10.0  # The temperature at t = 0, hot enough that every letter holds nearly the same mass.
-LAMBDA = 3.0  # The cooling rate, so tau(1) = 10 exp(-3) = 0.498.
-BETA = 2.0  # The noise scale; a larger beta makes the paths from one data token spread less.
-
-
-# %% 02. The Gumbel-max trick
-#
-# A Gumbel variable is the negative logarithm of an exponential clock, so the trick is Chapter 4's
-# race between clocks written with logarithms, where the winner is an argmax and the rates are
-# logits. All of the randomness sits in the noise, and the logits are ours to move.
-
-
-def gumbel(shape, generator=None):
-    """g ~ Gumbel(0, 1), as minus the log of minus the log of a uniform draw."""
-    u = torch.rand(shape, generator=generator).clamp(min=1e-12, max=1 - 1e-12)
-    return -torch.log(-torch.log(u))  # The distribution function is exp(-exp(-x)).
-
-
-def gumbel_max(probabilities, draws, generator=None):
-    """Frequencies of argmax_k [log pi_k + g_k] over ``draws`` samples. Returns [K]."""
-    logits = probabilities.log()[None, :] + gumbel((draws, len(probabilities)), generator)
-    counts = torch.bincount(logits.argmax(-1), minlength=len(probabilities)).double()
-    return counts / counts.sum()  # Approaches pi itself.
-
-
-# %% 03. The softmax relaxation and a cooling temperature
-#
-# A softmax of the noisy logits divided by a temperature is a point of the simplex. As the
-# temperature falls the point approaches the one-hot of the argmax, and as it grows the point
-# approaches the center. The temperature never reaches zero, so the path stays in the interior and
-# its density stays defined.
-
-
-def temperature(t, tau_max=TAU_MAX, rate=LAMBDA):
-    """tau(t) = tau_max exp(-lambda t), which is 10 at t = 0 and 0.498 at t = 1."""
-    return tau_max * math.exp(-rate * t)  # 2.231 at t = 0.5 and 0.672 at t = 0.9.
-
-
-# %% 04. The annealed path toward a data token
-#
-# Putting the data token into the logits and letting the cooling do the rest gives the conditional
-# path. The noise is drawn once and held for the whole trajectory, so the only quantity that changes
-# with time is the temperature and the point moves continuously toward the data corner.
-
-
-def annealed_path(logits, t):
-    """psi_t = softmax((e_k + g/beta)/tau(t)), with the logits already assembled. Returns [..., K]."""
-    return torch.softmax(logits / temperature(t), dim=-1)  # A point of the simplex at every time.
-
-
-def path_logits(targets, noise, K=DNA_SIZE, beta=BETA):
-    """e_k + g/beta at every position. targets: [B, L], noise: [B, L, K] -> [B, L, K]."""
-    return F.one_hot(targets, K) + noise / beta  # The data indicator plus the scaled Gumbel draw.
-
-
-# %% 05. Differentiating the path: the replicator equation
-#
-# The only time dependence sits in one over tau, which grows at relative rate lambda. Composing with
-# the softmax Jacobian gives the current mass times the gap between a letter's logit and the
-# mass-weighted average of all of them, which is the replicator equation with fitness
-# lambda ell_i/tau(t).
-
-
-def replicator_velocity(x, logits, t, rate=LAMBDA):
-    """u_{t,i} = (lambda/tau(t)) x_i (ell_i - sum_j x_j ell_j). x, logits: [..., K] -> [..., K]."""
-    average = (x * logits).sum(-1, keepdim=True)  # The mass-weighted average of the logits.
-    return (rate / temperature(t)) * x * (logits - average)  # Coordinates summing to zero.
-
-
-# %% 06. The inference-time velocity
-#
-# The sampler has only the current point, so the noise is set to zero. The logits become the data
-# indicator, the bracket becomes the gap to the data coordinate, and the velocity collapses to a
-# multiple of the direction toward the corner, with the factor x_k(1 - x_k) in the k-th component.
-
-
-def inference_velocity(x, k, t, rate=LAMBDA):
-    """u_t(x | x_1 = e_k) = (lambda/tau(t)) x_k (e_k - x). x: [..., K] -> [..., K]."""
-    corner = F.one_hot(torch.tensor(k), x.shape[-1]).to(x)  # The data corner.
-    return (rate / temperature(t)) * x[..., k, None] * (corner - x)  # Zero once x_k reaches zero.
-
-
-def marginal_velocity(x, p_hat, t, rate=LAMBDA):
-    """The posterior average over corners. x, p_hat: [..., K] -> [..., K]."""
-    weight = (rate / temperature(t)) * x * p_hat  # [..., K] the coefficient of each corner.
-    return weight - weight.sum(-1, keepdim=True) * x  # sum_k w_k (e_k - x), collected.
-
-
-# %% 07. Training and the marginal velocity
-#
-# Cross-entropy on the data token trains the denoiser, as it did for the mixture paths and for the
-# Dirichlet path, and the conditional velocity appears nowhere in the gradient. At sampling time the
-# marginal velocity is the posterior average, the fourth appearance of the same rule in this chapter.
-
-
-def relaxed_denoiser(length, width, K=DNA_SIZE, depth=2):
-    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token."""
-    model = SequenceDenoiser(K, length, width=width, depth=depth)
-    model.token_embedding = nn.Linear(K, width, bias=False)  # A point of the simplex, not an index.
-    return model
-
-
-@torch.no_grad()
-def generate(model, batch, length, steps, K=DNA_SIZE):
-    """Integrate the marginal velocity from the center of the simplex to t = 1. Returns [B, L]."""
-    x = annealed_path(gumbel((batch, length, K)) / BETA, 0.0)  # The source, with no data indicator.
-    h = 1.0 / steps
-    for k in range(steps):
-        t = k * h
-        p_hat = model(x, torch.full((batch,), t)).softmax(-1)  # [B, L, K] the denoiser.
-        x = x + h * marginal_velocity(x, p_hat, t)  # One Euler step along the marginal velocity.
-        x = x.clamp(min=1e-8)  # A large step can drive a coordinate below zero.
-        x = x / x.sum(-1, keepdim=True)  # Project back onto the simplex.
-    return x.argmax(-1)  # Read the token off the largest coordinate.
-
-
-# %% 08. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--draws", type=int, default=400_000, help="draws in the Gumbel-max check")
-    parser.set_defaults(steps=1000, width=64, samples=64)
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    pi = torch.tensor([0.5, 0.2, 0.2, 0.1])  # The categorical of the chapter's quick check.
-    report = {
-        "gumbel_max_frequencies": [round(float(v), 4) for v in gumbel_max(pi, args.draws)],
-        "target_probabilities": [round(float(v), 4) for v in pi],  # What the frequencies approach.
-        "tau_at_0": round(temperature(0.0), 4),  # 10.
-        "tau_at_1": round(temperature(1.0), 4),  # 0.498.
-    }
-
-    # Section 5.22.3: one base annealed toward G with the chapter's fixed noise.
-    scaled = torch.tensor([0.0, -0.45, 0.40, -0.05])  # g/beta over A, C, G, T.
-    logits = F.one_hot(torch.tensor(2), DNA_SIZE) + scaled  # The data indicator at G plus the noise.
-    report["path_logits"] = [round(float(v), 2) for v in logits]  # (0, -0.45, 1.40, -0.05).
-    for t in (0.0, 0.5, 0.9, 1.0):
-        point = annealed_path(logits, t)
-        report[f"path_at_{t}"] = [round(float(v), 4) for v in point]
-
-    # Section 5.22.4: the replicator velocity at the midpoint.
-    x = annealed_path(logits, 0.5)  # (0.2142, 0.1751, 0.4012, 0.2095).
-    report["tau_at_0.5"] = round(temperature(0.5), 4)  # 2.231.
-    report["rate_over_tau"] = round(LAMBDA / temperature(0.5), 3)  # 1.345.
-    report["average_logit"] = round(float((x * logits).sum()), 4)  # 0.4724.
-    noisy = replicator_velocity(x, logits, 0.5)
-    report["replicator_velocity"] = [round(float(v), 3) for v in noisy]  # (-0.136, -0.217, 0.5, -0.147).
-    report["replicator_sums_to"] = f"{float(noisy.sum()):.1e}"  # Zero, so the step stays on the simplex.
-
-    # Section 5.22.5: the same point with the noise set to zero.
-    clean = inference_velocity(x, 2, 0.5)
-    report["inference_velocity"] = [round(float(v), 3) for v in clean]  # (-0.116, -0.094, 0.323, ...).
-
-    # Section 5.22.6: one Euler step late in the trajectory.
-    late = torch.tensor([0.138, 0.084, 0.647, 0.131])  # The base at t = 0.8.
-    p_hat = torch.tensor([0.05, 0.10, 0.70, 0.15])  # The denoiser output there.
-    report["tau_at_0.8"] = round(temperature(0.8), 4)  # 0.9072.
-    coefficients = (LAMBDA / temperature(0.8)) * late * p_hat
-    report["corner_coefficients"] = [round(float(v), 4) for v in coefficients]  # (0.0228, ..., 1.4976).
-    v = marginal_velocity(late, p_hat, 0.8)
-    report["marginal_velocity"] = [round(float(c), 3) for c in v]  # (-0.200, -0.108, 0.454, -0.146).
-    report["euler_step_h_0.01"] = [round(float(c), 4) for c in late + 0.01 * v]  # (0.1360, ...).
-
-    # Training on the shared DNA, then generation.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    data = splits["train_tokens"]  # [N, L].
-    model = relaxed_denoiser(args.length, args.width)
-
-    def loss_fn(step):
-        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
-        x1 = data[rows]  # [B, L] the letters the denoiser has to recover.
-        t = torch.rand(args.batch_size)  # One time per sequence, uniform on [0, 1].
-        noise = gumbel((args.batch_size, args.length, DNA_SIZE))  # Held fixed along one trajectory.
-        cooled = path_logits(x1, noise) / torch.tensor(
-            [temperature(float(s)) for s in t])[:, None, None]  # Each sequence at its own temperature.
-        xt = cooled.softmax(-1)  # [B, L, K] a point on the conditional path.
-        return F.cross_entropy(model(xt, t).transpose(1, 2), x1)  # -log <x_theta, x_1>.
-
-    if not args.quiet:
-        print(f"parameters={count_parameters(model):,} tau_max={TAU_MAX} lambda={LAMBDA}")
-    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
-    tokens = generate(model, args.samples, args.length, args.sample_steps)
-
-    report["final_loss_nats"] = sum(losses[-20:]) / len(losses[-20:])
-    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
-    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
-    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
-    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
-    print_report("Gumbel-Softmax flow matching", report, decode_dna(tokens))
-    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))
-
-
-if __name__ == "__main__":
-    main()

lecture_5/kinetic_optimal.py

@@ -1,231 +0,0 @@
-"""Kinetic-optimal paths: choosing the velocity, and then the path, by minimizing an energy.
-
-Several velocities generate the same path and several paths connect the same endpoints, so a number
-is needed to choose among them. Written in the flux, both constraints are linear and the energy is
-convex, and the minimizer is the probability of the destination times the positive part of a
-difference of logarithmic growth rates. The mixture velocity of this chapter turns out to be that
-minimizer for the path it generates. Picking a distance on the tokens and cooling a softmax of
-negative distances then produces a path that moves probability only toward tokens closer to the data
-token.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.16.1 to 5.16.5 and 5.16.7, with the geodesic schedule of 5.16.6. The data marginal
-of the mask-path check is the letter frequency of the synthetic DNA of ``cis6270.data``. Prints the
-kinetic-optimal rates, the worked metric path for DNA, and the energy of two schedules.
-
-    python lecture_5/kinetic_optimal.py                     # 600 quadrature points per integral
-    python lecture_5/kinetic_optimal.py --steps 20 --quiet  # a coarse rule, visibly less accurate
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import DNA, DNA_SIZE, dna_splits
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-
-torch.set_default_dtype(torch.double)  # Everything here is arithmetic, so keep full precision.
-MASK = DNA_SIZE  # Index 4 is the mask symbol m, which the mask path visits.
-TRANSITION = torch.tensor([[0, 2, 1, 2], [2, 0, 2, 1], [1, 2, 0, 2], [2, 1, 2, 0]], dtype=torch.double)
-# d = 1 for a transition, A for G or C for T, and d = 2 for a transversion across the two classes.
-
-
-# %% 02. The flux and its energy
-#
-# The flux along an edge is the velocity times the probability of the state it leaves, and the master
-# equation is linear in it. The cost is a weight over the probability of the current state, times the
-# squared flux, which in the velocity reads as an expected squared rate under p_t.
-
-
-def flux(u, p):
-    """j_t(y, x) = u_t(y, x) p_t(x). u: [K, K], p: [K] -> [K, K]."""
-    return u * p[None, :]  # Each column scaled by the probability of its current state.
-
-
-def energy_density(u, p, weight=None):
-    """sum_{y != x} w_t(y, x) u_t(y, x)^2 p_t(x), the cost at one time written in the velocity."""
-    off = u - torch.diag(torch.diagonal(u))  # The diagonal holds no flux of its own.
-    weight = torch.ones_like(u) if weight is None else weight  # w = 1 gives the expected squared rate.
-    return float((weight * off ** 2 * p[None, :]).sum())
-
-
-# %% 03. The velocity of least energy
-#
-# Minimizing the energy edge by edge under the master equation, with the weight 1/p_t(y) that makes
-# the cost symmetric in the two ends of an edge, the multiplier comes out as the logarithmic growth
-# rate and the flux becomes p_t(y) p_t(x) times the positive part of its difference. Dividing by
-# p_t(x) returns the velocity.
-
-
-def kinetic_optimal(p, p_dot):
-    """u*_t(y, x) = p_t(y) [d/dt log p_t(y) - d/dt log p_t(x)]_+. p: [K], p_dot: [K] -> [K, K]."""
-    growth = torch.where(p > 0, p_dot / p.clamp(min=1e-300), torch.zeros_like(p))  # d/dt log p_t.
-    difference = growth[:, None] - growth[None, :]  # [K, K] destination minus current state.
-    u = p[:, None] * difference.clamp(min=0.0)  # Probability moves only toward the faster grower.
-    u = u * (p[None, :] > 0)  # A state with no probability sends none.
-    u.fill_diagonal_(0.0)
-    return u - torch.diag(u.sum(dim=0))  # The two rules fix the diagonal.
-
-
-# %% 04. Kinetic optimality of the mixture path
-#
-# On the conditional mixture path the two growth rates are kappa_dot/kappa at the data token and
-# minus kappa_dot/(1 - kappa) at the source token, and their difference collapses to the conditional
-# velocity built earlier in this chapter. The marginal mask path behaves the same way.
-
-
-def mask_path(q, kappa_t):
-    """p_t and p_dot_t over (A, C, G, T, m) for the mask source with data marginal q. Returns two [5]."""
-    p = torch.cat([kappa_t * q, torch.tensor([1.0 - kappa_t])])  # Tokens share kappa_t, m holds 1-kappa.
-    p_dot = torch.cat([q.clone(), torch.tensor([-1.0])])  # kappa_dot = 1 on the linear schedule.
-    return p, p_dot
-
-
-def mixture_rate(kappa_t, kappa_dot=1.0):
-    """kappa_dot_t/(1 - kappa_t), the conditional velocity out of the source token."""
-    return kappa_dot / (1.0 - kappa_t)  # 2 at t = 0.5 on the linear schedule.
-
-
-# %% 05. A distance on the tokens
-#
-# The formula for u* accepts any path, so we can design one. Cooling a softmax of negative distances
-# concentrates the conditional path on the data token, and the normalizing constant cancels between
-# the two logarithms, leaving the probability of the destination times beta_dot times the distance
-# the move gains.
-
-
-def metric_path(x1, beta_t, distance=TRANSITION):
-    """p_t(. | x_1) = softmax(-beta_t d(., x_1)). Returns [K]."""
-    return torch.softmax(-beta_t * distance[:, x1], dim=0)  # Uniform at beta = 0, a point mass later.
-
-
-def metric_velocity(x1, beta_t, beta_dot, distance=TRANSITION):
-    """u*_t(y, x | x_1) = p_t(y | x_1) beta_dot_t [d(x, x_1) - d(y, x_1)]_+. Returns [K, K]."""
-    p = metric_path(x1, beta_t, distance)  # [K] the path at this time.
-    d = distance[:, x1]  # [K] distance from each token to the data token.
-    gain = (d[None, :] - d[:, None]).clamp(min=0.0)  # [K, K] what the move from x to y gains.
-    u = p[:, None] * beta_dot * gain  # Faster when the step in distance is larger.
-    u.fill_diagonal_(0.0)
-    return u - torch.diag(u.sum(dim=0))  # Columns sum to zero.
-
-
-# %% 06. A worked metric path for DNA
-#
-# A transition swaps a purine for a purine or a pyrimidine for a pyrimidine, and a transversion
-# crosses between the two classes, so d = 1 and d = 2 record which mistakes are small. With
-# beta_t = t/(1 - t) and the data token G, the path at t = 0.5 favors the purine A over the two
-# pyrimidines, and the rates out of C run only toward A and G.
-
-
-def beta_schedule(t):
-    """beta_t = t/(1 - t) with beta_dot_t = 1/(1 - t)^2, which satisfies beta_0 = 0."""
-    return t / (1.0 - t), 1.0 / (1.0 - t) ** 2  # 1 and 4 at t = 0.5.
-
-
-def path_derivative(x1, t, h=1e-6, distance=TRANSITION):
-    """d p_t/dt by a central difference, for checking the velocity against the path it generates."""
-    forward, backward = metric_path(x1, beta_schedule(t + h)[0], distance), metric_path(
-        x1, beta_schedule(t - h)[0], distance)
-    return (forward - backward) / (2.0 * h)  # [K], compared against u* p_t.
-
-
-# %% 07. The energy of two schedules
-#
-# On a two-state conditional mixture path the energy is the time integral of kappa_dot squared over
-# kappa times one minus kappa. The linear schedule has a non-integrable singularity at each endpoint,
-# so its energy grows with the truncation and never settles. The geodesic schedule has the constant
-# integrand pi squared, which is the squared arc length of a quarter great circle.
-
-
-def schedule_energy(kappa, kappa_dot, steps, cut=1e-6):
-    """The integral of kappa_dot^2/(kappa (1 - kappa)) over [cut, 1 - cut], by a substituted midpoint.
-
-    The grid is uniform in s with t = 1/(1 + exp(-s)), whose Jacobian t(1 - t) cancels the endpoint
-    singularity exactly for the linear schedule and tames it for any other.
-    """
-    span = math.log((1.0 - cut) / cut)  # The s range that the truncation [cut, 1 - cut] maps to.
-    h, total = 2.0 * span / steps, 0.0
-    for k in range(steps):
-        s = -span + (k + 0.5) * h  # Midpoint of the k-th cell.
-        t = 1.0 / (1.0 + math.exp(-s))  # Back to the time variable.
-        k_t = min(max(kappa(t), 1e-12), 1.0 - 1e-12)  # Hold the schedule off both endpoints.
-        total += h * t * (1.0 - t) * kappa_dot(t) ** 2 / (k_t * (1.0 - k_t))
-    return total
-
-
-LINEAR = (lambda t: t, lambda t: 1.0)  # kappa_t = t.
-GEODESIC = (lambda t: math.sin(math.pi * t / 2) ** 2,  # The arc of Section 5.16.6 read as a schedule.
-            lambda t: (math.pi / 2) * math.sin(math.pi * t))
-
-
-def geodesic_schedule(t, source_probability=0.0):
-    """kappa_t(x_1) = 1 - sin^2((1 - t) Omega)/sin^2 Omega with Omega = arccos(sqrt(p(x_1)))."""
-    omega = math.acos(math.sqrt(source_probability))  # pi/2 under the mask source.
-    return 1.0 - math.sin((1.0 - t) * omega) ** 2 / math.sin(omega) ** 2
-
-
-# %% 08. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    # Section 5.16.3: the mask path at t = 0.5 with the chapter's data marginal.
-    q = torch.tensor([0.1, 0.6, 0.2, 0.1])  # [4] over A, C, G, T.
-    p, p_dot = mask_path(q, kappa_t=0.5)
-    u_star = kinetic_optimal(p, p_dot)
-    report = {
-        "mask_path_ACGTm": [round(float(v), 4) for v in p],  # (0.05, 0.30, 0.10, 0.05, 0.50).
-        "growth_at_mask": round(float(p_dot[MASK] / p[MASK]), 4),  # -2.
-        "growth_at_token": round(float(p_dot[1] / p[1]), 4),  # +2.
-        "rate_mask_to_C": round(float(u_star[1, MASK]), 4),  # 0.30 times 4, which is 1.2.
-        "mixture_rate_to_C": round(mixture_rate(0.5) * float(q[1]), 4),  # 2 times 0.6, the same.
-        "rate_between_tokens": round(float(u_star[1, 2]), 6),  # Zero, the no-revision mask path.
-        "generates_the_path": bool(torch.allclose(u_star @ p, p_dot, atol=1e-9)),
-    }
-
-    # The same check with the letter frequency of the synthetic DNA as the data marginal.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
-    frequency = counts / counts.sum()  # [4] the GC-rich background the lecture shares.
-    p_dna, p_dot_dna = mask_path(frequency, kappa_t=0.5)
-    u_dna = kinetic_optimal(p_dna, p_dot_dna)
-    report["dna_marginal"] = [round(float(v), 4) for v in frequency]
-    report["dna_rate_mask_to_G"] = round(float(u_dna[2, MASK]), 4)  # 2 times the frequency of G.
-    report["dna_energy_density"] = round(energy_density(u_dna, p_dna), 4)  # Expected squared rate.
-    report["dna_flux_mask_to_G"] = round(float(flux(u_dna, p_dna)[2, MASK]), 4)  # Rate times p_t(m).
-
-    # Section 5.16.5: the metric path toward G at t = 0.5.
-    g, c = DNA.index("G"), DNA.index("C")
-    beta_t, beta_dot = beta_schedule(0.5)  # 1 and 4.
-    p_metric = metric_path(g, beta_t)
-    u_metric = metric_velocity(g, beta_t, beta_dot)
-    report["metric_path_at_0.5"] = [round(float(v), 4) for v in p_metric]
-    report["rate_C_to_A"] = round(float(u_metric[0, c]), 3)  # 0.898.
-    report["rate_C_to_G"] = round(float(u_metric[g, c]), 3)  # 4.882.
-    report["rate_A_to_G"] = round(float(u_metric[g, 0]), 3)  # 2.441.
-    report["rate_out_of_C"] = round(float(-u_metric[c, c]), 3)  # 5.780.
-    step = torch.eye(DNA_SIZE)[c] + 0.05 * u_metric[:, c]  # One Euler step of size h = 0.05 from C.
-    report["euler_step_from_C"] = [round(float(v), 4) for v in step]
-    report["p_dot_from_velocity"] = [round(float(v), 4) for v in u_metric @ p_metric]
-    report["p_dot_from_path"] = [round(float(v), 4) for v in path_derivative(g, 0.5)]
-
-    # Section 5.16.7: the energy of the two schedules, and the geodesic schedule itself.
-    report["energy_linear_1e-6"] = round(schedule_energy(*LINEAR, args.steps, 1e-6), 2)  # 27.63.
-    report["energy_linear_1e-9"] = round(schedule_energy(*LINEAR, args.steps, 1e-9), 2)  # 41.45.
-    report["energy_geodesic"] = round(schedule_energy(*GEODESIC, args.steps, 1e-6), 4)  # 9.8696.
-    report["pi_squared"] = round(math.pi ** 2, 4)  # The squared arc length of a quarter circle.
-    report["geodesic_kappa_0.25"] = round(geodesic_schedule(0.25), 4)  # 0.1464.
-    report["geodesic_kappa_0.75"] = round(geodesic_schedule(0.75), 4)  # 0.8536.
-    report["geodesic_rate_0.9"] = round(math.pi * math.tan(math.pi * 0.9 / 2), 4)  # 19.84.
-    report["uniform_source_kappa_0.5"] = round(geodesic_schedule(0.5, 0.25), 4)  # 0.6667.
-    print_report("kinetic-optimal paths", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-if __name__ == "__main__":
-    main()

lecture_5/multi_objective.py

@@ -1,241 +0,0 @@
-"""Several objectives at once: MOG-DFM's cone and pCoMole's rollout.
-
-A peptide is scored on binding, solubility and half-life together, and no dataset labels one
-sequence for all of them, so the single value function of Discrete Guidance is out of reach.
-MOG-DFM drops the ratio form and builds a tilt out of a rank inside each objective and a direction
-against a weight vector, with a cone that removes the candidates pointing the wrong way. pCoMole
-keeps the ratio and replaces the property likelihood by an exponentiated worst-objective utility,
-whose value at a candidate is the average terminal reward over rollouts of the base flow.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.30 to 5.32. The single-objective tilt is in ``lecture_5/guidance.py``. The
-rollout runs on the synthetic DNA of ``cis6270.data``, scored by GC fraction against AT fraction.
-Prints the Pareto front, the chapter's MOG-DFM step with its cone, the chapter's pCoMole rollout,
-and a rollout on DNA whose value reverses the base rates.
-
-    python lecture_5/multi_objective.py                     # 600 rollouts per candidate on DNA
-    python lecture_5/multi_objective.py --steps 4000        # a tighter estimate of the second value
-    python lecture_5/multi_objective.py --steps 20 --quiet  # twenty rollouts, for a first look
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import DNA, DNA_SIZE, MOTIF, decode_dna, dna_splits
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-MASK = DNA_SIZE  # Index 4 is the mask symbol m, which a candidate still holds off its own position.
-OMEGA = torch.tensor([0.5, 0.5], dtype=torch.double)  # The trade-off the DNA rollout asks for.
-REFERENCE = torch.tensor([0.4, 0.3], dtype=torch.double)  # The levels r_k the objectives sit above.
-BETA = 2.0  # The inverse temperature of the terminal reward.
-
-
-# %% 02. Several objectives at once
-#
-# With several objectives there is no single best sequence. A candidate is on the Pareto front when
-# no other candidate matches it everywhere and beats it somewhere, and a weight vector picks one
-# point of that front by collapsing the scores to a single weighted number.
-
-
-def dominated(scores, row):
-    """True when some other candidate is at least as good everywhere and better somewhere."""
-    others = torch.cat([scores[:row], scores[row + 1:]])  # Every candidate except this one.
-    weak = (others >= scores[row]).all(dim=1)  # At least as good on every objective.
-    strict = (others > scores[row]).any(dim=1)  # Strictly better on at least one.
-    return bool((weak & strict).any())
-
-
-def pareto_front(scores):
-    """The indices of the candidates no other candidate dominates. scores: [M, N] -> list."""
-    return [row for row in range(len(scores)) if not dominated(scores, row)]
-
-
-# %% 03. MOG-DFM: rank, direction and a cone
-#
-# Without a joint likelihood there is no ratio to compute, so every candidate letter is scored twice.
-# Ranking inside each objective removes the units, the inner product with the weight vector measures
-# direction, and the two z-scored terms exponentiate into a tilt on the pretrained rate.
-
-
-def rank_scores(changes):
-    """I_n(y^i, x), the rank of each candidate's change within each objective, over the alphabet size."""
-    order = changes.argsort(dim=0).argsort(dim=0) + 1  # Rank one for the smallest change.
-    return order.to(changes.dtype) / changes.shape[0]  # [K, N] in {1/K, ..., 1}.
-
-
-def z_score(values):
-    """Norm[.], the z-score across the candidate letters at one position."""
-    spread = values.std(unbiased=False).clamp(min=1e-12)  # The population deviation, held off zero.
-    return (values - values.mean()) / spread  # Zero everywhere when no candidate stands out.
-
-
-def combined_score(changes, importance, weights, lam=1.0):
-    """Delta S = Norm[mean_n i_n I_n] + lambda Norm[Delta s . omega]. changes: [K, N] -> [K]."""
-    ranks = (rank_scores(changes) * importance).mean(dim=1)  # Importance-weighted average of ranks.
-    direction = changes @ weights  # [K] the directional score of each candidate.
-    return z_score(ranks) + lam * z_score(direction), ranks, direction
-
-
-# %% 04. Accepting a move only inside a cone
-#
-# A z-scored score puts about half the candidates above zero, so the tilt alone promotes the least
-# bad letter at a position where every letter moves away from the requested trade-off. The angle
-# between a candidate's score change and omega removes those, and a controller on the logarithm of
-# the angle holds the surviving fraction near one minus the target rejection rate.
-
-
-def cone_angles(changes, weights):
-    """alpha^i = arccos(Delta s . omega / (||Delta s|| ||omega||)), in degrees."""
-    cosine = (changes @ weights) / (changes.norm(dim=1) * weights.norm())  # [K].
-    return torch.rad2deg(torch.arccos(cosine.clamp(-1.0, 1.0))), cosine
-
-
-def update_cone(phi, rejected, target=0.2, gain=0.5, low=30.0, high=90.0):
-    """Phi_{t+h} = clip(Phi_t exp(eta (r - tau)), Phi_min, Phi_max), proportional control on log Phi."""
-    return min(max(phi * math.exp(gain * (rejected - target)), low), high)  # Wider when it rejects more.
-
-
-# %% 05. pCoMole: a reward in place of a likelihood
-#
-# Scoring a sequence by its weakest objective stops the sampler from trading one property away, and a
-# small multiple of the sum breaks ties. The exponentiated utility times a feasibility indicator is
-# the terminal reward, and the value of a candidate is the average reward over rollouts of the base
-# flow started from it.
-
-
-def utility(scores, weights, reference, rho=0.1):
-    """U(x) = min_k omega_k (f_k(x) - r_k) + rho sum_k omega_k (f_k(x) - r_k)."""
-    gaps = weights * (scores - reference)  # [N] the weighted gap on each objective.
-    return float(gaps.min() + rho * gaps.sum())  # The minimum, with the tie-breaking correction.
-
-
-def terminal_reward(scores, weights, reference, beta, feasible, rho=0.1):
-    """G(x) = exp(beta U(x)) 1[x in F], which is zero for an infeasible sequence."""
-    return math.exp(beta * utility(scores, weights, reference, rho)) * float(feasible)
-
-
-def rollout_value(rollouts, feasible, weights, reference, beta, rho=0.1):
-    """hat h_t(y_c), the average terminal reward over R continuations of the base flow."""
-    rewards = [terminal_reward(s, weights, reference, beta, f, rho) for s, f in zip(rollouts, feasible)]
-    return sum(rewards) / len(rewards), rewards  # A candidate with no feasible ending scores zero.
-
-
-# %% 06. A rollout on the synthetic DNA
-#
-# The base flow of a rollout is the mask-source sampler of this chapter with the letter frequency of
-# the data as its denoiser, so a continuation costs no network call. Two objectives that pull
-# against each other are the GC fraction and the AT fraction of the finished sequence, and the
-# feasible set is the sequences that carry the planted motif.
-
-
-def continuations(partial, frequency, rollouts, generator=None):
-    """Fill every masked position of a candidate from the base flow. partial: [L] -> [R, L]."""
-    grid = partial[None, :].repeat(rollouts, 1)  # [R, L] one copy of the candidate per rollout.
-    masked = grid == MASK  # [R, L] the positions the flow still owes a letter.
-    draws = torch.multinomial(frequency, int(masked.sum()), replacement=True, generator=generator)
-    return grid.masked_scatter(masked, draws)  # [R, L] the terminal sequences of the rollouts.
-
-
-def dna_objectives(sequences):
-    """The GC fraction and the AT fraction of each sequence, which sum to one. [R, L] -> [R, 2]."""
-    gc = ((sequences == 1) | (sequences == 2)).to(torch.get_default_dtype()).mean(dim=1)  # C and G.
-    return torch.stack([gc, 1.0 - gc], dim=1)  # [R, 2] raising one lowers the other.
-
-
-def carries_motif(sequences):
-    """One per sequence that holds the planted motif, which is the hard constraint here. [R] ints."""
-    return [int(MOTIF in sequence) for sequence in decode_dna(sequences)]
-
-
-def dna_rollout_value(partial, frequency, rollouts, weights, reference, beta, generator=None):
-    """hat h_t at one candidate, from ``rollouts`` continuations of the base flow. Returns a float."""
-    finished = continuations(partial, frequency, rollouts, generator)  # [R, L].
-    scores = dna_objectives(finished)  # [R, 2] the two objectives at the end of each rollout.
-    value, _ = rollout_value(list(scores), carries_motif(finished), weights, reference, beta)
-    return value  # Zero when no continuation of this candidate is feasible.
-
-
-# %% 07. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    # Section 5.30: four candidates on two objectives, and the two weight vectors.
-    candidates = torch.tensor([[0.8, 0.3], [0.6, 0.7], [0.5, 0.2], [0.6, 0.3]], dtype=torch.double)
-    front = pareto_front(candidates)
-    report = {
-        "pareto_front": front,  # Candidates 1 and 2; 3 and 4 are dominated.
-        "weighted_at_0.6_0.4": [round(float(s), 3) for s in candidates @ torch.tensor(
-            [0.6, 0.4], dtype=torch.double)],  # 0.60 for the first and 0.64 for the second.
-        "weighted_at_0.8_0.2": [round(float(s), 3) for s in candidates @ torch.tensor(
-            [0.8, 0.2], dtype=torch.double)],  # 0.70 and 0.62, which reverses the order.
-    }
-
-    # Section 5.31.2: one MOG-DFM step at a position holding A, with three candidate letters.
-    changes = torch.tensor([[0.0, 0.0], [0.20, -0.05], [0.05, 0.15], [-0.10, 0.08]], dtype=torch.double)
-    importance = torch.tensor([1.2, 0.8], dtype=torch.double)
-    weights = torch.tensor([0.6, 0.4], dtype=torch.double)
-    delta_s, ranks, direction = combined_score(changes, importance, weights)
-    report["rank_average"] = [round(float(r), 3) for r in ranks]  # (0.500, 0.700, 0.850, 0.450).
-    report["directional_scores"] = [round(float(d), 3) for d in direction]  # (0, 0.100, 0.090, -0.028).
-    report["combined_score"] = [round(float(s), 3) for s in delta_s]  # (-1.511, 1.541, 2.297, -2.327).
-    posterior = torch.tensor([0.3, 0.3, 0.3, 0.1], dtype=torch.double)  # The denoiser at t = 0.7.
-    base = posterior[1:] / (1.0 - 0.7)  # The pretrained rates out of A toward C, G and T.
-    report["base_rates"] = [round(float(r), 4) for r in base]  # (1.000, 1.000, 0.333).
-    tilted = base * delta_s[1:].exp()
-    report["guided_rates"] = [round(float(r), 3) for r in tilted]  # (4.667, 9.948, 0.033).
-    angles, cosines = cone_angles(changes[1:], weights)
-    report["cone_cosines"] = [round(float(c), 4) for c in cosines]  # (0.6727, 0.7894, -0.3032).
-    report["cone_angles_deg"] = [round(float(a), 1) for a in angles]  # (47.7, 37.9, 107.7).
-    inside = angles <= 60.0  # The cone test at Phi = 60 degrees.
-    survivors = tilted[inside]
-    report["after_cone"] = [round(float(r), 3) for r in survivors / survivors.sum()]  # (0.319, 0.681).
-    rejected = float((~inside).to(torch.get_default_dtype()).mean())  # One candidate of three.
-    report["next_cone_angle"] = round(update_cone(60.0, rejected), 1)  # 64.1.
-
-    # Section 5.32.2: three candidate edits, each rolled out four times.
-    omega, reference = torch.tensor([0.5, 0.5], dtype=torch.double), torch.zeros(2, dtype=torch.double)
-    edits = [
-        ([(0.8, 0.4), (0.6, 0.6), (0.9, 0.1), (0.5, 0.5)], [1, 1, 0, 1], 0.6),
-        ([(1.0, 0.0), (0.9, 0.1), (0.95, 0.05), (1.0, 0.1)], [1, 1, 1, 1], 1.2),
-        ([(0.4, 0.4)] * 4, [0, 0, 0, 0], 2.0),
-    ]
-    values = [rollout_value([torch.tensor(s, dtype=torch.double) for s in rollouts],
-                            feasible, omega, reference, 2.0)[0] for rollouts, feasible, _ in edits]
-    base_rates = torch.tensor([rate for _, _, rate in edits], dtype=torch.double)
-    products = torch.tensor(values) * base_rates  # The base rate times the continuation value.
-    report["rollout_values"] = [round(v, 3) for v in values]  # (1.390, 1.181, 0).
-    report["rollout_products"] = [round(float(p), 3) for p in products]  # (0.834, 1.417, 0).
-    report["rollout_selection"] = [round(float(p), 3) for p in products / products.sum()]
-    report["rollout_base_selection"] = [round(float(p), 3) for p in base_rates / base_rates.sum()]
-
-    # The same rollout on the shared DNA, from a partly revealed sequence with ACG already in place.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
-    frequency = counts / counts.sum()  # [4] the denoiser of the base flow, GC-rich like the data.
-    # The base rates below are kappa_dot/(1 - kappa) times that denoiser, read at t = 0.3.
-    revealed = torch.full((args.length,), MASK, dtype=torch.long)
-    revealed[2:5] = torch.tensor([0, 1, 2])  # Positions 3, 4 and 5 hold A, C and G.
-    finish = {letter: revealed.clone() for letter in "TC"}
-    for letter, candidate in finish.items():
-        candidate[5] = DNA.index(letter)  # Position 6 takes T, completing ACGT, or C, which does not.
-    seen = {letter: dna_rollout_value(candidate, frequency, args.steps, OMEGA, REFERENCE, BETA)
-            for letter, candidate in finish.items()}
-    rates = torch.tensor([float(frequency[DNA.index(letter)]) for letter in finish]) / (1.0 - 0.3)
-    guided = rates * torch.tensor([seen[letter] for letter in finish])  # The base rate times hat h.
-    report["dna_rollout_value_T"] = round(seen["T"], 4)  # Every continuation carries the motif.
-    report["dna_rollout_value_C"] = round(seen["C"], 4)  # Only the ones that plant it elsewhere.
-    report["dna_base_rates_T_C"] = [round(float(r), 4) for r in rates]  # The GC-rich data prefers C.
-    report["dna_base_selection"] = [round(float(r), 4) for r in rates / rates.sum()]
-    report["dna_guided_selection"] = [round(float(r), 4) for r in guided / guided.sum()]
-    print_report("several objectives at once", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-if __name__ == "__main__":
-    main()

lecture_5/rectify.py

@@ -1,293 +0,0 @@
-"""Rectification and AReUReDi: changing the coupling so that fewer steps suffice.
-
-A one-step sampler draws every position from its own marginal, so whatever structure the data has
-between positions is lost in that draw. The error it makes is the conditional total correlation of
-the coupling, which is a property of how the pairs were drawn and not of the data. Training on the
-current pairs, sampling from the model, and keeping the pairs it produces lowers that total
-correlation, and the discrete counterpart of the transport-cost bound says it never raises it.
-AReUReDi then composes a rectified flow with a Metropolis refinement whose proposal starts from the
-model's own per-position distribution.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.36 to 5.41. The second recoupling runs on adjacent base pairs of the synthetic
-DNA of ``cis6270.data``. Prints the total correlation of each round, the Hamming distance and the
-one-step success rate, and the AReUReDi proposal with its acceptance ratio.
-
-    python lecture_5/rectify.py                     # four rounds, 600 quadrature points each
-    python lecture_5/rectify.py --rounds 6          # closer to the limiting coupling
-    python lecture_5/rectify.py --steps 20 --quiet  # a coarse integration, visibly less accurate
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-
-from cis6270.data import DNA, DNA_SIZE, dna_splits
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-
-torch.set_default_dtype(torch.double)  # The recoupling is an exact integration, so keep precision.
-LENGTH = 2  # Two positions, so the state space has 16 sequences and can be enumerated.
-STATES = DNA_SIZE ** LENGTH
-LETTER = torch.stack([  # [L, 16] the letter each sequence index holds at each position.
-    torch.arange(STATES) // DNA_SIZE, torch.arange(STATES) % DNA_SIZE,
-])
-
-
-# %% 02. The pairing of source and data
-#
-# Since the first Euler step the joint distribution of the next state has been replaced by the
-# product of the per-position marginals. With small steps that costs almost nothing. One large step
-# draws every position at once from a product, and the divergence between the joint and that product
-# is the conditional total correlation of the coupling.
-
-
-def position_marginals(row):
-    """The per-position marginals of one distribution over sequences. row: [16] -> [L, K]."""
-    out = torch.zeros(LENGTH, DNA_SIZE)
-    for i in range(LENGTH):
-        out[i].index_add_(0, LETTER[i], row)  # Collect the mass of every sequence by its letter.
-    return out
-
-
-def total_correlation(coupling, source):
-    """TC_pi(X_1 | X_0), the expected divergence from the joint to the product of its marginals."""
-    total = 0.0
-    for s in range(STATES):
-        marginal = position_marginals(coupling[s])  # [L, K] given this source sequence.
-        product = marginal[0][LETTER[0]] * marginal[1][LETTER[1]]  # [16] the factorized draw.
-        live = coupling[s] > 0
-        total += float(source[s] * (coupling[s][live] * (coupling[s][live] / product[live]).log()).sum())
-    return total  # Zero exactly when the positions are independent given the source.
-
-
-def one_step_success(coupling, source, support):
-    """The chance a single factorized draw lands inside the data, averaged over source sequences."""
-    total = 0.0
-    for s in range(STATES):
-        marginal = position_marginals(coupling[s])
-        total += float(source[s] * sum(marginal[0][a] * marginal[1][b] for a, b in support))
-    return total  # 1/4 under the independent coupling on the complementary pairs.
-
-
-def hamming(coupling, source):
-    """The expected number of positions at which a source and its paired data sequence differ."""
-    distance = (LETTER[:, :, None] != LETTER[:, None, :]).sum(dim=0).to(coupling)  # [16, 16].
-    return float((source[:, None] * coupling * distance).sum())  # 1.5 under independent pairing.
-
-
-# %% 03. Total correlation on two bases
-#
-# Let the data be the four complementary pairs, each with probability one quarter, and the source be
-# uniform over all sixteen ordered pairs. Each position of the data alone is uniform over the four
-# bases, so the product of the two marginals spreads one sixteenth everywhere while the joint sits on
-# four states with one quarter apiece.
-
-
-def complementary_pairs():
-    """The data distribution over AT, TA, CG and GC, and the indices of those four sequences."""
-    support = [(DNA.index(a), DNA.index(b)) for a, b in ("AT", "TA", "CG", "GC")]
-    q = torch.zeros(STATES)
-    for a, b in support:
-        q[a * DNA_SIZE + b] = 0.25  # Each complementary pair with probability one quarter.
-    return q, support
-
-
-def adjacent_pairs(tokens):
-    """The empirical distribution of adjacent base pairs in the training DNA. Returns [16]."""
-    left, right = tokens[:, :-1].reshape(-1), tokens[:, 1:].reshape(-1)  # Every adjacent pair.
-    counts = torch.bincount(left * DNA_SIZE + right, minlength=STATES).to(torch.get_default_dtype())
-    return counts / counts.sum()  # [16] over the sixteen ordered pairs.
-
-
-# %% 04. ReDi: one recoupling round
-#
-# The model here is the entire many-step sampler, and the next coupling pairs a source sample with
-# the sequence the sampler produces from it. The velocity is the true marginal velocity of the
-# current coupling, and the master equation is integrated without discretization error, so the only
-# approximation left is the factorization across the two positions.
-
-
-def generator(joint, t):
-    """The factorized velocity at time t, in the variable u with t = 1 - exp(-u). Returns [16, 16]."""
-    eye = torch.eye(DNA_SIZE)
-    mix = (1 - t) * eye[:, :, None] + t * eye[:, None, :]  # [K, K, K] the mixture path at one position.
-    per = [mix[:, LETTER[i][:, None], LETTER[i][None, :]] for i in range(LENGTH)]  # [K, 16, 16] each.
-    weight = joint[None, None, :, :] * per[0][:, None, :, :] * per[1][None, :, :, :]
-    weight = weight.reshape(STATES, STATES, STATES)  # [x, x_0, x_1] the unnormalized posterior.
-    norm = weight.sum(dim=(1, 2)).clamp(min=1e-300)
-    posterior = weight.sum(dim=1) / norm[:, None]  # [x, x_1] the posterior over the data sequence.
-    rates, columns = torch.zeros(STATES, STATES), torch.arange(STATES)
-    for i in range(LENGTH):
-        denoiser = torch.zeros(STATES, DNA_SIZE)  # [x, a] the denoiser at position i.
-        denoiser.index_add_(1, LETTER[i], posterior)
-        for a in range(DNA_SIZE):
-            moved = LETTER.clone()
-            moved[i] = a  # [L, 16] every sequence with position i replaced by the letter a.
-            destination = moved[0] * DNA_SIZE + moved[1]  # [16] the index of that sequence.
-            keep = LETTER[i] != a  # A sequence already holding a at position i sends nothing.
-            rates[destination[keep], columns[keep]] += denoiser[keep, a]
-    return rates - torch.diag(rates.sum(dim=0))  # In u-time the schedule factor is one.
-
-
-def recouple(coupling, source, steps, horizon=14.0):
-    """Run the exact sampler from every source sequence and return the coupling it produces.
-
-    The substitution t = 1 - exp(-u) holds the schedule factor at one, so the integration stays
-    accurate as t approaches the horizon, where 1 - t is about 1e-6. Each cell takes a Heun step,
-    which is second order and reaches the chapter's figures at a few hundred cells.
-    """
-    joint = source[:, None] * coupling  # pi(x_0, x_1) over pairs.
-    probability = torch.eye(STATES)  # Row s is the state distribution started at the source s.
-    du = horizon / steps
-    for k in range(steps):
-        left = generator(joint, 1.0 - math.exp(-k * du))  # The velocity at the start of the cell.
-        right = generator(joint, 1.0 - math.exp(-(k + 1) * du))  # And at its end.
-        first = probability @ left.T  # The Euler slope.
-        second = (probability + du * first) @ right.T  # The slope the Euler step lands on.
-        probability = probability + 0.5 * du * (first + second)  # Their average.
-        probability = probability.clamp(min=0.0)
-        probability = probability / probability.sum(dim=1, keepdim=True)
-    return probability  # [16, 16] the next coupling, one row per source sequence.
-
-
-# %% 05. AReUReDi: the Tchebycheff score and the tilt
-#
-# Scoring a sequence by its weakest weighted objective stops the sampler from trading one property
-# away. The target is the base model times the exponentiated score, and the refinement anneals the
-# strength linearly over its iterations, so early ones explore and late ones commit.
-
-
-def tchebycheff(scores, weights):
-    """S_omega(x) = min_n omega_n s_n(x), the smallest weighted objective. scores: [M, N] -> [M]."""
-    return (scores * weights).min(dim=-1).values  # Raising a weight lowers the pressure on it.
-
-
-def anneal(iteration, total, low=1.0, high=20.0):
-    """eta_t = eta_min + (eta_max - eta_min) t/(T - 1), which is 1 + t at the chapter's settings."""
-    return low + (high - low) * iteration / (total - 1)  # 12 at the twelfth iteration of twenty.
-
-
-# %% 06. The proposal and its acceptance step
-#
-# The proposal starts from the model's own per-position distribution and reweights it by a balancing
-# function of how much the tilt would change. The identity g(u) = u g(1/u) collapses the acceptance
-# ratio to a ratio of two one-position normalizers, which costs one sum over the alphabet.
-
-
-def barker(u):
-    """g(u) = u/(1 + u), which satisfies g(u) = u g(1/u)."""
-    return u / (1.0 + u)  # The square root is the other common choice.
-
-
-def proposal(model_distribution, tilts, current):
-    """q_i(. | x) and its normalizer Z(x). tilts: [K] the value of W at each letter."""
-    ratios = tilts / tilts[current]  # r_i(y; x), the tilt at the proposed letter over the current one.
-    weighted = model_distribution * barker(ratios)  # The model's prior times the balancing function.
-    return weighted / weighted.sum(), float(weighted.sum()), ratios  # q, Z(x) and the ratios.
-
-
-def acceptance(z_current, z_proposed):
-    """alpha_i(x, x') = min(1, Z(x)/Z(x')), what the balancing identity leaves of the ratio."""
-    return min(1.0, z_current / z_proposed)  # One in both directions when the normalizers agree.
-
-
-# %% 07. Concentration on the maximizers
-#
-# Keeping only the maximizers in the denominator and bounding the numerator by the score gap gives an
-# exponential bound on the probability of anything outside the maximizing set, so the samples
-# concentrate there as the strength grows.
-
-
-def tilted_distribution(base, scores, eta):
-    """pi_{eta, omega}(x) proportional to p_1(x) exp(eta S_omega(x))."""
-    weighted = base * (eta * scores).exp()
-    return weighted / weighted.sum()  # [M] over the sequences under consideration.
-
-
-def concentration_bound(base, scores, eta):
-    """exp(-eta Delta) p_1(x) over the base mass of the maximizers, for a sequence outside the set."""
-    best = scores.max()
-    gap = float(best - scores[scores < best].max())  # Delta_omega, the gap to the next best score.
-    return math.exp(-eta * gap) / float(base[scores == best].sum()), gap
-
-
-# %% 08. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    parser.add_argument("--rounds", type=int, default=4, help="rectification rounds")
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    source = torch.full((STATES,), 1.0 / STATES)  # Uniform over all sixteen ordered pairs.
-    q, support = complementary_pairs()
-    coupling = q[None, :].repeat(STATES, 1)  # The independent coupling pairs every source with q.
-    report = {
-        "tc_round_0": round(total_correlation(coupling, source), 3),  # log 4 = 1.386.
-        "log_4": round(math.log(4), 3),  # What the enumeration of the four pairs gives.
-        "hamming_round_0": round(hamming(coupling, source), 3),  # 1.5.
-        "success_round_0": round(one_step_success(coupling, source, support), 4),  # 0.25.
-    }
-    for round_index in range(1, args.rounds + 1):
-        coupling = recouple(coupling, source, args.steps)  # One application of the operator.
-        report[f"tc_round_{round_index}"] = round(total_correlation(coupling, source), 3)
-        report[f"hamming_round_{round_index}"] = round(hamming(coupling, source), 3)
-        report[f"success_round_{round_index}"] = round(one_step_success(coupling, source, support), 4)
-        if round_index == 1:  # Where the first round's gain comes from, read at two sources.
-            columns = [a * DNA_SIZE + b for a, b in support]  # The four complementary pairs.
-            for name in ("AA", "AT"):  # One source outside the data and one inside it.
-                row = DNA.index(name[0]) * DNA_SIZE + DNA.index(name[1])
-                report[f"pairing_from_{name}"] = [round(float(coupling[row, c]), 3) for c in columns]
-    report["tc_limit"] = round(0.75 * math.log(2), 3)  # Twelve sources split in half, 0.520 nats.
-    report["success_limit"] = round(10.0 / 16.0, 4)  # 0.625, four exact sources and twelve halves.
-    report["hamming_limit"] = round(12.0 / 16.0, 4)  # 0.75, the least any coupling can reach.
-
-    # The same measurement on adjacent base pairs of the synthetic DNA, under independent pairing.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    dna = adjacent_pairs(splits["train_tokens"])  # [16] the empirical distribution of a base pair.
-    dna_coupling = dna[None, :].repeat(STATES, 1)  # Every source paired with the same distribution.
-    live = dna[dna > 0]
-    report["dna_pair_entropy"] = round(float(-(live * live.log()).sum()), 4)  # At most log 16 = 2.77.
-    report["dna_total_correlation"] = round(total_correlation(dna_coupling, source), 4)
-    report["dna_fraction_of_entropy"] = round(
-        total_correlation(dna_coupling, source) / float(-(live * live.log()).sum()), 4)
-
-    # Section 5.41.2: one refinement iteration at the twelfth of twenty, with two objectives.
-    scores = torch.tensor([[0.60, 0.80], [0.70, 0.76], [0.50, 0.90], [0.64, 0.62]])
-    weights = torch.tensor([0.5, 0.5])
-    eta = anneal(11, 20)  # The twelfth iteration, counting from zero, where 1 + t equals 12.
-    tchebycheff_scores = tchebycheff(scores, weights)
-    tilts = (eta * tchebycheff_scores).exp()  # W = exp(eta S) at that strength.
-    report["annealed_eta"] = round(eta, 4)  # 1 + t over twenty iterations.
-    report["tchebycheff"] = [round(float(s), 3) for s in tchebycheff_scores]  # (0.300, ...).
-    report["tilts"] = [round(float(w), 2) for w in tilts]  # (36.60, 66.69, 20.09, 41.26).
-    model_distribution = torch.tensor([0.40, 0.25, 0.20, 0.15])  # The rectified model at this position.
-    q_current, z_current, ratios = proposal(model_distribution, tilts, 0)  # The position holds A.
-    report["ratios_from_A"] = [round(float(r), 3) for r in ratios]  # (1.000, 1.822, 0.549, 1.127).
-    report["barker_from_A"] = [round(float(barker(r)), 3) for r in ratios]  # (0.500, 0.646, ...).
-    report["normalizer_at_A"] = round(z_current, 4)  # 0.5118.
-    report["proposal_at_A"] = [round(float(p), 3) for p in q_current]  # (0.391, 0.315, 0.138, 0.155).
-    _, z_proposed, _ = proposal(model_distribution, tilts, 1)  # The proposed sequence holds C.
-    report["normalizer_at_C"] = round(z_proposed, 4)  # 0.3704.
-    report["accept_A_to_C"] = round(acceptance(z_current, z_proposed), 3)  # One.
-    report["accept_C_to_A"] = round(acceptance(z_proposed, z_current), 3)  # 0.724.
-
-    # Section 5.41.3: the tilted distribution over four sequences, and the bound on a non-maximizer.
-    base = torch.tensor([0.50, 0.25, 0.15, 0.10])
-    objective = torch.tensor([0.20, 0.30, 0.45, 0.45])
-    for strength in (20.0, 60.0):
-        tilted = tilted_distribution(base, objective, strength)
-        bound, gap = concentration_bound(base, objective, strength)
-        report[f"mass_on_maximizers_eta_{strength:g}"] = f"{float(tilted[2:].sum()):.6f}"
-        report[f"second_sequence_eta_{strength:g}"] = f"{float(tilted[1]):.6f}"
-        report[f"bound_eta_{strength:g}"] = f"{bound * float(base[1]):.6f}"
-    report["score_gap"] = round(gap, 4)  # 0.15.
-    print_report("rectification", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-if __name__ == "__main__":
-    main()

lecture_5/velocity.py

@@ -1,212 +0,0 @@
-"""A jump rate written as a velocity, and the master equation it satisfies.
-
-Nothing here is learned. A velocity on a finite set is a rate matrix read from the destination's
-side, two rules fix which tables of numbers qualify, and the master equation says when such a table
-generates a probability path. The schedule factor of the mixture path turns out to be a hazard rate
-whose survival curve is the path itself, and the discrete velocity has the same shape as the
-continuous one of Chapter 2, with a posterior distribution in place of a posterior mean.
-
-Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
-u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).
-
-Chapter 5, Sections 5.2, 5.3, 5.6.3 and 5.8.1. The starting distribution of the long integration is
-the letter frequency of the synthetic DNA of ``cis6270.data``. Prints the chapter's worked velocity,
-its Euler step, the hazard at two times, and the two velocities side by side.
-
-    python lecture_5/velocity.py                     # 600 Euler steps in the long integration
-    python lecture_5/velocity.py --steps 20 --quiet  # a coarser grid, visibly less accurate
-"""
-
-# %% 01. Imports and configuration
-import math
-
-import torch
-import torch.nn.functional as F
-
-from cis6270.data import DNA, DNA_SIZE, dna_splits
-from cis6270.runner import common_parser, print_report, save_run, set_seed
-
-torch.set_default_dtype(torch.double)  # Every number here is arithmetic, so keep full precision.
-
-
-
-# %% 02. A jump rate written as a velocity
-#
-# Chapter 4 wrote R_t(x, y) for the rate of jumping from x to y. Discrete flow matching writes the
-# same number with the destination first, u_t(y, x) = R_t(x, y). Fixing x and letting y vary gives
-# one vector attached to the state x, with one component per destination.
-
-
-def velocity_from_rates(rates):
-    """Transpose a Chapter 4 rate matrix R_t(x, y) into a Chapter 5 velocity u_t(y, x)."""
-    return rates.transpose(0, 1)  # [K, K] with the destination now indexing the row.
-
-
-def jump_to_one_state(target, size=DNA_SIZE, rate=1.0):
-    """Every state other than ``target`` leaves toward it at ``rate``, and ``target`` never leaves."""
-    u = torch.zeros(size, size)  # [K, K], row = destination, column = current state.
-    for x in range(size):
-        if x == target:
-            continue
-        u[target, x] = rate  # The one positive off-diagonal entry in column x.
-        u[x, x] = -rate  # The diagonal is minus the total rate of leaving x.
-    return u
-
-
-# %% 03. The two rules a velocity obeys
-#
-# Off the diagonal the entries are nonnegative, since probability cannot flow out of a state at a
-# negative rate, and every column sums to zero, since probability leaving one state has to arrive at
-# another. The two rules fix the diagonal once the off-diagonal entries are chosen.
-
-
-def obeys_the_two_rules(u, tolerance=1e-9):
-    """True when every off-diagonal entry is nonnegative and every column sums to zero."""
-    off = u - torch.diag(torch.diagonal(u))  # The matrix with the diagonal removed.
-    return bool((off >= -tolerance).all() and (u.sum(dim=0).abs() < tolerance).all())
-
-
-def set_diagonal(u):
-    """u_t(x, x) = -sum_{y != x} u_t(y, x), the total rate of leaving x."""
-    u = u.clone()
-    u.fill_diagonal_(0.0)  # Keep the off-diagonal entries and discard whatever the diagonal held.
-    return u - torch.diag(u.sum(dim=0))  # Each column now sums to zero.
-
-
-# %% 04. The master equation in velocity form
-#
-# The rate of change of p_t(y) is the velocity into y times the probability of each current state,
-# summed over states. Written as a matrix equation it is p_dot = u p, with p a column vector and no
-# transpose, which is the column convention of this chapter.
-
-
-def master_equation(u, p):
-    """p_dot_t(y) = sum_x u_t(y, x) p_t(x). u: [K, K], p: [K] -> [K]."""
-    return u @ p  # The matrix multiplies the probability vector from the left.
-
-
-def arrivals_and_departures(u, p, y):
-    """The flux arriving at y from every other state, and the flux leaving y, as two numbers."""
-    others = [x for x in range(len(p)) if x != y]  # Every state except y itself.
-    arriving = sum(float(u[y, x] * p[x]) for x in others)  # u_t(y, x) p_t(x) along each edge in.
-    leaving = sum(float(u[x, y] * p[y]) for x in others)  # u_t(x, y) p_t(y) along each edge out.
-    return arriving, leaving  # Their difference is p_dot_t(y).
-
-
-def integrate(u, p, steps):
-    """Carry p forward by ``steps`` Euler steps over one unit of time under a constant velocity."""
-    h = 1.0 / steps  # The time increment of one step.
-    for _ in range(steps):
-        p = p + h * master_equation(u, p)  # p_{t+h} = p_t + h p_dot_t.
-    return p  # [K] after one unit of time.
-
-
-# %% 05. One velocity on DNA, with numbers
-#
-# Let the token be a single DNA base, let every base other than G jump to G at rate 1, and let G
-# never leave. Probability moves into G and the total stays at one, which the two rules guarantee.
-
-
-def worked_master_equation():
-    """Section 5.3.2 with p_t = (0.1, 0.2, 0.3, 0.4) and the all-to-G velocity."""
-    u = jump_to_one_state(DNA.index("G"))  # [4, 4] the matrix printed in the chapter.
-    p = torch.tensor([0.1, 0.2, 0.3, 0.4])  # [4] over A, C, G, T.
-    p_dot = master_equation(u, p)  # (-0.1, -0.2, 0.7, -0.4).
-    arriving, leaving = arrivals_and_departures(u, p, DNA.index("G"))  # 0.7 in and 0 out.
-    return u, p, p_dot, arriving, leaving
-
-
-# %% 06. Survival and hazard
-#
-# The schedule factor kappa_dot_t/(1 - kappa_t) is a hazard rate, and the probability that a position
-# still holds its source token is its survival curve. Separating variables gives S(t) = 1 - kappa_t,
-# the same weight the mixture path puts on the source token.
-
-
-def hazard(t):
-    """kappa_dot_t/(1 - kappa_t) on the linear schedule, which is 1/(1 - t)."""
-    return 1.0 / (1.0 - t)  # 1 at t = 0, 2 at t = 0.5, 10 at t = 0.9.
-
-
-def survival(steps, upper=0.9):
-    """Integrate dS/dt = -hazard(t) S(t) from S(0) = 1 and return S at ``upper``."""
-    h, s = upper / steps, 1.0  # One increment and the initial survival probability.
-    for k in range(steps):
-        s = s - h * hazard(k * h) * s  # Euler on the survival equation.
-    return s  # Exactly 1 - kappa_t = 1 - upper, since the two factors telescope at every grid.
-
-
-# %% 07. The continuous and discrete velocities side by side
-#
-# In continuous space the velocity is the schedule factor times the posterior mean minus the current
-# point. On tokens it is the same factor times the posterior distribution minus the point mass on the
-# current token. One step of size h displaces the point by h u_t(x); on tokens it redraws from the
-# denoiser with probability lambda_t, and the two agree in expectation on the one-hot embedding.
-
-
-def continuous_velocity(x, x1_hat, t):
-    """u_t(x) = kappa_dot/(1 - kappa) [x_hat_{1|t}(x) - x] at a point of R^d. x: [d] -> [d]."""
-    return hazard(t) * (x1_hat - x)  # The displacement that remains, divided by the time remaining.
-
-
-def discrete_velocity(x_i, p1, t, size=DNA_SIZE):
-    """u^i_t(., x) = kappa_dot/(1 - kappa) [p_{1|t}(. | x) - delta_{x^i}]. Returns [K]."""
-    return hazard(t) * (p1 - F.one_hot(torch.tensor(x_i), size).double())  # Signed, sums to zero.
-
-
-def side_by_side(x_i, p1, t, h):
-    """The continuous Euler step at the one-hot point and the discrete redraw distribution."""
-    x = F.one_hot(torch.tensor(x_i), len(p1)).double()  # The current token as a corner of R^K.
-    moved = x + h * continuous_velocity(x, p1, t)  # Displace by h times the velocity.
-    lam = h * hazard(t)  # lambda_t, the redraw probability of the two-part rule.
-    drawn = (1.0 - lam) * x + lam * p1  # Keep the current token, or redraw from the denoiser.
-    return moved, drawn, lam  # The two coincide, which is the last row of Table 5.1.
-
-
-# %% 08. The report
-def main():
-    parser = common_parser(__doc__.split("\n")[0])
-    args = parser.parse_args()
-    set_seed(args.seed)
-
-    u, p, p_dot, arriving, leaving = worked_master_equation()
-    report = {
-        "two_rules_hold": obeys_the_two_rules(u),  # Nonnegative off the diagonal, columns sum to zero.
-        "diagonal_recovered": bool(torch.allclose(u, set_diagonal(u))),  # The rules fix the diagonal.
-        "p_dot": [round(float(v), 4) for v in p_dot],  # (-0.1, -0.2, 0.7, -0.4).
-        "p_dot_sums_to": round(float(p_dot.sum()), 6),  # Zero, so probability is conserved.
-        "flux_into_G": round(arriving, 4),  # 0.7, the three incoming rates against their probabilities.
-        "flux_out_of_G": round(leaving, 4),  # 0, since G has no outgoing rate.
-        "euler_step_h_0.1": [round(float(v), 4) for v in p + 0.1 * p_dot],  # (0.09, 0.18, 0.37, 0.36).
-        "hazard_at_0.5": round(hazard(0.5), 4),  # 2.
-        "hazard_at_0.9": round(hazard(0.9), 4),  # 10.
-        "survival_at_0.9": f"{survival(args.steps):.6f}",  # 1 - 0.9 = 0.1 at any grid.
-        "survival_exact": 0.1,  # S(t) = 1 - kappa_t.
-    }
-
-    # Section 5.8.1: one position holding G with the denoiser of the chapter.
-    p1 = torch.tensor([0.1, 0.6, 0.2, 0.1])  # [4] over A, C, G, T.
-    moved, drawn, lam = side_by_side(DNA.index("G"), p1, t=0.5, h=0.1)
-    report["discrete_velocity_at_G"] = [round(float(v), 4) for v in discrete_velocity(2, p1, 0.5)]
-    report["continuous_step"] = [round(float(v), 4) for v in moved]  # (0.02, 0.12, 0.84, 0.02).
-    report["discrete_redraw"] = [round(float(v), 4) for v in drawn]  # The identical vector.
-    report["redraw_probability"] = round(lam, 4)  # lambda_t = 0.2.
-
-    # The same velocity run for a whole unit of time, starting from the data's letter frequency.
-    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
-    counts = torch.bincount(splits["train_tokens"].flatten(), minlength=DNA_SIZE).double()
-    start = counts / counts.sum()  # [4] the empirical frequency of A, C, G, T in the training DNA.
-    final = integrate(u, start, args.steps)  # Euler on the master equation, args.steps increments.
-    exact = start.clone()
-    exact[:] = start * math.exp(-1.0)  # Every base other than G decays as exp(-t).
-    exact[DNA.index("G")] = 1.0 - (start.sum() - start[DNA.index("G")]) * math.exp(-1.0)
-    report["dna_start"] = [round(float(v), 4) for v in start]  # The GC-rich background of the data.
-    report["dna_after_one_unit"] = [round(float(v), 4) for v in final]  # Euler at args.steps.
-    report["dna_exact"] = [round(float(v), 4) for v in exact]  # The closed-form solution.
-    report["euler_error"] = round(float((final - exact).abs().max()), 6)  # Shrinks with args.steps.
-    print_report("a jump rate written as a velocity", report)
-    save_run(args.out, config=vars(args), report=report)
-
-
-if __name__ == "__main__":
-    main()