Dynamic Search for Inference-Time Alignment in Diffusion Models

2025-03-03

AI Summary

Aligning diffusion model outputs with non-differentiable reward functions at inference time remains challenging. This paper formulates the problem as a dynamic search process and proposes a noise-aware adaptive beam search framework: it approximates intermediate-node rewards via subsampled denoising trajectories, designs a lookahead heuristic function, and dynamically adjusts beam width and tree expansion strategies according to noise level. To our knowledge, this is the first work to cast diffusion-based inference-time alignment as a dynamic search problem. Evaluated on biological sequence design, molecular optimization, and image generation, our method achieves substantial improvements in reward scores—ranging from +12.7% to +34.1%—while preserving generation quality, consistently outperforming existing gradient-free guidance approaches.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Constraint Satisfaction and Optimization: Search
• Computer Vision: Diffusion Models for Vision

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

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.

Problem

• Aligning diffusion model outputs with non-differentiable reward functions.
• Optimizing inference-time alignment using gradient-free guidance methods.
• Enhancing reward optimization in diverse domains like biology and image generation.

Innovation

• Dynamic Search for Diffusion (DSearch) optimizes alignment.
• DSearch dynamically adjusts beam width and tree expansion.
• Adaptive scheduling and lookahead heuristic refine decisions.