lecture_4/README.md

Lecture 4: Discrete Diffusion

Chapter 4 keeps every state inside a finite vocabulary and replaces the continuous displacement of Chapter 3 by a decision about which tokens to erase, reveal or replace. The chapter builds masked diffusion by hand from a survival schedule, then rebuilds it as one member of a family: a rate matrix fixes what a token can do in the next instant, the master equation follows the marginals, the matrix exponential gives the corruption kernel at any time, and time reversal turns the forward rates into the rates that generate. Seven published models come out of that family by changing the rate matrix, the time variable, what the network predicts, or the state space itself, and a further group keeps the rates and changes the order in which positions are filled, the weighting of the loss, or the distribution to sample from. Every file here runs on the shared synthetic DNA of cis6270.data, with t = 0 clean and t = 1 all masks, survival probability alpha_t, and rates R_t(x, y) in the row convention p_dot = p R.

The files

File	Command	What it trains and generates
ctmc.py	python lecture_4/ctmc.py	Trains no network, and works every number of the CTMC sections: the two rules, one row for DNA, waiting times, the master equation, the matrix exponential, the R_pi family, eigenvalues, detailed balance and time reversal.
mdlm.py	python lecture_4/mdlm.py	Trains the denoiser on the weighted cross-entropy at masked positions, generates DNA from all masks. The chapter's thirty-line model.
uniform.py	python lecture_4/uniform.py	Trains against the rate loss under uniform replacement, generates by tau-leaping from uniformly random tokens.
d3pm.py	python lecture_4/d3pm.py	Trains against the discrete-time variational bound plus the auxiliary cross-entropy, generates by T ancestral steps.
sedd.py	python lecture_4/sedd.py	Trains a network to output the concrete score through denoising score entropy, generates from the predicted ratios.
block.py	python lecture_4/block.py	Trains masked diffusion inside blocks with a clean prefix, generates block by block. --block moves the dial between MDLM and autoregression.
simplex.py	python lecture_4/simplex.py	Trains a denoiser that reads a Dirichlet belief at every position, generates by thinning the belief, drawing an innovation and mixing the two.
guidance.py	python lecture_4/guidance.py	Trains a label-dropout denoiser and a classifier on noisy sequences, then samples unguided, with exact classifier guidance, with its gradient approximation, and classifier free.
planner.py	python lecture_4/planner.py	Trains with the planner weights inside the bound, samples the same weights in a uniform order and in the planner's order with a remasking corrector.
search.py	python lecture_4/search.py	Trains masked diffusion, then searches the tree of partial sequences under the h-transform and keeps a Pareto archive over two objectives.

Two constructions of the chapter have no file here, and the chapter develops both in full. Per-token schedules, MD4 and GenMD4, give every clean token its own rate beta_a(t) and change nothing but the weight, which becomes w_a(t) = -alpha_dot_t(a)/(1 - alpha_t(a)) inside the sum over tokens. A2D2 lets the length change by adding an insertion jump beside the unmasking jump, which enlarges the state space and keeps the master equation, the reversal and the rate loss exactly as they are.

Every file takes the shared flags of cis6270.runner.common_parser, so --steps, --samples, --sample-steps, --seed and --quiet mean the same thing everywhere. ctmc.py works through rates the chapter writes down, so its --steps is the number of Euler steps in the matrix-exponential comparison and its --samples the number of simulated paths. d3pm.py takes T from --timesteps, since the same number fixes the corruption and the ancestral sampler. search.py reads --sample-steps as the length of one rollout, which is short because we score every leaf of the tree by its own rollouts.

What a default run produces

We took every number below from --seed 0 on a CPU. The worked numbers are exact and repeat on any machine; the trained numbers move a little with the hardware's floating point, and the sample statistics move with the seed. The training data have GC fraction 0.577 and motif fraction 0.777, and a sampler drawing its bases uniformly would land near 0.05 on the motif, since ACGT fits in a length-16 sequence at thirteen places and each one has probability 1/256.

File	Worked numbers it prints	Trained numbers
ctmc.py	exit rate 4 out of A, P(A, A) = 0.6 at h = 0.1, Euler step (0.16, 0.22, 0.28, 0.34), exp(tau R_pi) at 0.625 and 0.125, Euler products 0.480, 0.570, 0.616, reverse rates 2 and 0.5 out of C	none
mdlm.py	r = 0.5, reverse law (0.3, 0.2, 0, 0, 0.5), w(0.5) = 2, loss term 2.0996, rate error 1.0217	bound 1.25 nats per token, GC 0.56, motif 0.69
uniform.py	kernel 0.625 and 0.125, clean factor 5, posterior (0.375, 0.625), conditional rates 5 and 0.2, posterior average 2, model rates 2.76 and 1.16, rate error 0.754	rate loss 19.1 nats per sequence, GC 0.56, motif 0.50
d3pm.py	alpha_bar_2 = 0.56, Q_bar_2 at 0.67 and 0.11, the step-1 posterior (0.5795, 0.3523, 0.0341, 0.0341), Gaussian row 0.4088 and 0.0912, exp(0.3 R) first row (0.6747, 0.1964, 0.1122, 0.0167)	bound 18.2 nats per sequence, GC 0.57, motif 0.72
sedd.py	target ratio 2, score entropy 0.0107, 0.3863, 0.1891	score entropy 18.8 nats per sequence, mean predicted ratio 1.26 toward a conditional target of 5, GC 0.55, motif 0.33
block.py	worked bound 2.7567 nats on ACGTAC, integral of w(t)(1 - alpha_t) = 1	bound 1.21 nats per token at L' = 4 and 1.29 at L' = 1, GC 0.57, motif 0.67
simplex.py	c_t = 8 and c_s = 20, parameters (5, 1, 1, 1) and (17, 1, 1, 1), standard deviation 0.161 on A, increment (12, 0, 0, 0), W ~ Beta(8, 12) with mean 0.4, simulated standard deviation 0.0779 against the exact 0.0779	cross-entropy 0.42 nats per token, GC 0.57, motif 0.34
guidance.py	h = (2, 1, 0.5, 0.5), guided rates (1, 1.5, 0.5), at gamma = 2 rates (0.5, 4.5, 0.5) with exit rate 5.5 and share 0.8182 on G	GC 0.58 unguided, 0.68 with the exact classifier, 0.63 with the gradient approximation, 0.74 classifier free
planner.py	reveals positions 3 and 6, remasks position 5, both corrector fluxes 1, PAPL weights (0.4, 2.0, 0.6, 1.0), inner sums 2.9 and 2.12	motif 0.61 in a uniform order against 0.86 in the planner's order, GC 0.58 and 0.51
search.py	h(z) = 0.35, h(z') = 0.75, factor 2.143, a rate of 1.2 becomes 2.571, the archive keeps P1 and P2	144 rollouts scored, an archive of 3 sequences, best GC 0.88 and best motif match 1.0

On one core, we get about two seconds from a default run of ctmc.py, and between thirty and forty seconds of processor time from every other file. With --steps 20, each file runs in a few seconds.

mdlm.py --steps 3000 lowers the bound from 1.25 to 1.22 nats per token in about three times the time, and leaves the motif fraction inside the spread it already has across seeds, which at seeds 0, 1 and 2 is 0.69, 0.45 and 0.72. The one-position bound of block.py --block 1 is the exact autoregressive log-likelihood with no slack, and at this budget it still measures above the L' = 4 bound, because the estimator draws one block of sixteen and multiplies by sixteen, which is sixteen times noisier at the same number of draws, and because a denoiser restricted to a left prefix has a harder prediction problem than one that reads both sides. sedd.py sits below the other samplers on the motif because its sampler stops at the time floor, as the list below says. simplex.py --sample-steps 256 raises its motif fraction from 0.34 to 0.38, so the reverse grid accounts for only a small part of what holds that number down.

What is simplified

• The time weight w(t) = 1/t is unbounded at t = 0, so every file draws its times above a small floor. The chapter's integrand is correct as written; the floor keeps one draw from dominating a batch.
• The samplers that use tau-leaping rescale a step whose total jump probability would exceed one. uniform.py then finishes with the exact posterior step down to t = 0 that the chapter derives for the pi family. sedd.py predicts ratios, with no posterior over clean tokens to step through, so it stops at the time floor, where alpha_t = 0.92 leaves about eight percent of the positions redrawn from pi.
• uniform.py and sedd.py use the constant schedule beta(t) = 4, which is the schedule of the chapter's worked example and reaches alpha_t = 1/2 at t = (log 2)/4. Its survival probability at t = 1 is 0.018, and that residue is the KL(q_1 || p_1) term of the bound.
• d3pm.py samples one step index per sequence and multiplies by T, which estimates the sum over steps without evaluating all of it.
• block.py samples one block per sequence and multiplies by the number of blocks, for the same reason. To train a block of length L' we hide the suffix behind the mask in place of a block-causal attention mask.
• simplex.py holds the concentration constant at gamma = 1, realizes the sum over clean tokens by one categorical draw as the chapter does, and reads the final token as the corner the belief has reached. Its loss is the chapter's cross-entropy, so its nats per token are the prediction error of the denoiser at the sampled times, and the bounds above are on a likelihood.
• guidance.py uses a mean-pooling classifier on one-hot inputs, because we need a derivative with respect to the sequence to approximate the gradient, and the label of this data is a thresholded mean over positions.
• planner.py reveals the number of positions the schedule sets at each step, which is the deterministic version of the chapter's independent reveal draws, and remasks a fixed count.
• search.py scores --sample-steps rollout steps per completion and expands a fixed number of children, so the tree stays small enough to build in a few seconds.