Title: Test-Time Scaling of Diffusion Models via Noise Trajectory Search
Venue: NeurIPS 2025
Accepted: yes
Keywords: diffusion model

test-time scaling

Markov decision process

contextual bandit

optimization

search

MCTS

epsilon-greedy
Abstract: The iterative and stochastic nature of diffusion models enables *test-time scaling*, whereby spending additional compute during denoising generates higher-fidelity samples. Increasing the number of denoising steps is the primary scaling axis, but this yields quickly diminishing returns. Instead optimizing the *noise trajectory*—the sequence of injected noise vectors—is promising, as the specific noise realizations critically affect sample quality; but this is challenging due to a high-dimensional search space, complex noise-outcome interactions, and costly trajectory evaluations. We address this by first casting diffusion as a Markov Decision Process (MDP) with a terminal reward, showing tree-search methods such as Monte Carlo tree search (MCTS) to be meaningful but impractical. To balance performance and efficiency, we then resort to a relaxation of MDP, where we view denoising as a sequence of independent *contextual bandits*. This allows us to introduce an $\epsilon$-greedy search algorithm that *globally explores* at extreme timesteps and *locally exploits* during the intermediate steps where de-mixing occurs. Experiments on EDM and Stable Diffusion reveal state-of-the-art scores for class-conditioned/text-to-image generation, exceeding baselines by up to $164$% and matching/exceeding MCTS performance. To our knowledge, this is the first practical method for test-time noise *trajectory* optimization of *arbitrary (non-differentiable)* rewards.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis is a quite interesting and meaningful work, presenting a novel test-time scaling algorithm for diffusion models. The core idea is to treat each sampling step as an independent contextual bandit problem and then use an ϵ-greedy search algorithm to achieve a more comprehensive search. Compared to some previous test-time scaling algorithms, it is remarkably efficient and performs exceptionally well.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**:\n\n1. The problem of test-time scaling in diffusion models is notably challenging. The method proposed by the authors is highly novel, and I personally believe its exceptional performance stems from the $\\epsilon$-greedy search algorithm, which ensures that each sampling action is based on the previous iteration, thereby achieving a \"cumulative gain\" effect.\n\n2. The authors' proposed method is simple and efficient, which does not detract from its novelty or its beneficial impact on other works. For instance, this test-time scaling algorithm might enable step distillation algorithms for diffusion models to achieve superior results.\n\n3. The writing and logic of the article is clear and very engaging for the reader.\n\n**Weakness**:\n\nI believe this paper has two main shortcomings.\n\n1. The experiments in the paper are relatively limited. The authors might consider conducting further experiments on SDXL, SD3.5, and FLUX, evaluating them on benchmarks such as HPS v2, Image-Reward, GenEval, and T2I-CompBench.\n\n2. The approach of predicting $x_0$ solely using Tweedie's formula with a base model remains too crude, potentially yielding relatively low benefits, especially during the initial sampling stages. The authors could consider using a few-step model distilled from the corresponding base model to execute Tweedie's formula, which might lead to better results.\n\nQUESTIONS:\nNo.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper frames test time scaling of diffusion as Markov Decision Process (MDP) with a terminal reward. To avoid excessive computation costs of tree-search methods, they view denoising as a sequence of independent contextual bandits. Then, they introduce $\\epsilon$-greedy search algorithm.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. The paper tackles important problem, test-time scaling of diffusion models.\n\nWeakness\n1. Lack of novelty\n\t- The only novel component is $\\epsilon$-greedy noise search.\n\t- SVDD [1] has already proposed method to deal with non-differentiable rewards unlike the last sentence of the abstract.\n2. Lack of discussion on the proposed method.\n\t- The paper does not discuss why $\\epsilon$-greedy method improves the sampling.\n3. Unclear explanation of overall sampling process\n\t- In Algorithm 1, there is no explanation about how the denoising network $D_\\theta$ is used. It seems the written algorithm is incomplete.\n4. Limited and unconvincing experimental support\n\t- Insufficient scale: The valuation on ImageNet is based on only 36 images, which is too small sample size to support or claim meaningful conclusion.\n\t- Oversimplified tasks: The text prompts used for evaluation are very simple. To demonstrate the true value of a test-time scaling method, the experiments should involve more complex prompts which includes compositional information.\n\t- Quality degradation: The qualitative results (Example 2 of Figure 5, Example 3 of Figure 6) suggest, the proposed method may degrade image quality like color saturation. The authors should discuss regarding tradeoff between sample fidelity and text alignment. I suggest authors to include quantitative metrics regarding this. (measuring FID or IS along with Brightness/ Compressibility/ Classifier)\n[1] Li et al., Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding,\n\nQUESTIONS:\n1. Refer to weakness above.\n2. How many images are used for evaluation on Stable Diffusion?\n3. I suggest to show qualitative results in the main paper.\n\nLIMITATIONS:\n1. Refer to weakness above.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 1, "ORIGINALITY": 1, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work propose an efficient test-time ω-greedy search algorithm, inspired by the contextual bandit problem, for optimizing noisy trajectories in diffusion models under non-differentiable reward constraints  and for efficiently optimizing noisy trajectories in diffusion models by formulating denoising as a contextual bandit problem.  It empirically outperforms Monte Carlo Tree Search  and demonstrates global adaptive exploration and local exploitation  behaviors associated with specific diffusion steps, especially being more localized in intermediate steps involving demixing. This method achieves state-of-the-art performance in image generation and is also the first practical method to optimize non-differentiable rewards during test-time sampling.\n\nSTRENGTHS AND WEAKNESSES:\nThis study proposes an efficient test-time ω-greedy search algorithm inspired by contextual bandits for optimizing noisy trajectories in diffusion models under non-differentiable reward constraints. The algorithm empirically outperforms Monte Carlo Tree Search and demonstrates adaptive exploration  and local exploitation behaviors associated with specific diffusion steps, especially in the intermediate steps involving de-mixing. Through extensive comparisons with other search strategies, this study quantitatively demonstrates the critical importance of step-size adaptive search locality and shows the effectiveness of the method on EDM and stable diffusion, achieving sample quality improvement on various non-differentiable rewards such as classifier probability and VLM alignment. ​ \n\nWeaknesses: While I am not very familiar with this specific research area, the authors' presentation and experimental results seem methodologically sound and logically consistent. However, I was unable to find broader implications, limitations, and future work in the appendix. In addition, the authors only provide a small number of baselines to verify the effectiveness of the proposed testing-time strategy. I would like to know the impact of the authors' algorithm on a wider range of baselines.\n\nQUESTIONS:\nRefer to weaknesses part.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

comments: SUMMARY:
This paper presents a practical method for optimizing noise trajectories at test time for diffusion models. By framing sampling as an MDP and relaxing it to a sequence of contextual bandits, the authors propose an efficient epsilon greedy search algorithm that adapts global vs. local exploration based on the denoising step. The method shows improved generation quality on EDM and Stable Diffusion.

STRENGTHS AND WEAKNESSES:
Strength:

1. The empirical validations are provided for both class-conditional and text-to-image generation.

2. I find following the flow of this paper smooth, which is good.

Weakness:

1. No theoretically grounded analysis or guarantees for the search algorithm's performance.
2. To me, the compute cost of this method is not light, which is largely due to its MCTS nature. 
3. The related work part lacks many discussions of current literature on both the alignment of diffusion models and inference-time diffusion model scaling. I would encourage the authors to pay more attention to the literature. Besides, I find the experimental setting a bit weird. Many essential baselines are omitted in Table 2. There is a whole line of SMC-based algorithms that can effectively align DMs in inference-time (also requires some compute but can also scale well, for example https://arxiv.org/abs/2501.09685), whereas neither beam-search nor MCTS appears to be a classical baseline here.

QUESTIONS:
1. Can the authors comment on whether the effectiveness of this algorithm can be theoretically formalized?

2. How robust is performance to changes in guidance scale or step count T?

3. Can the authors explain the baseline selection and the lack of literature discussion?

LIMITATIONS:
NA
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-03