Dynamic Search for Inference-Time Alignment in Diffusion Models
This addresses the challenge of inference-time alignment in diffusion models for applications requiring non-differentiable rewards, representing a novel method for a known bottleneck.
The paper tackles the problem of aligning diffusion model outputs with non-differentiable reward functions by framing inference-time alignment as a search problem and proposing Dynamic Search for Diffusion (DSearch), which achieves superior reward optimization across domains like biological sequence design, molecular optimization, and image generation.
Diffusion models have shown promising generative capabilities across diverse domains, yet aligning their outputs with desired reward functions remains a challenge, particularly in cases where reward functions are non-differentiable. Some gradient-free guidance methods have been developed, but they often struggle to achieve optimal inference-time alignment. In this work, we newly frame inference-time alignment in diffusion as a search problem and propose Dynamic Search for Diffusion (DSearch), which subsamples from denoising processes and approximates intermediate node rewards. It also dynamically adjusts beam width and tree expansion to efficiently explore high-reward generations. To refine intermediate decisions, DSearch incorporates adaptive scheduling based on noise levels and a lookahead heuristic function. We validate DSearch across multiple domains, including biological sequence design, molecular optimization, and image generation, demonstrating superior reward optimization compared to existing approaches.