Information Estimation with Discrete Diffusion

AI summary

Overview

Research area: Information theory and machine learning — specifically, neural estimation of information-theoretic quantities (mutual information, Kullback–Leibler divergence, entropy) for high-dimensional discrete data.

Technical level: Advanced. The method is built on Continuous Time Markov Chain (CTMC) theory, time-reversed diffusion processes, Dynkin's formula, and score-based generative modeling, so the derivations assume familiarity with diffusion models and variational inference.

Scope: The paper introduces info-sedd, a discrete-diffusion-based family of estimators for mutual information, KL divergence, and entropy over discrete data, and evaluates it on synthetic benchmarks, text summarization, and genomics.

What This Paper Is About

Estimating mutual information (MI) on real-world discrete data — text, DNA — is hard: classical discrete estimators degrade quickly as dimensionality grows, and the common workaround is to embed discrete data into a continuous space and apply neural estimators designed for continuous distributions (the "embedding trick"), which requires careful engineering and loses the discrete structure. The authors propose info-sedd, which instead works directly in the discrete domain by adapting discrete diffusion (CTMC) models, whose learned score functions encode the probability ratios needed to compute KL divergences. The goal is an estimator that is lightweight, scalable, integrates with pretrained models, and is accurate in the high-MI, high-dimensional regimes where existing estimators fail.

Key Contributions

1. A CTMC-based KL estimator for discrete distributions. Using Dynkin's lemma, the authors convert the KL divergence between two discrete distributions into an expectation over a single forward diffusion process, where the unknown probability ratios are replaced by parametric score functions trained with the Diffusion Weighted Denoising Score Entropy (DWDSE) loss.
2. Two MI estimators plus an entropy estimator. info-sedd-j expresses MI as KL between the joint distribution and the product of marginals; info-sedd-c expresses MI as the expected KL between the conditional and the marginal; info-sedd-h computes entropy as log N − KL[p || u] against the uniform distribution.
3. A single-model design via an absorbing rate matrix. By choosing a sparse rate matrix with an absorbing state (each subcomponent can only jump to the absorbing state ∅, with a shared |χ| × |χ| local rate matrix), the marginal score ratios can be read off from a model trained only on the joint distribution, so MI across subsets of variables needs just one trained parametric model.
4. A theoretical error decomposition and a broad empirical evaluation. The authors bound the estimator's deviation from the true KL by an estimation-error term scaling with the score approximation errors and a truncation bias that decays exponentially in the probability of the absorbing state, and they benchmark against gan-dime, hd-dime, kl-dime, minde, mine, nwj, and smile.

Main Findings

• High-MI, high-dimensional synthetic benchmark: on Gaussian-style synthetic data with MI = 10, 20, 30, 40, 50 and dimensionality D = 10, 20, 30, 40, 50 respectively, info-sedd returned 9.92 ± 0.12, 20.02 ± 0.21, 29.83 ± 0.54, 39.11 ± 0.65, and 47.77 ± 1.18 (mean ± standard deviation over 10 seeds) — the closest estimates to ground truth, marked best in bold in the paper's Table 1. gan-dime overestimated at low MI (12.15 ± 0.89 at MI = 10) and collapsed at high MI (17.27 ± 1.46 at MI = 50); hd-dime, kl-dime, mine, and nwj increasingly underestimated; minde was consistently high but noisy (14.01 ± 2.91 at MI = 10, 32.60 ± 3.93 at MI = 50).
• Robustness and efficiency beyond accuracy: an ablation over support size |χ| (Section C.1.6) shows info-sedd is robust to changes in |χ| where competitors cannot provide accurate estimates; convergence experiments (Section C.1.3) show competitors such as gan-dime and smile need more epochs even when they eventually produce correct values; and sample-complexity experiments (Section C.1.5) show info-sedd remains accurate with only 10³ samples.
• Consistency on text summarization: on SummEval-derived text/summary pairs where BART summaries are paired with the original text with probability ρ (for ρ = 0.0, 0.1, …, 1.0), both info-sedd variants had the highest consistency among all estimators, closely matching the reference trend derived from entropy-rate estimates of 256 nats (Takahira et al., 2016) and 303 nats (Cover and King, 1978). info-sedd-c gave MI estimates closer to zero than the joint variant at ρ = 0.0. Variational estimators were capped by the logarithm of the batch size, and minde produced no meaningful estimates.
• MI as a model-selection signal: correlating MI estimates with SummEval human metrics (Pearson / Kendall's Tau), info-sedd-c reached 0.740 / 0.505 on consistency, 0.679 / 0.134 on fluency, 0.411 / 0.200 on relevance, 0.568 / 0.219 overall, and 0.209 / 0.105 on coherence — the strongest correlations overall. info-sedd-j followed (0.550 / 0.486 on consistency), then kl-dime (0.214 / 0.429 on consistency) and hd-dime (0.331 / −0.029); smile was negatively correlated on most metrics. The paper attributes the consistency result to the fact that consistency measures shared information, while coherence does not describe a source–summary relationship.
• Low-MI genomics consistency: on the "human vs. worm" dataset (10⁵ DNA sequences, 2 labels, median length 200 bases), info-sedd-c most closely tracked a classifier-based reference MI (built from H(Y) − H_b(Acc.)), while hd-dime performed relatively well only for ρ < 1.0 and the other competitors failed to keep MI estimates rising with ρ. The authors explain the gap between the two variants: info-sedd-c only parametrizes the label marginal and the label-given-DNA conditional, whereas info-sedd-j must diffuse over the full DNA sequence plus label.
• Ising-model entropy: the appendix reports that info-sedd produces accurate entropy estimates on Ising models.
• Motif discovery in promoters: the authors describe extending a randomized technique from Umarov and Solovyev (2017) to locate the TATA-box motif in Arabidopsis thaliana tata-promoter/non-promoter sequences (1497 and 2879 samples, each of fixed length 251 bases); the supplied excerpt ends before the numerical motif-location results are reported.

Methodology in Plain English

The starting point is a discrete diffusion process: take clean discrete data and gradually corrupt it — here, by letting each element of a sequence randomly jump to a special "absorbing" blank state — until the data becomes uninformative noise. Such a process is a Continuous Time Markov Chain, and the reverse process is also a Markov chain whose behaviour is governed by a "score," i.e., a set of learned probability ratios. The authors train a network to approximate those scores.

Their key move is to run two chains that share the same corruption mechanism but start from two different distributions, say p and q. Applying Dynkin's formula to the log-ratio of the two reverse-time distributions turns the KL divergence between p and q into an average, over the forward process, of quantities that depend only on the score functions. That gives a KL estimator that is computed by simple Monte Carlo: sample a time uniformly, simulate the corruption, and evaluate the network.

To make this tractable, they do not build a full transition matrix over the astronomically large state space. Instead they use a sparse design where only one subcomponent of the sequence can change at a time, governed by a small shared local rate matrix. Choosing this local matrix to be "absorbing" has a second payoff: the marginal score ratios can be recovered from a model trained only on the joint distribution, so MI between any subset of variables needs just one trained model. Because MI is a KL divergence, they obtain two flavours — joint versus product-of-marginals (info-sedd-j) and conditional versus marginal (info-sedd-c) — and entropy falls out by taking q to be uniform.

The theory section shows the estimator's error splits into a term proportional to the neural approximation error and a truncation bias that vanishes exponentially as the process approaches the absorbing state.

Why This Matters

The paper offers a way to compute information-theoretic quantities on exactly the kind of data where existing tools are weakest — long, high-dimensional discrete sequences — without having to design an embedding model, and it can be plugged into pretrained discrete diffusion models.

Real-world applications the paper identifies or demonstrates:

• Genomics and motif discovery: locating regulatory motifs such as the TATA-box in promoter sequences, and quantifying the dependency between motif presence and gene expression, a long-standing approach in transcription-regulation research.
• DNA and peptide sequencing analysis: the paper names sequencing as a domain that would benefit from scalable MI estimators.
• Text summarization model selection: using MI between source documents and generated summaries as a reference-free, human-aligned signal for choosing among summarization models.
• Neuroscience: the paper lists neuroscience among the fields needing scalable MI estimation on high-dimensional discrete data.
• Physics: entropy estimation in Ising models, used to characterize lattices of interacting spins.

Industry relevance: reference-free evaluation of generative text models is expensive and often requires learned metrics with large embedding models; an MI-based score that correlates with human consistency judgments and requires no elaborate embedding table is directly relevant to teams evaluating summarization or other conditional-generation systems. The method's compatibility with pretrained backbones (MDLM-small for text, CADUCEUS for DNA) means it can be added to existing checkpoints rather than trained from scratch.

Future Directions

• Extending beyond pairwise MI. The current estimators target MI between two variables, with a note that the absorbing-state construction allows MI across arbitrary subsets; systematic evaluation of multi-variable interactions is a natural next step.
• Closing the gap between the joint and conditional variants. The paper shows info-sedd-j lags info-sedd-c when one variable is high-dimensional (the DNA sequence), because it must diffuse over both modalities jointly; making the joint variant as easy to optimize is an open problem.
• Reducing the truncation bias / tuning the time horizon. The theoretical bound contains a truncation term controlled by how close the process gets to the absorbing state by time T, so the choice of schedule and horizon remains a tunable accuracy–cost trade-off.
• Broadening the downstream application set. The authors position their results as paving the way to use information measures as signals in other downstream tasks on high-dimensional discrete data, beyond the summarization and promoter-motif cases studied here.

Target Audience

Researchers and practitioners in machine learning and information theory who need to estimate MI, KL divergence, or entropy on high-dimensional discrete data (text, DNA, peptide sequences), and who are comfortable with diffusion models and score-based generative modeling. It is also relevant to applied scientists in computational biology and NLP evaluation who want a practical estimator that can be attached to a pretrained discrete diffusion backbone rather than requiring a bespoke embedding pipeline. Readers without a background in CTMCs or variational bounds will find the methodology and experimental sections accessible, but the theory in Section 2 and Appendix E assumes graduate-level probability.

Authors’ abstract

Information-theoretic measures, such as Mutual Information (MI), play a crucial role in understanding non-linear relationships between random variables and are widely used across scientific disciplines. Yet, their use on real-world discrete data remains challenging. Existing methods typically rely on embedding discrete data into a continuous space and apply neural estimators originally designed for continuous distributions. This process requires careful engineering for both the embedding model and estimator architecture, but suffers from issues related to high data dimensionality. In this work, we introduce InfoSEDD, a discrete diffusion-based approach that bridges information-theoretic estimation and generative modeling such that they can be used to compute Kullback-Leibler divergences. Backed by Continuous Time Markov Chains theory principles, the design of InfoSEDD is lightweight and scalable and allows seamless integration with pretrained models. We showcase the versatility of our approach through applications on motif discovery in genetic promoter data, semantic-aware model selection in text summarization, and entropy estimation in Ising models. Finally, we construct consistency tests on real-world textual and genomics data. Our experiments demonstrate that InfoSEDD outperforms alternatives that rely on the "embedding trick". Our results position InfoSEDD as a robust and scalable tool for information-theoretic analysis of discrete data.