Online Class Posteriors for Interpretable and Controllable Discrete Diffusion

Felix Koulischer; Gianni Van De Velde; François Remy; Chris Develder; Thomas Demeester

cs.LGNov 2, 2026

Online Class Posteriors for Interpretable and Controllable Discrete Diffusion

Authors: Felix Koulischer, Gianni Van De Velde, François Remy, Chris Develder, Thomas Demeester

Organizations: Ghent University - imec · Parallia Healthcare

Abstract

Discrete diffusion models remain far less understood than their continuous counterparts, despite recent breakthroughs across modalities, most notably in text. Leveraging tools from continuous diffusion, we formulate a theoretical framework to track the class-conditional posterior in an online manner by combining conditional and unconditional model predictions. In masked language modeling, this posterior offers highly interpretable dynamics: tokens aligned with a class boost the posterior, while contradictory ones suppress it. We also show that the posterior can be used for logical guidance operations (AND, OR, NOT) in discrete diffusion. These operations enable precise control over generation, such as avoiding specific classes. In addition, we discuss a mismatch in the formulation of classifier-free guidance between discrete and continuous settings: discrete models mix the data distributions, while continuous models mix the marginals. We validate our methods and discuss their potential applications on synthetic data, with all quantities analytically tractable, as well as on MNIST, text, and antibody sequences.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of the online posterior evaluation. In (a), evaluated on an analytically tractable setting in which data points consist of two tokens represented by blue and red with each class having an 80% majority of one of the two tokens, trajectories are shown above the posteriors. In (b), the posterior is evaluated on MNIST digit generation. On top for uniform, below for masked diffusion. The increase in the posterior occurs simultaneously with the emergence of the digit.

Panel of Figure 2 — Figure 2: In (a), the progress of the posterior approximation is shown throughout generation of four sentences on “ beverage ” with as class “ wine ”. Highlighted words show large jumps in the posterior, corresponding to those given in Table 1 . In (b), one-vs-rest AUC is evaluated throughout generation, using a modality-specific external classifier to determine the final class.

Topic	Sub-topic		1st		2nd		3rd
	(class c)		Token	$\Delta\log p_{c}$	Token	$\Delta\log p_{c}$	Token	$\Delta\log p_{c}$
beverage	wine	$+$	wine	+2.3	cab ernet	+1.3	red	+1.1
beverage	wine	$-$	syrup	-5.5	ginger	-3.9	espresso	-3.2
sweet	candy	$+$	candy	+2.8	lollip op	+0.8	are	+0.5
sweet	candy	$-$	mar zip an	-2.1	cookies	-0.7	sweet	-0.6
dairy product	cheese	$+$	cheddar	+2.4	Swiss	+1.4	cheese	+0.5

Table 1: Tokens pushing most towards (

+

) and away from (

-

) the class, estimated by our posterior estimation. (Sub)words in gray provide context for the token shown in black. Each gray context corresponds to the actual context for that token in all samples containing that token.

Panel of Figure 3 — Figure 3: Illustration of the not and or gates on MNIST under a uniform noise model. (a) The not gate, showing how unconditional samples containing the undesired number “1” are transformed into other, predominantly round numbers such as “0” and “9” . (b) The or gate evaluated on two distinct pairs, with corresponding samples shown above.

Figure 4: Pareto curves comparing the quality of the generated samples to the safety of the approach, respectively through the FID and the violation rate. Both negative CFG, as well as the theoretically justified dynamic not gate display a trade-off between quality and safety, with the dynamic approach outperforming the static one, albeit at the cost of an additional hyperparameter.

Method	Setting	Majority		Minority
		Viol. $\downarrow$	FPD $\downarrow$	Viol. $\downarrow$	FPD $\downarrow$
Unconditional		0.278	0.174	0.032	0.007
Negative CFG	majority-calibrated	0.013	0.104	0.002	0.326
Negative CFG	minority-calibrated	0.186	0.084	0.017	0.020
not gate	both	0.010	0.074	0.016	0.010

Table 2: not gate performance when removing majority and minority V-gene families. Static guidance is calibrated to match the dynamic method’s violation rate in the corresponding regime, then transferred across regimes. Gray entries show performance outside the calibration regime. The full Pareto front is given in Section F.5 .

Appendix figures & tables33 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Illustration of the non-commutativity of noising and mixing in the case of a multiplicative mixture. Continuous CFG first noises the data distribution and mixes the marginals, whereas discrete CFG first mixes the distributions and then only noises.

class	$p_{1}(x_{1}^{j}=A\|c)$ for $j\in\mathcal{H}_{1}$	$p_{1}(x_{1}^{j}=A\|c)$ for $j\in\mathcal{H}_{2}$
$A$	$p_{\mathrm{A}}$
$B$	$1-p_{\mathrm{A}}$
$A_{1}$	$p_{\mathrm{A}}$	$1/2$
$B_{1}$	$1-p_{\mathrm{A}}$	$1/2$
$A_{2}$	$1/2$	$p_{\mathrm{A}}$
$B_{2}$	$1/2$	$1-p_{\mathrm{A}}$

Appendix

Table 3: Synthetic data construction. The sequences are split into two halves

\mathcal{H}_{1},\mathcal{H}_{2}

of 10 tokens each, filled in with a majority token whose strength is set by a tunable parameter

p_{\mathrm{A}}

. For two classes both halves share the same majority likelihood. For four classes only one half is dominated by a token while the other is uniform, which increases the overlap between classes and makes logical operations such as and more interesting.

Panel of Figure 6 — Figure 6: Illustration of the online posterior evaluation under different samplers and evaluators. In (a), samples are taken from an ancestral sampler at the top and a dfm sampler at the bottom and the posterior likelihoods are updated using the ancestral transition likelihoods. In this case the tracked posterior is equal to the ground truth posterior. In (b), the exact same samples are shown but with this time the DFM transition likelihoods used for the posterior estimation, in which case the tracked quantity is not equal to the exact one anymore.

Figure 7: Evolution of the KL divergence between the distribution of

(n_{1},n_{2})

of the noisy state

x_{t}

and that of the clean target, for and

(A_{1},B_{2})

, or

(A_{1},B_{1})

and not

(A_{1})

, under masking (full) and uniform (dotted) noise. In all cases the derived gates sample from the desired target distribution, illustrated by the KL distance between generated and reference samples going to zero.

Architecture	DiT with AdaLN-Zero
Parameters	7.6M
Depth / hidden size / heads	6 / 256 / 4
Sequence length / vocabulary	784 / 256 (+ mask token)
Optimizer	AdamW
Learning rate	$4\cdot 10^{-4}$
Betas / weight decay	(0.9, 0.999) / 0.01

Appendix

Table 4: Hyperparameters of the MNIST models. The unconditional models are trained separately with label dropout 1.0, i.e., without ever seeing the label. The cosine decay is defined over a horizon longer than the training run, so that the learning rate only decays to about

3\cdot 10^{-4}

at the end of training.

Figure 8: Validation loss during training of the MNIST models, averaged over bins of noise levels

t\in[0,0.1],[0.1,0.2],\dots

, for masking (left) and uniform noise (right).

Figure 9: Sweep over the guidance scale for the MNIST-trained model. CFG is implemented as usual in the discrete case, i.e., a linear combination of logits. During inference 128 steps of the specified schedulers were taken. The stars correspond to the separately trained unconditional model.

Figure 10: Sweep over the number of sampling steps for the unconditional MNIST-trained model. Masking and uniform models are evaluated under three schedulers; sine (fast at first and then slow), linear (constant sped) and cosine (slow at first and then fast). From left to right, the FID, the Precision and Recall and the C2ST accuracy.

Figure 11: Sweep over the number of sampling steps for the conditional MNIST-trained model. Masking and uniform models are evaluated under three schedulers; sine (fast at first and then slow), linear (constant sped) and cosine (slow at first and then fast). From left to right, the FID, the Precision and Recall and the C2ST accuracy.

Panel of Figure 12 — Figure 12: Illustration of the online posterior evaluation. In (a-c), trajectories for which the online posterior estimate recovers the final class for both masking and uniform noise are shown. In (d), an example is shown for which the masking model does not recover the correct class label. All shown trajectories use $\tau=0.1$ , a value taken from a calibration run on the not gate.

Panel of Figure 13 — Figure 13: Evaluation of the online posterior estimation for different modalities during denoising. In (a), the OvR AUC. In (b), the accuracy.

Figure 14: Static negative guidance (negative prompting): FID and violation rate as a function of the constant guidance scale

\lambda

, for not (0), not (1) and not (2).

Panel of Figure 15 — Figure 15: FID and violation rate of the dynamic not gate as a function of the initial guidance scale and the inverse temperature $\tau$ , for not (0), not (1) and not (2) (top to bottom). Amber is best, blue is worst.

Figure 16: Dynamic negative guidance with

\tau=0.01

: FID and violation rate as a function of the initial guidance scale, for not (0), not (1) and not (2).

Figure 17: FID against violation rate for dynamic negative guidance (full) and the static baseline (dashed), for not (0), not (1) and not (2).

Figure 18: Guidance scale of the dynamic not gate over the sampling process (

\tau=0.01

, initial guidance scale 1): single trajectories (black) and their mean.

Panel of Figure 19 — Figure 19: Neutral, unconditional prompt

not gate (prob-space)	$p=p_{\text{neutral prompt}}-\lambda_{t}\cdot(p_{\text{neg prompt}}-p_{\text{neutral prompt}})$
Negative CFG (logit-space)	$l=l_{\text{neutral prompt}}-\lambda\cdot(l_{\text{neg prompt}}-l_{\text{neutral prompt}})$
Unconditional	$p=p_{\text{neutral prompt}}$
Positive prompted	$p=p_{\text{pos prompt}}$

Appendix

Table 5: Contrastive signals used in this paper.

Panel of Figure 20 — Figure 20: Negative prompt

Panel of Figure 21 — Figure 21: Posterior trajectories with different prompt types.

Panel of Figure 22 — Figure 22: ROC curves for classification of correct vs violated using different prompt types.

Figure 23: The posterior of multiple classes during generation of the prompt with ID 76 in the development set of Cognac. The answer of the model is “== Cabernet Sauvignon is a type of red wine made from the Cabernet Sauvignon grape variety. ==”. Some tokens are highlighted with a dashed line at the time step of generation. The blue line, which represents the wine class, clearly jumps on the words “wine”, “variety” and “red” being generated. On the other hand, the posteriors do not move much upon “a” or “the” being generated.

ID	Topic	Classes (%)	Accuracy
76	alcohol	wine 57%, brandy 18%, tequila 26%	99.0%
276	beverage	wine 55%, beer 28%, liqueur 17%	95.1%
353	liquor	gin 65%, brandy 23%, rum 12%	95.2%
418	domestic animal	cat 16%, dog 84%	95.6%

Appendix

Table 6: Accuracy of the posterior estimation, classifying which (manually defined) subtopic was generated using argmax of the posterior. The IDs refer to the Cognac dataset’s development set.

Topic	Sub-topic		1st		2nd		3rd
	(class c)		Token	$\Delta\log p_{c}$	Token	$\Delta\log p_{c}$	Token	$\Delta\log p_{c}$
sweet	candy	$+$	candy	+2.8	lollip op	+0.8	are	+0.5
sweet	candy	$-$	mar zip an	-2.1	cookies	-0.7	sweet	-0.6
dairy product	cheese	$+$	cheddar	+2.4	Swiss	+1.4	cheese	+0.5
dairy product	cheese	$-$	yogurt	-6.8	cream	-0.5	product	-0.5
dairy product	milk	$+$	milk	+0.8	with	+0.4	whole	+0.4

Appendix

Table 7: Tokens pushing most towards (

+

) and away from (

-

) the specified class, as estimated by our posterior estimation. (Sub)words in gray are context to put the token (black) in context for the reader. Each of there gray contexts were the actual context for that token for 100% of the samples with that token.

Method	Correct (Other classes)	Violated (Class)	Invalid
Unconditional	$0.70\,{\scriptstyle\pm 0.03}$	$0.23\,{\scriptstyle\pm 0.03}$	$0.07\,{\scriptstyle\pm 0.02}$
Prompted “ not ”	$0.79\,{\scriptstyle\pm 0.03}$	$0.10\,{\scriptstyle\pm 0.02}$	$0.11\,{\scriptstyle\pm 0.02}$
Negative CFG (logit-space)	$0.81\,{\scriptstyle\pm 0.02}$	$0.06\,{\scriptstyle\pm 0.01}$	$0.13\,{\scriptstyle\pm 0.02}$
not gate (prob-space)	$0.79\,{\scriptstyle\pm 0.03}$	$0.08\,{\scriptstyle\pm 0.02}$	$0.14\,{\scriptstyle\pm 0.02}$

Appendix

Table 8: Results of the not on the Cognac benchmark

Figure 24: Grid search on the development set of the Cognac benchmark. The optimal hyperparameters that were used for the final experiment are indicated with a bold box. We found that the text modality needs the $\delta$ from Equation 10 . We use $\Delta$ to indicate the cumulative $\delta$ over all steps. The first heatmap shows that $\tau$ is a little less sensitive to tune than the initial guidance scale. The second heatmap shows a downwards diagonal ridge, which seems to indicate that high initial guidance combined with high delta can make the guidance scale explode and yield poor results. The third heatmap again shows less sensitivity for $\tau$ , and larger sensitivity for $\Delta$ .

V-gene family	IGHV1	IGHV2	IGHV3	IGHV4	IGHV5	IGHV6	IGHV7
Share of sequences (%)	24.9	2.5	44.1	25.3	2.2	0.6	0.5

Appendix

Table 9: Class balance of the antibody dataset (human heavy chains of OAS95, measured on the held-out validation split of 1.1M sequences). The conditional model is trained on these seven classes.

Architecture	AMPLIFY with AdaLN-Zero
Parameters	50M
Depth / hidden size / heads	12 / 480 / 20
Sequence length / vocabulary	256 / 27 (incl. mask token)
Optimizer	AdamW
Learning rate	$1\cdot 10^{-4}$
Betas / weight decay	(0.9, 0.95) / 0.01

Appendix

Table 10: Hyperparameters of the antibody models.

Figure 25: Validation loss during training of the antibody models, averaged over bins of noise levels

t\in[0,0.1],[0.1,0.2],\dots

, for masking (left) and uniform noise (right).

Figure 26: Sweep over the number of sampling steps for the uniform-noise antibody model, for the unconditional (label dropped) and the conditional model, with a linear sampling schedule and 4096 sequences per point. The uniform model in full and the masking model in dashed lines. Shown are, from left to right, the FPD, the gap between the IgLM log-likelihood of the generated and of real sequences, and the C2ST accuracy. Class labels for the conditional model are drawn from the empirical class frequencies.

Figure 27: Negative CFG with static negative logit guidance. On the left, the Frechet Protein Distance measuring distributional coverage and sample quality, and on the right, the violation rate. Notice the strong differences between majority (IGHV1-4) and minority classes (IGHV5-7).

Figure 28: not gate with dynamic negative guidance. On the left, the Frechet Protein Distance measuring distributional coverage and sample quality, and on the rate the violation rate. Here, thanks to the prior weighted initial guidance scale, minority classes are much less harmed when applying guidance.

Figure 29: Pareto front between quality (FPD) and safety (violation rate) for majority and minority classes. Negative CFG with static logit guidance is compared to the not gate with dynamic negative guidance.

Discrete Diffusion Sampling

See all 20 papers

Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction
Jun 4, 2026Hongkun Dou, Zike Chen, Fengji Li +2Diffusion Model Guidance Molecular Generation
Beihang University, Beijing, China · Zhongguancun Academy, Beijing, China.
DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling
May 22, 2026Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3Diffusion Language Models Diffusion Language Model Distillation
NVIDIA
Interpolating Discrete Diffusion Models with Controllable Resampling
Apr 19, 2026Marcel Kollovieh, Sirine Ayadi, Stephan GünnemannDiffusion Model Sampling Molecular Graph Generation
School of Computation, Information and Technology, Technical University of Munich

Jun 4, 2026J/K move · Enter open · S save

Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction

Hongkun Dou, Zike Chen, Fengji Li, Hongjue Li+1

Beihang University, Beijing, China · Zhongguancun Academy, Beijing, China.

Diffusion Model Guidance Molecular Generation Diffusion Models

Open paper AskPDF arXiv

Controllable generation with discrete diffusion models is often hindered by high computational overhead or the need for retraining. In this paper, we present Gradient-Informed Logit Correction (GILC), a plug-and-play framework that efficiently estimates guidance signals by repurposing the pretrained denoising network as a variational proxy. To circumvent the gradient instability inherent in high-dimensional discrete spaces, we introduce a Jacobian-free mechanism that directly corrects the clean prediction logits, facilitating stable and effective guidance. Our method accommodates both differentiable and non-differentiable reward functions. Extensive experiments across DNA, protein sequence, and molecular generation tasks demonstrate that GILC achieves state-of-the-art performance without additional training, frequently outperforming fine-tuning approaches.

Abstract

Discrete Diffusion Sampling

Explore similar work

Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling

Interpolating Discrete Diffusion Models with Controllable Resampling