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

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

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

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

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

Hongkun Dou 1 Zike Chen 1 Fengji Li 1 2 Hongjue Li 1 Yue Deng 1 2 †

arXiv:2606.06303v1 [cs.LG] 4 Jun 2026

Abstract 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-ofthe-art performance without additional training, frequently outperforming fine-tuning approaches.

An intermediate step Sampling

M A

A C

Mask

T predictor M

A G C T

Clean prediction

Reward guidance

G A G C T

Guided results

Figure 1. Illustration of the guided discrete diffusion process by GILC. The reverse sampling process (top) iteratively denoises a DNA sequence from the fully masked state (t = 1) to the clean data (t = 0). The core correction mechanism (bottom) operates at each step: the mask predictor outputs the clean prediction xθ , which are then modified by the reward gradient, r(·), yielding the guided prediction xrθ . The next state zs is sampled from the transition distribution prθ (zs |zt ) = q(zs |zt , xrθ ).

1. Introduction Diffusion models (Ho et al., 2020; Song et al., 2021) have task objectives, moving beyond unconditional generation. become the gold standard for probabilistic modeling in conExamples include generating proteins with enhanced stabiltinuous domains, as powerfully demonstrated in image and ity (Widatalla et al., 2024) or novel molecules with specific video generation (Rombach et al., 2022; Esser et al., 2024; biochemical properties (Chang & Ye, 2024). Such tasks Liu et al., 2024). Correspondingly, recent studies (Austin necessitate conditional or controllable generation, where M M M M G G et al., 2021; Lou et al., 2023; Shi et al., 2024; Sahoo et al., classical approaches typically employ classifier/classifierM A A A A A 2024) indicate that diffusion models also hold significant free guidance (Schiff et al., 2025; Nisonoff et al., 2025) M M T T potential in discrete spaces, spanning domains suchTas lan- Tor extensive model fine-tuning (Venkatraman et al., 2024; M M C C C guage generation (Nie et al.,M2025), biological sequence Rector-Brooks et al., 2025; Wang et al., 2025). However, M M M M M synthesis (Campbell et al., 2024), and molecular design Cthese established methods demand additional training stages An intermediate step (Vignac et al., 2023). While discrete diffusion effectively (e.g., training a separate classifier or retraining the generacapture complex categorical data distributions, manySampling scien- A tive model), which fundamentally limit their generalizability M tific and Aindustrial applications require generated samples C and introduce significant computational overhead. Mask that are Tnotpredictor only realistic but also optimized for specific G An intriguing alternative is the plug-and-play paradigm, Reward A G C T A G C T M 1 guidance Beihang University, Beijing, China 2 Zhongguancun T which aims to guide generation using off-the-shelf reward Clean prediction Guided results Academy,M Beijing, China. Correspondence to: Yue Deng M functions, neural predictors, or property evaluators with<[email protected]>. out modifying or retraining the generative model (Chung et al., 2023; Yu et al., 2023; Song et al., 2023b; He et al., rd Proceedings of the 43 International Conference on Machine 2024). This flexible, training-free approach has achieved Learning, Seoul, South Korea. PMLR 306, 2026. Copyright 2026 by the author(s). notable results in continuous diffusion. While conceptually 1

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

elegant, these methods struggle to directly translate to discrete diffusion due to the inherent non-differentiability of categorical data. Existing efforts based on Sequential Monte Carlo (SMC) (Wu et al., 2023; Uehara et al., 2024) or importance sampling (Li et al., 2024; Lin et al., 2025; Ou et al., 2026) often suffer from prohibitively high computational costs and demonstrably limited steering effectiveness. For a more comprehensive discussion of related work, we refer the reader to Appendix A.

4. We empirically demonstrate state-of-the-art performance and computational efficiency for constrained generation across complex scientific domains.

2. Preliminaries This section provides the foundational concepts for our work, covering the core principles of discrete diffusion models and the specific task formulation we address for constrained generation.

This paper introduces Gradient-Informed Logit Correction (GILC), an efficient, principled, and training-free method for steering discrete diffusion (Fig. 1). We reframe the guidance task as estimating the gradient of a value function that measures the expected future reward from intermediate states. First, our framework employs a variational method where the pre-trained denoising network serves as a proxy for this value estimation. We then introduce two key practical insights for robust gradient computation: 1) combining the Gumbel-Softmax (GS) trick (Jang et al., 2017) and a Straight-Through (ST) estimator (Bengio et al., 2013) to maintain gradient flow in the discrete space, and 2) utilizing a Jacobian-free update that directly corrects the clean prediction logits for stable and effective guidance. Furthermore, we establish a formal connection between GILC and policy gradients (Williams, 1992) to universally handle nondifferentiable objectives. Crucially, GILC operates entirely on off-the-shelf objective functions, requiring no fine-tuning or auxiliary training.

2.1. Discrete Diffusion Models Diffusion models are a class of powerful generative models that define a forward process to progressively perturb data with noise, and a neural network is then trained to learn the reverse process to generate new data. In the context of discrete diffusion models, we operate on the data represented by one-hot encoded variables x ∈ {0, 1}K , where K is the vocabulary size. These models use the categorical distribution, denoted as Cat(·; π), to model the class probabilities, where π ∈ ∆K and ∆K is the K-simplex. While many discrete models define the forward process via a series of transition matrices (Austin et al., 2021; Lou et al., 2023), we adopt a simpler framework, proposed by Sahoo et al. (2024) and Shi et al. (2024), which directly interpolates between the clean data x and a fixed prior distribution π: q(zt |x) = Cat(zt ; αt x + (1 − αt )π)

We demonstrate the effectiveness of GILC across a broad set of scientific domains, including DNA sequence design, protein sequence engineering, and multimodal molecular generation. Empirically, GILC not only significantly outperforms popular training-free discrete diffusion guidance methods in both sample quality and computational efficiency, but also competitively matches the performance of fine-tuning-based approaches, achieving state-of-the-art results for controlled discrete generation. These results highlight the immense potential of GILC for advancing the field of discrete diffusion guidance.

(1)

where αt is a predefined noise schedule that decreases from 1 as t = 0 to 0 as t = 1. This process implies a tractable transition probability q(zt |zs ) = Cat(zt ; αt|s zs + (1 − αt|s )π) at each step, where αt|s = αt /αs . A key advantage of this design is that it yields a tractable posterior distribution q(zs |zt , x), which is crucial for training and inference: [αt|s zt +(1−αt|s )1π⊤ zt ]⊙[αs x+(1−αs )π] q (zs |zt , x) = Cat zs ; ⊤ ⊤ αt z x+(1−αt )z π t

(2) A particularly useful choice for the prior is an absorbing state, where all data eventually transitions to a single, special token (e.g., a [MASK] token). In this case, we set π = m, a one-hot vector corresponding to the mask token. With this choice, the posterior from the equation above simplifies into two distinct cases: ( Cat (zs ; zt ) zt ̸= m q (zs |zt , x) = (1−αs )m+(αs −αt )x Cat zs ; zt = m 1−αt (3) The first case, zt ̸= m, shows a simple identity transition when the token is unmasked. The second case, zt = m, demonstrates how the model can reverse the masking process to predict the original token. We use a neural network, parameterized by xθ (zt , t), to predict the clean data

In summary, our contributions are as follows: 1. We propose a training-free guidance framework by demonstrating the pre-trained diffusion network’s viability as a variational proxy for plug-and-play value function estimation. 2. We introduce logits correction guidance, which stably steers generation by computing the value gradient directly over the clean prediction logits, effectively resolving gradient instability issues in discrete spaces. 3. We establish a formal connection to the policy gradient formulation, extending the framework’s universality to encompass non-differentiable objectives. 2

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

and then substitute this prediction into the posterior, allowing us to perform the time-reversal sampling process via pθ (zs |zt ) = q (zs |zt , xθ ).

ation. This objective has a well-known closed-form solution: prθ (x) ∝ ppre θ (x) exp (r (x) /β) (Peters & Schaal, 2007; Peng et al., 2019).

When processing a sequence of discrete tokens x1:L (abbreviated as x) of length L, the denoising process is assumed to be factorizable across tokens conditioned on the noisy sequence z1:L (abbreviated as zt ). This yields the t full reverse process as a product of per-token distributions: QL pθ (zs |zt ) = ℓ=1 pθ zℓs |zt . To train the network xθ to predict the original tokens for each position ℓ, we minimize the following negative evidence lower bound (ELBO) objective:

The prevailing approaches to solving this objective typically involve fine-tuning the diffusion model or training a time-dependent classifier. However, these methods have significant drawbacks. Fine-tuning can suffer from reward hacking (Clark et al., 2024). Training a time-dependent classifier, while more flexible, still necessitates collecting new training data for each new target condition, which is both inconvenient and costly (Dhariwal & Nichol, 2021). In contrast to these established methods, our approach proposes a training-free strategy that can be seamlessly integrated with existing discrete diffusion and reward models. This allows for a flexible and efficient generation process without the need to modify model parameters or incur retraining.

Z 1

αt′ X log xℓθ (zt , t) , xℓ dt θ 0 1 − αt ℓ=1 (4) α′t Here, 1−αt is the weighting function, and the expectation is taken with respect to the forward process q(zt |x). This forward process is assumed to be factorizable across tokens, conditioned QL only on the clean token at each position: q(zt |x) = ℓ=1 q(zℓt |xℓ ). The term xℓθ (zt , t) represents the network’s prediction of the clean token at position ℓ, which is conditioned on the entire noisy sequence zt . Notably, this discrete diffusion formulation is equivalent to discrete flow models under the Continuous-Time Markov Chain (CTMC) framework (Campbell et al., 2024). Consequently, the proposed guidance method naturally extends to and can be directly applied within discrete flow models. θ∗ = arg minEx,zt

3. Method In this section, we propose Gradient-Informed Logit Correction (GILC), a general framework for guiding discrete diffusion in a plug-and-play fashion. We first demonstrate that solving the constrained generation objective (Eq. 5) relies on computing the gradient of a value function (Sec. 3.1). We then introduce a practical, training-free estimate for this value function via a variational proxy (Sec. 3.2). Finally, we present two robust approaches for calculating the necessary guidance gradient, including our key insight: omitting the Jacobian of the denoising network to achieve stable guidance by correcting the clean prediction logits (Sec. 3.3 and 3.4).

2.2. Task Formulation We consider a scenario where we have a pre-trained discrete diffusion model, ppre θ (x), trained using a standard objective like the ELBO. While such models excel at capturing the distribution of natural data, many real-world applications require generating data that satisfies specific properties.

3.1. Taylor Series Approximation of the Optimal Reverse Process The core idea behind guided generation is to replace the original, unguided reverse sampling step pθ (zs |zt ) with an optimal reverse process prθ (zs |zt ) that implicitly maximizes the target reward r(x). As shown in prior work (Uehara et al., 2024; Li et al., 2024), this optimal process is given by the following formulation:

Following prior work (Rafailov et al., 2023), we formulate this challenge as finding a new target distribution that balances adherence to a desired property with fidelity to the original pre-trained model. This is achieved by maximizing the following objective: prθ (x) = arg maxEx∼q [r (x)] −β DKL q (x) ||ppre θ (x) | {z } | {z } q Reward

pθ (zs |zt ) exp (v (zs ) /β) prθ (zs |zt ) = P zs pθ (zs |zt ) exp (v (zs ) /β)

(6)

where v(·) is the soft value function at state zs . It is rigorously defined as:

KL Regularization

(5) The first term, the Reward r(·), quantifies how well a data sample x complies with the desired conditional constraints. This term can be modeled by an external function or an auxiliary network. The second term, KL Regularization, ensures that the new distribution remains close to the original pre-trained distribution ppre θ (x). The hyperparameter β controls the trade-off between satisfying the reward and maintaining the quality of the original gener-

v (zs ) = β log Epθ (x|zs ) [exp (r(x)/β)]

(7)

For practical implementation, especially when β is small, this soft value function is commonly approximated by the expected reward: v (zs ) ≈ Epθ (x|zs ) [r(x)]. This approximation, justified in Li et al. (2024), simplifies the value function to the more intuitive concept of expected reward and helps mitigate numerical instability. 3

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

1.2

A significant challenge arises from the denominator in Eq. 6, which requires a summation over all possible states zs . For a sequence with length L, this involves a computationally intractable K L terms. In standard guidance frameworks (Schiff et al., 2025; Nisonoff et al., 2025), this obstacle is typically bypassed by training a parametric network vϕ (·) to explicitly estimate the expected target reward: min Eq(zt |x) (vϕ (zt ) − r(x))

L1 Error

1.0 MC5 MC10

0.6

MC20 MC50

DE (Baseline)

0.4 0.2

(8)

0.0

where vϕ (·) : RK×L → R is currently a continuously differentiable function and it optimal solution vϕ (zs ) ≈ Epθ (x|zs ) [r(x)]. By treating discrete states as a point constrained within the continuous domain of the function (Grathwohl et al., 2021), we can apply a first-order Taylor series expansion to approximate vϕ (zs ) around the current state zt :

Sampling Step

100

120

(a) Convergence of the value estimation error.

Mean Reward

vϕ (zs ) ≈ vϕ (zt ) + ⟨zs − zt , ∇zt vϕ (zt )⟩ L D E X =C+ zℓs − zℓt , ∇zℓt vϕ (zt )

0.8

SVDD (4.84)

The first term on the right-hand side is a constant with respect to zs and can therefore be absorbed into the normalization constant of Eq. 6. The estimation of the optimal policy thus hinges entirely on calculating the value function gradient, which can be calculated efficiently with a single forward and backward pass of vϕ (·). Moreover, the factorized form in the second line of Eq. 9 preserves the tractability of the reverse transition, enabling efficient sampling for high-dimensional discrete sequences. Additional details are provided in Appendix B.1.

Noisy State Guidance Clean Prediction Guidance

(9)

ℓ=1

4.1 + Improvement

6 DRAKES (5.61)

Pretrained (0.17)

Sampling Step

100

120

(b) Comparison of different guidance targets. Figure 2. Analysis of value estimation and guidance targets. (a) Convergence of the L1 error for the value function v(zt ) across varying Monte Carlo (MC) sample sizes. Increasing the sample size (from 5 to 50) consistently minimizes estimation error, significantly outperforming the Deterministic Estimation (DE) approach (Li et al., 2024). (b) Comparison of guidance targets during the denoising process. Guidance optimized in the clean-prediction logit (η) yields higher cumulative rewards and faster convergence compared to noisy-state guidance (zt ). Results are reported on a discrete DNA diffusion model using GILC-DB.

Although this approach is reasonable, it requires retraining the value network vϕ (·) for every new reward target. To overcome this constraint, the following section introduces a plug-and-play alternative that completely bypasses the retraining phase. Specifically, we demonstrate how this essential value gradient can be estimated directly by leveraging an off-the-shelf differentiable diffusion network. For notational brevity in subsequent derivations, we denote the differentiable value function vϕ (·) simply as v(·).

by minimizing the KL divergence between the two: arg minEx,zt [DKL (pθ (x|zt )∥p̃(x|zt ))] p̃

⇔ arg minEx,zt [− log p̃(x|zt )]

(10)

p̃

This equality shows that minimizing the KL divergence between our proxy and the true distribution is equivalent to minimizing the negative log-likelihood of the proxy. If we model our proxy using a mean-field assumption (Hoffman et al., 2013; Giordano et al., 2015), p̃(x|zt ) = Q ℓ ℓ ℓ Cat(x ; µ (zt , t)), the objective becomes the minimization of the negative log-likelihood of the per-token predictions,

3.2. A Variational Perspective on Value Function Estimation As concluded in the previous subsection, the value function v(zt ) requires computing an expectation Epθ (x|zt ) [·] over the unrolled multi-step reverse process trajectory leading to the final data x. Calculating or differentiating this expectation is computationally intractable. To overcome this, we introduce a variational approach by seeking a computationally efficient proxy distribution, p̃(x|zt ), that closely approximates the true distribution pθ (x|zt ). This is achieved

arg minEx,zt µ

L X

− log µℓ (zt , t) , xℓ

(11)

ℓ=1

We observe that this objective is exactly the negative loglikelihood loss (Eq. 4) used to train the discrete diffusion 4

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

Timestep 8 150

10 1

100

100 150

≈ 10 4

10 1 10 2 10 3 10 4

Timestep 16 150

100 150

100

00 100 10 1 10 2 10 3 10 4 10 5

Timestep 24

100

10 1

10 2 50

100

σk /σ1

100

150

Input Dimensions ≈ 10 5

100

150

≈ 10 4

10 1 10 2 10 3

100 150

Index (k)

100

150

Figure 3. Numerical instability in the model Jacobian. (Top) Visualization of the Jacobian ∂η/∂zt , aggregated over categorical dimensions using the Frobenius norm, for a discrete DNA diffusion model. (Bottom) Corresponding singular value spectrum. The high condition number (K ≈ 104 –105 ) signifies severe ill-conditioning, which causes the gradient flow to become numerically unstable.

model’s denoising network. This key insight means we do not need additional training and we can set the proxy distribution’s parameters to the readily available predictions from our pre-trained model: µ (zt , t) ← xθ (zt , t). This allows us to estimate the value function training-free using the Monte Carlo method through samples drawn from this variational proxy: n

v (zt ) ≈

1X r(x(i) ) n i=1

1) Enabling Gradient Flow Through Discrete Sampling. The sampling operation required to estimate the value funcP r(x(i) ), breaks the computational graph, tion, v(zt ) ≈ n1 preventing direct differentiation. To restore the flow of gradients, we employ a combination of the Gumbel-Softmax (GS) reparameterization trick (Jang et al., 2017) and a StraightThrough (ST) estimator (Bengio et al., 2013). The GS trick provides a differentiable, soft approximation of a categorical sample. Let η = η(zt , t) be the prediction logits from the denoising network, where the clean prediction is xθ = Softmax(η). The soft sample xsoft is calculated as: h i exp((ηK +gK )/τ ) 1 +g1 )/τ ) PK xsoft = PKexp((η , . . . , k=1 exp((ηk +gk )/τ ) k=1 exp((ηk +gk )/τ ) (13) where gk ∼ Gumbel(0, 1), and τ is a temperature coefficient controlling the sharpness of the approximation. However, many real-world reward models demand discrete or hard one-hot inputs. To ensure a consistent, accurate reward signal in the forward pass, we integrate the ST estimator. Specifically, we first obtain the hard, one-hot vector xhard by taking the argmax of the soft sample xhard = one-hot(arg maxk (xsoft )). We then compute the reward signal using the composite input x̂ = xhard − sg(xsoft ) + xsoft , where sg(·) is the stop-gradient operation. This formulation ensures that the reward function processes the discrete xhard in the forward pass, while the gradient flows backward robustly through the differentiable xsoft .

10 4 0

challenges of discrete sampling and network instability.

(12)

2) Logit Correction for the Clean Prediction. Following the reparameterization described above, the full gradient of the value function (Eq. 12) with respect to the noisy state zt can be expanded via the chain rule:

where samples x(1) , x(2) , · · · , x(n) ∼ p̃(x|zt ). Comparison to Deterministic Estimation: Note that prior work (Li et al., 2024) has proposed a deterministic value estimation method by directly using the clean prediction input to the reward model, i.e., v(zt ) ≈ r(xθ (zt , t)). In contrast, our approach utilizes stochastic sampling (n > 1) over the predicted distribution, which provides a better estimate of the expected future reward and captures the uncertainty inherent in the multi-step generation process.

∇zt v(zt ) ≈

1 X ∂r(x̂(i) ) ∂ x̂(i) n i=1 ∂ x̂(i) ∂η | {z } Logit Sensitivity

∂η ∂zt |{z}

(14)

Model Jacobian

While this formulation is mathematically rigorous, its direct computation is numerically unreliable. We observe that the model Jacobian ∂η/∂zt is often poorly conditioned (shown in Fig. 3), because denoising networks are trained to model clean data distributions, rather than to yield stable or smooth derivatives with respect to their inputs. As noted by Meng et al. (2021), a low training loss does not imply well-behaved Jacobian estimates. This instability is further exacerbated in discrete diffusion models, where the input zt inhabits a discrete token space. Differentiating through this non-smooth structure introduces substantial noise into the gradient, preventing the guidance signal from accumulating coherently across the denoising trajectory, consistent with the empirical behavior observed in Fig. 2b.

As illustrated in Fig. 2a, we validate our approximation against the ground-truth value function, derived from full reverse rollout samples. Our empirical results demonstrate that deterministic methods suffer from significant estimation bias even in the final stages of sampling, resulting in a performance plateau. In contrast, our proposed estimator achieves superior precision that scales predictably with the number of samples, consistently narrowing the gap to the ground-truth value. 3.3. Practical Designs for Calculating Gradients Our method requires calculating the gradient of the value function ∇zt v(zt ) to estimate the optimal reverse policy. This is achieved through two core designs that handle the

To circumvent this bottleneck, we draw inspiration from Score Distillation Sampling (SDS) and its variants (Poole 5

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

（a）

GS trick

Clean prediction

Logits

Reward function

Group computation

…

Variational proxy

Evaluate

Clean prediction

Reward function

…

（b）

… …

Variational proxy

Evaluate

Logits

Forward pass Backpropgation

ST estimator

Likelihood

Figure 4. Demonstration of gradient calculation for GILC. (a) Differentiable rewards: The gradient gη is calculated via direct backpropagation, utilizing the Gumbel-Softmax trick and the Straight-Through estimator to enable differentiation through the discrete samples x̂(i) . (b) Non-differentiable rewards: The gradient gη′ is estimated via the policy gradient formulation, where the rewards r(x(i) ) are converted into group relative advantages Ai to reduce variance. Both methods use the pre-trained network’s clean prediction xθ as the variational proxy.

3.4. Guidance with Non-Differentiable Objectives

et al., 2023; Hertz et al., 2023), which omit unstable Jacobian terms in continuous diffusion. Analogously, we propose to bypass the model Jacobian altogether and define a stable guidance correction directly in the logit space:

The gradient gη in Eq. 15 relies on the reward function r(·) being differentiable. While many rewards are provided by pre-trained differentiable networks, we propose an alternative, training-free approach for guiding systems that face a non-differentiable or black-box objective.

gη ≜

1 X ∂r(x̂(i) ) ∂ x̂(i) n i=1 ∂ x̂(i) ∂η

(15)

We address this using the concept of policy gradients (Williams, 1992). Our goal remains to calculate the gradient of the value function v (zs ) ≈ Epθ (x|zs ) [r(x)] with respect to the prediction logits η. The derivation is as follows:

This simplification is well-motivated: the clean-prediction logits η lie in a smooth, continuous space where gradient propagation is numerically stable, whereas zt resides in a discrete and highly non-smooth space. By extracting rewardinduced perturbations directly in logit space, we preserve the informative gradient direction while avoiding the instability caused by differentiating through zt .

∇η Epθ (x|zt ) [r(x)] =

X ∂p (x|z ) t

∂η

r(x)

pθ (x|zt )

∂pθ (x|zt ) 1 r(x) pθ (x|zt ) ∂η

As shown in Appendix Appendix B.2, under this approximation the optimal reverse process prθ (zs |zt ) (Eq. 6) admits a particularly interpretable form. Specifically, the guidance first modifies the clean-prediction logits as η r = η + gη /β, yielding a reward-corrected prediction xrθ = softmax(η r ). This corrected prediction is then incorporated into the reverse sampling step using the posterior from Eq. 3, namely q(zs |zt , xrθ ). We refer to the resulting algorithm as Gradient-Informed Logit Correction via Direct Backpropagation (GILC-DB). A complete description is provided in Algorithm 1.

h i θ (x|zt ) = Epθ (x|zt ) r(x) ∂ log p∂η

(16)

In this expression, pθ (x|zt ) acts as the policy, and the resulting expression facilitates estimation via Monte Carlo sampling. Recalling the insight from Sec. 3.2, we utilize the computationally tractable variational proxy p̃(x|zt ) to draw a group of samples x(i) . The expectation is then approximated empirically, and the log-likelihood term log pθ (x|zt ) is replaced by the log-likelihood of the proxy, log xθ (zt , t) , x(i) . 6

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

To further enhance gradient stability and reduce variance, we shift our focus from absolute rewards to relative advantages, similar to Group Relative Policy Optimization (GRPO) (Shao et al., 2024). The resulting gradient for correcting the logits is expressed as:

we use two reward oracles trained on disjoint subsets of the data: one used during guided generation and the other reserved exclusively for evaluation. Both oracles adopt the Enformer architecture (Avsec et al., 2021) to predict enhancer activity in the HepG2 cell line. To comprehensively evaluate enhancer generation performance, we generate 640 DNA samples and report metrics spanning three complementary aspects, following prior work (Wang et al., 2025): 1) Functional activity, including predicted activity (Pred-Activity) and binary classification of chromatin accessibility (ATAC-Acc); 2) Sequence fidelity, measuring similarity to real enhancer sequences via 3-mer Pearson correlation (3-mer Corr) and JASPAR motif enrichment correlation (JASPAR Corr); and 3) Distributional consistency, which assesses whether generated sequences remain within the data manifold using approximated loglikelihood (App-Log-Lik).

gη′ ≜

1 X ∂ log xθ (zt , t) , x(i) Ai n i=1 ∂η (i)

(1)

(17) (n)

)−mean({r(x ),··· ,r(x )}) where the advantage Ai = r(x std({r(x (1) ),··· ,r(x(n) )}) is computed using the rewards of the samples within the group. This formulation provides a robust, training-free mechanism for guidance that is entirely agnostic to the differentiability of the reward function, making it an efficient fallback for black-box guidance systems. We term this approach as Gradient-Informed Logit Correction via Policy Gradient (GILC-PG).

Results. Tab. 1 reports the quantitative results. Sequences generated by GILC-DB achieve the highest predicted activity in HepG2 cell lines. In particular, its Pred-Activity and ATAC-Acc metrics substantially outperform all other training-free baselines, and even surpass fine-tuned methods such as DRAKES. GILC-PG demonstrates superior similarity to natural enhancers, as evidenced by strong correlations with trinucleotide distributions and JASPAR motif patterns, together with favorable likelihood scores. In addition, Tab. 2 compares the computational efficiency of different methods. Best-of-N , SMC, and SVDD require maintaining multiple sampling trajectories, whereas the proposed GILC methods operate with a single trajectory. Notably, GILC-DB explicitly exploits the differentiability of the reward function r(·), resulting in significantly fewer reward evaluations than competing approaches.

4. Experiments In this section, we evaluate our guided generation methods, GILC-DB and GILC-PG, across three scientific domains: regulatory DNA sequence design, protein engineering, and small-molecule generation. In addition, we also conduct experiments on discrete diffusion in the image domain. A detailed ablation study of the key components is provided in Appendix D. We compare our approach against the following categories. 1) Standard guidance methods. Pretrained unconditional generation, Classifier Guidance (CG), and Classifier-Free Guidance (CFG) (Nisonoff et al., 2025). 2) Fine-tuning methods. DRAKES (Wang et al., 2025), a specialized approach for discrete diffusion and flows, along with its variant without KL regularization (DRAKES w/o KL). 3) Trainingfree guidance methods. Best-of-N (Beirami et al., 2025); Sequential Monte Carlo (SMC) (Wu et al., 2023); SVDD (Li et al., 2024); and TFG-Flow (Lin et al., 2025). Implementation details and hyperparameter settings are provided in Appendix C.

4.2. Protein Sequence Design Given a pre-trained inverse folding model that generates protein sequences conditioned on a fixed 3D backbone conformation, our objective is to guide the generation process toward sequences with enhanced thermodynamic stability. Experimental Setup. We utilize a pretrained discrete flow model (Campbell et al., 2024), using the ProteinMPNN architecture (Dauparas et al., 2022) as its backbone. The model is pretrained for inverse folding on the PDB dataset curated by Dauparas et al. (2022). To guide the generation process, we employ the reward oracle trained on the Megascale protein stability dataset (Tsuboyama et al., 2023), which comprises experimental stability measurements for approximately 1.8 million sequence variants across 983 protein domains. We maintain two distinct ProteinMPNN-based oracles: a guidance oracle used during sampling and an independent evaluation oracle reserved for final validation. Both oracles are trained to predict protein stability.

4.1. Regulatory DNA Sequence Design Here, our goal is to optimize regulatory DNA sequences to drive gene expression in a cell-type-specific manner (Taskiran et al., 2024). Experimental Setup. We conduct experiments on a largescale enhancer dataset introduced by Gosai et al. (2023), which measures enhancer activity across human cell lines for approximately 700K DNA sequences of length 200 bp using massively parallel reporter assays (MPRAs). For each sequence, the corresponding gene expression level is provided. We employ a masked discrete diffusion model (Sahoo et al., 2024) pretrained on the full sequence set. In addition, 7

Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction Table 1. Performance of various methods in regulating DNA sequence design. The table reports the mean and standard deviation of three random seeds.Bold: best, underline: second best. Pred-Activity ↑

ATAC-Acc ↑ (%)

3-mer Corr ↑

JASPAR Corr ↑

App-Log-Lik ↑

Pretrained CG (Nisonoff et al., 2025) CFG (Nisonoff et al., 2025)

0.17±0.04 3.30±0.00 5.04±0.06

1.5±0.2 0.0±0.0 92.1±0.9

−0.061±0.034 −0.065±0.001 0.746±0.001

0.249±0.015 0.212±0.035 0.864±0.011

−261±0.6 −266±0.6 −265±0.6

DRAKES w/o KL (Wang et al., 2025) DRAKES (Wang et al., 2025)

6.44±0.04 5.61±0.07

82.5±2.8 92.5±0.6

0.307±0.001 0.887±0.002

0.557±0.015 0.911±0.002

−281±0.6 −264±0.6

Best-of-N (Beirami et al., 2025) SMC (Wu et al., 2023) SVDD (Li et al., 2024) TFG-Flow (Lin et al., 2025) GILC-DB (Ours) GILC-PG (Ours)

3.73±0.41 4.15±0.33 4.84±0.28 3.48±0.21 7.04±0.26 5.21±0.18

35.8±7.6 39.9±8.7 51.9±6.3 21.2±5.4 95.2±2.1 84.0±1.5

0.813±0.037 0.840±0.045 0.870±0.061 0.262±0.058 0.900±0.044 0.910±0.027

0.671±0.071 0.756±0.068 0.826±0.057 0.566±0.044 0.935±0.039 0.937±0.031

−262±1.1 -259±2.5 −269±3.1 −263±1.5 −267±1.8 −270±2.4

Method

Table 2. Comparison of the number of calls to the denoising model and the reward function per step across different training-free guidance approaches in the DNA sequence design task. Method Best-of-N SMC SVDD TFG-Flow GILC-DB (Ours) GILC-PG (Ours)

Num-of-Diff ↓

Num-of-Reward ↓

20 20 20 1 1 1

20 20 20 20 5 20

4.3. Multimodal Molecule Generation We evaluate our method on the inverse design of molecules with targeted properties, a core challenge in computational chemistry (Hoogeboom et al., 2022; Gebauer et al., 2022). Experimental Setup. Following Lin et al. (2023), we adopt a pre-trained multimodal flow model trained on QM9 (Ramakrishnan et al., 2014) with an Equivariant Graph Neural Network (EGNN) backbone (Satorras et al., 2021), which jointly models discrete atomic types and continuous 3D coordinates. To avoid information leakage, the dataset is split into two disjoint subsets: one for training the guidance (reward) predictor and the other for training an independent evaluation predictor.

We employ the following three types of metrics to evaluate the stability of generated sequences and their ability to fold into target structures. 1) Predicted stability, evaluated using the independent evaluation oracle (Pred-ddG); 2) Self-consistency, measures the structural fidelity of the generated sequences. We utilize ESMFold (Lin et al., 2023) to predict the folded structures and calculate the Root Mean Square Deviation (RMSD) relative to the wild-type structure (scRMSD). 3) Finally, we define the success rate as the percentage of sequences that simultaneously satisfy Pred-ddG > 0 and scRMSD < 2 as in Campbell et al. (2024).

We apply GILC and training-free baselines to guide discrete atomic types toward following quantum properties: polarizability (α), dipole moment (µ), heat capacity (Cv ), HOMO energy (ϵHOMO ), LUMO energy (ϵLUMO ), and energy gap (∆ϵ). For each property, 4,096 samples are generated for evaluation. Guidance performance is assessed using mean absolute error (MAE) with the evaluation predictor. Results. Quantitative results summarized in Tab. 4. Both GILC-DB and GILC-PG consistently outperform all unsupervised baselines, including TFG-FLOW and SVDD. GILC-DB achieves the best overall performance by exploiting intrinsic reward gradients for more accurate guidance. Conversely, GILC-PG remained highly competitive despite treating the attribute predictor as a black box, offering superior flexibility. These results demonstrate the effectiveness and generality of the GILC framework for multimodal scientific data generation.

Results. As shown in Tab. 3, GILC-DB generates protein sequences with high structural stability, achieving the highest Pred-ddG scores while maintaining inverse folding success rates comparable to those of pre-trained models, measured by the percentage of samples with scRMSD < 2. When these criteria are considered jointly, GILC-DB substantially outperforms all baseline methods in terms of overall success rate, exceeding DRAKES by approximately 4 percentage points. GILC-PG also consistently outperforms other training-free baselines. By contrast, DRAKES is susceptible to reward hacking (Clark et al., 2024) and requires careful tuning of the KL-constraint strength, while CFG relies heavily on labeled data and consequently suffers from limited generalization. Overall, these results highlight the strong potential of GILC for protein sequence design.

4.4. Discrete image generation Experimental Setup. To evaluate the scalability and versatility of the GILC framework in larger discrete spaces, we conduct class-conditional and text-to-image generation

Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction Table 3. Model performance in the inverse protein folding task. The table reports the mean and standard deviation of three random seeds. Bold: best, underline: second best. The primary metric (overall success rate) is further highlighted in blue. Method

Pred-ddG ↑

%(ddG> 0) (%)↑

scRMSD ↓

%(scRMSD< 2)(%)↑

Success Rate (%)↑

Pretrained CG (Nisonoff et al., 2025) CFG (Nisonoff et al., 2025)

−0.544±0.037 −0.561±0.045 −1.186±0.035

36.6±1.0 36.9±1.1 11.0±0.4

0.849±0.013 0.839±0.012 3.146±0.062

90.9±0.6 90.9±0.6 29.4±1.0

34.4±0.5 34.7±0.9 1.3±0.4

DRAKES w/o KL (Wang et al., 2025) DRAKES (Wang et al., 2025)

1.108±0.004 1.095±0.026

100.0±0.0 86.4±0.2

7.307±0.054 0.918±0.006

34.1±0.2 91.8±0.5

34.1±0.2 78.6±0.7

Best-of-N (Beirami et al., 2025) SMC (Wu et al., 2023) SVDD (Li et al., 2024) TFG-Flow (Lin et al., 2025) GILC-DB (Ours) GILC-PG (Ours)

0.623±0.051 0.659±0.044 0.694±0.076 0.410±0.064 1.430±0.073 0.719±0.091

45.5±4.7 68.5±3.1 69.3±2.1 39.2±1.1 97.9±1.9 75.6±3.5

0.849±0.011 0.841±0.006 0.850±0.030 0.837±0.021 0.968±0.015 0.914±0.012

91.2±0.5 93.8±0.4 89.7±0.9 92.6±1.6 84.3±2.0 92.4±2.3

54.7±5.3 63.6±4.0 65.0±3.3 52.9±2.2 82.4±2.5 69.8±3.2

Table 4. Mean absolute error (MAE) of generated target quantum properties on the QM9 dataset, evaluated using the property predictor. Upper-bound, #Atoms, and lower-bound results are taken from the study (Bao et al., 2023). Mean and standard deviation are calculated over three random seeds. Among training-free methods, bold indicates the best performance, and underline indicates the second best. Method

∆ϵ

ϵHOMO

ϵLUMO

Upper bound #Atoms Lower bound

6.87 1.97 0.040

1.61 1.05 0.043

8.98 3.86 0.09

1464 886 65

645 426 39

1457 813 36

3.40±0.03 3.22±0.01 3.14±0.04 2.42±0.02 2.23±0.03 1.53±0.02

1.52±0.05 1.38±0.02 1.41±0.02 1.28±0.03 1.09±0.04 0.808±0.04

4.23±0.03 4.09±0.01 4.02±0.02 3.12±0.01 2.87±0.04 2.15±0.02

1213±6 1107±4 1140±5 902±3 810±8 783±5

614±5 561±3 557±2 476±5 448±4 374±7

1172±8 1103±4 1077±6 981±6 988±5 875±7

Best-of-N (Beirami et al., 2025) SMC (Wu et al., 2023) SVDD (Li et al., 2024) TFG-Flow (Lin et al., 2025) GILC-DB (Ours) GILC-PG (Ours)

tasks on natural image datasets.

5. Discussions and Limitations

1) Class-conditional image generation: We build upon the discrete diffusion model framework trained on CIFAR-10 (Campbell et al., 2022). The target class labels are utilized as guidance signals to steer the generation process.

This paper introduces the Gradient-Informed Logit Correction (GILC) framework, a versatile framework for guiding discrete diffusion in a plug-and-play manner. By employing a variational proxy to estimate the value function and utilizing gradient-informed logit correction, our method enables precise control over discrete data generation without requiring retraining. Through extensive evaluation, GILC demonstrates robust performance and high adaptability across diverse scientific domains.

2) Text-to-image generation: Following recent advancements, we adopt a text-to-image masked generative model termed Meissonic (Bai et al., 2024). To optimize the visual appeal of the generated outputs, we employ the aesthetic score (Pressman et al., 2022) as the reward function, which presents a challenging, non-differentiable optimization objective in high-dimensional pixel spaces.

While our approach reduces the computational overhead, specifically the number of calls to the diffusion network, it is not without limitations. Both GILC-PG and GILC-DB still necessitate multiple Monte Carlo samples to query the reward function. Further minimizing this sampling requirement while preserving guidance fidelity remains a promising direction for future research. Additionally, our current training-free variational proxy operates under a mean-field assumption of token independence. This assumption can introduce non-trivial errors in structured domains like natural language, where token dependencies are critical. Extending our framework to large-scale diffusion language models by relaxing this independence constraint, as in Xie et al. (2025), represents a significant next step.

Results. Both GILC-DB and GILC-PG successfully scale to these complex image domains without requiring any architectural modifications or additional training. For classconditional generation on CIFAR-10, qualitative evaluations demonstrate that GILC produces distinct, class-consistent images with high sample quality, confirming its efficacy in large-scale discrete state spaces (Fig. 8). Furthermore, in the more challenging text-to-image synthesis task with Meissonic, GILC remains highly effective, successfully generating high-resolution, visually appealing images aligned with the aesthetic score (Fig. 9).

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

Acknowledgements

Bengio, Y., Léonard, N., and Courville, A. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.

This work was supported by the National Natural Science Foundation of China under Grant 62405014 and 62325101.

Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S. Training diffusion models with reinforcement learning. In International Conference on Learning Representations, volume 2024, pp. 4965–4987, 2024.

Impact Statement Our proposed framework provides a training-free solution that enables discrete diffusion models to satisfy complex constraints. Owing to its versatility and computational efficiency, we anticipate that this framework will have a positive impact on related research domains, including gene therapy and drug design. Nevertheless, careful consideration is required to prevent potential misuse of the method, underscoring the need for robust ethical oversight.

Campbell, A., Benton, J., De Bortoli, V., Rainforth, T., Deligiannidis, G., and Doucet, A. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266–28279, 2022. Campbell, A., Yim, J., Barzilay, R., Rainforth, T., and Jaakkola, T. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein codesign. In International Conference on Machine Learning, pp. 5453–5512. PMLR, 2024.

References Austin, J., Johnson, D. D., Ho, J., Tarlow, D., and Van Den Berg, R. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:17981–17993, 2021.

Chang, J. and Ye, J. C. Bidirectional generation of structure and properties through a single molecular foundation model. Nature Communications, 15(1):2323, 2024.

Avsec, Ž., Agarwal, V., Visentin, D., Ledsam, J. R., GrabskaBarwinska, A., Taylor, K. R., Assael, Y., Jumper, J., Kohli, P., and Kelley, D. R. Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 18(10):1196–1203, 2021.

Chu, W., Wu, Z., Chen, Y., Song, Y., and Yue, Y. Split gibbs discrete diffusion posterior sampling. Advances in Neural Information Processing Systems, 38:145122– 145153, 2026. Chung, H., Kim, J., McCann, M. T., Klasky, M. L., and Ye, J. C. Diffusion posterior sampling for general noisy inverse problems. In 11th International Conference on Learning Representations, ICLR 2023, 2023.

Bai, J., Ye, T., Chow, W., Song, E., Chen, Q.-G., Li, X., Dong, Z., Zhu, L., and Yan, S. Meissonic: Revitalizing masked generative transformers for efficient highresolution text-to-image synthesis. In The Thirteenth International Conference on Learning Representations, 2024.

Clark, K., Vicol, P., Swersky, K., and Fleet, D. Directly fine-tuning diffusion models on differentiable rewards. In International Conference on Learning Representations, volume 2024, pp. 4793–4822, 2024.

Bajusz, D., Rácz, A., and Héberger, K. Why is tanimoto index an appropriate choice for fingerprint-based similarity calculations? Journal of cheminformatics, 7(1):20, 2015.

Dauparas, J., Anishchenko, I., Bennett, N., Bai, H., Ragotte, R. J., Milles, L. F., Wicky, B. I., Courbet, A., de Haas, R. J., Bethel, N., et al. Robust deep learning–based protein sequence design using proteinmpnn. Science, 378 (6615):49–56, 2022.

Bansal, A., Chu, H.-M., Schwarzschild, A., Sengupta, S., Goldblum, M., Geiping, J., and Goldstein, T. Universal guidance for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 843–852, 2023.

Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780–8794, 2021.

Bao, F., Zhao, M., Hao, Z., Li, P., Li, C., and Zhu, J. Equivariant energy-guided SDE for inverse molecular design. In The Eleventh International Conference on Learning Representations, 2023.

Esser, P., Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024.

Beirami, A., Agarwal, A., Berant, J., D’Amour, A., Eisenstein, J., Nagpal, C., and Suresh, A. T. Theoretical guarantees on the best-of-n alignment policy. In International Conference on Machine Learning, pp. 3580–3602. PMLR, 2025.

Gebauer, N. W., Gastegger, M., Hessmann, S. S., Müller, K.-R., and Schütt, K. T. Inverse design of 3d molecular 10

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

structures with conditional generative neural networks. Nature communications, 13(1):973, 2022.

Lin, H., Li, S., Ye, H., Yang, Y., Ermon, S., Liang, Y., and Ma, J. TFG-flow: Training-free guidance in multimodal generative flow. In The Thirteenth International Conference on Learning Representations, 2025.

Giordano, R. J., Broderick, T., and Jordan, M. I. Linear response methods for accurate covariance estimates from mean field variational bayes. Advances in neural information processing systems, 28, 2015.

Lin, Z., Akin, H., Rao, R., Hie, B., Zhu, Z., Lu, W., Smetanin, N., Verkuil, R., Kabeli, O., Shmueli, Y., et al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637): 1123–1130, 2023.

Gosai, S. J., Castro, R. I., Fuentes, N., Butts, J. C., Kales, S., Noche, R. R., Mouri, K., Sabeti, P. C., Reilly, S. K., and Tewhey, R. Machine-guided design of synthetic cell type-specific cis-regulatory elements. bioRxiv, 2023.

Liu, Y., Zhang, K., Li, Y., Yan, Z., Gao, C., Chen, R., Yuan, Z., Huang, Y., Sun, H., Gao, J., et al. Sora: A review on background, technology, limitations, and opportunities of large vision models. arXiv preprint arXiv:2402.17177, 2024.

Grathwohl, W., Swersky, K., Hashemi, M., Duvenaud, D., and Maddison, C. Oops i took a gradient: Scalable sampling for discrete distributions. In International Conference on Machine Learning, pp. 3831–3841. PMLR, 2021.

Lou, A., Meng, C., and Ermon, S. Discrete diffusion language modeling by estimating the ratios of the data distribution. 2023.

He, Y., Murata, N., Lai, C.-H., Takida, Y., Uesaka, T., Kim, D., Liao, W., Mitsufuji, Y., Kolter, Z., Salakhutdinov, R., et al. Manifold preserving guided diffusion. In International Conference on Learning Representations, volume 2024, pp. 44819–44850, 2024.

Meng, C., Song, Y., Li, W., and Ermon, S. Estimating high order gradients of the data distribution by denoising. Advances in Neural Information Processing Systems, 34: 25359–25369, 2021.

Hertz, A., Aberman, K., and Cohen-Or, D. Delta denoising score. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 2328–2337, 2023.

Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., and Li, C. Large language diffusion models. arXiv preprint arXiv:2502.09992, 2025.

Ho, J. and Salimans, T. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598, 2022.

Nisonoff, H., Xiong, J., Allenspach, S., and Listgarten, J. Unlocking guidance for discrete state-space diffusion and flow models. In International Conference on Learning Representations, volume 2025, pp. 36052–36106, 2025.

Ho, J., Jain, A., and Abbeel, P. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. Stochastic variational inference. the Journal of machine Learning research, 14(1):1303–1347, 2013.

Ou, Z., Pani, C., and Li, Y. Inference-time scaling of discrete diffusion models via importance weighting and optimal proposal design. In The Fourteenth International Conference on Learning Representations, 2026.

Hoogeboom, E., Satorras, V. G., Vignac, C., and Welling, M. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pp. 8867– 8887. PMLR, 2022.

Peng, X. B., Kumar, A., Zhang, G., and Levine, S. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177, 2019.

Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations, 2017.

Peters, J. and Schaal, S. Reinforcement learning by rewardweighted regression for operational space control. In Proceedings of the 24th international conference on Machine learning, pp. 745–750, 2007.

Langley, P. Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pp. 1207–1216, Stanford, CA, 2000. Morgan Kaufmann.

Poole, B., Jain, A., Barron, J. T., and Mildenhall, B. Dreamfusion: Text-to-3d using 2d diffusion. In The Eleventh International Conference on Learning Representations, 2023.

Li, X., Zhao, Y., Wang, C., Scalia, G., Eraslan, G., Nair, S., Biancalani, T., Ji, S., Regev, A., Levine, S., et al. Derivative-free guidance in continuous and discrete diffusion models with soft value-based decoding. arXiv preprint arXiv:2408.08252, 2024.

Prabhudesai, M., Goyal, A., Pathak, D., and Fragkiadaki, K. Aligning text-to-image diffusion models with reward backpropagation. arXiv preprint arXiv:2310.03739, 2023. 11

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

Pressman, J. D., Crowson, K., and Contributors, S. C. Simulacra aesthetic captions. Technical Report Version 1.0, Stability AI, 2022. url https://github.com/JDP/simulacra-aesthetic-captions .

Song, J., Vahdat, A., Mardani, M., and Kautz, J. Pseudoinverse-guided diffusion models for inverse problems. In International Conference on Learning Representations, 2023a.

Rafailov, R., Sharma, A., Mitchell, E., Manning, C. D., Ermon, S., and Finn, C. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems, 36: 53728–53741, 2023.

Song, J., Zhang, Q., Yin, H., Mardani, M., Liu, M.-Y., Kautz, J., Chen, Y., and Vahdat, A. Loss-guided diffusion models for plug-and-play controllable generation. In International Conference on Machine Learning, pp. 32483–32498. PMLR, 2023b.

Ramakrishnan, R., Dral, P. O., Rupp, M., and Von Lilienfeld, O. A. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1–7, 2014.

Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.

Rector-Brooks, J., Hasan, M., Peng, Z., Liu, C., Mittal, S., Dziri, N., Bronstein, M., Chatterjee, P., Tong, A., and Bose, J. Steering masked discrete diffusion models via discrete denoising posterior prediction. In International Conference on Learning Representations, volume 2025, pp. 25383–25414, 2025.

Taskiran, I. I., Spanier, K. I., Dickmänken, H., Kempynck, N., Pančı́ková, A., Ekşi, E. C., Hulselmans, G., Ismail, J. N., Theunis, K., Vandepoel, R., et al. Cell-type-directed design of synthetic enhancers. Nature, 626(7997):212– 220, 2024. Tsuboyama, K., Dauparas, J., Chen, J., Laine, E., Mohseni Behbahani, Y., Weinstein, J. J., Mangan, N. M., Ovchinnikov, S., and Rocklin, G. J. Mega-scale experimental analysis of protein folding stability in biology and design. Nature, 620(7973):434–444, 2023.

Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 10684–10695, 2022.

Uehara, M., Zhao, Y., Biancalani, T., and Levine, S. Understanding reinforcement learning-based fine-tuning of diffusion models: A tutorial and review. arXiv preprint arXiv:2407.13734, 2024.

Sahoo, S., Arriola, M., Schiff, Y., Gokaslan, A., Marroquin, E., Chiu, J., Rush, A., and Kuleshov, V. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 37:130136– 130184, 2024.

Venkatraman, S., Jain, M., Scimeca, L., Kim, M., Sendera, M., Hasan, M., Rowe, L., Mittal, S., Lemos, P., Bengio, E., et al. Amortizing intractable inference in diffusion models for vision, language, and control. Advances in neural information processing systems, 37:76080–76114, 2024.

Satorras, V. G., Hoogeboom, E., and Welling, M. E (n) equivariant graph neural networks. In International conference on machine learning, pp. 9323–9332. PMLR, 2021. Schiff, Y., Sahoo, S., Phung, H., Wang, G., Boshar, S., Dalla-Torre, H., Almeida, B., Rush, A., Pierrot, T., and Kuleshov, V. Simple guidance mechanisms for discrete diffusion models. In International Conference on Learning Representations, volume 2025, pp. 43776–43821, 2025.

Vignac, C., Krawczuk, I., Siraudin, A., Wang, B., Cevher, V., and Frossard, P. Digress: Discrete denoising diffusion for graph generation. In Proceedings of the 11th International Conference on Learning Representations, 2023.

Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Bi, X., Zhang, H., Zhang, M., Li, Y., Wu, Y., et al. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.

Wang, C., Uehara, M., He, Y., Wang, A., Lal, A., Jaakkola, T., Levine, S., Regev, A., Wang, H., and Biancalani, T. Fine-tuning discrete diffusion models via reward optimization with applications to dna and protein design. In International Conference on Learning Representations, volume 2025, pp. 47871–47899, 2025.

Shi, J., Han, K., Wang, Z., Doucet, A., and Titsias, M. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37:103131–103167, 2024.

Widatalla, T., Rafailov, R., and Hie, B. Aligning protein generative models with experimental fitness via direct preference optimization. bioRxiv, pp. 2024–05, 2024. 12

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

Williams, R. J. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning, 8(3):229–256, 1992. Wu, L., Trippe, B., Naesseth, C., Blei, D., and Cunningham, J. P. Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36:31372–31403, 2023. Xie, T., Xue, S., Feng, Z., Hu, T., Sun, J., Li, Z., and Zhang, C. Variational autoencoding discrete diffusion with enhanced dimensional correlations modeling. arXiv preprint arXiv:2505.17384, 2025. Yu, J., Wang, Y., Zhao, C., Ghanem, B., and Zhang, J. Freedom: Training-free energy-guided conditional diffusion model. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 23174–23184, 2023. Zekri, O. and Boullé, N. Fine-tuning discrete diffusion models with policy gradient methods. Advances in Neural Information Processing Systems, 38:152868–152906, 2026.

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

A. Related Works We review the landscape of diffusion models for constrained generation, focusing on the evolution from standard guidance to recent developments in diffusion with discrete state spaces. Classifier and Classifier-free Guidance. Classifier Guidance (CG) and Classifier-free Guidance (CFG) are the foundational frameworks for conditional diffusion generation (Dhariwal & Nichol, 2021; Song et al., 2021; Ho & Salimans, 2022). Recent efforts (Nisonoff et al., 2025; Schiff et al., 2025) have extended these paradigms to discrete diffusion and flow models. However, both face significant practical hurdles: CG necessitates training a dedicated noise-aware classifier, which precludes the use of pre-trained models optimized solely for clean data. Conversely, CFG requires paired datasets for joint training, a requirement that limits flexibility and often leads to suboptimal performance in data-scarce regimes. Training-free Guidance. To bypass the need for noise-aware classifiers, several variants have emerged that utilize approximations to estimate the guidance gradient. While these are effective for inverse problems and continuous conditional generation (Chung et al., 2023; Yu et al., 2023; Song et al., 2023a; Bansal et al., 2023; Song et al., 2023b), their application to discrete diffusion is non-trivial due to the discrete nature of the sampling state space. Alternative strategies, such as Sequential Monte Carlo (SMC) (Wu et al., 2023) and SVDD (Li et al., 2024), rely on particle filtering or candidate selection. However, these methods are inherently limited to the explored candidate space; as the sequence length expands, the number of particles required grows rapidly, leading to prohibitive sampling times. Furthermore, while TFG-low (Lin et al., 2025) was designed for multimodal flows, our empirical results suggest it struggles when guiding discrete diffusion in isolation. Other recent techniques involving variable splitting (Chu et al., 2026) show promise but are currently restricted to discrete diffusion with uniform transition matrices. Recently, Ou et al. (2026) suggests improving SMC by using reward gradient guidance as a better proposal distribution. In contrast, we directly derive our value function estimate from the variational objective and directly apply guidance to logits, which we found to be crucial for discrete diffusion. Fine-tuning and Reward Optimization. Conditional generation can also be framed as maximizing a reward function relative to a pre-trained model. While direct backpropagation (Prabhudesai et al., 2023; Clark et al., 2024) and reinforcement learning (Black et al., 2024) have been used to fine-tune continuous models, extensions to discrete states (Venkatraman et al., 2024; Rector-Brooks et al., 2025; Zekri & Boullé, 2026; Wang et al., 2025) are still gaining traction. Unlike these approaches, our method adopts a plug-and-play fashion. By utilizing reward functions without parameter updates, we avoid the heavy computational overhead of fine-tuning and mitigate the common pitfall of reward hacking. Tab. 5 shows a comparison between our proposed GILC framework and representative methods. Table 5. Comparison of guidance strategies for discrete diffusion in terms of plug-and-play capability (PnP), inference efficiency (NFE-Eff.), compatibility with non-differentiable objectives (Non-Diff.), and the ability to gradient utilization (Grad. Use). Methods marked († ) require fine-tuning of the diffusion network. Method

PnP NFE-Eff. Non-Diff. Grad. Use

GILC (Ours)

✓

CG/CFG (Nisonoff et al., 2025) SMC (Wu et al., 2023) SVDD (Li et al., 2024) TFG-Flow (Lin et al., 2025) DDPP† (Rector-Brooks et al., 2025) DRAKES† (Wang et al., 2025)

✗ ✓ ✓ ✓ ✗ ✗

✓ ✗ ✗ ✓ ✓ ✓

✗ ✓ ✓ ✓ ✓ ✗

✓ ✗ ✗ ✗ ✓ ✓

B. Additional Method Details B.1. Derivation of the Factorized Reverse Process In Sec. 3.1, we approximate the optimal reverse process prθ (zs |zt ) using a first-order Taylor expansion to ensure computational tractability. Taylor Expansion and Decomposition. The optimal reverse process is defined as: pθ (zs |zt ) exp(v(zs )/β) prθ (zs |zt ) = P zs pθ (zs |zt ) exp(v(zs )/β) 14

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

Applying a first-order Taylor expansion to v(zs ) around zt : v(zs ) ≈ v(zt ) +

L X ⟨zℓs − zℓt , ∇zℓt v(zt )⟩ ℓ=1

Substituting this into the exponential term and noting that exp(v(zt )/β) is a constant that cancels out in the normalization, we obtain: L Y 1 ℓ ℓ exp(v(zs )/β) ∝ exp ⟨z − zt , ∇zℓt v(zt )⟩ β s ℓ=1

ℓ ℓ=1 pθ (zs |zt ), the optimal

Factorization of the Partition Function. Assuming a factorized base transition pθ (zs |zt ) = process becomes: QL prθ (zs |zt ) ≈ P

ℓ ℓ=1 pθ (zs |zt ) exp

1 ℓ ℓ β ⟨zs − zt , ∇zℓt v(zt )⟩

ℓ ℓ=1 pθ (zs |zt ) exp

1 ℓ ℓ β ⟨zs − zt , ∇zℓt v(zt )⟩

P QL QL P Using the identity zs ℓ=1 fℓ (zℓs ) = ℓ=1 zℓ fℓ (zℓs ), the global normalization constant (the partition function) s factorizes into L independent local sums:

prθ (zs |zt ) ≈

pθ (zℓs |zt ) exp

L Y ℓ=1

1 ℓ ℓ β ⟨zs − zt , ∇zℓt v(zt )⟩

ℓ k=1 pθ (zs = k|zt ) exp

1 ℓ β ⟨k − zt , ∇zℓt v(zt )⟩

L Y

prθ (zℓs |zt )

ℓ=1

As a result, the final probability is fully decomposable across sequence positions. This allows each dimension ℓ to perform sampling independently, transforming an exponentially complex joint distribution into L parallelizable categorical distributions. B.2. Formal Derivation of Logit Guidance In this section, we demonstrate that for discrete diffusion models utilizing an absorbing state, the optimal reward-guided reverse process of GILC is exactly equivalent to applying a correction to the model’s prediction logits. Proof of Equivalence: prθ (zs | zt ) = q(zs | zt , xrθ ). In an absorbing state discrete diffusion model, let m denote the [MASK] state. When the current latent zt is masked, the base reverse transition for a single token is given by the categorical distribution: ( (αs − αt )xθ,k k ̸= m pθ (zs = k | zt = m) ∝ 1 − αs k=m k) where xθ,k = Pexp(η exp(ηj ) represents the model’s predicted probability for the clean-data class k. j

The optimal guided distribution prθ (Eq. 6) is defined by the reward-tilted density: prθ (zs = k | zt ) ∝ pθ (zs = k | zt ) exp

v(zs = ek ) β

where ek is the one-hot encoding for class k. We approximate the value function v(zs = ek ) via a first-order expansion around zt . Using the Jacobian-free correction vector gη to represent the gradient direction, and omitting terms independent of k (which vanish under the proportionality), the value term is approximated as: exp

v(zs = ek ) β

≈ exp

v(zt ) + ⟨gη , ek − zt ⟩ β 15

∝ exp

gη,k β

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

where gη,k = ⟨ek , gη ⟩ selects the k-th component of the correction vector gη . Substituting the Softmax form of xθ,k into the above equation for k ̸= m: ! gη,k g exp(η ) η,k k r exp pθ (zs = k | zt ) ∝ xθ,k · exp = P β β j exp(ηj ) gη,k ∝ exp ηk + β Normalizing this categorical distribution over the vocabulary j ∈ {1, . . . , K} yields: exp(ηk + gη,k /β) prθ (zs = k | zt ) = P j exp(ηj + gη,j /β) Conversely, we define the guided prediction as xrθ = Softmax(η + gη /β). Substituting this into the standard posterior q (Eq. 3) for the unmasking step: q(zs = k | zt , xrθ ) =

(αs − αt )xrθ,k + (1 − αs )I[k = m] 1 − αt

For any non-mask token k ̸= m, the term matches the result in prθ (zs = k | zt ) exactly (under the transition kernel’s normalization). Thus, the reward-tilted process is equivalent to a standard unmasking step using corrected clean prediction: prθ (zs | zt ) = q(zs | zt , xrθ ) This completes the proof. B.3. Pseudocodes This section presents the pseudocode implementation of the GILC framework through Algorithm 1. Algorithms 2 and 3 respectively detail the gradient computation for GILC-DB and CILC-PG. Algorithm 1 Gradient-Informed Logit Correction (GILC) 1: Input: Pre-trained discrete diffusion model pθ , reward function r(·), Monte Carlo number n, guidance scale β −1 2: Initialize fully masked sequence zT ← m ; 3: for t = T, . . . , 1 do 4: Predict clean data logits s ← t − 1, η ← pθ (zs |zt ) ; 5: Estimate correction gradient gη ← Correction-DB/PG(η, r, n) (Algorithms 2 or 3) ; 6: Apply logit guidance η r ← η + gη /β ; 7: Compute guided prediction xθr ← Softmax(η r ) ; 8: Sample the next state zs ∼ q(zs | zt , xθr ) (Eq. 3) ; 9: end for 10: Output: Generated sample x ← z0

Algorithm 2 Correction-DB (Direct Backpropagation Estimator) 1: Input: Logits η, reward function r(·), sample size n (1)

(n)

2: Sample soft samples xsoft , · · · , xsoft ∼ Gumbel-Softmax(η) ; (i)

(i)

3: Compute straight-through samples x̂(i) ← onehot arg max xsoft − sg xsoft + xsoft , i = 1, . . . , n ; 4: Evaluate rewards Ri ← r(x̂(i) ), i = 1, · · · , n ; 5: Estimate gradient with backpropgation gη ← n1 6: Output: Logit correction gη

i=1

∂r(x̂(i) ) ∂ x̂(i) ∂η ; ∂ x̂(i)

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

Algorithm 3 Correction-PG (Policy Gradient Estimator) 1: Input: Logits η, reward function r(·), sample size n 2: Sample a group of candidates x(1) , · · · , x(n) ∼ Cat(x; xθ (zt , t) = Softmax(η)) ; 3: Evaluate rewards Ri ← r(x(i) ), i = 1, · · · , n ; Ri −mean({R1 ,··· ,Rn }) ; std({R1 ,··· ,Rn } (i) P ∂ log x (z ,t),x ⟨ ⟩ t θ n Estimate policy gradient: gη′ ← n1 i=1 Ai ; ∂η ′ Output: Logit correction gη

4: Compute group relative advantages Ai ← 5: 6:

C. Experimental Details This section describes the implementation details and hyperparameter settings of all baseline methods. Quantitative results for classifier-based and classifier-free guidance, as well as fine-tuning methods, are primarily drawn from existing literature. For training-free baselines, we faithfully reproduced the original experimental setups to ensure fair and rigorous comparisons with our approach. All experiments were conducted on a single NVIDIA A100 GPU with 80 GB of memory. C.1. Implementation Details of Methods In this section, we describe the implementation details of the baseline methods used for comparison. Classifier Guidance (CG) (Nisonoff et al., 2025). Classifier Guidance (CG) steers generation toward target attributes during sampling by leveraging an auxiliary noise-aware classifier. For comparison, we implement CG following the procedures described in Wang et al. (2025) and Li et al. (2024). Specifically, the predictor is estimated using posterior mean estimation (Chung et al., 2023). This involves first extracting a denoised sequence from the noisy input zt using a pretrained diffusion model, and then evaluating the predicted clean sequence with the reward oracle. However, unlike in continuous diffusion settings, we observe that this posterior mean–based estimation is ineffective in discrete diffusion models. As a result, CG fails to provide reliable guidance signals in our setting, which explains its relatively poor performance reported in the experimental results. Classifier-free Guidance (CFG) (Nisonoff et al., 2025). In contrast to CG, Classifier-Free Guidance (CFG) trains conditional generation models from scratch (Ho & Salimans, 2022). We adopt the CFG implementation provided by Wang et al. (2025). To generate sequences with desired properties, CFG incorporates reward values as additional conditioning inputs to the diffusion model and encourages sampling toward high-reward regions. Concretely, reward values are binarized using the 95th percentile as the threshold, and sampling is performed conditioned on the high-reward label. We emphasize that CFG requires access to labeled training pairs with associated reward values. Consequently, its performance may degrade in regimes where labeled data are scarce. DRAKES (Wang et al., 2025). DRAKES is a method for directly fine-tuning discrete diffusion models via reward optimization. It enables backpropagation through the reward signal during sampling by maintaining differentiability using reparameterization techniques. In practice, DRAKES employs truncation strategies and observes that initiating backpropagation from intermediate diffusion time steps is often more effective than propagating gradients from the initial noise state. To further stabilize training and mitigate reward hacking, DRAKES introduces a KL-divergence regularization term that constrains the fine-tuned model to remain close to the pretrained distribution. In our experiments, we adopt the fine-tuned models released by Wang et al. (2025) and report their quantitative performance. Best-of-N (Beirami et al., 2025). Best-of-N is a simple yet effective baseline for guiding generative models, wherein multiple candidate samples are generated independently and the one with the highest reward is selected. In our experiments, we set the number of samples to N = 20. Sequential Monte Carlo (SMC) (Wu et al., 2023). SMC is a general-purpose sampling framework that maintains a population of particles and applies filtering operations during generation to approximate the target distribution. While SMC is theoretically exact in the limit of infinitely many particles, practical implementations necessarily operate with a finite particle budget. In our experiments, we set the number of particles to N = 20. SVDD (Li et al., 2024). SVDD is a derivative-free guidance method for diffusion models. At each diffusion time step, it samples multiple candidate states from the transition kernel, estimates the future value of each candidate using a deterministic 17

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

evaluator, and selects the candidate with the highest estimated value as the next state. In our experiments, we set the number of candidates to N = 20. TFG-Flow (Lin et al., 2025). TFG-Flow was originally proposed as a training-free guided multimodal flow model that jointly generates continuous and discrete components. When restricted to discrete guidance, it can also be applied as a discrete diffusion model. TFG-Flow requires estimating a guided rate matrix; in our experiments, this estimation is performed using N = 20 samples. All other hyperparameters are kept consistent with the official implementation. GILC-DB/PG (Ours). Our framework includes several tunable hyperparameters for guidance optimization. For GILC-DB, we use 20 Monte Carlo samples for value estimation in the protein unfolding task and 5 samples for DNA and small molecule generation, which we find sufficient to achieve state-of-the-art performance. The Gumbel-Softmax temperature τ is fixed at 1.0. For GILC-PG, the number of samples is set to 20 across all tasks. The guidance strength β is selected via grid search, yielding values of 10,000 for DNA, 1,000 for protein sequences, and 5,000 for small molecules. C.2. Use of Models All experiments were conducted using pre-trained models, with parameters kept frozen during guided sampling. For the DNA and protein sequence tasks, we employed a pre-trained discrete diffusion model trained on enhancer sequences and a pre-trained reverse folding model, respectively. The model architectures and associated reward functions were adopted directly from the repository provided by Wang et al. (2025). For the molecular generation task, we utilized a pre-trained multimodal flow model based on the EGNN architecture, as implemented by Lin et al. (2025). All reward functions and attribute predictors remained fixed throughout the experiments and were accessed solely during the inference phase. C.3. Evaluation metrics The following is a detailed introduction to the metrics used in our experiment. DNA Sequence Design. Following established conventions (Wang et al., 2025), we evaluate the generated enhancer sequences using the following metrics: • Predicted Activity (Pred-Activity): We measure the enhancer activity level in the HepG2 cell line using a reward oracle trained on a held-out evaluation subset. Crucially, the oracle used for guidance (or fine-tuning) is trained on a disjoint subset of data split by chromosomes, ensuring zero overlap with the evaluation oracle. • Chromatin Accessibility (ATAC-Acc): To validate whether the synthetic sequences correspond to accessible chromatin regions (a hallmark of active enhancers), we utilize an independent binary classification model trained on HepG2 chromatin accessibility data. While this metric is not used for optimization, it serves as an external validation of sequence biological plausibility. • 3-mer Pearson Correlation (3-mer Corr): We assess the distributional similarity between generated sequences and high-activity natural sequences. We calculate the Pearson correlation of 3-mer counts between the synthetic samples and the top 0.1% of sequences with the highest HepG2 activity in the reference dataset. • JASPAR Motif Analysis (JASPAR Corr): We analyze the biological relevance of generated motifs using JASPAR transcription factor binding profiles. We calculate the Spearman correlation of motif frequencies between the generated sequences and the top 0.1% of natural high-activity sequences, verifying if the model captures the motif patterns driving enhancer activity. • Approximated Log-Likelihood (App-Log-Lik): To quantify how natural the generated sequences appear to the pre-trained base model, we compute the approximate log-likelihood using the Evidence Lower Bound (ELBO) of the discrete diffusion model. Lower likelihoods indicate out-of-distribution sequences that may result from over-optimizing the reward oracle. Protein Inverse Folding. We evaluate the stability of generated sequences and their structural consistency with the target backbone. Note that for all evaluations, we condition on protein backbone conformations from the test set that were unseen during fine-tuning. 18

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

• Predicted Stability (Pred-ddG): We use an evaluation oracle trained on the full Megascale dataset (train, validation, and test splits) to predict protein stability (∆∆G). In contrast, the oracle used for guidance during sampling is trained strictly on the Megascale training set, preventing information leakage. • Self-consistency RMSD (scRMSD): To verify if the generated sequence folds into the target structure, we predict the structure of the generated sequence using ESMFold (?) and calculate the Root Mean Square Deviation (RMSD) between the predicted structure and the ground-truth wild-type backbone. • Success Rate: Following Campbell et al. (2024), we also define the success rate as the proportion of generated sequences that satisfy both stability and structural constraints: specifically, Pred-ddG > 0 and scRMSD < 2 Å. Molecular Generation. • Mean Absolute Error (MAE): To evaluate the alignment between generated molecules and the target property, we follow the protocol from (Bao et al., 2023). The MAE is calculated as: M

MAE =

1 X |ϕp (xi ) − ci | M i=1

where ϕp is the evaluation predictor, xi represents a generated molecule and ci denotes its target property.

4.31

2 0

90 80 70

DE MC3 MC5 MC8 MC10

73.0 DE MC3 MC5 MC8 MC10

1.1 1.0 0.93 0.88 0.90 0.91 0.9 0.8 0.7 0.62 0.6 0.5 DE MC3 MC5 MC8 MC10

1.0

0.93 0.94 0.94

0.9 0.8 0.7

0.95

0.72

App-Log-Lik ( )

93.9 95.2 95.3 96.3

100

3-mer Corr ( )

6.70

7.03 7.29 7.51

ATAC-Acc (\%) ( )

Pred-Activity ( )

110 8

JASPAR Corr ( )

D. Additional Results 255

265

-265

270

-267 -272

275

DE MC3 MC5 MC8 MC10

Methods of Value Estimation

-260 -261

260

DE MC3 MC5 MC8 MC10

0.5 0.0

0.55 DE MC5 MC10 MC15 MC20

0.93 0.95 0.96 0.98

1.0 0.8 0.6 0.51 0.4 0.2 0.0

DE MC5 MC10 MC15 MC20

1.02 1.00

1.01

1.02 0.99 0.98

0.98

0.97

0.96 DE MC5 MC10 MC15 MC20

Methods of Value Estimation

0.85

0.85 0.84 0.83 0.82 0.81

0.84 0.82

0.84

0.83

DE MC5 MC10 MC15 MC20

Success Rate ( )

1.0

1.2

%(scRMSD < 2) ( )

1.43

scRMSD ( )

1.38 1.23 1.30

1.5

%(ddG > 0) ( )

Pred-ddG ( )

Figure 5. Performance comparison on DNA sequence design. The deterministic estimation (DE) baseline is compared with Monte Carlo (MC) estimation using varying sample sizes (n = 3 ∼ 10). Metrics include predicted activity, chromatin accessibility (ATAC-Acc), motif correlations (3-mer, JASPAR), and log-likelihood. Red stars (⋆) indicate the best performance for each metric. 0.9

0.76 0.79

0.8

0.82 0.82

0.7 0.6 0.5 0.4

0.52 DE MC5 MC10 MC15 MC20

Figure 6. Performance comparison on protein sequence design. The deterministic estimation (DE) baseline is compared with Monte Carlo (MC) estimation using varying sample sizes (n = 5 ∼ 20). Metrics evaluate stability (Pred-ddG, positive proportion) and structural self-consistency (scRMSD) and success rate. Red stars (⋆) indicate the best performance for each metric.

D.1. Result of Molecular Generation In Tab. 4, the Upper bound corresponds to a diagnostic baseline obtained by randomly shuffling attribute labels within the unseen half of the QM9 training set, thereby removing correlations between molecular structures and target properties. The MAE is then computed using this shuffled dataset. If a method performs better than this upper bound, it indicates that the model successfully incorporates conditional attribute information into the generated molecules rather than relying on spurious correlations. The #Atoms baseline predicts molecular properties solely based on the number of atoms in the molecule. Performance surpassing this baseline suggests that the model captures conditional information beyond simple 19

82.6

7.08

0 0 0 0 0 0 0 700 800 900 1000 1100 1200 1300

82.50%

82.4 82.2 82.0 81.8

700

800

900

Polarizability MAE ( )

7.10 7.05 7.00 6.95 6.90 6.85

Success Rate (%)

Pred-Activity ( )

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

0 0 0 0 100 110 120 130

Guidance Strength

0.825 0.820 0.815 0.810 0.805

0.808 0 0 0 0 0 0 0 200 300 400 500 600 700 800

Figure 7. Ablation study on guidance strength β on GILC-DB. The impact of varying guidance strength is evaluated across three tasks. Red stars (⋆) indicate the optimal performance for each metric.

size-related cues in the generated molecular structures. Finally, the Lower bound represents the MAE achieved by directly predicting target properties using the pretrained property predictor itself, serving as an approximate oracle reference (Bao et al., 2023). To evaluate framework generality under more challenging scenarios, we incorporate a structure-guided molecular generation task constrained by non-differentiable reward functions. Here, molecular structures are characterized using discrete molecular fingerprints (Gebauer et al., 2022), and the Tanimoto coefficient (Bajusz et al., 2015) is utilized to measure structural similarity against target molecules. As reported in Tab. 6, GILC-PG significantly outperforms all baselines, achieving the highest similarity score. Overall, these findings highlight the effectiveness and generality of the GILC framework in guiding discrete and multimodal flows toward both differentiable and non-differentiable targets without requiring additional training. Table 6. Structural similarity results on the QM9 dataset under non-differentiable fingerprint rewards.

Method

Similarity ↑

Best-of-N (Beirami et al., 2025) SMC (Wu et al., 2023) SVDD (Li et al., 2024) TFG-Flow (Lin et al., 2025) GILC-PG (Ours)

0.182±0.016 0.178±0.002 0.234±0.011 0.271±0.006 0.308±0.004

(a) GILC-DB

(b) GILC-PG

Figure 8. Class-conditional images generated by GILC-DB and GILC-PG on CIFAR-10.

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

GILC-PG

GILC-DB

4.23

5.04

5.96

5.07

6.08

6.30

5.53

6.15

6.26

Pretrain

GILC-PG

GILC-DB

5.69

6.15

6.27

5.46

5.73

6.20

6.03

6.16

6.47

ant

goat

lizard

monkey

cat

lion

Pretrain

Figure 9. Improving aesthetic score for text-to-image generation. The aesthetic score is shown at the bottom of the image.

D.2. Result of Image Discrete Diffusion Figs. 8 and 9 present qualitative visualizations to intuitively demonstrate the generation quality and guidance efficacy of our proposed framework on two distinct image discrete diffusion tasks: class-conditional generation on CIFAR-10 (Campbell et al., 2022) and high-resolution text-to-image synthesis via Meissonic (Bai et al., 2024). D.3. Ablation Study We first perform a series of ablation studies to rigorously validate the effectiveness of our framework’s core components. Subsequently, we investigate the sensitivity of GILC-DB and GILC-PG to various hyperparameter configurations. Value Function Estimation. To validate the efficacy of our proposed variational proxy method for value function estimation, we conduct an ablation study by substituting this component with deterministic estimation methods from the literature (Li et al., 2024) while holding all other experimental factors constant. Evaluations are performed on DNA sequence design and protein inverse folding tasks. As illustrated in Figs. 5 and 6, our variational proxy consistently yields more effective guidance than deterministic estimators. Notably, increasing the number of Monte Carlo samples further enhances estimation accuracy, directly translating to superior performance. This suggests that the variational approach captures the underlying distribution of the discrete state space more effectively than point-based estimates. Impact of the Model Jacobian. We further evaluate the effect of excluding the discrete diffusion model Jacobian. Specifically, we compare generation results when the corrected gradient is computed with respect to both the clean prediction logits and the noisy state, versus when the model Jacobian term is omitted. As shown in Tabs. 7 and 8, removing the Jacobian term consistently improves guidance performance across tasks. This empirical observation is consistent with our theoretical analysis. In high-dimensional discrete settings, the model Jacobian is often severely ill-conditioned, which can introduce numerical instability into gradient-based guidance. By excluding this term, we obtain a more stable and reliable guidance signal, enabling more robust convergence toward the desired target attributes. Effect of Guidance Strength β. Our method relies on gradient-based corrections, making the guidance strength β a critical hyperparameter. As derived in Eq. 6, β controls the magnitude of the guidance signal. We visualize its effect on DNA enhancer optimization, protein stability enhancement, and molecular polarizability α guidance. As shown in Fig. 7, an appropriately chosen β leads to strong performance. In contrast, overly large values of β can reduce sample diversity and hinder exploration during early sampling, while overly small values provide insufficient guidance, both resulting in degraded performance. 21

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

Constant Linear Decay Exponential Decay

Early Stage (High Uncertainty)

Number of Calls

40 30

Fixed Total Budget

1280

20 10 0

100

120

Sampling Time Step

Figure 10. Illustration of constant, linear, and exponential decay strategies normalized to a fixed budget (N ≈ 128 × 10). The exponential decay schedule prioritizes early-stage evaluations to counter high initial uncertainty.

Allocation of Reward Function Calls. As shown in Fig. 2a, estimation errors of the value function are typically larger during the early stages of sampling. This suggests that, under a fixed budget of reward function evaluations, allocating more calls to early time steps is more effective than distributing them uniformly. We evaluate several scheduling strategies, including constant, linear decay, and exponential decay schedules (shown in Fig. 10). Our results in Tabs. 9 and 10 show that both linear and exponential decay schedules, which emphasize early-stage evaluations, substantially improve performance. Accordingly, we adopt exponential decay scheduling as the default in all experiments. Table 7. Ablation study on logit correction by omitting model Jacobian on regulatory DNA sequence design. Bold indicates the best performance within the same method. Method

Pred-Activity ↑

ATAC-Acc ↑ (%)

3-mer Corr ↑

JASPAR Corr ↑

App-Log-Lik ↑

GILC-DB

w/ Jacobian w/o Jacobian

4.18 7.04

48.8 95.2

0.816 0.900

0.904 0.935

−263 −267

GILC-PG

w/ Jacobian w/o Jacobian

3.44 5.21

34.7 84.0

0.518 0.910

0.673 0.937

−259 −270

Table 8. Ablation study on logit correction by omitting model Jacobian on protein sequence design. Bold indicates the best performance within the same method. Method

Pred-ddG ↑

%(ddG> 0) (%)↑

scRMSD ↓

%(scRMSD< 2)(%)↑

Success Rate (%)↑

GILC-DB

w/ Jacobian w/o Jacobian

0.809 1.430

79.5 97.9

0.928 0.968

87.6 84.3

70.1 82.4

GILC-PG

w/ Jacobian w/o Jacobian

0.528 0.719

44.7 75.6

0.941 0.914

89.2 92.4

48.9 69.8

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

Table 9. Ablation study of reward call schedules on DNA sequence generation using GILC-DB. The best performance is highlighted in bold, and the second best is underlined. Schedule Constant Linear Decay Exponential Decay

Pred-Activity ↑

ATAC-Acc ↑ (%)

3-mer Corr ↑

JASPAR Corr ↑

App-Log-Lik ↑

6.20 6.62 7.04

87.3 95.0 95.2

0.806 0.853 0.900

0.890 0.912 0.935

-274 -270 -267

Table 10. Ablation study of reward call schedules on protein backbone design using GILC-DB. The best performance is highlighted in bold, and the second best is underlined. Method Constant Linear Decay Exponential Decay

Pred-ddG ↑

%(ddG> 0) (%) ↑

scRMSD ↓

%(scRMSD< 2) (%) ↑

Success Rate (%) ↑

1.33 1.35 1.43

93.2 95.8 97.9

0.967 1.000 0.968

83.0 85.4 84.3

76.9 81.6 82.4

Record · ID 259490 · SHA-256 8f522a806a2ba07e
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.