Generalized Discrete Diffusion from Snapshots

Introduction

• Existing dLLMs mostly pair masking or uniform replacement with mean parametrization, limiting neighborhood-aware corruption and flexible reverse uncertainty at LLM scale.The forward process ignores discrete-space neighborhoods, while mean parametrization restricts how denoising uncertainty becomes reverse dynamics.Source: §1 · pp. 1–2
• GDDS provides a generalized interpolating discrete-diffusion framework for arbitrary Markovian noising processes that encompasses existing approaches.The framework extends beyond the restricted subclasses explored in prior work.Source: §1 · p. 2
• Its forward process enables fast arbitrary corruption for large vocabularies using only column access to the noising rate matrix.Source: §1 · p. 2
• GDDS parameterizes reverse transitions and trains with a simple snapshot-based ELBO, predicting clean-token posteriors from exact forward-CTMC snapshots rather than full paths.Source: §1 · p. 2, Figure 2 · p. 2
• GDDS is the first fully general and computationally efficient discrete-diffusion framework, with state-of-the-art modeling and generation quality on large-scale language modeling tasks.The experiments demonstrate state-of-the-art modeling and generation quality.Source: §1 · p. 2

2.1. Discrete Diffusion

• Discrete diffusion models token evolution as a continuous-time Markov chain over a finite vocabulary, with forward transitions mapping q0 to qt and allowing noisy tokens to be sampled directly.The transition matrices Kt act on marginals as qt = Ktq0.Source: §2.1 · pp. 2–3
• The reverse process is defined by time reversal of the forward dynamics, and a neural network learns to simulate it to reconstruct clean data from a fully noised quantity.The section sets T = 1, noting that any bounded time interval can be rescaled to [0, 1].Source: §2.1 · p. 3
• Forward and reverse rate matrices have nonnegative off-diagonal entries and mass-conserving diagonals, and can be factored into exit rates and jump kernels.Exit rates control how often the chain leaves a state; jump kernels specify where it moves.Source: §2.1 · p. 3

2.2. Designs of the rate matrix

• At m = 50,257, storing dense F requires more than 2.5 × 10^9 parameters (≈20GB in double precision), and matrix-vector products cost O(m^2), making it impractical.Common designs therefore use structured kernels such as uniform or masking corruption. Equal exit rates and a time-independent kernel yield a closed-form Kt, but these structures impose rigid corruption patterns.Source: §2.2 · p. 3

3.1. Generalized interpolating discrete diffusion

• The mixing rate α_t controls corruption intensity while Π_t determines token-mixing structure, recovering masked and uniform schemes and covering Zhou et al. (2025) and GenMD4 (Shi et al., 2024), unlike GIDD (von R¨utte et al., 2025a).More generally, the framework allows structured and token-dependent corruption mechanisms.Source: §3.1 · pp. 3–4

3.2. Efficient forward noising through uniformization

• Uniformization represents exact noising as a Poisson number of transitions through F_t, preserving its transition structure; non-shared exit rates extend via Poisson thinning.The exposition focuses on shared exit rates to simplify implementation.Source: §3.2 · p. 4
• Uniformization yields exact samples for arbitrary continuous-time noising using Poisson sampling and rate-matrix column access rather than matrix exponentiation.The procedure also generalizes to a parallel sequence-level algorithm.Source: §3.2 · pp. 4–5

4.1. The core mismatch in reverse parametrization

• Using the mean denoiser through the reverse-time model does not generally train it to estimate the clean-token posterior, because it controls both reverse jump timing and destinations.GDDS instead aligns the mean network with a snapshot generative model and uses a separate jump network for path-wise reverse modeling.Source: §4.1 · p. 5

4.2. Path-wise model and loss function

• The jump-states parametrization learns the reverse jump kernel while fixing the true reverse exit rates, yielding a weighted cross-entropy ELBO with exit-rate weights for any forward rate matrix.Even with shared forward exit rates, reverse exit rates need not be shared.Source: §4.2 · p. 5
• The path-wise loss can require tractable forward conditionals, so the Campbell estimator rewrites it as a sum over the uniformization path using only network outputs and that path.Path-wise training evaluates many snapshot-wise conditionals and calls for a two-stream mechanism that is not naturally aligned with standard transformers and empirically underperforms them.Source: §4.2 · p. 6

4.3. Snapshot model and loss function

• The snapshot-latent ELBO trains the mean predictor on clean tokens and noisy snapshots, without requiring explicit knowledge of the forward conditional, and is compatible with standard architectures.Replacing full-path latents with snapshots trades some information for an objective better aligned with the architecture and easier to optimize; minimizing it enforces the intended clean-token posterior and empirically produces better samples.Source: §4.3 · pp. 6–7

5. Experiments

• The experiments cover Text8 character-level language modeling and BPE-tokenized OWT language modeling and generation, comparing retrained models under matched compute.Text8 models train for 1M steps and OWT models for 500k; retrained models share a Transformer backbone, with diffusion using bidirectional attention and AR using causal attention.Source: §5 · p. 7

5.1. Language modeling

• On Text8, GDDS Absorb outperforms the AR baseline; on OWT, GDDS achieves tighter variational bounds than retrained UDLM/MDM and can outperform matched AR.The comparisons are under matched compute; Table 1 reports Text8 BPC, while Table 2 and Figure 4 report OWT validation perplexity and training curves.Source: Table 1 · p. 7, §5.1 · p. 8, Table 2 · p. 8
• GDDS Gauss outperforms all baselines across every OOD dataset, while GDDS Absorb improves on retrained MDM on most datasets and GDDS Uniform on retrained UDLM across all but one.The authors suggest semantically structured corruptions may provide a generalization advantage.Source: Table 3 · p. 8, §5.1 · p. 9

5.2. Language generation

• GDDS improves the generation quality-diversity tradeoff: GDDS Absorb matches MDM's entropy with lower Gen-PPL at 64 steps, even when MDM uses up to 1024 steps.UDLM has low Gen-PPL but lower entropy, while MDM's increased entropy at larger decoding budgets comes with a steep Gen-PPL degradation.Source: Figure 5 · p. 9, §5.2 · p. 9
• GDDS Uniform increases Distinct-1/2/3 over UDLM and matches the strong diversity of masked methods, without the large high-budget Gen-PPL penalties observed for MDM.Distinct-n measures the fraction of unique n-grams in generated samples.Source: §5.2 · p. 9, Table 4 · p. 9

• GDDS beats previous discrete diffusion and autoregressive models on language modeling tasks at this scale.Source: §6 · p. 9
• GDDS supports efficient noising processes with arbitrary rate matrices.Source: §6 · p. 9
• Its simple training loss uses snapshot samples instead of the entire noising path and is compatible with standard architectures.Source: §6 · p. 9
• Future work could develop Semantic-Informed Kernels to make noising processes meaningful based on similarities between words.Source: §6 · p. 9

Impact Statement

• The paper frames its contribution as advancing discrete diffusion-based generative modeling and notes that it has potential societal consequences.Source: Impact Statement · p. 10

A.1. Transition operator for time-inhomogeneous CTMCs

• For a time-varying generator, the transition operator solves a matrix-valued ODE and maps an initial distribution to its future CTMC law.It has a time-ordered exponential representation, which reduces to an ordinary exponential when the generators commute pairwise.Source: §A.1 · p. 13

A.2. Coordinate-wise rate matrix and product-form marginals

• When jumps modify one token at a time under coordinate-wise generators, the sequence-level conditional marginal factorizes across coordinates.Each coordinate evolves as an independent time-inhomogeneous CTMC.Source: §A.2 · pp. 13–14

B.1. CTMC rate matrices

• For an invertible interpolating kernel Kt, Qt = K̇tK_t^-1 is a valid CTMC rate matrix.The proof establishes nonnegative off-diagonal entries and zero column sums.Source: §B.1 · p. 14

B.2. Uniformization and exact sampling

• Uniformization samples the noised state exactly using a Poisson jump count, ordered jump times, and transitions from the time-specific matrices Ft.The appendix explicitly states that this sampling procedure is exact.Source: §B.2 · pp. 15–16

B.3.1. REVERSE PROCESS

• When μθ matches the exact forward posterior q0|t, the plug-in reverse transition equals the true reverse kernel.The paper identifies q0|t as the training target for μθ.Source: §B.3.1 · p. 17
• The true reverse generator separates jump timing from destination, but the plug-in construction learns both and complicates path-wise ELBO optimization.Its event-level objective has gradients through both the destination cross-entropy and the rate/normalization term.Source: §B.3.1 · p. 18
• The reverse generator learns jump destinations but prescribes exit rates using a nonnegative schedule r[x](t) independent of θ.This parameterization is intended to mirror the true CTMC factorization.Source: §B.3.1 · p. 18
• Its off-diagonal rates are r[x](t)jθ(x, t)[y], while its diagonal rate is −r[x](t).Source: §B.3.1 · p. 18
• Unlike the plug-in Bayes transition in Eq. (16), the network determines where the chain jumps, while the schedule determines when it jumps.Source: §B.3.1 · p. 18

B.3.2. PROOF OF PROPOSITION 4.1

• Proposition B.3 states a general path-wise ELBO for any forward CTMC and its parametrized time reversal, restating a result established in prior work.The passage cites Campbell et al. 2022; Shi et al. 2024; von Rütte et al. 2025a; Zhou et al. 2025.Source: §B.3.2 · pp. 18–19
• The jump-state parametrization makes the only θ-dependent contribution a weighted cross-entropy over reverse jump matrices, keeping the reverse rate fixed.This concentrates learning on predicting jump destinations rather than fitting a separate factor controlling the time-change.Source: §B.3.2 · p. 20
• Unlike the jump-state objective, SEDD has a linear score and cross-entropy term, while mean parametrization has a θ-dependent reverse-rate term alongside cross-entropy.Both alternatives require fitting jump destinations and an additional θ-dependent factor that controls time-change, adding training burden and variance.Source: §B.3.2 · p. 20
• The jump-state approach is closely related to Amin et al. (2025), but does not condition the path-wise generative model on an explicit event schedule.It retains the same joint distribution defining the generative model as prior work.Source: §B.3.2 · p. 20

B.3.3. PROOF OF PROPOSITION 4.2

• The forward jump random measure has intensity measure Λ_x0, so Campbell’s formula converts sums over forward jumps into expectations under the jump measure.The proof verifies the intensity on rectangles by calculating expected jump counts for transitions x → y over a time interval.Source: §B.3.3 · pp. 20–21
• At sequence level, the jump measure is the superposition of independent token jump measures, and linearity of expectation yields the sequence-level expression.The forward marginal factorizes over tokens in the stated sequence setup.Source: §B.3.3 · p. 21

B.3.4. RECOVERING THE MASKED DIFFUSION LOSS

• The appendix says this reverse-rate independence is specific to masked diffusion; other generalizations of Austin et al. (2021) and Campbell et al. (2022) yield more complicated losses.Those losses are contrasted with the weighted cross-entropy from the jump-states parametrization.Source: §B.3.4 · p. 21
• In masked diffusion, the mean and jump-states parametrizations coincide because the conditional reverse rate is independent of θ.In the absorbing-mask setup, a token is either still clean or [MASK], while the schedule is fixed by α_t.Source: §B.3.4 · pp. 21–22

B.4. Snapshot parametrization and Evidence Lower Bound (ELBO)

• The snapshot latent is s = (x_t, t), sampled by drawing t from ρ and extracting x_t from the forward process; the predictor uses the mean network’s output.For the proof, ρ is uniform on [0, 1].Source: §B.4 · p. 23
• The snapshot-to-path expected NLL gap is nonnegative because the full path contains at least as much information about x_0 as its snapshot.The appendix expresses the gap using the conditional-entropy difference H(x_0 | s) − H(x_0 | ω).Source: §B.4 · p. 24
• The snapshot NLL’s conditional-entropy term depends on the data and forward process, not θ, and its time integral is weighted by ρ(t) when ρ is nonuniform.The appendix identifies this as the unconditional expected snapshot NLL after averaging over the data distribution.Source: §B.4 · p. 24
• Unlike the snapshot ELBO, the path-wise ELBO includes a θ-dependent diffusion term and trains local previous-state conditionals; its minimizer need not match the snapshot predictor’s.The path-wise objective predicts z_{k−1} from (z_k, T_k), whereas calibration of x_0 given the full path is a different conditional task.Source: §B.4 · pp. 24–25

C.1. Experimental setting

• Under matched protocols, GDDS yields tighter likelihood bounds and stronger generation quality despite sensitivity of retrained baseline results to implementation and system conditions.On Text8, retrained AR, MDM, and UDLM scored 1.35, 1.58, and 1.67 BPC, respectively; OWT training used 500k steps at sequence length 1024, totaling 262B tokens.Source: §C.1 · pp. 25–26

C.2. Semantic-Informed Kernel (SIK)

• The Semantic-Informed Kernel uses embedding-space affinities with Gaussian or cosine distances, implemented either as sparse KNN sampling or dense lazy KEOPS evaluation.KNN samples from top-64 neighbors; KEOPS computes scores through blockwise GPU reductions without materializing the full transition matrix.Source: §C.2 · p. 26

C.3. Metrics

• Evaluation reports likelihood-based metrics alongside unconditional-sample metrics for generative quality and diversity.The metrics include BPC or perplexity, Generative Perplexity, sequence entropy, and Distinct-n.Source: §C.3 · p. 27

C.3.1. NEGATIVE LOG-LIKELIHOOD (NLL), BPC, AND PPL

• For likelihood evaluation, Text8 uses BPC and OWT uses perplexity; when likelihood is intractable, the reported NLL is a variational upper bound.Lower NLL, BPC, or perplexity indicates better likelihood-based performance.Source: §C.3.1 · pp. 27–28

C.3.2. GENERATIVE PERPLEXITY

• Generative Perplexity evaluates unconditional samples under GPT2-large, with lower values indicating greater predictability under that reference model and empirically correlating with perceived fluency.Likelihood metrics do not directly assess sample quality when generation uses an approximate sampler.Source: §C.3.2 · p. 28

C.3.3. SEQUENCE ENTROPY

• Sequence entropy measures average within-sample unigram entropy to quantify token diversity and is plotted against Generative Perplexity to show the quality–diversity tradeoff.Higher entropy indicates broader, less repetitive token use.Source: §C.3.3 · p. 28

C.3.4. DISTINCT-n

• Distinct-k measures corpus-level lexical diversity as the proportion of unique k-grams, with higher values indicating fewer repeated k-grams.The paper reports Distinct-n statistics.Source: §C.3.4 · p. 28

C.4. Qualitative samples on Text8

• The appendix presents two unconditional Text8 samples each from GDDS Absorb and GDDS Uniform models.The samples are shown in Tables 7 and 8, respectively.Source: §C.4 · p. 28, Table 7 · p. 29, Table 8 · p. 29

C.5. Qualitative samples on OpenWebText

• Tables 9–11 present unconditional OpenWebText samples from the GDDS Absorb, Uniform, and Gauss models after lightweight rendering cleanup.The cleanup replaces paragraph markers with line breaks, decodes common Unicode escapes, and renders special tokens in typewriter formatting.Source: §C.5 · p. 29, Table 9 · p. 30, Table 10 · p. 31

D. Sampling procedures

• Sampling uses a discrete-time grid with decoding budgets K ∈ {32, 64, 128, 256, 512, 1024}; formulas run token-wise and in parallel across sequence positions.At sampling time, the mean network provides a clean-token distribution conditioned on the current token and time.Source: D. Sampling procedures · p. 33

D.1. Ancestral sampling

• The ancestral reverse kernel uses the standard plug-in approximation of the clean token by the mean network’s predicted distribution.The kernel combines the forward transition with the forward marginal induced by that prediction.Source: §D.1 · p. 33
• On a discretized grid, the sampler applies the plug-in Bayes/ancestral kernel between each pair of adjacent times in reverse order.The kernel uses the forward transition and the predictor-induced mixture forward marginal at the preceding time.Source: §D.1 · p. 33
• The appendix gives closed forms for Uniform and Absorb sampling and an operator form for SIK.It identifies the Uniform and Absorb forms as already known in the literature (Sahoo et al., 2024; Shi et al., 2024; Ou et al., 2025; Schiff et al., 2025).Source: §D.1 · p. 33
• Algorithm 4 initializes the terminal state from its forward marginal, then iterates backward, predicts clean-token distributions, computes marginals, and samples the preceding state.The procedure is token-wise and parallel over sequence positions.Source: §D.1 · p. 33

D.2. Instantiations of ancestral sampling

• Unlike the closed-form uniform and absorbing kernels, SIK ancestral updates require approximating bridge terms and forward marginals, making decoding slower and more delicate.The implementation uses cached uniformization-based matrix-vector products to approximate these quantities.Source: §D.2 · pp. 34–35
• At K=128 and 256, GDDS-SIK reaches the target natural-OWT entropy regime, but Gen-PPL worsens beyond K=64 and its sampler does not beat uniform or absorbing samplers.SIK sampling approximates the time-ordered exponential with uniformization-based cached matrix-vector products; accumulated approximation error may explain the worsening quality.Source: §D.2 · pp. 34–35, Table 12 · p. 35
• Adaptive or confidence-based schedules are a proposed direction for avoiding repeated explicit SIK operator approximations in sampling.The stated goal is samplers and snapshot-ELBO training objectives that do not require approximating forward transition operators in closed form.Source: §D.2 · p. 35

E. Architectural details for Campbell

• Along the sequence-level diffusion path, each jump changes one coordinate, so the reverse kernel concerns only the predecessor token at that updated position.The observed transition pair is associated with the coordinate updated at the jump.Source: §E.1 · p. 35

E.1. Campbell objective at the sequence-level : an any-order autoregressive objective

• Jensen’s inequality yields a sequence-level ELBO whose parameter-dependent term is Campbell’s objective, an event-wise cross-entropy sum up to an additive constant.The derivation extends the token-level ELBO beyond sequences of length one.Source: §E.1 · p. 36
• Unlike XLNet’s random-permutation factorization, Campbell’s ordering comes from Poisson jump events and uses noised contexts like MPNet, while applying beyond masked diffusion.The objective trains over all jumps along the forward path.Source: §E.1 · p. 36

Two-stream architecture

• The architecture handles the whole path in an average of E[N_ℓ] NFEs, independent of sequence length n, while masked diffusion requires exactly one pass.Naive neural implementations may require an enormous number of NFEs to evaluate Eq. (31).Source: §E.2 · p. 36
• The Campbell objective factorizes over observed jumps, and a shared network jθ assigns predecessor-token probabilities conditioned on each post-jump state.Each observed jump contributes a cross-entropy loss term.Source: §E.2 · p. 36
• Predicting the pre-jump token requires conditioning on clean tokens with later masking times, creating a sample-dependent, non-causal factorization that fixed left-to-right masks cannot enforce.The rank r_ℓ orders positions by decreasing masking time, so τ_j > τ_ℓ exactly when r_j < r_ℓ.Source: §E.2 · p. 36
• The encoder builds contextual representations, while the decoder uses masked queries and attends only to keys with lower rank, preventing leakage from the target or unrevealed tokens.The decoder is conditioned on continuous time through Adaptive LayerNorm, while the encoder remains time-agnostic.Source: §E.2 · p. 36

E.3. Empirical results

• The two-stream architecture was trained on Text8 and OWT using the same experimental setup as previously, with absorbing forward noising.Table 13 reports BPC and validation perplexity under the setup from Section C.Source: §E.3 · p. 37, Table 13 · p. 37
• Training-loss standard deviations were 3.08×10−2 for AR, 1.63×10−1 for MDM, 1.65×10−1 for GDDS Absorb, and 2.92×10−2 for Campbell.These values cover the last 300k steps on Text8; Campbell’s curve was markedly smoother than snapshot-based training and comparable to AR.Source: §E.3 · p. 37, Figure 6 · p. 37
• The two-stream architecture clearly underperforms on likelihood under the tested architectural constraints, despite Campbell’s smoother optimization trajectory.Table 13 reports validation perplexity for the two-stream architecture; the prose provides no specific table values.Source: §E.3 · p. 37, Table 13 · p. 37
• Campbell sums per-jump cross-entropies along the uniformization path, providing at least n supervision terms per clean sequence versus one for snapshot-based training.The paper attributes Campbell’s lower iteration-to-iteration objective variance to this additional supervision, while noting that Campbell conditions on full path information ω.Source: §E.3 · p. 37