A toy diffusion model for text gen using Karpathy's shakespear data

A developer built a discrete diffusion model for text generation using Andrej Karpathy's Shakespeare dataset, implementing a continuous-time Markov chain (CTMC) with a rate matrix to handle discrete token transitions. The model, trained at the character level, demonstrates how diffusion techniques typically used for images can be adapted to generate Shakespeare-like text.

The code below is trimmed to the essentials; the full, runnable version lives in the original notebook: github.com/litlig/notebooks/discrete diffusion.ipynb Diffusion models are largely used in image and video generation. They operate in a high-dimensional continuous space, where we convert noise to an image by following a denoising process. Text generation, on the other hand, is discrete in nature. For a text sequence, we tokenize it and get a vector of token indices. The numerical difference between tokens has no natural meaning — a smaller gap between two token ids does not mean the tokens are semantically closer. Discrete diffusion is designed for this discrete setting. Here we use the Shakespeare dataset from Andrej Karpathy’s LLM series, build a discrete DiT model from scratch, and see if it can generate Shakespeare-like text. We work at the character level, so the “vocabulary” is just the set of distinct characters in the corpus: vocab = sorted list set text vocab size = len vocab 65 stoi = { ch:i for i,ch in enumerate vocab } itos = { i:ch for i,ch in enumerate vocab } encode = lambda s: stoi c for c in s decode = lambda l: ''.join itos i for i in l CTMC and the rate matrix ctmc-and-the-rate-matrix To recall see the earlier 2-D diffusion post https://blog.strayforge.com/posts/diffusion-2d-notebook/ for the continuous case in full , the continuous diffusion model is done by flow matching, where we define a vector field \ u t\ that tells the direction and velocity a noised image \ x t\ should move at time \ t\ . Since we do not know the distribution of the image manifold, \ u t\ is not tractable. However, if we know the final state a sample image , \ u t\ can be computed analytically — one solution is to linearly connect the noised point \ x t\ and the final state \ z\ . What we actually need is not the vector field for a given \ z\ , but the average vector field over all \ z\ that follow the posterior distribution \ p z\,|\,x t \ . In neural net training the loss is always an average over samples, so by training a network to minimize the average loss against the conditional vector field, we also recover the marginal vector field we need. For a text sequence, the transition from noise to a meaningful sequence is a series of jumps across discrete states. Instead of a vector field, we define a rate matrix \ Q t y\,|\,x \ , the rate of jumping from state \ x\ to state \ y\ at time \ t\ . If we are at state \ x\ at time \ t\ , then over a very small interval \ h\ the probability we jump to \ y\ with \ y \neq x\ is \ h\, Q t y\,|\,x \ . This is a continuous-time Markov chain CTMC . The number of states is exponential in the sequence length — \ V^d\ , where \ V\ is the vocab size and \ d\ is the sequence length. To keep \ Q\ manageable, we set the rate to zero whenever more than one position changes. For \ X t = x\ , a rate matrix \ Q t v, j \ then defines the rate of changing position \ j\ to token \ v\ , and we can sample with an Euler approximation: python def euler step x t, rates, h : delta {v,x}: B, d, V one-hot at the current token delta = F.one hot x t, num classes=rates.size -1 .to rates.dtype off diag = h rates 1.0 - delta h q v for v = x, else 0 stay = 1.0 - off diag.sum -1, keepdim=True 1 - h sum {v =x} q v probs = off diag + delta stay return torch.distributions.Categorical probs=probs .sample def sample model, n seq, n steps, noise gen : ts = torch.linspace 0.0, 1.0, n steps + 1, device=device x = noise gen.gen 1, n seq for i in range n steps : s, t = ts i , ts i + 1 rate mtx = model x, s n seq, n vocab x = euler step x, rate mtx, t - s return x Before training anything, we can sanity-check the machinery with a mock model that returns random rates. Starting from uniform noise and running the sampler produces exactly what you’d expect — noise: python class MockModel nn.Module : def forward self, x, t : logits = torch.randn x.shape -1 , vocab size, device=device return F.softmax logits, dim=-1 - F.one hot x, num classes=vocab size .to device nwAlQgrsSo 'bmm,DjUmT?hkdwWpui:&N iVXMEEqDm'zkh 3FfL?EK,oAvGkdLAK x,cX.bG Factorized mixture path factorized-mixture-path Given an initial noise distribution \ p {\mathrm{init}}\ and a final data distribution \ p {\mathrm{data}}\ , a discrete probability path is a family \ p t\ with \ p 0 \sim p {\mathrm{init}}\ and \ p 1 \sim p {\mathrm{data}}\ . The most commonly used discrete probability path is the factorized mixture path , which defines the conditional path as: where \ \kappa t\ is the noise schedule. Each token position is treated independently. The rate matrix conditioned on \ z\ is: \ Q^z t i, v\,|\,x i = \frac{\dot{\kappa} t}{1 - \kappa t}\big \delta {z i} v - \delta {x i} v \big \ and the marginal rate matrix is: \ Q t i, v\,|\,x i = \sum z Q^z t i, v\,|\,x i \, p z\,|\,x = \frac{\dot{\kappa} t}{1 - \kappa t}\big p z j = v\,|\,x - \delta {x i} v \big \ So the problem reduces to a categorization: predict \ z\ given \ x\ and \ t\ . The rate matrix is just a reparameterization of that prediction. The model: a discrete DiT the-model-a-discrete-dit The network is a transformer that takes noised tokens \ x t\ and a scalar time \ t\ , and outputs logits over the vocabulary at each position — its job is to predict the clean tokens \ z\ . Time is injected through AdaLN-Zero conditioning, as in the DiT architecture: a sinusoidal time embedding is passed through an MLP to produce a conditioning vector, which modulates each transformer block. The final projection is zero-initialized so each block starts as an identity map, which stabilizes training at initialization. class DiscreteDiffusionTransformer nn.Module : """ Input: x t B, L token ids, t B, Output: logits B, L, vocab size predicting the original tokens """ def forward self, x t, t, padding mask=None : B, L = x t.shape x = self.token emb x t + self.pos emb :, :L, : c = self.time mlp t B, cond dim conditioning vector shared across layers for block in self.blocks: x = block x, c, key padding mask=padding mask return self.final x, c Training to predict z training-to-predict-z Training follows the factorized mixture path directly. For each batch we sample a time \ t\ , set the noise level \ \kappa t\ here the schedule is simply \ \kappa t = t\ , and construct \ x t\ by keeping each clean token with probability \ \kappa t\ and replacing it with uniform noise otherwise. The network then predicts the clean sequence \ z\ , and the loss is a plain cross-entropy over positions: python def schedule t : return t def get batch : t = torch.rand n sample, device=device kappa = schedule t seq idx = torch.randint 0, len text - n seq + 1, n sample, , device=device .unsqueeze 1 \ + torch.arange n seq, device=device z = data seq idx sample, seq masks = torch.bernoulli kappa.unsqueeze 1 .expand -1, n seq .long keep clean where 1 noise = noise gen.gen n sample, n seq x = masks z + 1 - masks noise return x, t, z def loss fn z pred, z : return F.cross entropy z pred.view -1, vocab size , z.view -1 Sampling with the trained model sampling-with-the-trained-model At sampling time we turn the model’s clean-token prediction back into a rate matrix using the marginal formula above, then take an Euler step: python @torch.no grad def sample model, n seq, n steps, noise gen : model.eval ts = torch.linspace 0.0, 1.0, n steps + 1, device=device x = noise gen.gen 1, n seq for i in range n steps : s, t = ts i , ts i + 1 logits = model x, s.expand x.shape 0 rate mtx = F.softmax logits, dim=-1 - F.one hot x, num classes=vocab size .to device / 1 - s x = euler step x, rate mtx, t - s return x After a short training run, the samples are no longer random noise — the model has picked up Shakespeare’s shape: capitalized speaker names, line breaks, and vaguely English words. d did neyce? HONRY VIA: sutur selon'sl mave, Myoury abenly tuke It is far from coherent — this is a tiny model trained for a few thousand steps on characters — but the discrete diffusion process clearly works: starting from pure noise, a sequence of denoising jumps lands us in something that looks like a play.