Adaptive Inference-Time Scaling via Cyclic Diffusion Search

TL;DR

Adaptive Bi-directional Cyclic Diffusion (ABCD) rethinks inference-time compute by casting diffusion-model inference as a flexible search. It combines Cyclic Diffusion Search, Automatic Exploration-Exploitation Balancing, and Adaptive Thinking Time to allocate computation adaptively per instance and terminate when further gains are unlikely. Across planning, maze solving, Sudoku, molecule generation, and text-to-image tasks, ABCD yields stronger performance under comparable or lower compute than fixed-schedule baselines. This work demonstrates that instance-aware inference-time scaling can significantly improve both efficiency and outcome quality in diffusion-based reasoning and generation.

Abstract

Diffusion models have demonstrated strong generative capabilities across domains ranging from image synthesis to complex reasoning tasks. However, most inference-time scaling methods rely on fixed denoising schedules, limiting their ability to allocate computation based on instance difficulty or task-specific demands adaptively. We introduce the challenge of adaptive inference-time scaling-dynamically adjusting computational effort during inference-and propose Adaptive Bi-directional Cyclic Diffusion (ABCD), a flexible, search-based inference framework. ABCD refines outputs through bi-directional diffusion cycles while adaptively controlling exploration depth and termination. It comprises three components: Cyclic Diffusion Search, Automatic Exploration-Exploitation Balancing, and Adaptive Thinking Time. Experiments show that ABCD improves performance across diverse tasks while maintaining computational efficiency.

Figures (21)

• Figure 1: Overview of our method. (1) Start with $N$ particles, use jumpy denoising $x_T \rightarrow x_{T- j}\rightarrow \cdot\cdot\cdot \rightarrow x_{0}$. (2) Replicate each particles $J$ times and send each particles to multiple noise levels (AEEB). (3) Re-denoise particles. (4) Select Top $K$ particles. (5) We repeat steps (2)–(4) until the adaptive terminal condition is satisfied (CDS). The number of cycles executed in this way defines the Adaptive Thinking Time (ATT).
• Figure 2: Evaluated Tasks. (a) Mixture of Gaussian for Proof-of-concept that we need multiple go back temperature. (b) Sudoku puzzle completion demanding logical consistency across multiple constraints simultaneously; (c) Pixel Maze Path Finding testing generalization to novel environmental structures not encountered during training; (d) Molecular structure prediction requiring physically and chemically valid 3D conformations. (e) OGBench PointMaze navigation requiring complex planning over 1000+ steps (f) Text-to-image generation requiring complex, high-dimensional outputs that are well aligned with the given textual descriptions;
• Figure 3: Mean accuracy across four difficulty levels—Very Hard (17-19 provided), Hard (20-21 provided), Medium (22-24 provided), and Easy (25-27 provided).
• Figure 4: Left: Sudoku Puzzle Completion result. Mean accuracy on Harder dataset (17-28 entities provided) by giving more computational budget (sec.). Middle: Pixel maze path finding result. Success rate on the OOD Pixel Maze size-15 test set. Right: Molecular 3D structure prediction task result. Molecular Stability rate on the QM9 dataset. For each method, inference-time computation was scaled as much as possible along its core controllable axis for inference time scaling to ensure a fair comparison under expanded budgets.
• Figure 5: Pixel Maze path finding results on multiple maze size. Success rate comparison across maze sizes in the Pixel Maze task. All settings involve OOD evaluation, where models are trained on smaller mazes (sizes 4$\sim$6) and tested on larger ones. As maze size increases, the search space grows and the performance gap between other methods widens, highlighting the importance of adaptive inference-time exploration in complex scenarios.

Theorems & Definitions (7)

• Theorem 1
• Lemma E.1
• proof
• Lemma E.2
• proof
• Lemma E.3
• proof