Inference-time Scaling of Diffusion Models through Classical Search

2025-05-29

AI Summary

This work addresses the challenge of flexibly adapting diffusion models to diverse downstream test objectives during inference—without retraining or architectural modification. We propose the first inference-time expansion framework grounded in classical search paradigms. Methodologically, we uniquely integrate annealed Langevin MCMC (for local search) with breadth-first or depth-first tree search (for global exploration) directly into the diffusion sampling process, enabling dynamic, differentiable, and objective-driven control over generation trajectories. Our framework significantly improves generation quality and objective satisfaction rates across planning, offline reinforcement learning, and image synthesis tasks, while reducing computational overhead. Extensive experiments demonstrate that classical search serves as a general, practical, and effective mechanism for inference-time control in diffusion models, establishing a novel paradigm for controllable generation.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Computer Vision: Diffusion Models for Vision
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: ML for personalized search and recommendations
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Classical search algorithms have long underpinned modern artificial intelligence. In this work, we tackle the challenge of inference-time control in diffusion models -- adapting generated outputs to meet diverse test-time objectives -- using principles from classical search. We propose a general framework that orchestrates local and global search to efficiently navigate the generative space. It employs a theoretically grounded local search via annealed Langevin MCMC and performs compute-efficient global exploration using breadth-first and depth-first tree search. We evaluate our approach on a range of challenging domains, including planning, offline reinforcement learning, and image generation. Across all tasks, we observe significant gains in both performance and efficiency. These results show that classical search provides a principled and practical foundation for inference-time scaling in diffusion models. Project page at diffusion-inference-scaling.github.io.

Problem

• Inference-time control in diffusion models for diverse objectives
• Efficient generative space navigation via local and global search
• Performance and efficiency gains in planning, RL, and image generation

Innovation

• Uses classical search for diffusion model control
• Combines local and global search strategies
• Applies annealed Langevin MCMC and tree search