Mean-to-Score Discrete Diffusion: Posterior-Mean Denoisers for Score Entropy

2026-07-23

AI Summary

This work addresses a fundamental limitation in discrete diffusion models, where conventional score estimation violates Bayes realizability, leading to negative sampling weights and degraded generation quality. To resolve this, the authors propose the Mean-to-Score (M2S) method, which first predicts the posterior mean of the clean token and then applies a kernel-dependent linear mapping to project it precisely onto the bridge polytope. This ensures that the resulting score function always corresponds to a valid posterior distribution, inherently satisfying the Bayes realizability constraint across a broad class of continuous-time Markov chain (CTMC) noise processes. Empirical results demonstrate significant improvements: on CIFAR-10, the method achieves a bits-per-dim (BPD) of 3.129 and markedly better FID scores; on 262B OpenWebText tokens, it attains a perplexity (PPL) of 143.3 with only 128 sampling steps, outperforming existing baselines.

Abstract

Score Entropy Discrete Diffusion (SEDD) parameterizes discrete reverse processes with unconstrained positive score ratios. While positivity guarantees nonnegative reverse jump rates, it does not ensure Bayes realizability: ratios at a noisy state need not be jointly induced by any clean-token posterior under the forward kernel. The score-entropy loss has the correct population optimum but does not enforce this constraint away from it. In a trained pure-uniform SEDD checkpoint, roughly one quarter of complete score vectors violate the coordinate box, while more than half lie inside it yet remain materially incompatible with any valid posterior. Such violations can produce negative pre-normalization weights in finite-step sampling. Projecting raw scores onto the bridge polytope removes all observed negative weights and improves external generative PPL from $203.6$ to $175.1$ without changing the sampler. We introduce \emph{mean-to-score} (M2S), which predicts a clean-token posterior mean and converts it to the score through an exact kernel-dependent linear map. The construction applies to any known coordinate-wise continuous-time Markov chain (CTMC) satisfying a mild support condition. For uniform corruption, it maps the probability simplex onto the bridge polytope; for absorbing-mask corruption, the resulting objective recovers MD4 exactly. In a controlled 28.4M-parameter CIFAR-10 comparison, M2S lowers test BPD from $3.173$ to $3.129$ and FID-50k from $\CifarSEDDFID$ to $\CifarMtwoSFID$. A 170M-parameter M2S model trained on about 262B OpenWebText token slots outperforms the evaluated pure-uniform SEDD, GIDD, and Neural CTMC checkpoints at every tested sampling budget, reaching generative PPL $143.3$ at 128 steps versus $183.6$ for the strongest pure-uniform baseline.

Problem

• discrete diffusion
• score entropy
• Bayes realizability
• posterior consistency
• reverse process

Innovation

• Mean-to-Score
• Discrete Diffusion
• Bayes Realizability
• Posterior-Mean Denoiser
• Score Entropy