Information-Theoretic Discrete Diffusion

AI summary

Overview

Research area: Machine learning — discrete diffusion models, information theory, and likelihood estimation for generative sequence models.

Technical level: Advanced. The paper relies on continuous-time Markov chains, information-theoretic identities (I-MMSE style arguments), and score-matching objectives.

Scope (one sentence): The paper derives discrete analogues of the I-MMSE identity — the I-MDSE and I-MDCE relations — proving that standard denoising score entropy (DSE) and denoising cross-entropy (DCE) losses are exact estimators of data log-likelihood rather than merely variational bounds, and it develops practical likelihood, conditional likelihood, and likelihood-ratio estimators from those identities.

What This Paper Is About

Continuous (Gaussian) diffusion models already have a well-established information-theoretic foundation: the I-MMSE identity connects mutual information between clean and noisy variables to minimum mean squared error, and a pointwise generalization decomposes the data log-likelihood exactly in terms of denoising error. The discrete case had not been studied from this angle. This paper builds the analogous framework for discrete diffusion — including masked/absorbing diffusion, the setting used by large masked diffusion language models such as LLaDA — and shows that the losses practitioners already train with are exact, principled likelihood estimators.

Key Contributions

1. I-MDSE relation. The authors introduce the Information–Minimum Denoising Score Entropy identity, which links mutual information decay in the forward discrete diffusion process to the minimum denoising score entropy (DSE) loss, and derive a time-integral decomposition of the negative log-likelihood (NLL) in terms of that minimum loss.

2. I-MDCE relation for masked diffusion. They formulate the pointwise denoising cross-entropy (DCE) loss for masked (absorbing) diffusion, prove its exact pointwise and time-integrated equivalence to the DSE loss under a time reparameterization, and use this to establish the Information–Minimum Denoising Cross-Entropy identity and a parallel NLL decomposition.

3. Practical estimator extensions. Building on I-MDCE they derive a time-free reformulation of the log-likelihood expressed as an expectation over randomly selected unmasked token subsets, a conditional likelihood estimator for structured settings such as prompt–response modeling, and a coupled Monte Carlo estimator for likelihood ratios.

4. Empirical validation. They test the estimators on synthetic data with known ground-truth distributions, verify variance reduction against baselines, and demonstrate real-world utility through auditing experiments on LLaDA (out-of-distribution detection and model influence analysis).

Main Findings

• Information decays in proportion to the minimum DSE loss. The pointwise I-MDSE relation states that the derivative of the KL divergence between the forward conditional and marginal distributions equals the negative pointwise minimum DSE, and the marginal form states that the derivative of mutual information between the clean and diffused variables equals the negative minimum DSE. The negative sign reflects that information decays over time, since DSE is always nonnegative.

• The DSE loss is an exact likelihood estimator, not just a bound. I-MDSE yields an NLL decomposition: negative log-likelihood equals the time integral of the pointwise minimum DSE plus a KL residual term that vanishes in the full-noise limit. Practically, replacing the optimal score with a learned score network gives the previously known time-integrated DSE objective as an estimator of NLL — showing no higher-order corrections are needed.

• DSE and DCE losses are exactly equivalent. Lemma 3.3 shows the pointwise DSE loss equals a time-dependent factor times the pointwise DCE loss when the score and conditional predictors are related by the time-free reparameterization; Theorem 3.4 upgrades this to exact equality of the time-integrated training losses. This extends prior work that had shown only asymptotic equivalence in the full-noise limit.

• The DCE loss is also an exact likelihood estimator. The I-MDCE relation gives NLL equal to an integral over the noise level of the minimum DCE divided by the noise level, plus a KL residual that vanishes at full noise. Replacing the optimal predictor with a learned one recovers the standard DCE training objective as the estimator.

• A time-free estimator has substantially lower variance. Theorem 4.1 expresses NLL as the L-th harmonic number times an expectation over randomly selected unmasked index sets, weighted by a Beta-function-based sampling distribution. Empirically this time-free form achieves substantially lower Monte Carlo variance than the time-integral form.

• Conditional likelihoods extend the framework. The paper gives both an integral form and a time-free form for conditional negative log-likelihood over disjoint target and context index sets, enabling prompt–response style estimation; the sampling distribution for the time-free conditional case again involves the Beta function and harmonic number.

• Coupling reduces likelihood-ratio variance. A coupled Monte Carlo formulation using a shared mask for both sequences is unbiased and substantially reduces variance compared with decoupled estimation. The authors note this helps explain the empirical stability of recent alignment methods based on likelihood ratios in masked diffusion language models.

• Synthetic experiments recover ground-truth likelihoods. On synthetic DNA sequences over the alphabet {A, T, G, C}, a ground-truth distribution over 128 sequences of length 8 was used, with one million samples drawn to train a RADD model; Figure 1(a) compares true versus estimated likelihoods on 64 sequences and shows strong agreement. In a conditional setting, a length-five-million DNA sequence generated by a 4th-order Markov chain provided the ground truth; subsequences of length 32 were sampled for training, and a held-out sequence was split into a prompt (first 16 bases) and response (remaining 16 bases). Figure 1(b) shows estimated conditional likelihoods closely match the Markov ground truth.

• Variance reduction confirmed on real models and datasets. Using LLaDA, the authors compared Monte Carlo variance of conditional log-likelihood estimates across HellaSwag, ARC-hard, and PIQA for the time-free estimator versus the time-integral baseline, with variance measured over 15 independent samples, across various numbers of Monte Carlo samples. They also measured log-likelihood ratio variance on the BeaverTails dataset. The full numeric tables are not included in the provided text.

• Real-world auditing results. Conditional likelihood estimates detected out-of-distribution inputs and revealed distributional shifts in LLaDA.

Methodology in Plain English

The authors start from an analogy. In Gaussian diffusion, a well-known identity says the rate at which information about the clean data is lost as noise increases equals the best achievable denoising error — the minimum mean squared error. The team asks what the discrete equivalent looks like, where "error" is not squared distance but a score-ratio mismatch measured by the denoising score entropy loss.

They formalize discrete diffusion as a continuous-time Markov chain with a rate matrix, define the minimum achievable DSE loss (both globally and for a single data point), and prove a differential identity: the rate of change of a KL divergence or mutual information equals the negative of that minimum loss. Integrating this identity over time produces an exact decomposition of the negative log-likelihood into an integral of the loss plus a residual KL term that disappears as the process runs to the stationary distribution.

For masked diffusion they define a pointwise cross-entropy loss over masked token positions, then show algebraically that this loss equals the DSE loss up to a time-dependent factor once the score network and the conditional predictor are related by the known time-free reparameterization. That lets them carry the whole information-theoretic argument over to masked models, producing the I-MDCE relations and matching NLL decomposition.

To make this usable, they remove the time integral entirely. Instead of integrating over noise levels, they sample a random subset of tokens to leave unmasked and average a cross-entropy over the remaining positions, scaled by the harmonic number of the sequence length. The same trick extends to conditional likelihoods given a context, and a shared-mask coupling turns the difference of two log-likelihoods into a low-variance ratio estimate.

They then validate empirically in three tiers: toy problems with analytically known distributions, variance comparisons against baseline estimators, and auditing experiments on the open-source LLaDA model.

Why This Matters

Impact on research. The paper repositions two widely used discrete diffusion training losses — DSE and DCE — from "convenient upper bounds" to "exact likelihood estimators." That change in status matters for any downstream method that depends on comparing or subtracting likelihoods, such as likelihood-ratio-based alignment and reward modeling, where bound-based estimates give no guarantees. It also supplies a discrete counterpart to a body of Gaussian diffusion theory that has driven practical likelihood estimation there.

Real-world applications:

• Language model auditing. Detecting out-of-distribution inputs and distributional shifts in deployed masked diffusion language models, as demonstrated on LLaDA.
• Model influence analysis. Quantifying how training data affects model behavior, which the paper demonstrates on LLaDA.
• Prompt–response scoring. The conditional likelihood estimator targets structured generative settings where a response must be scored given a context, relevant to ranking, filtering, and evaluation.
• Sequence modeling in scientific domains. The toy experiments use DNA sequences with Markov-chain ground truth, suggesting use in genomic or other symbolic sequence analysis where calibrated likelihoods matter.

Industry relevance. Masked diffusion language models are an active alternative to autoregressive models, and practitioners need cheap, reliable likelihood estimates for evaluation, safety filtering, and preference optimization. A time-free estimator that avoids diffusion-time integration and a coupled ratio estimator that reduces variance translate directly into lower compute for the same accuracy. The release of code at https://github.com/Dongjae0324/infodis lowers the barrier to adoption.

Future Directions

1. Scaling the estimator to large models and long sequences. The reported real-world experiments use LLaDA; whether the variance and accuracy advantages hold at much longer sequence lengths and in larger-scale deployment settings is not established in the provided text.

2. Tighter characterization of the residual KL term. The NLL decompositions carry a KL residual that vanishes only in the full-noise limit (or at the full-noise level in the masked case). Quantifying how large this term is at finite noise levels, and when it can be safely dropped, remains an open question.

3. Extending beyond absorbing masked processes. The I-MDCE analysis is developed specifically for absorbing/ masked diffusion. Whether analogous identities hold for the uniform transition matrix, which the paper notes is the other case with closed-form transition matrices, is not addressed.

4. Connecting more tightly to alignment and preference optimization. The authors point to recent alignment methods based on likelihood ratios in masked diffusion language models and argue their exact decomposition explains empirical stability. Turning that observation into a rigorous theory of when ratio-based training is guaranteed to behave well would be a natural next step.

Target Audience

This paper is aimed at machine learning researchers working on diffusion models, discrete generative models, and language model theory — particularly those interested in the theoretical foundations of training objectives and in likelihood-based evaluation. It will also be useful to practitioners building or auditing masked diffusion language models who need calibrated likelihoods or likelihood ratios, and to information theorists interested in discrete analogues of I-MMSE-style identities. The mathematical content assumes comfort with continuous-time Markov chains, KL divergence, mutual information, and score-based generative modeling, so readers without that background will find the appendices and proofs (referenced as Appendix C and D, though not included in the provided text) necessary for full understanding.

Authors’ abstract

We present an information-theoretic framework for discrete diffusion models that yields principled estimators of log-likelihood using score-matching losses. Inspired by the I-MMSE identity for the Gaussian setup, we derive analogous results for the discrete setting. Specifically, we introduce the Information-Minimum Denoising Score Entropy (I-MDSE) relation, which links mutual information between data and its diffused version to the minimum denoising score entropy (DSE) loss. We extend this theory to masked diffusion and establish the Information-Minimum Denoising Cross-Entropy (I-MDCE) relation, connecting cross-entropy losses to mutual information in discrete masked processes. These results provide a time-integral decomposition of the log-likelihood of the data in terms of optimal score-based losses, showing that commonly used losses such as DSE and DCE are not merely variational bounds but tight and principled estimators of log-likelihood. The I-MDCE decomposition further enables practical extensions, including time-free formula, conditional likelihood estimation in prompt-response tasks, and coupled Monte Carlo estimation of likelihood ratios. Experiments on synthetic and real-world data confirm the accuracy, variance stability, and utility of our estimators. The code is publicly available at https://github.com/Dongjae0324/infodis.