Guided Trajectory Optimization with Sparse Scaling for Test-Time Diffusion

The paper introduces RTS, a novel Reward-guided Trajectory Scaling method designed to enhance diffusion model generation performance during inference by efficiently navigating the high-dimensional noise space. This method addresses limitations in existing Test-Time Scaling (TTS) techniques by integrating a reward-guided optimization strategy with a sparse test-time scaling framework, aiming to unlock the generative potential of diffusion models while maintaining computational efficiency.

Problem Statement and Goal

The core problem is to synthesize an optimal sample, denoted as the optimal noise trajectory Z∗, that maximizes a given reward paradigm R during inference: F(Dθ, c, R, Z) → x∗0 (Equation 1). The main challenge is the high dimensionality of the noise trajectory Z ∈ R L×d, which leads to an exponential expansion of the search space. To tackle this, RTS employs a dual strategy: first, a sparse optimization strategy to identify key denoising steps, and second, a reward-guided strategy to flexibly explore high-reward noise regions at each key step.

Sparse Test-Time Scaling Framework

The framework is designed to optimize both the initial and intermediate noises while ensuring computational efficiency. It achieves this through two main components:

1. A sparse optimization strategy using PCA-based dimensionality reduction to analyze each denoising step and identify a sparse set of key inflection points. This involves collecting high-dimensional latent vectors, constructing a data matrix X, performing Singular Value Decomposition (SVD), and projecting the latents into a 3D subspace.

2. A reward-guided noise optimization strategy that provides gradient-driven exploration by alternating between coarse and fine search rounds to navigate the search process effectively.

Reward-Guided Noise Optimization Strategy

This strategy utilizes a coarse-to-fine alternating mechanism to accelerate the search process and actively direct the search toward promising regions. The mechanism alternates between:

1. Coarse-grained search, which aims to stochastically explore the surrounding region to estimate a surrogate gradient. This involves generating neighbor candidates M = 3, where N is the neighbor count, by introducing random perturbations to the base noise z and estimating a proxy gradient gt based on differential signals of these candidates.

2. Fine-grained search, which is conducted for targeted refinement guided by the estimated gradient. This round reuses the base noise from the previous round and generates new neighbors by integrating the estimated gradient with a perturbation component, balancing historical gradients and stochasticity.

Key Components of Optimization

The RTS framework explicitly targets both initial and intermediate noises:

1. For the initial noise, it employs multi coarse-to-fine alternating search cycles to identify an optimal starting point.

2. For each point from the sparse set identified by PCA, a single coarse-to-fine cycle is executed to obtain an optimal intermediate noise state, ensuring comprehensive coverage of the entire denoising space.

Results and Contributions

Experiments demonstrate that RTS outperforms baselines by 15.6% across GenEval Score and a 60.4% enhancement in ImageReward score, setting a new SOTA. The method achieves substantial gains over search baselines like Best-of-N (BoN), Zeroth-Order (ZO), and Feynman-Kac (FK) steering, demonstrating superior performance on both generation quality and efficiency across various benchmarks. Ablation studies confirm that both the initial and intermediate optimization stages secure substantial performance gains, showing a strong synergy between the two components. The paper concludes that RTS provides a practical solution for efficient test-time scaling by flexibly allocating compute along the denoising trajectory.

Limitations and Future Work

The effectiveness of RTS is bounded by the quality and granularity of the adopted reward model; if signals are weak, it may fail to resolve detailed semantic errors. A promising direction for future work is obtaining more reliable and fine-grained reward signals, such as region-aware semantic evaluators or predictive reward estimation methods. The broader impact lies in promoting the accessibility of advanced AI by reducing the computational energy typically required for extensive model training, facilitating deployment on computationally constrained devices.

Key Phrases Highlighted:

Reward-guided Trajectory Scaling method

reward-guided noise optimization strategy to proactively explore high-reward regions within the noise space.

sparse test-time scaling framework that enhances diffusion model performance by flexibly allocating compute along the denoising trajectory.

coarse-to-fine alternating mechanism to accelerate the search process.

PCA-driven curvature analysis to identify a sparse set of key inflection points."

jointly optimizing initial and intermediate noises, our framework achieves comprehensive coverage of the entire denoising space."

**"full scaling degrades performance (e.g., lowering ImgReward). We attribute this to the phase-sensitive nature of diffusion models...

As a fastidious researcher, I have analyzed the proposed methodology of Guided Trajectory Optimization with Sparse Scaling for Test-Time Diffusion (RTS). This framework introduces a novel, two-pronged approach—reward-guided noise optimization and sparse test-time scaling—to enhance diffusion models during inference.

Here are the specific improvements that can be made to AI systems using this paper's concepts, and what these improved systems can achieve:

)1. Enhanced Inference Efficiency via Sparse Trajectory Search

The RTS framework replaces exhaustive search (like Best-of-N or full trajectory sampling) with a computationally efficient method by employing PCA-driven curvature analysis.

• Specific Improvement: Implement a Sparse Test-Time Scaling Framework that uses Principal Component Analysis (PCA) to project the high-dimensional noise trajectory into a low-dimensional subspace, identifying only the key inflection points (top 3 principal components) where the latent space exhibits maximum variance.

• What it achieves: This drastically reduces computational cost during inference by focusing optimization efforts only on critical steps, allowing for faster generation while maintaining or improving quality (as shown in Table 6 and Table 7).

)2. Proactive, Gradient-Driven Noise Exploration

The paper proposes a coarse-to-fine alternating mechanism guided by a surrogate gradient estimation based on relative reward signals between neighboring noise samples.

• Specific Improvement: Integrate the Reward-Guided Noise Optimization Strategy (Algorithm 1) into the inference pipeline. During generation, use this mechanism to iteratively refine the initial Gaussian noise and intermediate noise states by estimating a surrogate gradient that points toward regions with higher predicted reward scores (using ImageReward as a primary verifier).

• What it achieves: This moves beyond static or purely stochastic sampling. The system proactively steers the denoising trajectory toward high-reward areas, leading to superior final image quality and better semantic alignment (GenEval and ImageReward gains shown in Table 1 and 2).

)3. Unified Initial and Intermediate Noise Optimization

The RTS framework treats the entire denoising process as a single trajectory, optimizing both the starting point (initial noise) and the path taken through intermediate steps.

• Specific Improvement: Develop a unified test-time scaling module that applies coarse-to-fine alternating search to optimize the initial latent vector, followed by targeted refinement of intermediate noise states only at the identified key inflection points.

• What it achieves: This ensures comprehensive coverage of the entire denoising manifold, preventing trajectory collapse or error accumulation that occurs when optimizing only one component (as shown in Table 3 and Table 9).

)4. Versatility Across Diffusion Solvers

The methodology is designed to be robust across different underlying diffusion model architectures (ODE-based vs. SDE-based solvers).

• Specific Improvement: Design the scaling framework to dynamically adapt its search parameters based on whether the model utilizes a deterministic ODE solver or a stochastic SDE solver. For ODEs, prioritize initial noise refinement; for SDEs, focus on intermediate noise exploration.

• What it achieves: The resulting AI system gains versatility, enabling high performance across various diffusion architectures without requiring separate model fine-tuning for each solver paradigm (as noted in Section 3).

)5. Adaptive Resource Allocation and Performance Scaling

The framework allows researchers to tune the complexity of the search process by adjusting parameters like the number of key steps and neighborhood size.

• Specific Improvement: Implement a dynamic control mechanism that monitors real-time computational budget (e.g., NFE constraints) and adjusts search parameters (e.g., number of rounds, neighborhood count, or key step selection criteria) to find the optimal trade-off between generation quality and inference time.

• What it achieves: The system becomes highly resource-aware, enabling deployment on computationally constrained devices while maximizing performance gains (as evidenced by the Scaling Behavior analysis in Section 14).

In summary, the improved AI system will be a diffusion model inference engine capable of generating high-fidelity images with significantly reduced latency and increased semantic accuracy by intelligently navigating the noise space based on reward signals and curvature analysis.

Test-Time Scaling (TTS) paradigm offers a promising perspective for enhancing the generation performance of diffusion models. However, current solutions largely restrict their search to predefined noise candidates or suffer from inflexible exploration across the denoising trajectory. To bridge this gap, we propose RTS, a novel Reward-guided Trajectory Scaling method to fully unlock the generative potential of diffusion models. Unlike existing methods, RTS facilitates the synthesis of refined, high-fidelity images via two core innovations: 1) a coarse-to-fine noise optimization mechanism that exploits historical search experience to actively steer the exploration toward high-reward regions and 2) a unified sparse test-time scaling framework featuring PCA-driven curvature analysis, which eliminates temporal redundancy by flexiblely allocating compute to a sparse set of key timesteps that represent critical shifts in the denoising direction. Extensive experiments across SD v3, FLUX, and Qwen-Image architectures demonstrate that RTS outperforms baselines, improving the GenEval score by 20.7%, 15.6%, and 12.2%, respectively. Notably, empirical findings indicate that these key points primarily cluster in the mid-stage of the trajectory, distinct from the structure-sensitive early phases and the late attribute refinement phases.