Reward-Guided Discrete Diffusion via Clean-Sample Markov Chain for Molecule and Biological Sequence Design

TL;DR

This paper introduces the CSMC sampler, a novel method for reward-guided discrete diffusion that improves molecule and biological sequence generation by effectively performing local search without relying on noisy intermediate rewards.

Contribution

The paper proposes the Clean-Sample Markov Chain (CSMC) sampler, enabling effective test-time reward-guided sampling for discrete diffusion models without using intermediate rewards.

Findings

CSMC outperforms prior reward-guided methods in molecule and sequence generation.

The method effectively constructs a Markov chain with the target distribution as stationary.

Experiments demonstrate improved sample quality with various reward functions.

Abstract

Discrete diffusion models have recently emerged as a powerful class of generative models for chemistry and biology data. In these fields, the goal is to generate various samples with high rewards (e.g., drug-likeness in molecules), making reward-based guidance crucial. Most existing methods are based on guiding the diffusion model using intermediate rewards but tend to underperform since intermediate rewards are noisy due to the non-smooth nature of reward functions used in scientific domains. To address this, we propose Clean-Sample Markov Chain (CSMC) Sampler, a method that performs effective test-time reward-guided sampling for discrete diffusion models, enabling local search without relying on intermediate rewards. CSMC constructs a Markov chain of clean samples using the Metropolis-Hastings algorithm such that its stationary distribution is the target distribution. We design a proposal distribution by sequentially applying the forward and backward diffusion processes, making the acceptance probability tractable. Experiments on molecule and biological sequence generation with various reward functions demonstrate that our method consistently outperforms prior approaches that rely on intermediate rewards.