Plug-and-Play Guidance for Discrete Diffusion Models via Gradient-Informed Logit Correction
About
Controllable generation with discrete diffusion models is often hindered by high computational overhead or the need for retraining. In this paper, we present \underline{\textbf{G}}radient-\underline{\textbf{I}}nformed \underline{\textbf{L}}ogit \underline{\textbf{C}}orrection (\textbf{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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| DNA sequence design | DNA sequence design | Pred-Activity7.04 | 23 | |
| Inverse protein folding | Inverse protein folding | Pred-ddG1.43 | 11 | |
| Molecular property generation | QM9 | Heat Capacity (Cv)2.23 | 9 | |
| Structure-guided molecular generation | QM9 | Similarity Score30.8 | 5 |