Upload 87 files

0ba9d09 verified about 1 month ago

16.2 kB

"""Flow Map Language Models, the composition identity in denoiser coordinates.
We lift each token to a one-hot vector, the source is Gaussian over that space, and the state moves
through R^{L x K}, so no atom has to be split. If we step by the full remaining time 1 - s in place
of the elapsed time t - s, the average velocity becomes a two-time denoiser that lands on the
simplex. The map is then a convex combination of staying put and jumping to the prediction, and the
cocycle is a convex combination of two shorter predictions. A KL divergence measures the gap.
Chapter 6, Sections 6.25 to 6.27 of the notes. Trained on the synthetic DNA of
``cis6270.data.synthetic_dna``. Prints the convex weights, the composition weight gamma = 0.2 at
(0.2, 0.5, 0.8), the worked KL of 0.030747 nats, the two gradients that decide the loss, the
decoding clock of a 30,522-token vocabulary and the four steps it produces, and samples from the
trained map.
    python lecture_6/fmlm.py                 # 1500 steps, about half a minute
    python lecture_6/fmlm.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import TwoTimeMLP, count_parameters
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train


torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
# than it saves. ``tests/conftest.py`` sets the same thing for the test session, and this line
# acts when a reader runs the script directly.

BOOK_VOCAB = 30522  # The vocabulary the chapter puts on the clock, a standard tokenizer.
ENDPOINT_FRACTION = 0.25  # Share of each batch drawn with t = 1, the slice the sampler reads.


# %% 02. The convex form and the composition weight
#
# X_{s,t}(x) = (1-t)/(1-s) x + (t-s)/(1-s) delta_{s,t}(x). Both weights are non-negative for
# s <= t <= 1 and they sum to one, so each step is a convex combination of staying put and jumping
# to the predicted endpoint. At t = s the pair is (1, 0) and the identity axiom holds.
def convex_weights(s, t):
    """The pair ((1-t)/(1-s), (t-s)/(1-s)), the weights on the state and on the prediction."""
    remaining = (1.0 - s).clamp(min=1e-6)
    return (1.0 - t) / remaining, (t - s) / remaining


def composition_weight(s, u, t):
    """gamma = (1-t)(u-s)/((1-u)(t-s)), the share of the long prediction made over the first leg."""
    return (1.0 - t) * (u - s) / (((1.0 - u) * (t - s)).clamp(min=1e-6))


# %% 03. The cocycle in denoiser coordinates, checked on the scalar flow
#
# delta_{s,t}(x) = gamma delta_{s,u}(x) + (1-gamma) delta_{u,t}(X_{s,u}(x)). On dx/dt = x the
# two-time denoiser is x[1 + (1-s)(e^{t-s}-1)/(t-s)], so every term is available in closed form.
def scalar_denoiser(x, s, t):
    """delta_{s,t}(x) on the flow dx/dt = x."""
    return x * (1.0 + (1.0 - s) * math.expm1(t - s) / (t - s))


def scalar_cocycle(x: float = 1.0, s: float = 0.2, u: float = 0.5, t: float = 0.8):
    """Every term of the denoiser cocycle at the chapter's triple of times."""
    gamma = (1.0 - t) * (u - s) / ((1.0 - u) * (t - s))
    first = scalar_denoiser(x, s, u)  # The prediction over the first leg.
    reached = x * math.exp(u - s)  # X_{s,u}(x), the state the first leg arrives at.
    second = scalar_denoiser(reached, u, t)  # The prediction over the second leg.
    long = scalar_denoiser(x, s, t)  # The prediction over the whole interval.
    keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
    return {"gamma": gamma, "delta_s,u": first, "delta_u,t": second, "delta_s,t": long,
            "two legs combined": gamma * first + (1.0 - gamma) * second,
            "convex weights": (round(keep, 4), round(move, 4)),
            "convex form": keep * x + move * long, "X_s,t(x)": x * math.exp(t - s)}


# %% 04. One term of the loss, on the DNA alphabet
#
# The composed target is a convex combination of two simplex points, so it is a distribution and
# the divergence against it needs no projection. Its gradient pushes mass onto the token with more
# weight in the two-leg decomposition than in the long prediction.
def worked_kl():
    """The chapter's four-token example, with the target composed from two student predictions."""
    first = torch.tensor([0.10, 0.20, 0.60, 0.10])  # delta_{0.2,0.5} over (A, C, G, T).
    second = torch.tensor([0.05, 0.10, 0.80, 0.05])  # delta_{0.5,0.8} at the state the leg reaches.
    gamma = (1.0 - 0.8) * (0.5 - 0.2) / ((1.0 - 0.5) * (0.8 - 0.2))
    target = gamma * first + (1.0 - gamma) * second  # [4] the composed target.
    student = torch.tensor([0.10, 0.15, 0.65, 0.10])  # The student's long prediction.
    terms = target * (target / student).log()  # [4] the four terms of the divergence.
    return {"worked gamma": gamma, "composed target": [round(v, 4) for v in target.tolist()],
            "target mass": float(target.sum()), "worked KL": float(terms.sum()),
            "the four terms": [round(v, 6) for v in terms.tolist()]}


# %% 05. The two gradients of a softmax head
#
# With a softmax head the KL gradient in the logits is q - p, so the size of the correction in
# coordinate i is the error in coordinate i. The squared error passes through the softmax Jacobian
# q_j(delta_ij - q_i) and carries a factor q_i in every term, so a token the model has learned to
# rule out gets almost no correction. Here we take both gradients by automatic differentiation.
# That checks the two expressions of Section 6.27 as well as evaluating them.
def logit_gradients(target, student):
    """The gradient of KL(p || q) and of half the squared error in the logits of q. Two [K]."""
    for_kl = student.log().clone().requires_grad_(True)  # Logits whose softmax is the student.
    divergence = (target * (target / torch.softmax(for_kl, -1)).log()).sum()
    kl_gradient, = torch.autograd.grad(divergence, for_kl)  # [K] equal to q - p.
    for_squared = student.log().clone().requires_grad_(True)  # A second graph on the same logits.
    squared = 0.5 * ((torch.softmax(for_squared, -1) - target) ** 2).sum()
    squared_gradient, = torch.autograd.grad(squared, for_squared)  # [K] with a q_i throughout.
    return kl_gradient, squared_gradient


def gradient_comparison(floor: float = 1e-6):
    """The two gradients at the chapter's student, and again once that student has saturated."""
    target = torch.tensor([0.06, 0.12, 0.76, 0.06])  # p, the composed target of Section 6.27.
    student = torch.tensor([0.10, 0.15, 0.65, 0.10])  # q, the student's long prediction.
    kl, squared = logit_gradients(target, student)
    ratio = (squared / kl).abs()  # Between a quarter and a half, coordinate by coordinate.
    rest = student[1:] * (1.0 - floor) / student[1:].sum()  # The other three keep their shares.
    saturated = torch.cat([torch.full_like(student[:1], floor), rest])  # [K], summing to one.
    kl_low, squared_low = logit_gradients(target, saturated)
    return {
        "KL gradient": [round(v, 5) for v in kl.tolist()],
        "squared-error gradient": [round(v, 5) for v in squared.tolist()],
        "squared over KL": [round(v, 4) for v in ratio.tolist()],
        "saturated student": [round(v, 6) for v in saturated.tolist()],
        "saturated KL gradient": float(kl_low[0]),
        # The surviving squared-error gradient is of order 1e-08, below what four printed
        # decimals show, so the report carries it in units of 1e-08.
        "saturated squared-error gradient, in 1e-08": 1e8 * float(squared_low[0]),
        "orders of magnitude between them": math.log10(
            abs(float(kl_low[0])) / abs(float(squared_low[0]))),
    }


# %% 06. The decoding clock
#
# A position decodes correctly when its data coordinate is the largest of the |V| coordinates,
# a shifted standard normal against the maximum of the rest. The reparameterized time tau is the
# fraction of the initial decoding error already removed, and the sampler runs on that clock.
def decoding_error(t: float, vocab_size: int, nodes: int = 4001) -> float:
    """P_e(t), the probability that one position decodes to the wrong token at time t."""
    if t >= 1.0:
        return 0.0
    shift = t / (1.0 - t)  # a = t/(1-t), the gap the data coordinate opens.
    z = torch.linspace(-12.0, 14.0, nodes, dtype=torch.float64)  # Quadrature grid over the draw.
    density = torch.exp(-0.5 * z * z) / math.sqrt(2.0 * math.pi)  # phi(z).
    cumulative = 0.5 * torch.erfc(-(z + shift) / math.sqrt(2.0))  # Phi(z + a).
    log_all = (vocab_size - 1) * cumulative.clamp(min=1e-300).log()  # All the others fall short.
    return 1.0 - float(torch.trapz(density * log_all.exp(), z))


def clock(vocab_size: int):
    """tau(t) = (P_e(0) - P_e(t))/P_e(0), as a callable, with the source error rate beside it."""
    start = decoding_error(0.0, vocab_size)
    return (lambda t: (start - decoding_error(t, vocab_size)) / start), start


def even_in_tau(vocab_size: int, pieces: int = 4):
    """Solve tau(t) = k/pieces by bisection and return the grid with its two endpoints."""
    tau, _ = clock(vocab_size)
    grid = [0.0]
    for k in range(1, pieces):
        low, high = 0.0, 1.0 - 1e-6
        for _ in range(60):  # Bisection; tau is increasing in t.
            middle = 0.5 * (low + high)
            low, high = (middle, high) if tau(middle) < k / pieces else (low, middle)
        grid.append(0.5 * (low + high))
    return grid + [1.0]


# %% 07. The model, a two-time denoiser with a softmax head
def predict(model, x, s, t, length, vocab):
    """delta_{s,t}(x) at every position, a point of the simplex produced by a softmax head."""
    logits = model(x.flatten(1), s, t).view(-1, length, vocab)  # [B, L, K] the network output.
    return torch.softmax(logits, dim=-1)  # [B, L, K] on the simplex at every position.


def step_map(x, prediction, s, t):
    """The convex form, the flow map written in denoiser coordinates."""
    keep, move = convex_weights(s, t)
    return keep[:, None, None] * x + move[:, None, None] * prediction  # [B, L, K].


def make_loss(model, data, length, vocab, batch_size, device):
    eye = torch.eye(vocab, device=device)  # The one-hot lift of the alphabet.

    def loss_fn(step: int) -> torch.Tensor:
        index = torch.randint(0, len(data), (batch_size,), device=device)
        x1 = eye[data[index]]  # [B, L, K] the data, lifted.
        x0 = torch.randn_like(x1)  # [B, L, K] the Gaussian source.
        t_diag = torch.rand(batch_size, device=device)  # [B] a time for the diagonal term.
        xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1  # [B, L, K].
        diagonal = -(predict(model, xt, t_diag, t_diag, length, vocab).clamp_min(1e-9).log()
                     .gather(-1, data[index][..., None])).mean()  # Cross-entropy to the token.

        s = torch.rand(batch_size, device=device)  # [B] start time.
        t = s + (1.0 - s) * torch.rand(batch_size, device=device)  # [B] end time.
        # The slice t = 1 is the one-step case the sampler reads, and a continuous draw never
        # lands on it, so we put a fixed share of the batch there on purpose (Section 6.26).
        endpoint = torch.rand(batch_size, device=device) < ENDPOINT_FRACTION  # [B] which rows.
        t = torch.where(endpoint, torch.ones_like(t), t)
        u = s + (t - s) * torch.rand(batch_size, device=device)  # [B] the split point.
        xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1  # [B, L, K] the interpolant.
        with torch.no_grad():  # Both legs and the intermediate state are frozen targets.
            first = predict(model, xs, s, u, length, vocab)
            second = predict(model, step_map(xs, first, s, u), u, t, length, vocab)
            gamma = composition_weight(s, u, t)[:, None, None]
            target = gamma * first + (1.0 - gamma) * second  # [B, L, K], still on the simplex.
        student = predict(model, xs, s, t, length, vocab).clamp_min(1e-9)
        off_diagonal = (target * (target.clamp_min(1e-9) / student).log()).sum(-1).mean()
        return diagonal + off_diagonal

    return loss_fn


# %% 08. Sampling on a grid of times
@torch.no_grad()
def sample(model, count, length, vocab, grid, device):
    """Walk the grid with the convex form, and read the alphabet off the last prediction."""
    x = torch.randn(count, length, vocab, device=device)  # [N, L, K] the Gaussian source.
    for a, b in zip(grid[:-1], grid[1:]):
        s = torch.full((count,), a, device=device)
        t = torch.full((count,), b, device=device)
        x = step_map(x, predict(model, x, s, t, length, vocab), s, t)
    return x.argmax(-1)  # [N, L] token ids, read by an argmax at every position.


# %% 09. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed)
    data = tokens.to(device)  # [1024, L] the training sequences.

    report = scalar_cocycle()
    report.update(worked_kl())
    report.update(gradient_comparison())
    if not args.quiet:
        print(f"at the saturated student the KL gradient is "
              f"{report['saturated KL gradient']:.4f} and the squared-error gradient "
              f"{1e-8 * report['saturated squared-error gradient, in 1e-08']:.3e}")
    book_tau, book_start = clock(BOOK_VOCAB)
    for t in (0.60, 0.75, 0.80, 0.85):
        report[f"tau({t:.2f}) at |V|=30522"] = book_tau(t)
    book_grid = even_in_tau(BOOK_VOCAB)
    report["four steps on the clock"] = [round(t, 6) for t in book_grid]
    rounded = [round(t, 3) for t in book_grid]  # The chapter quotes the grid to three decimals.
    report["convex weights on that grid"] = [
        (round((1 - b) / (1 - a), 4), round((b - a) / (1 - a), 4))
        for a, b in zip(rounded[:-1], rounded[1:])]
    report["gamma on the last three"] = float(composition_weight(
        *(torch.tensor([v]) for v in rounded[1:4])))
    report["gamma over the long first leg"] = float(composition_weight(
        *(torch.tensor([v]) for v in (0.0, rounded[1], rounded[2]))))
    report["tau on a uniform grid"] = [round(book_tau(t), 5) for t in (0.25, 0.5, 0.75)]
    if not args.quiet:
        print(f"clock at |V|=30522: P_e(0) = {book_start:.6f}, four steps at "
              f"{[round(t, 3) for t in book_grid]}")

    model = TwoTimeMLP(dim=args.length * DNA_SIZE, width=args.width).to(device)
    if not args.quiet:
        print(f"\ntraining the two-time denoiser  parameters={count_parameters(model):,}")
    losses = train(model, make_loss(model, data, args.length, DNA_SIZE, args.batch_size, device),
                   steps=args.steps, lr=args.lr, quiet=args.quiet)

    dna_grid = even_in_tau(DNA_SIZE)  # The same clock, read at the four-letter alphabet.
    report["DNA clock grid"] = [round(t, 4) for t in dna_grid]
    samples: list[str] = []
    for name, grid in (("one step", [0.0, 1.0]), ("four, uniform", [0.0, 0.25, 0.5, 0.75, 1.0]),
                       ("four, on the clock", dna_grid)):
        torch.manual_seed(args.seed + 1)  # The same source noise for every schedule.
        drawn = sample(model, args.samples, args.length, DNA_SIZE, grid, device)
        report[f"{name}: GC fraction"] = gc_fraction(drawn)
        report[f"{name}: motif fraction"] = motif_fraction(drawn)
        report[f"{name}: token entropy"] = token_entropy(drawn, DNA_SIZE)
        report[f"{name}: unique fraction"] = unique_fraction(drawn)
        samples.append(f"{name:<20s} {decode_dna(drawn[:1])[0]}  GC {gc_fraction(drawn):.3f}")
    report["training GC fraction"] = gc_fraction(data)
    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])

    print_report("Flow Map Language Models, Chapter 6", report, samples)
    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)


if __name__ == "__main__":
    main()

	"""Flow Map Language Models, the composition identity in denoiser coordinates.

	We lift each token to a one-hot vector, the source is Gaussian over that space, and the state moves
	through R^{L x K}, so no atom has to be split. If we step by the full remaining time 1 - s in place
	of the elapsed time t - s, the average velocity becomes a two-time denoiser that lands on the
	simplex. The map is then a convex combination of staying put and jumping to the prediction, and the
	cocycle is a convex combination of two shorter predictions. A KL divergence measures the gap.

	Chapter 6, Sections 6.25 to 6.27 of the notes. Trained on the synthetic DNA of
	``cis6270.data.synthetic_dna``. Prints the convex weights, the composition weight gamma = 0.2 at
	(0.2, 0.5, 0.8), the worked KL of 0.030747 nats, the two gradients that decide the loss, the
	decoding clock of a 30,522-token vocabulary and the four steps it produces, and samples from the
	trained map.

	python lecture_6/fmlm.py # 1500 steps, about half a minute
	python lecture_6/fmlm.py --steps 20 # a few seconds, for a first look
	"""

	# %% 01. Imports and configuration
	import math

	import torch

	from cis6270.data import DNA_SIZE, decode_dna, synthetic_dna
	from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
	from cis6270.nets import TwoTimeMLP, count_parameters
	from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train


	torch.set_num_threads(1) # These networks are small, and splitting one matrix product costs more
	# than it saves. ``tests/conftest.py`` sets the same thing for the test session, and this line
	# acts when a reader runs the script directly.

	BOOK_VOCAB = 30522 # The vocabulary the chapter puts on the clock, a standard tokenizer.
	ENDPOINT_FRACTION = 0.25 # Share of each batch drawn with t = 1, the slice the sampler reads.


	# %% 02. The convex form and the composition weight
	#
	# X_{s,t}(x) = (1-t)/(1-s) x + (t-s)/(1-s) delta_{s,t}(x). Both weights are non-negative for
	# s <= t <= 1 and they sum to one, so each step is a convex combination of staying put and jumping
	# to the predicted endpoint. At t = s the pair is (1, 0) and the identity axiom holds.
	def convex_weights(s, t):
	"""The pair ((1-t)/(1-s), (t-s)/(1-s)), the weights on the state and on the prediction."""
	remaining = (1.0 - s).clamp(min=1e-6)
	return (1.0 - t) / remaining, (t - s) / remaining


	def composition_weight(s, u, t):
	"""gamma = (1-t)(u-s)/((1-u)(t-s)), the share of the long prediction made over the first leg."""
	return (1.0 - t) * (u - s) / (((1.0 - u) * (t - s)).clamp(min=1e-6))


	# %% 03. The cocycle in denoiser coordinates, checked on the scalar flow
	#
	# delta_{s,t}(x) = gamma delta_{s,u}(x) + (1-gamma) delta_{u,t}(X_{s,u}(x)). On dx/dt = x the
	# two-time denoiser is x[1 + (1-s)(e^{t-s}-1)/(t-s)], so every term is available in closed form.
	def scalar_denoiser(x, s, t):
	"""delta_{s,t}(x) on the flow dx/dt = x."""
	return x * (1.0 + (1.0 - s) * math.expm1(t - s) / (t - s))


	def scalar_cocycle(x: float = 1.0, s: float = 0.2, u: float = 0.5, t: float = 0.8):
	"""Every term of the denoiser cocycle at the chapter's triple of times."""
	gamma = (1.0 - t) * (u - s) / ((1.0 - u) * (t - s))
	first = scalar_denoiser(x, s, u) # The prediction over the first leg.
	reached = x * math.exp(u - s) # X_{s,u}(x), the state the first leg arrives at.
	second = scalar_denoiser(reached, u, t) # The prediction over the second leg.
	long = scalar_denoiser(x, s, t) # The prediction over the whole interval.
	keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
	return {"gamma": gamma, "delta_s,u": first, "delta_u,t": second, "delta_s,t": long,
	"two legs combined": gamma * first + (1.0 - gamma) * second,
	"convex weights": (round(keep, 4), round(move, 4)),
	"convex form": keep * x + move * long, "X_s,t(x)": x * math.exp(t - s)}


	# %% 04. One term of the loss, on the DNA alphabet
	#
	# The composed target is a convex combination of two simplex points, so it is a distribution and
	# the divergence against it needs no projection. Its gradient pushes mass onto the token with more
	# weight in the two-leg decomposition than in the long prediction.
	def worked_kl():
	"""The chapter's four-token example, with the target composed from two student predictions."""
	first = torch.tensor([0.10, 0.20, 0.60, 0.10]) # delta_{0.2,0.5} over (A, C, G, T).
	second = torch.tensor([0.05, 0.10, 0.80, 0.05]) # delta_{0.5,0.8} at the state the leg reaches.
	gamma = (1.0 - 0.8) * (0.5 - 0.2) / ((1.0 - 0.5) * (0.8 - 0.2))
	target = gamma * first + (1.0 - gamma) * second # [4] the composed target.
	student = torch.tensor([0.10, 0.15, 0.65, 0.10]) # The student's long prediction.
	terms = target * (target / student).log() # [4] the four terms of the divergence.
	return {"worked gamma": gamma, "composed target": [round(v, 4) for v in target.tolist()],
	"target mass": float(target.sum()), "worked KL": float(terms.sum()),
	"the four terms": [round(v, 6) for v in terms.tolist()]}


	# %% 05. The two gradients of a softmax head
	#
	# With a softmax head the KL gradient in the logits is q - p, so the size of the correction in
	# coordinate i is the error in coordinate i. The squared error passes through the softmax Jacobian
	# q_j(delta_ij - q_i) and carries a factor q_i in every term, so a token the model has learned to
	# rule out gets almost no correction. Here we take both gradients by automatic differentiation.
	# That checks the two expressions of Section 6.27 as well as evaluating them.
	def logit_gradients(target, student):
	"""The gradient of KL(p \|\| q) and of half the squared error in the logits of q. Two [K]."""
	for_kl = student.log().clone().requires_grad_(True) # Logits whose softmax is the student.
	divergence = (target * (target / torch.softmax(for_kl, -1)).log()).sum()
	kl_gradient, = torch.autograd.grad(divergence, for_kl) # [K] equal to q - p.
	for_squared = student.log().clone().requires_grad_(True) # A second graph on the same logits.
	squared = 0.5 * ((torch.softmax(for_squared, -1) - target) ** 2).sum()
	squared_gradient, = torch.autograd.grad(squared, for_squared) # [K] with a q_i throughout.
	return kl_gradient, squared_gradient


	def gradient_comparison(floor: float = 1e-6):
	"""The two gradients at the chapter's student, and again once that student has saturated."""
	target = torch.tensor([0.06, 0.12, 0.76, 0.06]) # p, the composed target of Section 6.27.
	student = torch.tensor([0.10, 0.15, 0.65, 0.10]) # q, the student's long prediction.
	kl, squared = logit_gradients(target, student)
	ratio = (squared / kl).abs() # Between a quarter and a half, coordinate by coordinate.
	rest = student[1:] * (1.0 - floor) / student[1:].sum() # The other three keep their shares.
	saturated = torch.cat([torch.full_like(student[:1], floor), rest]) # [K], summing to one.
	kl_low, squared_low = logit_gradients(target, saturated)
	return {
	"KL gradient": [round(v, 5) for v in kl.tolist()],
	"squared-error gradient": [round(v, 5) for v in squared.tolist()],
	"squared over KL": [round(v, 4) for v in ratio.tolist()],
	"saturated student": [round(v, 6) for v in saturated.tolist()],
	"saturated KL gradient": float(kl_low[0]),
	# The surviving squared-error gradient is of order 1e-08, below what four printed
	# decimals show, so the report carries it in units of 1e-08.
	"saturated squared-error gradient, in 1e-08": 1e8 * float(squared_low[0]),
	"orders of magnitude between them": math.log10(
	abs(float(kl_low[0])) / abs(float(squared_low[0]))),
	}


	# %% 06. The decoding clock
	#
	# A position decodes correctly when its data coordinate is the largest of the \|V\| coordinates,
	# a shifted standard normal against the maximum of the rest. The reparameterized time tau is the
	# fraction of the initial decoding error already removed, and the sampler runs on that clock.
	def decoding_error(t: float, vocab_size: int, nodes: int = 4001) -> float:
	"""P_e(t), the probability that one position decodes to the wrong token at time t."""
	if t >= 1.0:
	return 0.0
	shift = t / (1.0 - t) # a = t/(1-t), the gap the data coordinate opens.
	z = torch.linspace(-12.0, 14.0, nodes, dtype=torch.float64) # Quadrature grid over the draw.
	density = torch.exp(-0.5 * z * z) / math.sqrt(2.0 * math.pi) # phi(z).
	cumulative = 0.5 * torch.erfc(-(z + shift) / math.sqrt(2.0)) # Phi(z + a).
	log_all = (vocab_size - 1) * cumulative.clamp(min=1e-300).log() # All the others fall short.
	return 1.0 - float(torch.trapz(density * log_all.exp(), z))


	def clock(vocab_size: int):
	"""tau(t) = (P_e(0) - P_e(t))/P_e(0), as a callable, with the source error rate beside it."""
	start = decoding_error(0.0, vocab_size)
	return (lambda t: (start - decoding_error(t, vocab_size)) / start), start


	def even_in_tau(vocab_size: int, pieces: int = 4):
	"""Solve tau(t) = k/pieces by bisection and return the grid with its two endpoints."""
	tau, _ = clock(vocab_size)
	grid = [0.0]
	for k in range(1, pieces):
	low, high = 0.0, 1.0 - 1e-6
	for _ in range(60): # Bisection; tau is increasing in t.
	middle = 0.5 * (low + high)
	low, high = (middle, high) if tau(middle) < k / pieces else (low, middle)
	grid.append(0.5 * (low + high))
	return grid + [1.0]


	# %% 07. The model, a two-time denoiser with a softmax head
	def predict(model, x, s, t, length, vocab):
	"""delta_{s,t}(x) at every position, a point of the simplex produced by a softmax head."""
	logits = model(x.flatten(1), s, t).view(-1, length, vocab) # [B, L, K] the network output.
	return torch.softmax(logits, dim=-1) # [B, L, K] on the simplex at every position.


	def step_map(x, prediction, s, t):
	"""The convex form, the flow map written in denoiser coordinates."""
	keep, move = convex_weights(s, t)
	return keep[:, None, None] * x + move[:, None, None] * prediction # [B, L, K].


	def make_loss(model, data, length, vocab, batch_size, device):
	eye = torch.eye(vocab, device=device) # The one-hot lift of the alphabet.

	def loss_fn(step: int) -> torch.Tensor:
	index = torch.randint(0, len(data), (batch_size,), device=device)
	x1 = eye[data[index]] # [B, L, K] the data, lifted.
	x0 = torch.randn_like(x1) # [B, L, K] the Gaussian source.
	t_diag = torch.rand(batch_size, device=device) # [B] a time for the diagonal term.
	xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1 # [B, L, K].
	diagonal = -(predict(model, xt, t_diag, t_diag, length, vocab).clamp_min(1e-9).log()
	.gather(-1, data[index][..., None])).mean() # Cross-entropy to the token.

	s = torch.rand(batch_size, device=device) # [B] start time.
	t = s + (1.0 - s) * torch.rand(batch_size, device=device) # [B] end time.
	# The slice t = 1 is the one-step case the sampler reads, and a continuous draw never
	# lands on it, so we put a fixed share of the batch there on purpose (Section 6.26).
	endpoint = torch.rand(batch_size, device=device) < ENDPOINT_FRACTION # [B] which rows.
	t = torch.where(endpoint, torch.ones_like(t), t)
	u = s + (t - s) * torch.rand(batch_size, device=device) # [B] the split point.
	xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1 # [B, L, K] the interpolant.
	with torch.no_grad(): # Both legs and the intermediate state are frozen targets.
	first = predict(model, xs, s, u, length, vocab)
	second = predict(model, step_map(xs, first, s, u), u, t, length, vocab)
	gamma = composition_weight(s, u, t)[:, None, None]
	target = gamma * first + (1.0 - gamma) * second # [B, L, K], still on the simplex.
	student = predict(model, xs, s, t, length, vocab).clamp_min(1e-9)
	off_diagonal = (target * (target.clamp_min(1e-9) / student).log()).sum(-1).mean()
	return diagonal + off_diagonal

	return loss_fn


	# %% 08. Sampling on a grid of times
	@torch.no_grad()
	def sample(model, count, length, vocab, grid, device):
	"""Walk the grid with the convex form, and read the alphabet off the last prediction."""
	x = torch.randn(count, length, vocab, device=device) # [N, L, K] the Gaussian source.
	for a, b in zip(grid[:-1], grid[1:]):
	s = torch.full((count,), a, device=device)
	t = torch.full((count,), b, device=device)
	x = step_map(x, predict(model, x, s, t, length, vocab), s, t)
	return x.argmax(-1) # [N, L] token ids, read by an argmax at every position.


	# %% 09. The report
	def main() -> None:
	parser = common_parser(__doc__.split("\n")[0])
	parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16)
	args = parser.parse_args()
	set_seed(args.seed)
	device = pick_device(args.device)
	tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed)
	data = tokens.to(device) # [1024, L] the training sequences.

	report = scalar_cocycle()
	report.update(worked_kl())
	report.update(gradient_comparison())
	if not args.quiet:
	print(f"at the saturated student the KL gradient is "
	f"{report['saturated KL gradient']:.4f} and the squared-error gradient "
	f"{1e-8 * report['saturated squared-error gradient, in 1e-08']:.3e}")
	book_tau, book_start = clock(BOOK_VOCAB)
	for t in (0.60, 0.75, 0.80, 0.85):
	report[f"tau({t:.2f}) at \|V\|=30522"] = book_tau(t)
	book_grid = even_in_tau(BOOK_VOCAB)
	report["four steps on the clock"] = [round(t, 6) for t in book_grid]
	rounded = [round(t, 3) for t in book_grid] # The chapter quotes the grid to three decimals.
	report["convex weights on that grid"] = [
	(round((1 - b) / (1 - a), 4), round((b - a) / (1 - a), 4))
	for a, b in zip(rounded[:-1], rounded[1:])]
	report["gamma on the last three"] = float(composition_weight(
	*(torch.tensor([v]) for v in rounded[1:4])))
	report["gamma over the long first leg"] = float(composition_weight(
	*(torch.tensor([v]) for v in (0.0, rounded[1], rounded[2]))))
	report["tau on a uniform grid"] = [round(book_tau(t), 5) for t in (0.25, 0.5, 0.75)]
	if not args.quiet:
	print(f"clock at \|V\|=30522: P_e(0) = {book_start:.6f}, four steps at "
	f"{[round(t, 3) for t in book_grid]}")

	model = TwoTimeMLP(dim=args.length * DNA_SIZE, width=args.width).to(device)
	if not args.quiet:
	print(f"\ntraining the two-time denoiser parameters={count_parameters(model):,}")
	losses = train(model, make_loss(model, data, args.length, DNA_SIZE, args.batch_size, device),
	steps=args.steps, lr=args.lr, quiet=args.quiet)

	dna_grid = even_in_tau(DNA_SIZE) # The same clock, read at the four-letter alphabet.
	report["DNA clock grid"] = [round(t, 4) for t in dna_grid]
	samples: list[str] = []
	for name, grid in (("one step", [0.0, 1.0]), ("four, uniform", [0.0, 0.25, 0.5, 0.75, 1.0]),
	("four, on the clock", dna_grid)):
	torch.manual_seed(args.seed + 1) # The same source noise for every schedule.
	drawn = sample(model, args.samples, args.length, DNA_SIZE, grid, device)
	report[f"{name}: GC fraction"] = gc_fraction(drawn)
	report[f"{name}: motif fraction"] = motif_fraction(drawn)
	report[f"{name}: token entropy"] = token_entropy(drawn, DNA_SIZE)
	report[f"{name}: unique fraction"] = unique_fraction(drawn)
	samples.append(f"{name:<20s} {decode_dna(drawn[:1])[0]} GC {gc_fraction(drawn):.3f}")
	report["training GC fraction"] = gc_fraction(data)
	report["final loss"] = sum(losses[-20:]) / len(losses[-20:])

	print_report("Flow Map Language Models, Chapter 6", report, samples)
	save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)


	if __name__ == "__main__":
	main()