Unlocking Guidance for Discrete State-Space Diffusion and Flow Models

2024-06-03

arXiv.org

AI Summary

Discrete-state-space generative models—e.g., for small molecules, DNA, and protein sequences—lack principled, controllable guidance mechanisms. Existing continuous-domain guidance paradigms do not generalize to discrete spaces, hindering attribute-controlled generation. Method: This paper introduces the first general, differentiable guidance framework for discrete generative models based on continuous-time Markov processes, unifying discrete diffusion and flow-matching architectures. It overcomes the fundamental incompatibility of continuous guidance with discrete state spaces by establishing a theoretically grounded guidance theory for discrete domains. The framework enables arbitrary differentiable guidance objectives without model retraining, leveraging probability path reweighting and gradient-driven discrete sampling. Contribution/Results: Experiments demonstrate substantial improvements in target property satisfaction rates and sample diversity across diverse biomolecular generation tasks, while maintaining flexibility and strong generalization across guidance objectives and model architectures.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Mixed Discrete/Continuous Search
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Generative models on discrete state-spaces have a wide range of potential applications, particularly in the domain of natural sciences. In continuous state-spaces, controllable and flexible generation of samples with desired properties has been realized using guidance on diffusion and flow models. However, these guidance approaches are not readily amenable to discrete state-space models. Consequently, we introduce a general and principled method for applying guidance on such models. Our method depends on leveraging continuous-time Markov processes on discrete state-spaces, which unlocks computational tractability for sampling from a desired guided distribution. We demonstrate the utility of our approach, Discrete Guidance, on a range of applications including guided generation of small-molecules, DNA sequences and protein sequences.

Problem

• Extending guidance to discrete state-space diffusion models
• Enabling controllable generation in discrete state-spaces
• Applying guidance to molecular and biological sequences

Innovation

• General method for discrete state-space guidance
• Leverages continuous-time Markov processes
• Enables guided generation of biomolecules