Upload 87 files

0ba9d09 verified about 1 month ago

15.1 kB

"""Simplex diffusion on DNA, a belief at every position in place of a token.
Every model before this one draws a token at each reverse step and discards the rest of the
denoiser's distribution, so a position the model was unsure about looks the same as one it was
certain about. Here the state at a position is a probability vector over the four bases, and the
corruption is a Dirichlet around the uniform-replacement kernel. One reverse step thins the
current belief, draws an innovation and mixes the two.
Chapter 4, "From tokens to beliefs" and "Thinning, an innovation, and a mixture". Time runs from
t = 0 clean to t = 1 noise, the survival probability is alpha_t = 1 - t, the concentration is
c_t = gamma K / (1 - alpha_t), and the rates R_t(x, y) act on probability rows, p_dot = p R.
Trained on the synthetic DNA of cis6270.data. Prints the worked Dirichlet parameters, the worked
thinning weight and its Monte Carlo check, then the sample statistics.
    python lecture_4/simplex.py                    # train and generate
    python lecture_4/simplex.py --mixture-draws 10000  # a coarser check of the mixture
    python lecture_4/simplex.py --steps 20         # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import torch  # Tensors, gradients and the Dirichlet, Beta and categorical draws.
import torch.nn as nn  # The denoiser, where only the input layer changes.
import torch.nn.functional as F  # Cross-entropy at every position.

from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import SequenceDenoiser, count_parameters  # The transformer the belief feeds.
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

GAMMA = 1.0  # The constant of the concentration schedule c_t = gamma K / (1 - alpha_t).
TIME_FLOOR = 1e-3  # The concentration grows like 1/(1 - alpha_t), so the grid stops above zero.


# %% 02. The forward marginal as a Dirichlet
#
# Fact 2 requires a forward marginal in closed form, and we want the belief to fluctuate around
# the marginal the token models already use,
# E[P_t | x_0] = alpha_t onehot(x_0) + (1 - alpha_t) pi. The concentration c_t sets the spread.
def concentration(alpha: float | torch.Tensor) -> float | torch.Tensor:
    """c_t = gamma K / (1 - alpha_t), the schedule that makes the increment nonnegative."""
    return GAMMA * DNA_SIZE / (1.0 - alpha)  # 8 at alpha_t = 0.5 and 20 at alpha_s = 0.8.


def belief_mean(alpha: float, x0: int) -> torch.Tensor:
    """E[P_t | x_0] over (A, C, G, T), the uniform-replacement kernel of the token models. [K]."""
    mean = torch.full((DNA_SIZE,), (1.0 - alpha) / DNA_SIZE)  # [K] the share from pi.
    mean[x0] = mean[x0] + alpha  # The clean token keeps the survival mass on top of its share.
    return mean  # [K] a distribution, since alpha + (1 - alpha) = 1.


def dirichlet_parameters(alpha: float, x0: int) -> torch.Tensor:
    """c_t E[P_t | x_0], the Dirichlet parameters at one position. Returns [K]."""
    return concentration(alpha) * belief_mean(alpha, x0)  # (5, 1, 1, 1) at alpha_t = 0.5.


def parameter_field(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
    """The same parameters for a whole batch. x0: [B, L], alpha: [B, 1] -> [B, L, K].
    Every coordinate other than the clean token has parameter gamma, constant in time, and the
    clean coordinate has gamma (1 + K alpha_t/(1 - alpha_t)), the factor of the uniform rates.
    """
    off = torch.full(x0.shape + (DNA_SIZE,), GAMMA, device=x0.device)  # [B, L, K] the pi share.
    clean = GAMMA * DNA_SIZE * alpha / (1.0 - alpha)  # [B, 1] the survival mass times c_t.
    return off.scatter_add(2, x0[:, :, None], clean[..., None].expand_as(x0[:, :, None]))


def corrupt(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
    """Draw P_t ~ Dir(c_t E[P_t | x_0]) at each position. x0: [B, L], alpha: [B, 1] -> [B, L, K]."""
    return torch.distributions.Dirichlet(parameter_field(x0, alpha)).sample()  # [B, L, K] beliefs.


# %% 03. Thinning, an innovation, and a mixture
#
# A Dirichlet with parameter a + b equals W V + (1 - W) U for V ~ Dir(a), U ~ Dir(b) and
# W ~ Beta(sum a, sum b). With this concentration schedule the increment b sits entirely on the
# clean coordinate, so the innovation U is the one-hot vector of the clean token.
def increment(alpha_t: float, alpha_s: float, x0: int) -> torch.Tensor:
    """b = c_s E[P_s | x_0] - c_t E[P_t | x_0], nonnegative for alpha_s > alpha_t. Returns [K]."""
    return dirichlet_parameters(alpha_s, x0) - dirichlet_parameters(alpha_t, x0)  # (12, 0, 0, 0).


def thinning_law(alpha_t: float, alpha_s: float) -> tuple[float, float]:
    """The two Beta parameters, sum a = c_t and sum b = c_s - c_t, or 8 and 12 here."""
    return float(concentration(alpha_t)), float(concentration(alpha_s) - concentration(alpha_t))


def reverse_step(belief: torch.Tensor, clean: torch.Tensor,
                 alpha_t: float, alpha_s: float) -> torch.Tensor:
    """P_s = W P_t + (1 - W) onehot(x_0), one Beta draw per position. belief: [B, L, K]."""
    first, second = thinning_law(alpha_t, alpha_s)  # The two parameters of the thinning weight.
    w = torch.distributions.Beta(  # [B, L] one weight per position, with mean c_t / c_s.
        torch.tensor(first, device=belief.device),
        torch.tensor(second, device=belief.device)).sample(belief.shape[:2])
    innovation = F.one_hot(clean, DNA_SIZE).to(belief.dtype)  # [B, L, K] the corner it moves to.
    return w[..., None] * belief + (1.0 - w[..., None]) * innovation  # [B, L, K] the mixture.


# %% 04. The denoiser that reads a belief
#
# Only the input layer changes. We look a token up in the embedding table, and a belief is the
# same K rows of that table combined in its own proportions, so a one-hot input recovers the
# denoiser of every other file in this lecture exactly.
class BeliefDenoiser(nn.Module):
    """d_theta(. | P_t, t), the shared sequence denoiser reading K numbers at every position."""

    def __init__(self, length: int, width: int = 128):
        super().__init__()
        self.inner = SequenceDenoiser(DNA_SIZE, length, width=width)  # Embeddings and transformer.

    def forward(self, beliefs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """beliefs: [B, L, K], t: [B] -> logits over the clean tokens [B, L, K]."""
        inner = self.inner
        positions = torch.arange(beliefs.shape[1], device=beliefs.device)  # [L].
        h = (beliefs @ inner.token_embedding.weight  # [B, L, D] the belief mixes the K embeddings.
             + inner.position_embedding(positions)[None, :, :]  # [1, L, D] where it sits.
             + inner.time_embedding(t)[:, None, :])  # [B, 1, D] how far along the path we are.
        return inner.head(inner.transformer(h))  # [B, L, K] every position attends to every other.


# %% 05. The loss
#
# Fact 5 is the cross-entropy whose optimum is the posterior over clean tokens. The reverse step
# moves the denoiser's answer through the mixture of Section 03, and its weight already comes
# from the concentration schedule, so the time derivative of alpha_t multiplies nothing here.
def loss(model: nn.Module, x0: torch.Tensor) -> torch.Tensor:
    """One Monte Carlo sample of E[-log d_theta(x_0 | P_t, t)], summed over positions."""
    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
    alpha = 1.0 - t  # [B, 1] the linear schedule, alpha_0 = 1 and alpha_1 = 0.
    beliefs = corrupt(x0, alpha)  # [B, L, K] one Dirichlet draw per position.
    logits = model(beliefs, t[:, 0])  # [B, L, K] the denoiser reads the whole field of beliefs.
    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
    return ce.sum(-1).mean()  # Nats per sequence, with every position counted.


# %% 06. The sampler
#
# Generation starts from Dir(gamma, ..., gamma), the law the corruption ends at when alpha_1 = 0.
# Each step draws one clean token from the denoiser, thins the current belief and mixes the two,
# so the belief moves through the interior and reaches a corner at the last step.
@torch.no_grad()  # Generation needs no gradient tracking.
def sample(model: nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
    """Generate from beliefs scattered around pi down to clean DNA. Returns [B, L] over ACGT."""
    device = next(model.parameters()).device
    start = torch.full((batch, length, DNA_SIZE), GAMMA, device=device)  # [B, L, K] at alpha_1 = 0.
    belief = torch.distributions.Dirichlet(start).sample()  # [B, L, K] scattered around pi.
    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid, t down to 0.
    for t, s in zip(ts[:-1], ts[1:]):  # One interval [s, t] at a time.
        logits = model(belief, t.expand(batch))  # [B, L, K] the clean-token prediction.
        clean = torch.distributions.Categorical(logits=logits).sample()  # [B, L] the innovation.
        belief = reverse_step(belief, clean, float(1.0 - t), float(1.0 - s))  # [B, L, K].
    return belief.argmax(-1)  # [B, L] the corner each position has reached.


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--mixture-draws", type=int, default=400_000,
                        help="draws in the check of the thinning representation")
    # The belief holds the denoiser's uncertainty between steps. Enough reverse steps have to run
    # for it to appear. A default run trains 1000 steps and measures 64 sequences, inside about a
    # minute.
    parser.set_defaults(steps=1000, sample_steps=64, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)  # The same seed gives the same training batches and the same samples.
    device = pick_device(args.device)
    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.

    # Section 02, the worked Dirichlet at the two noise levels of the chapter, with a clean A.
    at_t = dirichlet_parameters(0.5, 0)  # [K] (5, 1, 1, 1) at the later of the two times.
    at_s = dirichlet_parameters(0.8, 0)  # [K] (17, 1, 1, 1) at the earlier one.
    total_t, total_s = at_t.sum(), at_s.sum()  # c_t and c_s, the totals that set every variance.
    spread_t = (at_t * (total_t - at_t) / (total_t ** 2 * (total_t + 1))).sqrt()  # [K] exact.
    spread_s = (at_s * (total_s - at_s) / (total_s ** 2 * (total_s + 1))).sqrt()  # [K] exact.
    talk(f"alpha_t = 0.5       c_t {concentration(0.5):.0f}  mean "
         f"{[round(float(v), 3) for v in belief_mean(0.5, 0)]}  parameters "
         f"{[round(float(v), 0) for v in at_t]}  std on A {float(spread_t[0]):.3f}")
    talk(f"alpha_s = 0.8       c_s {concentration(0.8):.0f}  mean "
         f"{[round(float(v), 2) for v in belief_mean(0.8, 0)]}  parameters "
         f"{[round(float(v), 0) for v in at_s]}")

    # Section 03, the increment, the thinning weight, and the mean the mixture lands on.
    step = increment(0.5, 0.8, 0)  # [K] all of the increase sits on the clean coordinate.
    first, second = thinning_law(0.5, 0.8)  # Beta(8, 12), with mean 0.4.
    mean_weight = first / (first + second)  # c_t / c_s, the average share the old belief keeps.
    mixed = mean_weight * belief_mean(0.5, 0) + (1.0 - mean_weight) * torch.eye(DNA_SIZE)[0]  # [K].
    talk(f"increment           {[round(float(v), 0) for v in step]}  W ~ Beta({first:.0f}, "
         f"{second:.0f})  mean {mean_weight:.1f}  mixture mean "
         f"{[round(float(v), 2) for v in mixed]}")

    # The same mixture by simulation, against the exact standard deviations of Dir(17, 1, 1, 1).
    thinned = torch.distributions.Dirichlet(at_t).sample((args.mixture_draws,))  # [N, K] draw V.
    weight = torch.distributions.Beta(  # [N] the Beta draw of the representation.
        torch.tensor(first), torch.tensor(second)).sample((args.mixture_draws,))
    drawn = weight[:, None] * thinned + (1.0 - weight[:, None]) * torch.eye(DNA_SIZE)[0]  # [N, K].
    talk(f"mixture check       {args.mixture_draws} draws  std "
         f"{[round(float(v), 3) for v in drawn.std(0)]}  against the exact "
         f"{[round(float(v), 3) for v in spread_s]}")

    # Train the denoiser on the shared synthetic DNA. We reseed here so the training batches stay
    # independent of how many draws the check above consumed.
    set_seed(args.seed)
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the loss below.
    model = BeliefDenoiser(args.length, width=args.width).to(device)
    talk(f"device={device} parameters={count_parameters(model):,} "
         f"train={len(train_tokens)} length={args.length}")

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

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

    # Generate, then measure the structure the data was built with.
    torch.manual_seed(args.seed + 1)  # Fix the Beta and categorical draws, so the samples repeat.
    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
    report = {
        "c_t at alpha_t = 0.5": float(concentration(0.5)),
        "Dirichlet parameter on A at 0.5": float(at_t[0]),
        "standard deviation on A at 0.5": float(spread_t[0]),
        "Dirichlet parameter on A at 0.8": float(at_s[0]),
        "increment on A": float(step[0]),
        "mean thinning weight": mean_weight,
        "simulated standard deviation on A": float(drawn.std(0)[0]),
        "exact standard deviation on A": float(spread_s[0]),
        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "validation loss, nats/seq": held_out,
        "loss, nats per token": held_out / args.length,
        "GC fraction": gc_fraction(samples),
        "GC fraction in the data": gc_fraction(train_tokens),
        "motif fraction": motif_fraction(samples),
        "motif fraction in the data": motif_fraction(train_tokens),
        "distinct samples": unique_fraction(samples),
        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
    }
    print_report("simplex: beliefs in place of tokens", report, decode_dna(samples))
    save_run(args.out, model=model, config=vars(args), report=report,
             losses=losses, samples=decode_dna(samples))


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

	"""Simplex diffusion on DNA, a belief at every position in place of a token.

	Every model before this one draws a token at each reverse step and discards the rest of the
	denoiser's distribution, so a position the model was unsure about looks the same as one it was
	certain about. Here the state at a position is a probability vector over the four bases, and the
	corruption is a Dirichlet around the uniform-replacement kernel. One reverse step thins the
	current belief, draws an innovation and mixes the two.

	Chapter 4, "From tokens to beliefs" and "Thinning, an innovation, and a mixture". Time runs from
	t = 0 clean to t = 1 noise, the survival probability is alpha_t = 1 - t, the concentration is
	c_t = gamma K / (1 - alpha_t), and the rates R_t(x, y) act on probability rows, p_dot = p R.
	Trained on the synthetic DNA of cis6270.data. Prints the worked Dirichlet parameters, the worked
	thinning weight and its Monte Carlo check, then the sample statistics.

	python lecture_4/simplex.py # train and generate
	python lecture_4/simplex.py --mixture-draws 10000 # a coarser check of the mixture
	python lecture_4/simplex.py --steps 20 # a few seconds, for a first look
	"""

	# %% 01. Imports and configuration
	import torch # Tensors, gradients and the Dirichlet, Beta and categorical draws.
	import torch.nn as nn # The denoiser, where only the input layer changes.
	import torch.nn.functional as F # Cross-entropy at every position.

	from cis6270.data import DNA_SIZE, decode_dna, dna_splits # ACGT, K = 4 and the shared sequences.
	from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
	from cis6270.nets import SequenceDenoiser, count_parameters # The transformer the belief feeds.
	from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

	GAMMA = 1.0 # The constant of the concentration schedule c_t = gamma K / (1 - alpha_t).
	TIME_FLOOR = 1e-3 # The concentration grows like 1/(1 - alpha_t), so the grid stops above zero.


	# %% 02. The forward marginal as a Dirichlet
	#
	# Fact 2 requires a forward marginal in closed form, and we want the belief to fluctuate around
	# the marginal the token models already use,
	# E[P_t \| x_0] = alpha_t onehot(x_0) + (1 - alpha_t) pi. The concentration c_t sets the spread.
	def concentration(alpha: float \| torch.Tensor) -> float \| torch.Tensor:
	"""c_t = gamma K / (1 - alpha_t), the schedule that makes the increment nonnegative."""
	return GAMMA * DNA_SIZE / (1.0 - alpha) # 8 at alpha_t = 0.5 and 20 at alpha_s = 0.8.


	def belief_mean(alpha: float, x0: int) -> torch.Tensor:
	"""E[P_t \| x_0] over (A, C, G, T), the uniform-replacement kernel of the token models. [K]."""
	mean = torch.full((DNA_SIZE,), (1.0 - alpha) / DNA_SIZE) # [K] the share from pi.
	mean[x0] = mean[x0] + alpha # The clean token keeps the survival mass on top of its share.
	return mean # [K] a distribution, since alpha + (1 - alpha) = 1.


	def dirichlet_parameters(alpha: float, x0: int) -> torch.Tensor:
	"""c_t E[P_t \| x_0], the Dirichlet parameters at one position. Returns [K]."""
	return concentration(alpha) * belief_mean(alpha, x0) # (5, 1, 1, 1) at alpha_t = 0.5.


	def parameter_field(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
	"""The same parameters for a whole batch. x0: [B, L], alpha: [B, 1] -> [B, L, K].

	Every coordinate other than the clean token has parameter gamma, constant in time, and the
	clean coordinate has gamma (1 + K alpha_t/(1 - alpha_t)), the factor of the uniform rates.
	"""
	off = torch.full(x0.shape + (DNA_SIZE,), GAMMA, device=x0.device) # [B, L, K] the pi share.
	clean = GAMMA * DNA_SIZE * alpha / (1.0 - alpha) # [B, 1] the survival mass times c_t.
	return off.scatter_add(2, x0[:, :, None], clean[..., None].expand_as(x0[:, :, None]))


	def corrupt(x0: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
	"""Draw P_t ~ Dir(c_t E[P_t \| x_0]) at each position. x0: [B, L], alpha: [B, 1] -> [B, L, K]."""
	return torch.distributions.Dirichlet(parameter_field(x0, alpha)).sample() # [B, L, K] beliefs.


	# %% 03. Thinning, an innovation, and a mixture
	#
	# A Dirichlet with parameter a + b equals W V + (1 - W) U for V ~ Dir(a), U ~ Dir(b) and
	# W ~ Beta(sum a, sum b). With this concentration schedule the increment b sits entirely on the
	# clean coordinate, so the innovation U is the one-hot vector of the clean token.
	def increment(alpha_t: float, alpha_s: float, x0: int) -> torch.Tensor:
	"""b = c_s E[P_s \| x_0] - c_t E[P_t \| x_0], nonnegative for alpha_s > alpha_t. Returns [K]."""
	return dirichlet_parameters(alpha_s, x0) - dirichlet_parameters(alpha_t, x0) # (12, 0, 0, 0).


	def thinning_law(alpha_t: float, alpha_s: float) -> tuple[float, float]:
	"""The two Beta parameters, sum a = c_t and sum b = c_s - c_t, or 8 and 12 here."""
	return float(concentration(alpha_t)), float(concentration(alpha_s) - concentration(alpha_t))


	def reverse_step(belief: torch.Tensor, clean: torch.Tensor,
	alpha_t: float, alpha_s: float) -> torch.Tensor:
	"""P_s = W P_t + (1 - W) onehot(x_0), one Beta draw per position. belief: [B, L, K]."""
	first, second = thinning_law(alpha_t, alpha_s) # The two parameters of the thinning weight.
	w = torch.distributions.Beta( # [B, L] one weight per position, with mean c_t / c_s.
	torch.tensor(first, device=belief.device),
	torch.tensor(second, device=belief.device)).sample(belief.shape[:2])
	innovation = F.one_hot(clean, DNA_SIZE).to(belief.dtype) # [B, L, K] the corner it moves to.
	return w[..., None] * belief + (1.0 - w[..., None]) * innovation # [B, L, K] the mixture.


	# %% 04. The denoiser that reads a belief
	#
	# Only the input layer changes. We look a token up in the embedding table, and a belief is the
	# same K rows of that table combined in its own proportions, so a one-hot input recovers the
	# denoiser of every other file in this lecture exactly.
	class BeliefDenoiser(nn.Module):
	"""d_theta(. \| P_t, t), the shared sequence denoiser reading K numbers at every position."""

	def __init__(self, length: int, width: int = 128):
	super().__init__()
	self.inner = SequenceDenoiser(DNA_SIZE, length, width=width) # Embeddings and transformer.

	def forward(self, beliefs: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
	"""beliefs: [B, L, K], t: [B] -> logits over the clean tokens [B, L, K]."""
	inner = self.inner
	positions = torch.arange(beliefs.shape[1], device=beliefs.device) # [L].
	h = (beliefs @ inner.token_embedding.weight # [B, L, D] the belief mixes the K embeddings.
	+ inner.position_embedding(positions)[None, :, :] # [1, L, D] where it sits.
	+ inner.time_embedding(t)[:, None, :]) # [B, 1, D] how far along the path we are.
	return inner.head(inner.transformer(h)) # [B, L, K] every position attends to every other.


	# %% 05. The loss
	#
	# Fact 5 is the cross-entropy whose optimum is the posterior over clean tokens. The reverse step
	# moves the denoiser's answer through the mixture of Section 03, and its weight already comes
	# from the concentration schedule, so the time derivative of alpha_t multiplies nothing here.
	def loss(model: nn.Module, x0: torch.Tensor) -> torch.Tensor:
	"""One Monte Carlo sample of E[-log d_theta(x_0 \| P_t, t)], summed over positions."""
	t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR) # [B, 1] one time each.
	alpha = 1.0 - t # [B, 1] the linear schedule, alpha_0 = 1 and alpha_1 = 0.
	beliefs = corrupt(x0, alpha) # [B, L, K] one Dirichlet draw per position.
	logits = model(beliefs, t[:, 0]) # [B, L, K] the denoiser reads the whole field of beliefs.
	ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none") # [B, L].
	return ce.sum(-1).mean() # Nats per sequence, with every position counted.


	# %% 06. The sampler
	#
	# Generation starts from Dir(gamma, ..., gamma), the law the corruption ends at when alpha_1 = 0.
	# Each step draws one clean token from the denoiser, thins the current belief and mixes the two,
	# so the belief moves through the interior and reaches a corner at the last step.
	@torch.no_grad() # Generation needs no gradient tracking.
	def sample(model: nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
	"""Generate from beliefs scattered around pi down to clean DNA. Returns [B, L] over ACGT."""
	device = next(model.parameters()).device
	start = torch.full((batch, length, DNA_SIZE), GAMMA, device=device) # [B, L, K] at alpha_1 = 0.
	belief = torch.distributions.Dirichlet(start).sample() # [B, L, K] scattered around pi.
	ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device) # The reverse grid, t down to 0.
	for t, s in zip(ts[:-1], ts[1:]): # One interval [s, t] at a time.
	logits = model(belief, t.expand(batch)) # [B, L, K] the clean-token prediction.
	clean = torch.distributions.Categorical(logits=logits).sample() # [B, L] the innovation.
	belief = reverse_step(belief, clean, float(1.0 - t), float(1.0 - s)) # [B, L, K].
	return belief.argmax(-1) # [B, L] the corner each position has reached.


	# %% 07. The report
	def main() -> None:
	parser = common_parser(__doc__.split("\n")[0])
	parser.add_argument("--mixture-draws", type=int, default=400_000,
	help="draws in the check of the thinning representation")
	# The belief holds the denoiser's uncertainty between steps. Enough reverse steps have to run
	# for it to appear. A default run trains 1000 steps and measures 64 sequences, inside about a
	# minute.
	parser.set_defaults(steps=1000, sample_steps=64, samples=64)
	args = parser.parse_args()
	set_seed(args.seed) # The same seed gives the same training batches and the same samples.
	device = pick_device(args.device)
	talk = (lambda *a: None) if args.quiet else print # --quiet keeps only the report block.

	# Section 02, the worked Dirichlet at the two noise levels of the chapter, with a clean A.
	at_t = dirichlet_parameters(0.5, 0) # [K] (5, 1, 1, 1) at the later of the two times.
	at_s = dirichlet_parameters(0.8, 0) # [K] (17, 1, 1, 1) at the earlier one.
	total_t, total_s = at_t.sum(), at_s.sum() # c_t and c_s, the totals that set every variance.
	spread_t = (at_t * (total_t - at_t) / (total_t ** 2 * (total_t + 1))).sqrt() # [K] exact.
	spread_s = (at_s * (total_s - at_s) / (total_s ** 2 * (total_s + 1))).sqrt() # [K] exact.
	talk(f"alpha_t = 0.5 c_t {concentration(0.5):.0f} mean "
	f"{[round(float(v), 3) for v in belief_mean(0.5, 0)]} parameters "
	f"{[round(float(v), 0) for v in at_t]} std on A {float(spread_t[0]):.3f}")
	talk(f"alpha_s = 0.8 c_s {concentration(0.8):.0f} mean "
	f"{[round(float(v), 2) for v in belief_mean(0.8, 0)]} parameters "
	f"{[round(float(v), 0) for v in at_s]}")

	# Section 03, the increment, the thinning weight, and the mean the mixture lands on.
	step = increment(0.5, 0.8, 0) # [K] all of the increase sits on the clean coordinate.
	first, second = thinning_law(0.5, 0.8) # Beta(8, 12), with mean 0.4.
	mean_weight = first / (first + second) # c_t / c_s, the average share the old belief keeps.
	mixed = mean_weight * belief_mean(0.5, 0) + (1.0 - mean_weight) * torch.eye(DNA_SIZE)[0] # [K].
	talk(f"increment {[round(float(v), 0) for v in step]} W ~ Beta({first:.0f}, "
	f"{second:.0f}) mean {mean_weight:.1f} mixture mean "
	f"{[round(float(v), 2) for v in mixed]}")

	# The same mixture by simulation, against the exact standard deviations of Dir(17, 1, 1, 1).
	thinned = torch.distributions.Dirichlet(at_t).sample((args.mixture_draws,)) # [N, K] draw V.
	weight = torch.distributions.Beta( # [N] the Beta draw of the representation.
	torch.tensor(first), torch.tensor(second)).sample((args.mixture_draws,))
	drawn = weight[:, None] * thinned + (1.0 - weight[:, None]) * torch.eye(DNA_SIZE)[0] # [N, K].
	talk(f"mixture check {args.mixture_draws} draws std "
	f"{[round(float(v), 3) for v in drawn.std(0)]} against the exact "
	f"{[round(float(v), 3) for v in spread_s]}")

	# Train the denoiser on the shared synthetic DNA. We reseed here so the training batches stay
	# independent of how many draws the check above consumed.
	set_seed(args.seed)
	splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
	train_tokens = splits["train_tokens"].to(device) # [N, L] clean sequences over ACGT.
	valid_tokens = splits["valid_tokens"].to(device) # [M, L] held out for the loss below.
	model = BeliefDenoiser(args.length, width=args.width).to(device)
	talk(f"device={device} parameters={count_parameters(model):,} "
	f"train={len(train_tokens)} length={args.length}")

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

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

	# Generate, then measure the structure the data was built with.
	torch.manual_seed(args.seed + 1) # Fix the Beta and categorical draws, so the samples repeat.
	samples = sample(model, args.samples, args.length, args.sample_steps) # [B, L] over ACGT.
	report = {
	"c_t at alpha_t = 0.5": float(concentration(0.5)),
	"Dirichlet parameter on A at 0.5": float(at_t[0]),
	"standard deviation on A at 0.5": float(spread_t[0]),
	"Dirichlet parameter on A at 0.8": float(at_s[0]),
	"increment on A": float(step[0]),
	"mean thinning weight": mean_weight,
	"simulated standard deviation on A": float(drawn.std(0)[0]),
	"exact standard deviation on A": float(spread_s[0]),
	"training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
	"validation loss, nats/seq": held_out,
	"loss, nats per token": held_out / args.length,
	"GC fraction": gc_fraction(samples),
	"GC fraction in the data": gc_fraction(train_tokens),
	"motif fraction": motif_fraction(samples),
	"motif fraction in the data": motif_fraction(train_tokens),
	"distinct samples": unique_fraction(samples),
	"letter entropy, nats": token_entropy(samples, DNA_SIZE),
	}
	print_report("simplex: beliefs in place of tokens", report, decode_dna(samples))
	save_run(args.out, model=model, config=vars(args), report=report,
	losses=losses, samples=decode_dna(samples))


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