Language generation with discrete diffusion · StrayForge↓<br>Skip to main content<br>StrayForge
Table of Contents<br>Table of Contents
The code below is trimmed to the essentials; the full, runnable version lives in the original notebook:<br>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.<br>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.<br>We work at the character level, so the “vocabulary” is just the set of distinct characters in the corpus:<br>vocab = sorted(list(set(text)))<br>vocab_size = len(vocab) # 65
stoi = { ch:i for i,ch in enumerate(vocab) }<br>itos = { i:ch for i,ch in enumerate(vocab) }<br>encode = lambda s: [stoi[c] for c in s]<br>decode = lambda l: ''.join([itos[i] for i in l])
CTMC and the rate matrix<br>To recall (see the earlier 2-D diffusion post 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.<br>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).<br>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:<br>def euler_step(x_t, rates, h):<br># delta_{v,x}: (B, d, V) one-hot at the current token<br>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<br>stay = (1.0 - off_diag.sum(-1, keepdim=True)) # 1 - h*sum_{v!=x} q(v)
probs = off_diag + delta * stay<br>return torch.distributions.Categorical(probs=probs).sample()
def sample(model, n_seq, n_steps, noise_gen):<br>ts = torch.linspace(0.0, 1.0, n_steps + 1, device=device)<br>x = noise_gen.gen(1, n_seq)<br>for i in range(n_steps):<br>s, t = ts[i], ts[i + 1]<br>rate_mtx = model(x, s) # (n_seq, n_vocab)<br>x = euler_step(x, rate_mtx, t - s)<br>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:<br>class MockModel(nn.Module):<br>def forward(self, x, t):<br>logits = torch.randn(x.shape[-1], vocab_size, device=device)<br>return F.softmax(logits, dim=-1) - F.one_hot(x, num_classes=vocab_size).to(device)
nwAlQgrsSo 'bmm,DjUmT?hkdwWpui:&N<br>iVXMEEqDm'zkh<br>3FfL?EK,oAvGkdLAK!x,cX.bG
Factorized mixture path<br>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}}\).<br>The most commonly used discrete probability path is the factorized mixture path , which defines the conditional path as:<br>\[p_t(x\,|\,z) = \prod_{j=1}^d \big[(1 - \kappa_t)\, p_{\mathrm{init}}^{(j)}(x_j) + \kappa_t\, \delta_{z_j}(x_j)\big]\]where \(\kappa_t\) is the noise schedule. Each token position is treated independently. The rate matrix conditioned on \(z\) is:<br>\[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...