Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding (2408.08252v5)

Abstract: Diffusion models excel at capturing the natural design spaces of images, molecules, DNA, RNA, and protein sequences. However, rather than merely generating designs that are natural, we often aim to optimize downstream reward functions while preserving the naturalness of these design spaces. Existing methods for achieving this goal often require ``differentiable'' proxy models (\textit{e.g.}, classifier guidance or DPS) or involve computationally expensive fine-tuning of diffusion models (\textit{e.g.}, classifier-free guidance, RL-based fine-tuning). In our work, we propose a new method to address these challenges. Our algorithm is an iterative sampling method that integrates soft value functions, which looks ahead to how intermediate noisy states lead to high rewards in the future, into the standard inference procedure of pre-trained diffusion models. Notably, our approach avoids fine-tuning generative models and eliminates the need to construct differentiable models. This enables us to (1) directly utilize non-differentiable features/reward feedback, commonly used in many scientific domains, and (2) apply our method to recent discrete diffusion models in a principled way. Finally, we demonstrate the effectiveness of our algorithm across several domains, including image generation, molecule generation, and DNA/RNA sequence generation. The code is available at \href{https://github.com/masa-ue/SVDD}{https://github.com/masa-ue/SVDD}.

• The paper presents SVDD, a novel derivative-free optimization approach that bypasses the need for fine-tuning or differentiable proxies.
• It integrates soft value functions with iterative sampling to achieve superior performance in tasks like image aesthetics and molecule docking.
• The method reduces computational costs while maintaining sample naturalness across continuous and discrete domains such as biological sequences.

An Overview of Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding

The paper "Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding" introduces a novel methodology for optimizing diffusion models, which are popular for generating realistic samples across domains such as image, molecule, and biological sequence generation. The focus is on enhancing the capability of these models to optimize downstream reward functions while preserving sample naturalness, without relying on fine-tuning or differentiated proxy models.

Key Contributions

The authors propose Soft Value-Based Decoding in Diffusion Models (SVDD), an iterative sampling method that integrates soft value functions into diffusion models for improved performance on downstream tasks. Unlike existing methods, SVDD does not require constructing differentiable proxy models or involve heavy computational fine-tuning, which are common challenges when leveraging pre-trained diffusion models for real-world applications.

Key aspects of SVDD include:

• Derivative-Free Optimization: SVDD eliminates the need to construct differentiable proxy models, allowing direct optimization based on non-differentiable features or reward feedback common in scientific applications, thus broadening applicability.
• Non-Fine-Tuning Approach: By decoupling the optimization process from the fine-tuning of generative models, SVDD significantly cuts down on computational resources, making it scalable for large models often referred to as "foundation models."
• Algorithmic Efficiency: Through a combination of importance sampling (IS) and resampling based on soft value functions, SVDD efficiently approximates a value-weighted policy, showing strong empirical performance across domains including image aesthetics and molecule docking score optimization.

Numerical Results and Implications

The paper substantiates SVDD's effectiveness through experiments in diverse generative tasks:

• In image generation, specifically for optimizing aesthetic scores and compressibility, SVDD outperformed other inference-time methods, maintaining high fidelity to the natural image space while producing samples with higher aesthetic and compression scores.
• In molecule generation, SVDD achieved superior scores across metrics like QED and docking performance against several proteins, demonstrating its potential utility in drug discovery.
• For biological sequences, where discrete diffusion models were used, SVDD showcased its flexibility and robustness in optimizing DNA enhancer and RNA 5’UTR activity levels.

The authors compare SVDD to traditional approaches such as classifier guidance, noting the latter's inherent limitation of requiring differentiable gradients. SVDD's derivative-free nature and adaptive sampling approach make it particularly suited for scientific domains where simulation-based feedback is prevalent and non-differentiable.

Extensions and Future Work

Several extensions to SVDD are proposed, such as combining it with proposal distributions different from pre-trained models, or even using it for model distillation when fine-tuning is necessary to reduce inference latency. Moreover, SVDD's robustness against reward exploitation—a common issue in RL-based fine-tuning—enhances its utility as a reliable optimization tool.

This approach opens pathways for future research into more generalizable algorithms applicable across diverse downstream tasks without heavily relying on fine-tuning adjustments. Potential developments include enhancing the method for further application in domains like protein design and real-world problem-solving in complex multi-objective environments where trade-offs between competing objectives are often necessary.

In summary, "Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding" contributes a significant stride in diffusion model optimization, making them more applicable across various disciplines while ensuring operational efficiency and adherence to natural data characteristics.