Dynamic Search for Diffusion

• DSearch is a dynamic family of techniques that reformulates diffusion model inference as a reward-driven search, integrating methods such as Monte Carlo Tree Search, differentiable solver search, and neural architecture search.
• It enables adaptive computation allocation through strategies like dynamic beam search and cyclic inference to boost sample efficiency and quality in applications from image generation to game planning.
• DSearch optimizes the balance between exploration and exploitation by leveraging state-of-the-art planning and search techniques, reducing compute costs while achieving superior alignment and fidelity.

Dynamic Search for Diffusion (DSearch) refers to a broad family of techniques and frameworks that transform inference in diffusion models from a fixed, monolithic process into a flexible, data- or reward-driven search problem. DSearch approaches—spanning Monte Carlo Tree Search (MCTS)–inspired inference, differentiable solver search, segment-wise neural architecture search, beam and cyclic search during generation, and even model-parameter search over entire diffusion processes—share the principle of adaptively allocating computation to maximize generation quality under a fixed compute budget. These methods have set state-of-the-art sample efficiency and alignment on a variety of tasks, ranging from image and molecule generation to long-horizon game planning, sequence alignment, and video synthesis.

1. Formalization: Inference-Time Alignment and Search

DSearch methods recast inference in a pretrained diffusion model as a reward-guided search through the reverse diffusion chain. The canonical generative model is

$$p_\theta(x_0, \ldots, x_T) = p(x_T) \prod_{t=1}^T p_\theta(x_{t-1} \mid x_t)$$

with inference-time alignment seeking to sample from the reward-tilted distribution

$$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$

Marginalizing latents yields the posterior for $x_0$:

$\pi^*(x_0) \propto p_\theta(x_0) \exp r(x_0).$

This generalizes class/prompt-conditional sampling, reward-guided alignment, and non-differentiable acceptability constraints. The key DSearch insight is that this sampling can be formulated as optimal control, stochastic planning, or search—where the reverse denoising process is a multi-step decision trajectory, and optimal samples correspond to high-reward leaves in a search tree or trajectory beam (Jain et al., 25 Jun 2025, Li et al., 3 Mar 2025).

2. Monte Carlo Tree Search and Tree-Based Policy Improvement

The core instance of DSearch is Diffusion Tree Sampling (DTS) (Jain et al., 25 Jun 2025), which instantiates the denoising chain as a finite-horizon MCTS on a depth-$T$ tree. States correspond to partially denoised samples at timestep $t$; transitions correspond to sampling $x_{t-1}$ from $p_\theta(\cdot \mid x_t)$. Each full rollout $$x_T \rightarrow x_{T-1} \rightarrow \cdots \rightarrow x_0$$ terminates in a reward $r(x_0)$.

A crucial mechanism is soft value backup along the tree, mirroring the soft Bellman recursion:

$$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$0

with $$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$1. Empirical value estimates $$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$2 are updated via log-sum-exp over descendant nodes, leading to refined sampling of $$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$3 weighted by exponentiated value. DTS thus reuses all intermediate partial rollouts, increasing estimator accuracy at high noise levels. Sampling from the tree with Boltzmann (softmax) branching converges asymptotically to the reward-aligned posterior (Jain et al., 25 Jun 2025).

The greedy variant, DTS*, employs value-maximizing policy selection at each node, returning global best leaves. Progressive widening, visit count–based exploration bonuses, and anytime sampling (immediate output at any search budget) provide scalability and controllable compute-quality trade-offs. This tree search approach achieves state-of-the-art FID with up to $$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$4 compute reduction and supports reward functions unavailable to gradient-based guidance (Jain et al., 25 Jun 2025).

3. Differentiable Solver and Architecture Search

DSearch also denotes trainable, data-driven search over algorithmic scheduling of diffusion steps and solvers:

• Differentiable Solver Search (Wang et al., 27 May 2025): The reverse diffusion is treated as an ODE integration problem. Classical multistep solvers (Adams–Bashforth) use hand-derived, timestep-only weights. DSearch instead directly optimizes solver update times $$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$5 and multistep coefficients $$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$6 via differentiable loss, minimizing the trajectory error vs. a high-fidelity Euler baseline. This approach yields simple closed-form updates:

$$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$7

with learned, tabulated coefficients that transfer seamlessly across models, architectures, and resolutions. Results include state-of-the-art FID at 10 sampling steps on ImageNet256 ($$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$8 for DiT-XL/2, vs. $$p^*(x_0, \ldots, x_T) \propto p_\theta(x_0, \ldots, x_T) \exp R(x_0).$$9/$x_0$0 for baseline solvers).

• Segment-Wise NAS (Flexiffusion) (Huang et al., 3 Jun 2025): Instead of uniform or step-wise NAS—which is combinatorially intractable ($x_0$1 search space)—DSearch decomposes the denoising sequence into $x_0$2 segments, each segment dynamically choosing between full, partial (cache-reused), and null (skipped) computation. Segmentwise evolutionary search is conducted using an rFID metric computed against teacher model outputs, and architectural candidates are mutated with per-segment granularity. This reduces search space exponentially, yielding 2–5.1$x_0$3 speedups at $x_0$4 FID/CLIP degradation.

4. Dynamic Beam, Lookahead, and Cyclic Inference

DSearch methods extend beyond tree search to robust inference-time alignment mechanisms:

• Dynamic Beam Search and Lookahead (Oshima et al., 31 Jan 2025, Li et al., 3 Mar 2025): In text-to-video, image, or language sequence generation, DSearch instantiates a beam of hypotheses at each denoising step, expanding each via the pretrained kernel, and scoring either via one-step or $x_0$5-step lookahead, using either reconstructed sample predictions or short deterministic DDIM rollouts for reward estimation. Adaptive beam scheduling (size, width by noise), search scheduling (aggressive expansion near low noise), and per-beam pruning boost reward optimization and diversity, outperform static best-of-$x_0$6 and SMC across text, molecule, and image tasks (Li et al., 3 Mar 2025).
• Cyclic Diffusion Search (ABCD) (Lee et al., 20 May 2025): CDS employs alternating reverse (jumpy denoising) and forward (re-noising back to higher $x_0$7) traversals, governed by a temperature pool and exploration-exploitation balancing across cycles. Adaptive thinking time terminates the search once all top-$x_0$8 survivors originate from fine-scale refinements (no go-back), implementing instance-adaptive compute allocation. ABCD matches or exceeds existing search and SMC baselines in broad domains while using per-instance dynamic compute.
• Confidence-Switched Search in DLMs (Cao et al., 11 Feb 2026): For diffusion LLM decoding, dynamic search interleaves position beam search (exploration at low confidence) and parallel decoding (acceleration at high confidence), using per-token model uncertainty as the switch. This improves code generation and math reasoning benchmarks, closing much of the quality gap to static beam methods at near-greedy speed.

5. Model and Scheduling Search in Training

DSearch is additionally realized as automatic search over the diffusion process itself:

• Joint Optimization of Diffusion Schedules and SDEs (Singhal et al., 2023): Automatic Multivariate Diffusion Training (AMDT) treats the design of multivariate diffusion processes as a parameter optimization problem. By parameterizing both the drift (skew-symmetric and symmetric positive semi-definite blocks) and the noise schedule within a general linear SDE, the algorithm optimizes the Evidence Lower Bound (ELBO) over this process family. The result is a fully automated search over inference diffusions, producing new high-order stochastic processes (e.g., ALDA, MALDA) that outperform fixed diffusions at matched model size and architecture.

6. Connections to Planning, Search, and Computational Efficiency

DSearch generalizes well-known planning and search frameworks from reinforcement learning (Monte Carlo Tree Search, bandit strategies, lookahead planning, optimism in face of uncertainty) to the domain of diffusion models—where the stochastic chain structure and sample efficiency constraints are critical. By recycling partial rollouts, dynamically adjusting search width and depth, and exploiting bandit-style allocation of exploration vs. exploitation, DSearch circumvents the exponential search cost of naïve tree search.

In cost- and exploration-aware settings (e.g., active search over action sequences) (Banerjee et al., 23 Feb 2026), DSearch provides a scalable alternative to rollout-based MCTS, using batch diffusion sampling and gradient-based reward guidance to recover near-optimal long-horizon plans. Empirically, these DSearch-based planners match or beat tree search at substantially reduced inference cost.

7. Limitations and Future Directions

DSearch frameworks introduce complexity in rollout management, value estimation, and hyperparameter tuning (beam width, scheduling, search budgets). Theoretical guarantees (asymptotic convergence to reward-tilted posterior (Jain et al., 25 Jun 2025)) exist for the principal tree-based methods; other search variants (segment NAS, cyclic search) are mostly substantiated by empirical efficiency and ablation studies. Diversity can collapse under aggressive beam-based or batch-resample search; maintaining balance between diversity and reward remains an explicit research target (Li et al., 3 Mar 2025). Scaling DSearch to high-dimensional latent models, aligning reward approximation with human or model preferences (via calibrated metrics), integrating value networks, and leveraging asynchronous distributed rollouts are active areas of development (Jain et al., 25 Jun 2025).

---

Key References

Approach	Main Contribution	arXiv ID
Diffusion Tree Sampling	Reward-aligned MCTS tree search for inference	(Jain et al., 25 Jun 2025)
Differentiable Solver	Solver schedule search for ODE integration	(Wang et al., 27 May 2025)
Flexiffusion	Segment-wise, cache-adaptive NAS for inference	(Huang et al., 3 Jun 2025)
Dynamic Beam Search	Adaptive beam + lookahead in sequence/video diff.	(Oshima et al., 31 Jan 2025, Li et al., 3 Mar 2025)
Cyclic Diffusion Search	Bidirectional, adaptive cycles for instance scaling	(Lee et al., 20 May 2025)
AMDT (Diff SDE Search)	Joint ELBO optimization over SDE/diffusion process	(Singhal et al., 2023)

DSearch thus provides a unified paradigm for efficient, reward-driven, and scalable inference in diffusion generative models, subsuming tree-based, architectural, algorithmic, and optimizer-level dynamic search innovations across modalities and tasks.