Context-Guided Diffusion for Out-of-Distribution Molecular and Protein Design

TL;DR

This paper introduces context-guided diffusion (CGD), a plug-and-play method that enhances out-of-distribution sampling in generative models for molecular and protein design, leveraging unlabeled data and smoothness constraints.

Contribution

The authors propose a novel context-guided diffusion approach that improves out-of-distribution generation in guided diffusion models for molecular and protein design tasks.

Findings

Significant performance improvements in out-of-distribution sampling.

Effective across continuous, discrete, and graph-structured diffusion processes.

Applicable to drug discovery, materials science, and protein design.

Abstract

Generative models have the potential to accelerate key steps in the discovery of novel molecular therapeutics and materials. Diffusion models have recently emerged as a powerful approach, excelling at unconditional sample generation and, with data-driven guidance, conditional generation within their training domain. Reliably sampling from high-value regions beyond the training data, however, remains an open challenge -- with current methods predominantly focusing on modifying the diffusion process itself. In this paper, we develop context-guided diffusion (CGD), a simple plug-and-play method that leverages unlabeled data and smoothness constraints to improve the out-of-distribution generalization of guided diffusion models. We demonstrate that this approach leads to substantial performance gains across various settings, including continuous, discrete, and graph-structured diffusion processes with applications across drug discovery, materials science, and protein design.