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

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.

Many important scientific problems—such as designing new proteins, optimizing DNA sequences, or discovering drug molecules—require computers to generate new candidates with specific desired properties. Current AI systems based on discrete diffusion models can produce realistic biological sequences, but controlling them to meet precise requirements (e.g., a protein that is more stable) usually requires expensive retraining or additional models. This paper introduces GILC, a lightweight method that guides existing discrete diffusion models toward desired goals without any retraining. The key idea is to reuse the model's own predictions to estimate how to adjust the output in a desired direction. By applying a simple correction at each step of the generation process, GILC steers the model toward sequences that satisfy user-specified criteria. The method works for both numerical targets and black-box quality scores, such as an expert lab test result. Experiments on DNA, protein, and molecular generation tasks show that GILC matches or outperforms approaches that require full retraining, while being plug-and-play—meaning it can be combined with any existing discrete diffusion model immediately.