Distributional Multi-objective Black-box Optimization for Diffusion-model Inference-time Multi-Target Generation

TL;DR

This paper introduces the IMG algorithm, which optimizes diffusion models at inference-time for efficient multi-objective generation, outperforming traditional methods in molecule design tasks.

Contribution

The paper proposes a novel inference-time optimization method for diffusion models that directly generates multi-objective samples, improving efficiency and performance over existing approaches.

Findings

IMG achieves higher hypervolume with a single diffusion pass.

The multi-target Boltzmann distribution has a log-likelihood interpretation.

The method can be integrated into existing diffusion-based algorithms.

Abstract

Diffusion models have been successful in learning complex data distributions. This capability has driven their application to high-dimensional multi-objective black-box optimization problem. Existing approaches often employ an external optimization loop, such as an evolutionary algorithm, to the diffusion model. However, these approaches treat the diffusion model as a black-box refiner, which overlooks the internal distribution transition of the diffusion generation process, limiting their efficiency. To address these challenges, we propose the Inference-time Multi-target Generation (IMG) algorithm, which optimizes the diffusion process at inference-time to generate samples that simultaneously satisfy multiple objectives. Specifically, our IMG performs weighted resampling during the diffusion generation process according to the expected aggregated multi-objective values. This weighted resampling strategy ensures the diffusion-generated samples are distributed according to our desired multi-target Boltzmann distribution. We further derive that the multi-target Boltzmann distribution has an interesting log-likelihood interpretation, where it is the optimal solution to the distributional multi-objective optimization problem. We implemented IMG for a multi-objective molecule generation task. Experiments show that IMG, requiring only a single generation pass, achieves a significantly higher hypervolume than baseline optimization algorithms that often require hundreds of diffusion generations. Notably, our algorithm can be viewed as an optimized diffusion process and can be integrated into existing methods to further improve their performance.