Paper deep dive
Trust-Region Noise Search for Black-Box Alignment of Diffusion and Flow Models
Niklas Schweiger, Daniel Cremers, Karnik Ram
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 3/22/2026, 5:10:01 AM
Summary
The paper introduces Trust-Region Noise Search (TRS), a black-box optimization algorithm for aligning pre-trained diffusion and flow-based generative models to target rewards at inference time. By treating the generative and reward models as black boxes and optimizing only the source noise, TRS balances global exploration and local exploitation using adaptive trust regions. It demonstrates superior performance in text-to-image, molecule, and protein design tasks compared to existing gradient-based and search-based alignment methods.
Entities (5)
Relation Signals (3)
Trust-Region Noise Search → optimizes → source noise
confidence 98% · TRS... only optimizes the source noise.
Trust-Region Noise Search → aligns → Diffusion Models
confidence 95% · We introduce trust region search (TRS), a simple approach for inference-time reward alignment of black-box diffusion and flow models.
ImageReward → evaluates → Diffusion Models
confidence 90% · We use two generative models... with ImageReward and HPSv2 as the reward models.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Optimizing the noise samples of diffusion and flow models is an increasingly popular approach to align these models to target rewards at inference time. However, we observe that these approaches are usually restricted to differentiable or cheap reward models, the formulation of the underlying pretrained generative model, or are memory/compute inefficient. We instead propose a simple trust-region based search algorithm (TRS) which treats the pre-trained generative and reward models as a black-box and only optimizes the source noise. Our approach achieves a good balance between global exploration and local exploitation, and is versatile and easily adaptable to various generative settings and reward models with minimal hyperparameter tuning. We evaluate TRS across text-to-image, molecule and protein design tasks, and obtain significantly improved output samples over the base generative models and other inference-time alignment approaches which optimize the source noise sample, or even the entire reverse-time sampling noise trajectories in the case of diffusion models. Our source code is publicly available.
Tags
Links
- Source: https://arxiv.org/abs/2603.14504v1
- Canonical: https://arxiv.org/abs/2603.14504v1
Trouble viewing inline? Open PDF directly →
Full Text
83,191 characters extracted from source content.
Expand or collapse full text
Trust-Region Noise Search for Black-Box Alignment of Diffusion and Flow Models Niklas Schweiger 1 Daniel Cremers 1,† Karnik Ram 1,† 1 TU Munich niklas.schweiger@tum.de Abstract Optimizing the noise samples of diffusion and flow models is an increasingly popular approach to align these models to target rewards at inference time. However, we observe that these approaches are usually restricted to differentiable or cheap reward models, the formulation of the underlying pre- trained generative model, or are memory/compute inefficient. We instead propose a simple trust-region based search algorithm (TRS) which treats the pre-trained generative and reward models as a black-box and only optimizes the source noise. Our approach achieves a good balance between global exploration and local exploitation, and is versatile and easily adaptable to various generative settings and reward models with minimal hyperparameter tuning. We evaluate TRS across text-to-image, molecule and protein design tasks, and obtain significantly improved output samples over the base generative models and other inference-time alignment approaches which optimize the source noise sample, or even the entire reverse-time sampling noise trajectories in the case of diffusion models. Our source code is publicly available ∗ . 1 Introduction Generative models such as diffusion and flow-based models have revolutionized diverse domains, including high-fidelity image and video synthesis, molecule generation, and language modeling, by learning complex data distributions from large-scale pretraining data (Ho et al., 2020; Lipman et al., 2023). While scaling model capacity and training compute has significantly improved sample quality, pretrained models often fall short of specific, fine-grained requirements. This includes generating molecules with precise binding affinities (Guan et al., 2023) or images that strictly align with complex, multi-attribute prompts. This gap has motivated the emergence of inference-time alignment for diffusion and flow- based models, an exciting new paradigm where the quality of generated samples is optimized post-training using feedback from target reward models (Bansal et al., 2024; Eyring et al., 2024). Unlike fine-tuning, this approach requires no additional training data and instead invests additional compute at inference time to steer samples toward desired properties. Several approaches for inference-time alignment have been proposed. Gradient-based methods back-propagate through the entire (often ODE-based) iterative process of the model to adjust the initial noise sample (Ben-Hamu et al., 2024; Wang et al., 2025). But these are prone to high GPU-memory costs and risk drifting off the training-data manifold. Sequence-based methods, typically used for SDE samplers, range from filtering approaches that iteratively reweigh and resample (Singhal et al., 2025; Kim et al., 2025) to tree-search based methods (Jain et al., 2025; Li et al., 2025a). But these often require a high number of expensive reward calls or rely on accurate value estimates for terminal rewards, which are not always available or reliable. In contrast, black-box search methods which are relative underexplored, treat the † Equal supervision. ∗ Project page: https://niklasschweiger.github.io/trust-region-noise-search/ 1 arXiv:2603.14504v1 [cs.LG] 15 Mar 2026 Aesthetic Score (↑) 6.97.57.97.78.18.2 Property Alignment (↓) 0.75 0.65 0.63 0.43 0.41 0.35 Designability (↑) 0.041 0.057 0.0590.2010.3460.506 Figure 1: Progression of output samples produced by our algorithm (TRS) across iterations i ∈ 0,...,15for text-to-image, molecules, and proteins. Top: Aesthetic reward alignment for the prompt “Animated movie poster...”. Middle: Molecule property alignment. Bottom: Protein designability. The horizontal arrow indicates the direction of optimization, while the arrows↑ /↓ indicate whether the objective is maximized or minimized. generator and reward models as a black box and apply search heuristics for the source noise sample (Ma et al., 2025; Tan et al., 2025). While versatile and applicable to any generative architecture and reward model, we notice that existing approaches often struggle to find a good balance between global exploration and local exploration, tending toward one extreme or the other. Inspired by Bayesian optimization algorithms (Eriksson et al., 2019), which also optimize expensive black-box functions, we propose a simple and effective trust region search (TRS) approach that balances between global exploration and local exploitation in a structured manner. By controlling only the source noise, TRS is readily applicable to a wide range of generative models and reward functions without internal modifications. TRS begins by exploring multiple seed noise samples, which are then pruned and iteratively refined with local perturbations. Crucially, these perturbations are adaptively controlled in both magnitude and direction based on observed reward values, ensuring the search remains within the data manifold to produce stable and high-quality samples. The contributions of our work are as follows: 1.We introduce trust region search (TRS), a simple approach for inference-time reward alignment of black-box diffusion and flow models via adaptive source noise control. 2.We provide an extensive evaluation on text-to-image generation, demonstrating that TRS yields significantly more aligned and higher-quality samples compared to existing search heuristics and even full-noise sequence search baselines under identical compute budgets. 3.We demonstrate the versatility of our approach through an extended evaluation on small molecule and protein design tasks, where TRS proves effective even with expensive reward functions and requires minimal hyperparameter tuning. 2 2 Background and Related Work 2.1 Diffusion and Flow-based Models Diffusion and flow-based generative models aim to transport a simple noise distribution p noise =p 0 to a complex data distributionp data =p 1 . A sample x 0 ∼ p 0 is gradually transformed into a data sample x 1 ∼ p 1 through a continuous-time processx t t∈[0,1] , typically discretized intoTsteps witht k =k/T. In continuous data regimes,p 0 is usually white Gaussian noise N (0, I). Diffusion models are characterized by a forward noising process that gradually transforms data into noise. This process is governed by a stochastic differential equation (SDE): dx t = f (t)x t dt + g(t)dw t , where f (t) andg(t) are the drift and diffusion coefficients that determine the noising schedule (Song et al., 2021). To generate new samples, this process must be reversed. The generative (reverse) process relies on the score functionε(x t ,t)≜∇ x logp t (x t ), which represents the gradient of the log-density of the noisy sample. Once the score is estimated via a neural network, sampling can be performed using either the deterministic or stochastic formulations shown in Table 1. Flow matching models instead learn a continuous vector field v θ (x t ,t) =dx t /dt(Lipman et al., 2023) that defines the velocity of the samples. Unlike diffusion, the paths are not necessarily tied to a corruption process; a common choice is the optimal transport (OT) path (Liu et al., 2023), defined by linear interpolation x t = (1− t)x 0 + tx 1 , which yields a constant target velocity v t = x 1 −x 0 . Standard flow matching integrates the learned ODE for sampling, while stochastic variants (Bose et al., 2024) can introduce Brownian motion and optional score corrections as detailed in Table 1. Unified view. Both paradigms can be expressed in a general formdx t = f (x t ,t )dt+g(t)dw t . As summarized in Table 1, the two frameworks differ primarily in how the drift f (x t ,t ) is parameterized, either through the score functionε(x t ,t) to reverse a corruption process, or through a velocity field v θ (x t ,t) to match a probability flow. Table 1: Generative sampling formulations (t= 0 as noise,t= 1 as data). Note that for diffusion,ε(x t ,t) denotes the score function ∇ x logp t (x). FormulationODESDE Diffusiondx t = [f (t)x t − 1 2 g 2 (t)ε(x t ,t)]dt dx t = [f (t)x t −g 2 (t)ε(x t ,t)]dt + g(t)dw t Flow Matching dx t = v θ (x t ,t)dtdx t = v θ (x t ,t)dt + g(t)dw t 2.2 Reward alignment of diffusion and flow models. While pre-trained diffusion and flow models generate high-quality samples, they often fail to meet specific downstream criteria, such as aesthetics Wu et al. (2023) in the case of image generation, or specific properties such as stability or binding affinity in the case of molecule (Guan et al., 2023; Xu et al., 2022) and protein design (Ingraham et al., 2023; Watson et al., 2023). This alignment can be achieved by fine-tuning the model’s parameters (Clark et al., 2024; Black et al., 2024) but this requires collecting new samples for every new target property, and is incompatible with new properties that are presented only at inference-time. A promising alternative approach is reward alignment where the noise samples during generation are aligned to meet a target reward objective. Conceptually, these aligned or optimized noise 3 Table 2: Feature comparison across noise optimization methods. We compare methods by their ability to support gradient-free objectives, black-box models, parallelized batch efficiency, and global or local optimization. MethodGrad-free Black-box Batch Eff. Global Local OC-Flow (Tang et al., 2025)✘✔ Fast Direct (Tan et al., 2025)✔✘ DTS* (Jain et al., 2025)✔✘✔✘ Random Search (Ma et al., 2025)✔✘ Zero-Order (Ma et al., 2025)✔✘✔ TRS (Ours)✔ samples create a new target distribution that balances two competing objectives: staying true to the structural distribution learned during pretraining while shifting the probability mass towards target samples that maximize the reward. We categorize these noise optimization or search methods as follows. Gradient-based guidance. One approach to noise optimization involves back-propagating the gradients from a reward function through the entire iterative sampling process to update the source (or intermediate) noise samples (Ben-Hamu et al., 2024; Wang et al., 2025; Guo et al., 2024; Tang et al., 2025). While effective, these methods require differentiable reward functions and incur substantial GPU memory and computational overhead, as they necessitate storing or recomputing the full diffusion trajectory during inference. This limitation becomes particularly severe in high-dimensional settings, such as image or 3D generation, or when many solver steps are required. Furthermore, pure gradient-based refinement often shifts the generation off the natural data manifold, necessitating additional regularization or alignment terms to preserve sample quality (Wang et al., 2025). An exception is direct noise optimization (DNO) (Tang et al., 2025), which additionally supports gradient approximations for non-differentiable rewards and an SDE-based mode that can refine intermediate noises along the trajectory. Despite these extensions, the same disadvantages of high runtime and GPU memory consumption remain. Noise sequence search. Another popular approach is to guide the generation process throughout the whole or parts of the sampling trajectoryx 0 ,...,x 1 . This includes sequential Monte Carlo (SMC) approaches using resampling (Singhal et al., 2025; Kim et al., 2025), tree search methods such as DSearch and DTS (Li et al., 2025a; Jain et al., 2025), and Fast Direct (Tan et al., 2025) which uses Gaussian process-based black-box optimization over the full noise sequence (Tan et al., 2025). While these methods can be effective for stochastic sampling, they typically rely on intermediate reward approximations (Li et al., 2025b) or are difficult to utilize with batched evaluations. Black-box search. Most versatile are source noise optimization algorithms which treat the generative model and the reward model as a black-box. They can be applied to any generative model, that maps a source noise distribution to some data distribution, and any reward model which can be differentiable or non-differentiable. They are conceptually much simpler and easy to implement, while also leading to surprisingly good results. This category of approaches is relatively underexplored but includes some recent work by Ma et al. Ma et al. (2025) where they apply random search and zero-order search for the source noise. Some additional recent (concurrent) work is by Jajal and Eliopoulos et al. Jajal et al. (2025) where they apply evolutionary and genetic search algorithms for the source noise. Our work belongs to this category and we apply trust-region search, inspired by Bayesian optimization algorithms Eriksson et al. (2019), and show strong results across data modalities and reward models. 4 3 Methodology 3.1 Problem Statement LetF:R M → R D be a pre- trained generative diffusion or flow-based model which maps source noise x 0 to a data sam- ple x 1 andR:R D → Ra re- ward model processing x 1 and returning a scalar rewardr. Note that generative models typically operate in a x 0 Generative Model F (x 0 ) Reward Function R(x 1 ) r x 1 Black-box: x 0 → r Figure 2: Illustration of the black-box, defined by the generative modelFand the reward functionR(x 1 ), which connect noise samples x 0 with scalar rewards r. compressed latent space Rombach et al. (2022), leading toM < D. We do not intervene with the model in any other way, treating it as a black-box, as illustrated in Figure 2. The resulting objective is: x ⋆ 0 = arg max x 0 ∈R M R F (x 0 ) .(1) We treat the entire mappingR(F(x 0 )) as a computationally expensive black-box, as illustrated in Figure 2, either due to the generative model, the reward model, or the use of large batch sizes. Consequently, we focus on methods that can effectively steer generation under a strict evaluation budget. 3.2 Trust-Region Search (TRS) TRS is a trust-region search algorithm with structured noise sampling designed for inference-time alignment with expensive black-box evalua- tions. Our approach is in- spired by TuRBO (Eriksson et al., 2019) for Bayesian opti- mization, but introduces sev- eral important modifications for inference-time steering of large generative models. In particular, unlike surrogate- based Bayesian optimization, TRS relies purely on struc- tured sampling, since we find that surrogates do not con- tribute meaningfully to per- formance, due to the highly non-linear noise space, which is difficult to predict under Algorithm 1 Trust-Region Search (TRS) Require:Total/Warmup BudgetN total /N warm , Batch sizeB, Regions k, Model F , Reward function R 1:Warm-up: Samplex 0,i N warm i=1 ∼ p 0 ; evaluater i = R(F (x 0,i )) in batches of size B. 2: Init: Set centersx c 0,j k j=1 to topk; set lengthsℓ j ← ℓ init ; set remaining budget N ← N total − N warm 3: while budget N ≥ B do 4: S batch ←∅ 5:for region T j and sample b∈1..B/k do 6:p j,b ∼ U(p min ,p max ); m j,b ∼ Ber(p j,b ) M 7: ̃ x 0,j,b ∼ Perturb(ℓ j ) ▷Gauss or Sobol 8:x 0,j,b ← x c 0,j + ( ̃ x 0,j,b ⊙ m j,b ) 9: S batch ←S batch ∪x 0,j,b 10:end for 11:Evaluate: r j,b ← R(F (S batch )) 12:Adapt: Update lengths ℓ j 13:Shift: Re-center x c 0,j to the global top k. 14: N ← N − B 15: end while 16: return best F (x 0 ),r constrained budgets. Additionally, we use a different center selection scheme, by always choosing the top-k observed noises, which is crucial for the performance. An overview of the method is given in Algorithm 1. Warm-up. We begin with a short warm-up phase for bootstrapping the search. We sample N warm initial noise samples from the model prior p 0 , x 0,i ∼ p 0 ,x 1,i =F (x 0,i ), r i = R(x 1,i ),(2) withi= 1,...,N warm , and select the top-performingkpoints as initial trust-region centers x c 0,j k j=1 . All trust regions are initialized with the same side lengthℓ j =ℓ init . In practice, this warm-up is implemented using standard iterations with batch sizeB. We find allocating approximately 20% of the total evaluation budgetN total to this phase to be a robust heuristic across tasks and budgets; an ablation is provided in Section D.3. 5 Noise Space R M T 1 T 2 Model F Data Manifold R D 1 2 Legend T j : Trust Region x c 0,j : Center x 0,j,b : Candidates F: Mapping x 1,j,b : Data Points k Top-k Samples m j,b : Locked Dim. ̃ x 0,j,b : Perturbation New Candidates x 0,j,b ← x c 0,j + m j,b ⊙ ̃ x 0,j,b x 0,j,b = x c 0,j + ̃ x 0,j,b ⊙ m j,b (c) Perturbation Sobol ℓ j ̃ x 0,j,b = a j + (b j − a j )⊙ u j,b , u j,b ∈ [0, 1] M Gaussian ̃ x 0,j,b ∼ N(0,σ 2 j I), σ j = ℓ j / √ 12 (d) Region Updates Success (ℓ 1 ↑) Failure (ℓ 2 ↓) T 1 new T 2 newT 1 old T 2 old Recentering (a) (b)(c)(d) Figure 3: Illustration of the trust region search algorithm with a two-region example for the prompt “A panda making latte art” from DrawBench (Saharia et al., 2022). (a) Samples from the noise spaceR M are mapped to the data manifold inR D via the generative model F. Generated samples from the same region (T j ) exhibit visual similarity; here,T 1 shows markedly better prompt alignment thanT 2 . (b) New candidates x 0,j,b are generated by adding masked (m j,b ) relative perturbations ( ̃ x 0,j,b ) to the current center x c 0,j . (c) Comparison of Sobol and Gaussian perturbation schemes used to fill the trust-region hypercube. (d) The update and shift logic:T 1 expands upon identifying a top-ksample, while the underperforming T 2 is re-centered (shifted) to a more promising region. Trust-region iterations. After warm-up, the algorithm maintainskhypercubic trust regionsT j ⊂ R M , each defined by a center x c 0,j and side lengthℓ j . In each iteration, we perform the following steps, where a global batch setS batch of sizeBacross all regions is maintained. 1.Propose: For each regionj, generateB/kcandidate noise vectors x 0,j,b by perturbing the center x c 0,j . Following (Eriksson et al., 2019), perturbations ̃ x 0,j,b are generated within an axis-aligned hypercube with side lengthℓ j . Additionally, we combine these perturbations with a stochastic coordinate mask. This mask is built by drawing a per- turbation probabilityp j,b ∼ Uniform(p min ,p max ) and applying m j,b ∼ Bernoulli (p j,b ) M . The resulting candidates are: x 0,j,b = x c 0,j + ̃ x 0,j,b ⊙ m j,b .(3) 2.Evaluate: AllBcandidates across regions are aggregated intoS batch and evaluated in parallel to obtain the rewards r j,b = R(F (S batch )). 3.Update: After batch evaluation, we adapt trust-region side lengthsℓ j using success- based rules per region. A key distinction from vanilla TuRBO (Eriksson et al., 2019) is that trust regions are not treated independently. While exploring multiple regions is beneficial early on, we observe that only a subset typically remains promising. Accord- ingly, after each batch iteration we re-center all trust regionsx c 0,j at the globally best kpoints observed so far. This mechanism naturally shifts computation from exploration toward exploitation, reallocating evaluation budget to promising regions over time. An ablation study on different re-centering strategies is provided in Section D.1. Perturbations. For generating the perturbations in the proposal step we use two different schemes. Deterministic low-discrepancy samplers such as Sobol (Sobol’, 1967) are efficient to fill search spaces, but can only be computed only up to 21k dimensions (Joe & Kuo, 2008). For higher dimensional search spaces, like in our experiment with SDXL-Lightning in 6 Section 4.1 with noise dimension 65,536, we use a Gaussian perturbation scheme, that is designed to follow the trust-region hypercube exploration closely. For Sobol-based proposals, a point u j,b ∈[0,1] M is sampled and mapped affinely into the trust region. Let a j =− 1 2 ℓ j 1 M and b j = 1 2 ℓ j 1 M , where 1 M denotes aM-dimensional vector of ones. The resulting proposal is ̃ x 0,j,b = a j + (b j − a j )⊙ u j,b . (4) For Gaussian perturbations, we sample ̃ x 0,j,b ∼N (0,σ 2 j I), σ j = ℓ j / √ 12,(5) where I∈ R M×M is the identity matrix. The standard deviationσ j =ℓ j / √ 12is chosen to match the variance of a uniform distributionU[−ℓ j /2,ℓ j /2], ensuring the Gaussian proposals cover the trust-region hypercube with equivalent spread. Trust-region adaptation. The update rules for the trust-region side length largely follow the strategy introduced in TuRBO (Eriksson et al., 2019). For each region, we maintain success and failure counters based on whether newly evaluated candidates improve upon the best value observed within that region. If a candidate yields an improvement, the success counter is incremented and the failure counter is reset. Otherwise, the failure counter is incremented and the success counter is reset. Two thresholds,c succ andc fail , govern the adaptation of the trust-region side lengthℓ j . When the success counter reachesc succ , the trust region is expanded according to ℓ new j = min(ℓ j · α ℓ ,ℓ max j ),(6) whereα ℓ is an expansion factor for the length (we set it to 1.5 by default). Conversely, when the failure counter reaches c fail , the trust region is contracted as ℓ new j = max(ℓ j /α ℓ ,ℓ min j ).(7) This mechanism enables adaptive control over the exploration scale based on recent optimiza- tion progress. We further illustrate the core steps of TRS in Figure 3 and provide additional details in the appendix. 4 Experiments We evaluate our trust-region search (TRS) method across three diverse generative settings: text-to-image diffusion models in Section 4.1, ODE-based molecular flow matching in Sec- tion 4.2, and protein backbone design in Section 4.3. These experiments span different modalities, reward models, and sampling procedures, allowing us to assess both the effective- ness and robustness of TRS under varying optimization budgets and reward costs. Across all settings, we compare against representative gradient-based, noise-sequence, and black-box search baselines under matched compute constraints. Overall, the results demonstrate that TRS consistently achieves stronger alignment with target objectives while maintaining sample quality and requiring minimal task-specific tuning. 4.1 Text-to-Image Setup. We evaluate on the DrawBench (Saharia et al., 2022) benchmark which comprises 200 prompts across diverse categories (e.g., counting, composition). We use two generative models: SD1.5 (Rombach et al., 2022) (50 steps) and SDXL-Lightning (Lin et al., 2024a) (8 steps), with ImageReward (Xu et al., 2023) and HPSv2 (Wu et al., 2023) as the reward models. TRS is compared against gradient-based OC-Flow (Wang et al., 2025), state-of-the- art Diffusion Tree Sampling (DTS)* (Jain et al., 2025), Fast Direct (FD) (Tan et al., 2025), and black-box methods including random and zero-order search (RS, ZO) (Ma et al., 2025). The number of function evaluations (NFE) is fixed across methods. Further implementation details for the baselines are in Section F. 7 Prompt TRS (Ours)RSZODTSFDOC-Flow Three cats and two dogs sitting on the grass. A storefront with ’Google Brain Toronto’ written on it. A zebra underneath a broccoli. Figure 4: Examples of optimized samples from different algorithms for all methods in Section 4.1, with SDXL-Lightning (Lin et al., 2024a). The first row is optimized with ImageReward (Xu et al., 2023) and the lower two rows with HPSv2 (Wu et al., 2023). All prompts are from DrawBench (Saharia et al., 2022). Outputs from TRS adhere to the prompt more closely in terms of specified animal count, text, and relative positions. Further examples (incl. randomized) are given in the appendix. Metrics. We report the mean of the best per-prompt reward Table 3 and average best rewards across varying NFE budgets in Figure 5. Additionally, we report the number of reward oracle evaluations induced by the fixed NFE budget (Table 3). Direct NFE comparison for backpropagation-based OC-Flow is provided via runtime analysis in Section H. Note that the ImageReward outputs are designed to roughly follow a normal distribution, so outputs beyond the range of [−2,2] are rare. HPSv2 shows differences in a smaller scale where typical high-quality images fall within the range of [0.25, 0.35]. Results. Table 3 demonstrates that TRS is consistently better than all baselines across generative models and reward functions. On both models, TRS exceeds the state-of-the-art DTS* performance with up to 4×reduction in wall-clock time and fewer reward evaluations. OC-Flow and Fast Direct underperform relative to random search, highlighting the difficulty of surrogate modeling and gradient-based optimization in high-dimensional noise spaces, while TRS shows the best results. Figure 4 show some qualitative results and Figure 5 further illustrates the scaling trends from small to high budgets, where TRS generally shows the best trends, where else the other algorithms saturate earlier. Full hyperparameter configurations of all methods are displayed in Section B. 4.2 Molecule Generation Setup. In this experiment, we move to ODE-based flow matching and consider a different data modality. Our goal is to generate small molecules with specified target values of chemical properties (Ramakrishnan et al., 2014). We use EquiFM (Song et al., 2023) with 50 integration steps, operating on a joint continuous Gaussian coordinate- and encoded feature space (291 total dimensions), with pretrained chemical property prediction models. This setting is a standard benchmark for evaluating classifier-free guidance (Hoogeboom et al., 2022) and was also used by OC-Flow (Wang et al., 2025) for noise optimization toward single-property targets. However, we observe that random search based methods already achieve very strong performance in this regime and therefore extend the benchmark to multi-property target matching. Specifically, we define a distance-based loss as the mean absolute deviation between predicted and target values across multiple chemical properties. Since our framework is formulated as reward maximization, we maximize the negative of this 8 Table 3: Comparison of noise optimization algorithms on DrawBench. We report the mean best rewards (IR, HPSv2) for SD1.5 and SDXL alongside the computational cost in terms of number of reward calls.↑ /↓indicate the direction of improvement. Best values are in bold, second best are underlined. SD1.5SDXL AlgorithmIR ↑HPS ↑IR ↑HPS ↑Reward calls Base-0.160.2460.500.2621 Gradient-based guidance OC-Flow (Wang et al., 2025)0.420.2770.850.287– Noise sequence search Fast Direct (Tan et al., 2025)1.350.2961.500.314420 DTS* (Jain et al., 2025)1.590.3031.620.327503 Black-box search Random (Ma et al., 2025)1.440.3021.540.324400 Zero-order (Ma et al., 2025)1.500.3131.590.332400 TRS (Ours)1.620.3221.660.340400 5k10k15k20k 1.3 1.4 1.5 1.6 NFE Mean best reward SD1.5 (IR) (↑) 5k10k15k20k 0.3 0.31 0.32 NFE SD1.5 (HPS) (↑) 0.8k1.6k2.4k3.2k 1.45 1.55 1.65 NFE SDXL (IR) (↑) 0.8k1.6k2.4k3.2k 0.32 0.33 0.34 NFE SDXL (HPS) (↑) TRS (Ours) RandomZero-order DTS* Figure 5: Here we plot the mean best rewards for SD1.5/SDXL optimizing for HPSv2 and ImageRe- ward across different NFE budgets. TRS shows the best scaling performance among all methods. loss. We consider two reward functions,R 3 andR 6 , involving 3 and 6 chemical properties, respectively and additional property details are provided in Section C.2. As baselines, we exclude noise-sequence search methods, as they are incompatible with ODE-based flow matching, and compare against OC-Flow, random search, and zero-order search. Metrics. As our primary metric, we report the mean best lossesL 3 (x) =−R 3 (x) and L 6 (x) =−R 6 (x) (mean distance to the target properties) over 200 molecules, with the target property values sampled randomly. In addition, we track other relevant metrics such as the molecule stability percentage (MSP), and valid and unique percentage (VUP). Results. The quantitative results are summarized in Figure 7. Additionally, we visualize some randomly chosen optimized samples from different solversFigure 6. We observe that TRS consistently achieves the lowest losses, indicating the highest alignment with the specified multi-property targets. Importantly, this improvement does not come at the expense of other quality metrics: molecule stability and novelty remain comparable to those of the base model. In contrast, the gradient-based OC-Flow method exhibits degraded stability and novelty, even with regularization. This suggests that gradient-based optimization tends to drift off the data manifold in this setting. We find that sampling-based approaches are better suited for this task; notably, even random search outperforms OC-Flow. This performance gap likely stems from conflicting gradients between different property classifiers, which makes local optimization difficult. As illustrated in the bottom row of Figure 6, OC-Flow often fails 9 (a) Molecules OC-FlowRandom Zero-orderTRS 1.040.630.600.37 0.930.620.590.55 (b) Proteins Random Zero-orderTRS 0.530.590.65 0.330.390.45 Figure 6: Optimized molecule and protein samples. We visualize randomly selected samples produced by different optimization methods in the molecule and protein experiments. (a) In the top row optimization is conducted with three target properties (L 3 ), while the bottom row with the six (L 6 ). Lower is better. (b) In the top rown res =50 and in the bottom row n res =100. The proteins are optimized for the designability reward. to escape unfavorable initial samples, whereas sampling-based methods explore the space more effectively. Algorithm L 3 (↓) L 6 (↓) MSP VUP Base1.211.1588.388.3 OC-Flow0.830.9368.068.0 Random0.470.6284.284.1 Zero-order0.430.5984.484.3 TRS (Ours) 0.39 0.5585.785.6 (a) 25k50k75k100k 0.4 0.5 0.6 NFE Loss (Property dist.) L 3 25k50k75k100k 0.5 0.6 0.7 NFE L 6 TRS (Ours) RandomZO (b) Figure 7: Molecule optimization results. (a) Distance to the combined targets of 3 and 6 chemical property values, together with MSP stability and validity VUP metrics which we do not explicitly optimize. The best alignemnt values are in bold, second best areunderlined. (b) Scaling trends over the NFE budget. 4.3 Protein Design Setup. Protein design is another challenging data modality with expensive reward models and is relatively underexplored by inferenece-time alignment methods. State of the art 3D protein generation models, like Proteina (Geffner et al., 2025), are flow matching models that are typically trained via ODE integration, but sampling is often performed using an SDE scheduler with noise reduction to improve designability (Bose et al., 2024; Lin et al., 2024b). However, this procedure alters the target distribution and no longer samples from the full flow-matching distribution, which can negatively impact other metrics such as diversity and novelty. In this experiment, we therefore focus on steering ODE-based sampling while preserving the full distribution. For completeness, we provide a comparison to SDE-based sampling in Section E. We further note that SDE-based noise-sequence optimization methods are not applicable in this setting: they either rely on value estimation functions that are unavailable for 3D protein design, or are computationally infeasible due to the high cost of 10 reward evaluation. We design the experiment by fixing the number of residuesn res = 50 and n res = 100. For the reward function, we use the computationally expensive designability reward based on large structure extraction and protein folding models (see Section C for details). We compare against purely black-box baselines, including random search and zero-order search (Ma et al., 2025). Metrics. As alignment metric we report the mean best rewards for each run. Additionally we show the cluster diversity and the pairwise TM-score and a novelty metric, which compares the proteins to the entire PDB dataset (Berman et al., 2000). Further details of these are given in Section G.2. n res MethodDes.↑ Clu.↑ TM↓ Nov.↓ 50 Base0.05 0.71 0.50 0.70 Random0.53 0.27 0.68 0.85 Zero-order0.590.39 0.62 0.85 TRS(Ours) 0.65 0.30 0.67 0.89 100 Base0.02 0.71 0.56 0.86 Random0.33 0.49 0.55 0.84 Zero-order0.39 0.59 0.54 0.84 TRS(Ours) 0.45 0.44 0.57 0.86 (a) 16k32k48k64k 0.4 0.5 0.6 NFE Mean Best Reward n res = 50 16k32k48k64k 0.2 0.3 0.4 NFE n res = 100 TRS (Ours) RSZO (b) Figure 8: Protein optimization results. (a) Optimized mean designability rewards at 64k NFE, together with diversity and novelty metrics which we do not explicity optimize. The best values are in bold, second best are underlined. (b) Reward improvement across NFE budgets. Results. In terms of the designability steering, we see in Figure 8, that we significantly improve over other search algorithms in both settings. We see that generally the rewards are higher when we optimize proteins with 50 residues, which is expected, but also their diversity and novelty metrics decrease compared to the base model. We assume that, since the solution space is rather small, it is likely that well designable proteins share similar features. However, this is not comparable to the mode collapse, which we observe for SDE noise reduction (see Section E). Naturally, due to a larger solution space, optimizing proteins with 100 residues leads to better diversity and novelty, while still gaining major improvements over the base model. This counts for all optimization methods, but TRS achieves the best reward alignment again. 5 Conclusion In this work we investigate inference-time scaling and preference alignment, where our simple trust-region source noise search achieves state-of-the-art performance across text-to-image, molecule, and protein design tasks. Our approach is model and reward agnostic, making it particularly suited for real-world settings where reward functions are often expensive or unknown. It also offers good balance between exploration and exploitation by searching multiple noise regions early and refining promising ones, and does not drift off the data manifold and remains stable as we observe in Section 4.2. While all methods are limited by the accuracy of the reward models (Section I), scaling improvements (Wu et al., 2025) suggest this limitation will diminish and our efficient source noise optimization is particularly well-suited for this development. Future work includes exploring the geometry of the source noise space further and developing improved perturbation schemes that adhere more to this geometry. 11 References Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Roni Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal Guidance for Diffusion Models. In The Twelfth International Conference on Learning Representations, 2024. Heli Ben-Hamu, Omri Puny, Itai Gat, Brian Karrer, Uriel Singer, and Yaron Lipman. D- Flow: Differentiating through Flows for Controlled Generation. In Forty-first International Conference on Machine Learning, 2024. Helen M Berman, John Westbrook, Zukang Feng, Gary Gilliland, Talapady N Bhat, Helge Weissig, Ilya N Shindyalov, and Philip E Bourne. The protein data bank. Nucleic acids research, 28(1):235–242, 2000. Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training Diffusion Models with Reinforcement Learning. In The Twelfth International Conference on Learning Representations, 2024. Joey Bose, Tara Akhound-Sadegh, Guillaume Huguet, Kilian FATRAS, Jarrid Rector-Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael M. Bronstein, and Alexander Tong. SE(3)-Stochastic Flow Matching for Protein Backbone Generation. In The Twelfth International Conference on Learning Representations, 2024. Kevin Clark, Paul Vicol, Kevin Swersky, and David J. Fleet. Directly Fine-Tuning Diffusion Models on Differentiable Rewards. In The Twelfth International Conference on Learning Representations, 2024. J. Dauparas, I. Anishchenko, N. Bennett, H. Bai, R. J. Ragotte, L. F. Milles, B. I. M. Wicky, A. Courbet, R. J. de Haas, N. Bethel, P. J. Y. Leung, T. F. Huddy, S. Pellock, D. Tischer, F. Chan, B. Koepnick, H. Nguyen, A. Kang, B. Sankaran, A. K. Bera, N. P. King, and D. Baker. Robust deep learning–based protein sequence design using ProteinMPNN. Science, 378(6615):49–56, 2022. Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion Models Beat GANs on Image Synthesis. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. David Eriksson, Michael Pearce, Jacob Gardner, Ryan D Turner, and Matthias Poloczek. Scalable Global Optimization via Local Bayesian Optimization. In Advances in Neural Information Processing Systems, p. 5496–5507, 2019. Luca Eyring, Shyamgopal Karthik, Karsten Roth, Alexey Dosovitskiy, and Zeynep Akata. ReNO: Enhancing One-step Text-to-Image Models through Reward-based Noise Optimiza- tion. In Neural Information Processing Systems (NeurIPS), 2024. Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, and Karsten Kreis. Proteina: Scaling Flow-based Protein Structure Generative Models. In The Thirteenth International Conference on Learning Representations, 2025. Jiaqi Guan, Wesley Wei Qian, Xingang Peng, Yufeng Su, Jian Peng, and Jianzhu Ma. 3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction. In The Eleventh International Conference on Learning Representations, 2023. Xiefan Guo, Jinlin Liu, Miaomiao Cui, Jiankai Li, Hongyu Yang, and Di Huang. Initno: Boosting Text-to-Image Diffusion Models via Initial Noise Optimization. In CVPR, p. 9380–9389, 2024. Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS ’20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546. 12 Emiel Hoogeboom, V ́ıctor Garcia Satorras, Cl ́ement Vignac, and Max Welling. Equivariant Diffusion for Molecule Generation in 3D. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, p. 8867–8887. PMLR, 17–23 Jul 2022. John Ingraham, Max Baranov, Zak Costello, Karl Barber, Wujie Wang, Ahmed Ismail, Vincent Frappier, Dana Lord, Christopher Ng-Thow-Hing, Erik Vlack, Shan Tie, Vincent Xue, Sarah Cowles, Alan Leung, Jo ̃ao Rodrigues, Claudio Morales-Perez, Alex Ayoub, Robin Green, Katherine Puentes, and Gevorg Grigoryan. Illuminating protein space with a programmable generative model. Nature, 623:1–9, 11 2023. Vineet Jain, Kusha Sareen, Mohammad Pedramfar, and Siamak Ravanbakhsh. Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025. Purvish Jajal, Nick John Eliopoulos, Benjamin Shiue-Hal Chou, George K Thiruvathukal, James C Davis, and Yung-Hsiang Lu. Inference-Time Alignment of Diffusion Models with Evolutionary Algorithms. arXiv preprint arXiv:2506.00299, 2025. Stephen Joe and Frances Y. Kuo. Constructing Sobol Sequences with Better Two-Dimensional Projections. SIAM Journal on Scientific Computing, 30(5):2635–2654, 2008. Sunwoo Kim, Minkyu Kim, and Dongmin Park. Test-time Alignment of Diffusion Models without Reward Over-optimization. In The Thirteenth International Conference on Learning Representations, 2025. Greg Landrum, Paolo Tosco, Brian Kelley, Ricardo Rodriguez, David Cosgrove, Riccardo Vianello, sriniker, Peter Gedeck, Gareth Jones, Eisuke Kawashima, NadineSchneider, Dan Nealschneider, tadhurst cdd, Andrew Dalke, Matt Swain, Brian Cole, Samo Turk, Aleksandr Savelev, Niels Maeder, Alain Vaucher, Maciej W ́ojcikowski, Hussein Faara, Ichiru Take, Rachel Walker, Vincent F. Scalfani, Daniel Probst, Kazuya Ujihara, Axel Pahl, guillaume godin, and Juuso Lehtivarjo. rdkit/rdkit: 2025095 (Q3 2025) Release, January 2026. Xiner Li, Masatoshi Uehara, Xingyu Su, Gabriele Scalia, Tommaso Biancalani, Aviv Regev, Sergey Levine, and Shuiwang Ji. Dynamic Search for Inference-Time Alignment in Diffusion Models. arXiv preprint arXiv:2503.02039, 2025a. Xiner Li, Yulai Zhao, Chenyu Wang, Gabriele Scalia, G ̈okcen Eraslan, Surag Nair, Tommaso Biancalani, Shuiwang Ji, Aviv Regev, Sergey Levine, and Masatoshi Uehara. Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-based Decoding. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025b. Shanchuan Lin, Anran Wang, and Xiao Yang. SDXL-Lightning: Progressive Adversarial Diffusion Distillation, 2024a. Yeqing Lin, Minji Lee, Zhao Zhang, and Mohammed AlQuraishi. Out of Many, One: Designing and Scaffolding Proteins at the Scale of the Structural Universe with Genie 2, 2024b. Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637):1123–1130, 2023. Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow Matching for Generative Modeling. In The Eleventh International Conference on Learning Representations, 2023. Xingchao Liu, Chengyue Gong, and qiang liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. In The Eleventh International Conference on Learning Representations, 2023. 13 Nanye Ma, Shangyuan Tong, Haolin Jia, Hexiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang, Yandong Li, Tommi Jaakkola, Xuhui Jia, and Saining Xie. Scaling Inference Time Compute for Diffusion Models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 2523–2534, June 2025. Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning Transferable Visual Models From Natural Language Super- vision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, p. 8748–8763. PMLR, 18–24 Jul 2021. Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1): 140022, 2014. Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bj ̈orn Ommer. High-Resolution Image Synthesis With Latent Diffusion Models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), p. 10684– 10695, June 2022. Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Raphael Gontijo-Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic Text-to-Image Diffusion Models with Deep Language Understanding. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmar- czyk, and Jenia Jitsev. LAION-5B: An open large-scale dataset for training next-generation image-text models. arXiv preprint arXiv:2210.08402, 2022. Raghav Singhal, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. A General Framework for Inference-time Scaling and Steering of Diffusion Models. In Forty-second International Conference on Machine Learning, 2025. I.M Sobol’. On the distribution of points in a cube and the approximate evaluation of integrals. USSR Computational Mathematics and Mathematical Physics, 7(4):86–112, 1967. ISSN 0041-5553. Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representations, 2021. Yuxuan Song, Jingjing Gong, Minkai Xu, Ziyao Cao, Yanyan Lan, Stefano Ermon, Hao Zhou, and Wei-Ying Ma. Equivariant Flow Matching with Hybrid Probability Transport for 3D Molecule Generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. Kim Yong Tan, Yueming Lyu, Ivor Tsang, and Yew-Soon Ong. Fast Direct: Query-Efficient Online Black-box Guidance for Diffusion-model Target Generation. In The Thirteenth International Conference on Learning Representations, 2025. Zhiwei Tang, Jiangweizhi Peng, Jiasheng Tang, Mingyi Hong, Fan Wang, and Tsung-Hui Chang. Inference-Time Alignment of Diffusion Models with Direct Noise Optimization, 2025. Michel van Kempen, Stephanie S. Kim, Charlotte Tumescheit, Milot Mirdita, Jeongjae Lee, Cameron L. M. Gilchrist, Johannes S ̈oding, and Martin Steinegger. Fast and accurate protein structure search with Foldseek. Nature Biotechnology, 42(2):243–246, 2024. 14 Luran Wang, Chaoran Cheng, Yizhen Liao, Yanru Qu, and Ge Liu. Training Free Guided Flow-Matching with Optimal Control. In The Thirteenth International Conference on Learning Representations, 2025. Joseph L. Watson, David Juergens, Nathaniel R. Bennett, Brian L. Trippe, Jason Yim, Helen E. Eisenach, Woody Ahern, Andrew J. Borst, Robert J. Ragotte, Lukas F. Milles, et al. De novo design of protein structure and function with RFdiffusion. Nature, 620 (7976):1089–1100, 2023. Jie Wu, Yu Gao, Zilyu Ye, Ming Li, Liang Li, Hanzhong Guo, Jie Liu, Zeyue Xue, Xiaoxia Hou, Wei Liu, Yan Zeng, and Weilin Huang. RewardDance: Reward Scaling in Visual Generation, 2025. Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human Preference Score v2: A Solid Benchmark for Evaluating Human Preferences of Text-to-Image Synthesis, 2023. Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. ImageReward: Learning and Evaluating Human Preferences for Text- to-Image Generation. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. GeoDiff: A Geometric Diffusion Model for Molecular Conformation Generation. In International Conference on Learning Representations, 2022. 15 A Appendix This appendix provides supplementary material to support the results presented in the main text. We begin by detailing the hyperparameters and implementation of TRS in Section B, followed by expanded experimental settings in Section C. We then present additional ablation studies (Section D) and discuss the comparative impact of ODE versus SDE sampling (Section E). Furthermore, we provide a comprehensive overview of all baselines (Section F), offer deeper insights into the evaluation metrics (Section G), and report the runtime for all methods (Section H). Finally, we address the limitations of our approach in Section I and showcase random samples from the image experiments in Section J. B Algorithm Details and Hyperparameters In this section we explain the remaining details about our method of Section 3.2 and show all hyperparameter details from the experiments in Section 4 for TRS and across all baselines. Restart logic. When the failure tolerance thresholdc fail is triggered while the trust-region side length is already at its minimum valueℓ min j , a region restart is triggered. In this case, the side length is reset, and the region is subsequently re-centered at one of the globally best points observed so far. This restart mechanism allows the algorithm to escape from locally saturated regions and to continue allocating evaluations toward more promising areas of the search space. Length-dependent perturbation constraints. To maintain stable behavior across different trust-region scales, we apply simple constraints that couple the side lengthℓ j with the perturbation probabilityp j,b used for coordinate masking m j,b , by rejection sampling. These constraints limit the number of perturbed dimensions when the trust region is large, which is particularly important in high-dimensional noise spaces. • If ℓ j ≥ 2.0→ p j ≤ 0.2. • If ℓ j ≥ 1.6→ p j ≤ 0.5. • If ℓ j ≥ 1.2→ p j ≤ 0.7. These constraints were found to improve robustness without introducing additional tuning complexity. We provide an intuitive visualization for those rules in Figure 10, showing under which combinations of mask and region length the perturbations fail. Hyperparameters. While TRS contains several hyper-parameters, we find them to be robust across all the experiments we conduct in Section 4. We show this in Table 4, where we can see that the number of regions is the only one we change, which is due to the different batch sizes we use for efficiency. But even this shows good robustness, which we show in an ablation in Figure 11. 16 Table 4: Hyperparameters for all our experiments in section 4. When entries are separated with a slash / , we refer to the difference between SD1.5 / SDXL. The dagger † means that is the setting the entry refers to the highest NFE setting in the respective experiment in section 4. The asterix * excludes OC-Flow and DTS. Hyperparameter T2IMoleculsProtein All algorithms Noise Dimension D16384 / 655362913× n res Batchsize B*201008 Inference steps50/850400 OC-Flow (Wang et al., 2025) Batch size B11– Step size η0.25– Weight decay0.9980.995– Weight constraint0.40.4 Number of steps5050– OptimizerSGDL-BFGS– Learning rate α1.01.0– DTS* (Jain et al., 2025) Expansion steps [50, 40, 30, 20, 10] / [8, 6, 4, 2] – Exploration Constant λ0.1 (IR), 0.01 (HPS)– Progressive Width Constant C2.0– Progressive Width α0.4– Exploration typeUCB– Zero-Order Search (Ma et al., 2025) Added noise ε0.10.10.1 Fast Direct (Tan et al., 2025) Step size α808080 Total steps T6 † 6 † 6 † Noise sigma (GP) σ n 0.10.10.1 TRS (Ours) PerturbationsSobol / GaussianSobolSobol # Trust regions k15205 Initial Trust Region length l init 0.80.80.8 Minimal Length l min 0.050.050.05 Maximal Lenght l max 2.42.42.4 Length Update factor α ℓ 1.51.51.5 Success counter threshold c succ 333 Failure counter threshold c fail 333 Fraction warm-up iterations20%20%20% Min. prob. for masks p min 0.10.10.1 Max. prob. for masks p max 0.90.90.9 C Experiment Details In this section, we provide additional details of the experiments in Section 4 with a specific focus on the generative models and the reward functions that are applied. We provide the links to all external resources in Table 5. C.1 Text-to-Image Generative models. We use Stable Diffusion (Rombach et al., 2022) v1.5(∼859.5M parameters), which generates 3×512×512 images from a noise space of 4×64×64, and a distilled version of the larger SDXL Lin et al. (2024a) (∼2.57B parameters), which generates 17 Table 5: Summary of generative models and reward functions for the experiments in Section 4. DomainModel / ToolParamsSource / Repository T2ISD v1.5∼860MHF: SD-v1.5 SDXL-Lightning ∼2.6BHF: SDXL-Lightning ImageReward∼446MGitHub: ImageReward HPSv2∼986MGitHub: HPSv2 Aesthetic Pred. ∼428MGitHub: Aesthetic MoleculesEquiFM∼22MGitHub: EquiFM / OC-Flow ProteinsProteina (M small FS ) ∼60MGitHub: Proteina ProteinMPNN∼1.7MGitHub: ProteinMPNN ESMFold-v1∼3BGitHub: ESMFold 3×1024×1024 images from a noise space of 4×128×128. Both models encode text prompts using CLIP-based text encoders (Radford et al., 2021). Classifier-free guidance (Dhariwal & Nichol, 2021) is applied during sampling by linearly combining conditional and unconditional model predictions with conditioning signalc. For image sampling, we use the DDIM scheduler. We setη= 0 (deterministic sampling) for black-box source noise optimization methods andη= 1.0 (stochastic sampling) for noise sequence optimization methods. The number of inference steps is set to 50 for SD1.5 and 8 for SDXL, and all experiments are conducted in float16 precision. Reward functions. We use three reward functions for the text-to-image experiments, each defined by the scalar output of a pretrained model. ImageReward (Xu et al., 2023) and HPSv2 (Wu et al., 2023) evaluate image–prompt alignment and are defined asR(x 1 ,c), where x 1 denotes the generated image andcthe text prompt. In contrast, the Aesthetic predictor (Schuhmann et al., 2022) evaluates only the generated image and is defined as R(x 1 ). C.2 Molecule Generation Generative model. We use EquiFM (Song et al., 2023) (∼22M parameters) to generate small QM9 molecules and base our implementation on the code of OC-Flow (Wang et al., 2025). EquiFM is an ODE-based flow matching model that operates on a combined continuous Gaussian noise representation for both atomic coordinates and atom types, with total dimensionalityM= 291. The number of atoms per molecule ranges from 3 to 29. To accommodate variable-sized molecules while keeping a fixed input dimension, unused atom entries are masked to zero. In the experiments described in Section 4.2, the number of atoms is sampled according to the empirical distribution in the QM9 training set, which most frequently yields molecules with 15 to 20 atoms. Reward functions. The reward is defined as the negative sum of absolute deviations between target property values and the corresponding predictions of pretrained regression models. For theR 6 reward, we consider the propertiesα,μ,∆ε,ε HOMO ,ε LUMO ,c v , cor- responding to isotropic polarizability, dipole moment, HOMO–LUMO gap, HOMO energy, LUMO energy, and heat capacity. For theR 3 reward, we restrict the objective to the subset α,μ,ε LUMO . C.3 Protein Design Generative model. For protein design, we use theM small FS variant of Proteina (Geffner et al., 2025) with 400 integration steps, which consists of approximately 60M transformer parameters and does not include triangle layers. The source noise is drawn from a contin- uous Gaussian distributionN(0,I) with dimensionalityn res ×3, corresponding to the 3D coordinates of n res residues. 18 Once the model has learned the ODE dynamics, sampling can be performed either via deterministic ODE integration, dx t = v θ t (x t , ̃c) dt,or using stochastic SDE sampling, dx t = v θ t (x t , ̃c)− γg(t)s θ t (x t , ̃c) dt+ p 2γg(t)dw t .We discuss the implications of these sampling choices in more detail in Section E. In practice, setting γ < 1 improves designability at the expense of diversity and novelty. Reward functions. Following (Geffner et al., 2025), we optimize the designability of generated 3D protein backbones. Given a generated structure, a pretrained inverse folding model, ProteinMPNN (Dauparas et al., 2022) (∼1.7M parameters), is used to design a compatible amino acid sequence. This sequence is subsequently processed by a folding model, for which we use the transformer-based ESMFold-v1 (Lin et al., 2023) (∼3B parameters), to predict its 3D structure. We compare the predicted structure to the original generated backbone by computing the self-consistency root mean square deviation (scRMSD). The scalar reward is then defined asR(x 1 ) =exp(−scRMSD(x 1 )),which maps the score to the range (0,1], with values closer to 1 indicating higher designability. Note that usually 8 sequences are extracted and folded and the best of those is chosen, but for computational efficiency we only extract 1 per protein backbone. D Ablations Here we perform further ablation studies to justify the design choices in Section 3.2 and hyperparameters in Section B. Most importantly, Section D.1 emphasises the importance of our top-kcenter selection instead of keeping them strictly apart. Section D.2 provides visual and quantitative insights into the relationship of the trust-region lengthsℓ j and the mask probabilitiesp j , while Section D.3 shows that 20% is a good heuristic for the fraction of warm-up budgets across multiple total budgets. Finally, Section D.4 shows that even the number of regionsk, is a robust hyperparameter, which is the only one that we change across our experiments in Section 4. D.1 Center Selection Setup. We investigate four center selection strategies for text-to-image generation using Stable Diffusion 1.5 (SD1.5) and ImageReward as reward function. LocalBest and Lo- calLastIter enforce strict region separation, updating centers using either the historical best or the most recent best sample within each regionT j , respectively. In contrast, GlobalTopk and GlobalLastIter allow interaction by re-centering regions globally based on the overall top-ksamples observed so far or in the latest iteration. All experiments use a batch size B = 20 and 20 iterations (20k NFE), consistent with the configuration in Section 4.1. Metrics. Performance is evaluated on the full DrawBench benchmark. We use the mean best ImageReward as the primary metric to rank the generated noise samples and determine center updates. Random search is included as baseline as a reference to compare the relative improvement of each selection strategy. Results. As shown in Figure 9, GlobalTopk consistently outperforms all other strategies. The results indicate that enforcing strict separation between regions (as in the original TuRBO algorithm (Eriksson et al., 2019)) limits performance in this setting. Conversely, global re-centering enables a more effective allocation of the evaluation budget toward the most promising areas of the search space, significantly improving generation quality. D.2 Trust Region Dynamics: Interaction of Length and Masks Setup. We perform a grid ablation across varying initial trust region lengthsℓ init and masking sparsity ranges defined byp min andp max . First, we do this quantitatively by investigating different combinations of initial trust-region lengthsℓ init and set the mask probabilities in different ranges, with different min and max probabilitiesp min ,p max . This experiment is conducted using SD1.5 on a representative subset of 55 prompts from Draw- Bench, covering all categories of the original benchmark to ensure broad coverage of the 19 1.41.451.51.551.6 RandomSearch LocalLastIter LocalBest GlobalLastIter GlobalTopk Mean best reward Center Selection Strategies (SD1.5 + IR) Figure 9: Effect of center selection strategy on DrawBench using SD1.5 and ImageReward. We compare the mean best rewards across the prompts in the benchmark. latent space. Additionally, we visualize the relationship of the trust-region lengthℓ j and sampled mask probability p j in Figure 10. Metrics. In Figure 10 we aim to track the point at which the samples begin to exhibit structural degradation and visualize these noising effects, while we show the mean best ImageReward for the quantitative experiment in Table 6. Results. Figure 10 reveals a clear stability threshold at approximatelyℓ j ≈1.6, beyond which images show visible noise. However, we find that low-probability masks (e.g.,p j = 0.05) act as a regularizer and they allow for significantly higher exploration lengths, up toℓ j = 6.4, while maintaining valid, coherent structures. Table 6 reveals that the exact masking range is not particalarely important, it works well for different ranges. The initial shows also good robustness, even when changing it to very high or low numbers, due to the adaptive nature of the algorithm. Table 6: Mean best rewards across varying trust-region init lengths (ℓ init ) and [p min ,p max ] settings for SD1.5 and ImageReward. TR init lengthMask probabilities [p min , p max ] (ℓ init )[0.05, 0.50][0.10, 0.90][0.30, 0.70][0.50, 1.00] 0.2 (Low)1.62231.64301.63531.6564 0.8 (Medium)1.69281.67611.68201.6543 2.0 (Large)1.64041.61971.63061.6185 D.3 Warm-Up Setup. We investigate the impact of the warm-up phase on performance to determine the optimal allocation of the total compute budget between initialization and the main opti- mization process. This experiment uses Stable Diffusion 1.5 on the subset of 55 DrawBench prompts, equally distributed across all categories. We evaluate three total evaluation budgets (N total ∈120, 360, 720) with a fixed batch size of B = 24. Metrics. To assess the effectiveness of different budget allocations, we utilize ImageReward as the reward function. We report the final mean best rewards achieved across all 55 prompts, tracking how the optimization efficiency changes as the ratio of warm-up to main optimization varies. Results. As shown in Figure 11a, TRS achieves peak performance when the warm-up phase accounts for 10% to 20% of the total compute budget across all tested settings. We 20 (a) SD1.5 (D = 16384)(b) SDXL (D = 65536) 0.050.20.40.60.81.00.050.20.40.60.81.0 0.4 0.8 1.6 3.2 6.4 Figure 10: TRS perturbation landscape comparison: Interaction between perturbation length (rows) and masking probability (columns) for SD1.5 and SDXL. For both we used the prompt: ”A breathtaking view from behind the Cristo Redentor (Christ the Redeemer) statue in Rio de Janeiro, Brazil, with layered mountains stretching into the distance and the sparkling Atlantic Ocean clearly visible below with blue color; warm golden-hour light, atmospheric haze, ultra-detailed, cinematic, wide-angle landscape, beautiful and serene and the city below.” 020406080 1.4 1.6 1.8 Warmup Fraction (%) Mean Best ImageReward (a) Warm-up Fraction 1248121624 1.55 1.6 1.65 Image (SD1.5 & IR) Budget 120Budget 360Budget 720 Figure 11: Ablation Studies. (a) Impact of the warm-up phase fraction across different compute budgets. (b) Sensitivity of TRS to the number of regionsk. Performance remains stable across a wide range of k values. note that while performance is sensitive to very small warm-up budgets, it degrades only slowly when the warm-up phase is slightly larger than the optimal range. This suggests that the exact choice is not overly critical, provided the initialization period is sufficient; consequently, we set the warm-up fraction to 20% for all main experiments. D.4 Number of Regions Setup. We ablate the impact of the number of trust regionsk ∈1,2,4,8,16,24on the optimization performance for text-to-image generation. These experiments use SD1.5 on a 55-prompt subset of DrawBench, utilizing a fixed batch size ofB= 24. The optimization process is partitioned into 3 warm-up iterations followed by 12 trust-region iterations. Metrics. Consistent with our previous ablations, we align the optimization to ImageReward. We evaluate the final mean best rewards across the prompt subset to identify the optimal range forkand to assess how the number of regions interacts with the total evaluation budget and batch size. 21 Results. As illustrated in Figure 11b, using a very small number of regions leads to sub-optimal performance. However, we find that selectingkin a medium range between 1 and the batch sizeBconsistently yields strong results. These findings suggest thatkis not a highly sensitive hyperparameter, as performance remains stable across a wide range of values, indicating that it does not require extensive per-task tuning. Sincekis the only hpyerparameter that was changed between the main experiments in Section 4, we conclude that TRS is a highly robust algorithm for noise optimization of flow and diffusion models. E ODE vs. SDE Sampling Dynamics TRS is a black-box optimization method that relies on local exploration to find better samples. This approach is more effective with ODE sampling because the deterministic paths provide a clear sense of locality. However, since SDE sampling with noise reduction leads to higher designability in modern generative 3D protein models (Geffner et al., 2025; Bose et al., 2024), we want to see how our algorithm handles the trade-off between the designability gains of SDEs and the diversity preserved by ODEs. We also include experiments for text-to-Image (T2I) generation to see if these sampling dynamics remain consistent across different model types. Setup. We compare TRS against standard random search (RS) in both deterministic and stochastic settings. For protein generation, we generate 100 proteins of length 50, similar to Section 4.3 optimizing specifically for designability. For T2I generation, we use the same setup as in Section 4.1, utilizing the DDIM scheduler with both a stochastic setting (η= 1.0) and a deterministic setting (η= 0.0). A key difference between these setups is the inference budget: proteina uses 400 steps, while the T2I model uses 50. Metrics. We evaluate the protein results based on designability, diversity, and novelty and the Rank-1 percentage and show the results in Table 7. With this last metric, we want to show not only into how many clusters the proteins fall in this analysis, but also what the distribution between those clusters is. Rank-1 is the percentage of proteins that fall into the most dominant cluster. In more detail Figure 12 shows the distribution across the 10 most dominant clusters using rank–size and cumulative mass plots to show how concentrated the generated samples are. Results and Discussion. As shown in Table 7, SDE sampling does lead to higher designability, but it significantly diminishes diversity and novelty. When we compare TRS with ODE to random search and TRS with SDE (γ= 0.6), we observe that the designability metrics are on par, while TRS with ODE shows much better diversity and novelty metrics. This becomes especially clear when observing the Rank-1 metric, where we see that SDE- optimized samples, for which 54 or even 68% of the samples belong to the same cluster. Figure 12 further illustrates the mode collapse of SDE-based sampling in protein design, where we can see that in SDE-based optimization, the designable samples are highly concentrated at the few most dominant clusters. ODE-based optimization instead shows a healthier distribution, creating structures distributed at different clusters more equally. We also find that most current SDE-based scaling methods are not practical for proteina, either because they lack necessary value estimations (Li et al., 2025b; Singhal et al., 2025) or because they are too slow without efficient batch utilization (Jain et al., 2025), especially, when many inference steps are necessary and the reward function is expensive. Interestingly, TRS shows a smaller advantage over random search in proteina when using SDEs compared to T2I. This is likely because the high number of inference steps in proteina allows the stochastic noise to eventually override the local search signal provided by the trust region. 22 Table 7: Comparison of ODE and SDE variants for random search (RS) and our trust-region search (TRS) in protein design and text-to-image (T2I). ImageReward and designability are the respective alignment rewards and Rank-1, Cluster Divergence, pairwise TM-Score and PDB Novelty are additional metrics to uncover the effects of SDE vs. ODE based alignment with the 400-step proteina model. In T2I we use the 50-step SD1.5 as generative model. T2IProtein (SDE with γ = 0.6) AlgorithmImageReward↑ Design.↑ Rank-1 ↓ Clust. Div.↑ TM Div.↓ Nov.↓ RS + ODE1.430.53370.270.680.85 RS + SDE1.450.66680.190.730.94 TRS (Ours) + ODE1.620.65180.300.670.89 TRS (Ours) + SDE1.520.66540.270.700.92 12345678910 1 2 5 10 20 50 100 Cluster rank Cluster size Cluster Rank–Size Distribution RS + ODERS + SDE TRS + ODETRS + SDE 00.20.40.60.81 0 0.25 0.5 0.75 1 Fraction of clusters Cumulative fraction Cumulative Cluster Mass Figure 12: Protein cluster distributions. The rank–size plot (left) shows that ODE methods maintain higher diversity, whereas SDE methods suffer from mode-collapse. The cumulative mass (right) confirms that TRS + ODE provides the most balanced distribution of samples. F Baselines In this section, we provide more detail about the baselines and explain their hyperparameters. The exact numbers used in the experiments are shown in Table 4. F.1 Gradient-based guidance Optimal Control Flow (OC-Flow). OC-Flow (Wang et al., 2025) provides a theoretically grounded, training-free framework for guided flow matching by framing generation as an optimal control problem. It augments pre-trained flow dynamics with a time-dependent control term u t : ̇ x t = v t (x t ) + u t .The framework seeks to minimize a cost functional J(u) =R(x 1 ) + R 1 0 1 2λ ∥u t ∥ 2 dt, comprising a terminal reward lossR(x 1 ) and a quadratic running cost that regulates the trajectory’s deviation from the prior distribution. Leveraging Pontryagin’s Maximum Principle, the control trajectory is optimized iteratively over a fixed number of steps using either SGD or L-BFGS. In each iteration, gradient information is propagated backward via a co-state flowμ t to update the control parameters with a step sizeηand weight decay. To ensure trajectory regularity and stable convergence, a weight constraint is enforced on the magnitude of the control term throughout the optimization process. F.2 Noise sequence search Diffusion Tree Sampling (DTS). DTS (Jain et al., 2025) frames the inference-time alignment of diffusion models as a tree-structured optimization problem over the denoising 23 sequence. The framework employs a recursive value-based search guided by a soft value functionV(x t ), which is estimated via a soft-Bellman backup with an exploration constant λ: V (x t ) = 1 λ log E x t−1 ∼p θ (·|x t ) [exp(λV (x t−1 ))] At each node, the search proceeds for a fixed number of expansion steps, where the selection of trajectories is governed by a specific exploration type (e.g., UCT). To effectively navigate the continuous branching space of the diffusion process, the algorithm utilizes progressive widening to determine the number of childrenkfor a node x t based on its visit countN(x t ): k(x t ) =⌈C· N (x t ) α ⌉ whereCis the progressive width constant andαis the progressive width exponent. A stochastic rollout parameterρserves as a decision gate: with probabilityρ, the algorithm performs a full rollout to the terminal statet= 0 to obtain an exact reward, while with probability 1− ρ, it continues recursive tree expansion. Fast Direct. Fast Direct (Tan et al., 2025) is a trajectory-level optimization method designed for black-box guidance and optimizes the complete noise sequenceε t T t=1 simulta- neously. The method identifies a pseudo-target x ∗ on the manifold via Gaussian Process (GP) regression over previous evaluations and computes a universal direction to update the entire sequence: ε new t = Norm(ε old t + α(x ∗ − x 0 )) whereαis the step size. This global refinement allows it to converge to high-reward regions in fewer iterations than per-step filtering methods. F.3 Source noise search Random Search. Random search is the simplest search-based method one can apply, whereNsource noises x 0 are randomly sampled and evaluated. Thus, it is often referred to as best-of-N in the literature Jain et al. (2025). This technique can be used both with ODE- and SDE-based samplers. Zero-Order Search. Zero-order Search (Ma et al., 2025) can be seen as a special case of our algorithm. Like ours, it iteratively refines the initial latent noise to maximize a target reward. The algorithm begins by sampling an initial set ofBGaussian noise vectorsx j 0 B j=1 and selecting the candidate that yields the highest reward as the initial center x c 0 . This is equivalent to using only one warm-up iteration in TRS and setting the number of regions k= 1. In each subsequent search iteration,Bnew candidates are generated by perturbing the current center: x new 0 = x c 0 + ε·δ,δ ∼N (0, I) whereε(scalar) controls the radius of the local neighborhood search. This perturbation scheme is similar to using TRS without probability masks and choosing a fixed region length. Similar to ours, this technique works best for ODE samplers. G Metrics In this section, we describe all additional evaluation metrics used in this work that are not employed as reward functions. The reward-based metrics are detailed in Section C. We consider evaluation criteria for both molecule generation and protein design. G.1 Molecule Generation For the evaluation of 3D molecule generation, we follow the methodology described by (Song et al., 2023), which assesses the physical and chemical plausibility of the generated Cartesian coordinates and atom types. 24 Stability. We evaluate the structural integrity of the generated samples using Atom Stability and Molecule Stability. Following the convention of Hoogeboom et al. (Hoogeboom et al., 2022), chemical bonds are inferred based on the Euclidean distances between atoms using a threshold-based lookup table of covalent radii. An individual atom is considered stable if its inferred bond count matches its expected chemical valency (e.g., 4 for carbon, 1 for hydrogen). Molecule Stability is then defined as the percentage of generated molecules for which all constituent atoms are stable. This metric serves as a proxy for the geometric consistency of the generated 3D structures. Validity, Uniqueness, and Novelty. The chemical validity of a molecule is determined by its ability to be successfully parsed into a molecular graph using RDKit (Landrum et al., 2026). We report the Valid Fraction, the proportion of samples satisfying fundamental valency constraints. The Unique Fraction indicates the percentage of non-redundant molecules among generated compounds. We also report Novelty, defined as the fraction of generated molecules not present in the training set. To provide a stringent quality assessment, we track the Molecule Stability Percentage (MSP) and the Valid and Unique Percentage (VUP). The former assesses the geometric consistency of the 3D coordinates, while the latter ensures the model explores the chemical space without collapsing into a few valid structures. G.2 Protein Design To evaluate the structural quality and variability of the generated protein backbones, we adopt the metric suite introduced by (Geffner et al., 2025). These metrics assess two complementary aspects: the uniqueness of the generated structures relative to known proteins (Novelty) and the structural variability within the generated set (Diversity). Following (Geffner et al., 2025), we compute these scores only for designable samples, which means that scRMSD ¿ 2.0 ̊ A (equivalently designability ¿ 0.1353). Novelty. Novelty measures the extent to which the model generates protein folds that differ from those observed in existing experimental and predicted structure databases. For each generated, designable backbone structure, we compute the maximum TM-score against all entries in a reference set using Foldseek (van Kempen et al., 2024). We then report the average of these maximum TM-scores across the entire sample set. Lower average maximum TM-scores indicate higher structural novelty. We evaluate novelty with respect to the PDB dataset (Berman et al., 2000). Diversity. We quantify the internal diversity of the generated, designable backbone samples using two complementary metrics: 1.Average Pairwise TM-score: We compute the mean pairwise TM-score between all designable backbone samples for each generated protein length. These values are subsequently aggregated to obtain a global average. Since the TM-score measures structural similarity on a scale from 0 to 1, lower values correspond to higher diversity. 2.Cluster Ratio: We cluster the generated backbones using Foldseek with a TM-score threshold of 0.5. The diversity score is defined as the ratio of the number of unique clusters to the total number of designable samples. Because a more diverse set of samples results in a higher number of clusters, higher ratios indicate greater structural variety within the generated set. H Runtime Comparison We analyze the runtime of the experiments described in Section 4 for all considered methods on a single NVIDIA A100-SXM4-40GB. The quantitative comparison of computation time and memory usage is visualized in Figure 13. For these benchmarks, we use the settings of ImageReward for T2I,R 6 for molecule generation, and protein designability for sequences of 50 residues. 25 We observe that DTS * generally exhibits the highest compute time. This arises because the algorithm’s sequential nature is difficult to combine with fixed-batch-size calls of the generative model. In our benchmarks, this results in DTS * being approximately 4×slower than TRS, which is naturally parallelizable. For OC-Flow, which operates with a batch size of 1, the computation time is high for SD1.5 due to the requirement of back-propagating through 50 integration steps. Conversely, for the distilled SDXL Lightning, OC-Flow is faster but incurs substantially higher memory usage. In the molecule generation task, OC-Flow is the slowest method, while its memory footprint in this domain is lower than that of TRS and other sampling-based methods, which leverage a larger batch size of 100. Regarding Fast Direct, we employ slightly higher NFE budgets because their iterative step- increasing mechanism makes it difficult to match an exact NFE while maintaining a constant batch size. Beyond this, their computation times and memory usage are only marginally higher than our baselines, as the overhead from their Gaussian Process (GP) is minimal. Finally, Random Search, Zero-Order Search, and TRS are the most efficient overall, as their runtime consists almost entirely of the generative model forward passes and reward function evaluations. 0 500 1,000 Time (s) SD1.5 + IR Runtime (↓) 0 200 400 600 800 SDXL + IR Runtime (↓) 0 50 100 EquiFM + R 6 Runtime (↓) 0 200 400 600 800 Prot. n res = 50 Runtime (↓) TRS (Ours) RSZOGNSOOC-Flow DTS* 0 10 20 30 VRAM (GB) SD1.5 + HPS Mem. (↓) 0 10 20 30 SDXL + HPS Mem. (↓) 0 5 10 EquiFM + R 3 Mem. (↓) 0 5 10 Prot. n res = 50 Mem. (↓) Figure 13: Comparison of runtime and memory consumption across methods on a single NVIDIA A100-SXM4-40GB. I Limitations Reward Functions. A primary limitation of these methods lies in the reliability of reward functions; when derived from neural network predictors, performance is inherently constrained by these models. In our experiments, reliance on pretrained networks for feedback makes optimization susceptible to model bias and failure modes. Figure 14 illustrates this via TRS-optimized SDXL images that receive high rewards despite failing to capture the core prompt intent. This highlights a fundamental challenge: reward models may score samples highly that exploit shortcuts in the learned signal rather than satisfying the true objective. To address this, recent work has focused on larger, more accurate reward models (Wu et al., 2025). TRS is well-positioned to integrate with these advances, offering a flexible and scalable framework as the field progresses toward more expressive models. Diversity. While TRS begins with a global exploration phase across multiple promising regions, it ultimately converges to and exploits the most dominant region. Our experimental results demonstrate that this is an effective strategy for identifying a single optimal sample that maximizes the reward objective. However, TRS may be less suited for tasks requiring a diverse ensemble of samples. A direction for future work would be extending the framework to optimize for a diverse set of high-reward samples rather than a single point estimate. For 26 instance, diversity could be explicitly enforced in TRS by incorporating a cosine similarity constraint on the center noise vectors of the trust regions. Two cats and two dogs sitting on the grass A carrot on the left of a broccoliA cat on the right of a tennis racket 1.861.901.531.571.831.89 Four cars on the streetA painting by Grant Wood of an astronaut couple, american gothic style A banana on the left of an apple 0.3170.3290.3030.3090.3030.308 Figure 14: Visualizing reward misalignment: cases where higher reward values are assigned even though the model fails to correctly generate the specific conceptual or spatial re- quirements of the prompt. All images are generated with SDXL. The upper row shows missalignment to ImageReward and the lower one to HPSv2. J Optimized samples We randomly sample 10 prompts out of DrawBench and show the qualitative results of the six algorithms we compare in Section 4.1 including our TRS. The images are the best images per prompt from experiment Section 4.1. We provide in total four figures, including Figure 15 for SD1.5 and ImageReward, Figure 16 for SD1.5 and HPSv2, Figure 17 for SDXL-Lightning and ImageReward and Figure 18 for SDXL-lightning optimized for HPSv2. 27 Prompt TRS (Ours)RSZODTS*FDOC-Flow A keyboard made of water, the water is made of light, the light is turned off. A red colored dog. Hyper- realistic photo of an abandoned industrial site during a storm. A cat on the left of a dog. A giraffe underneath a microwave. A baby fennec sneezing onto a strawberry, detailed, macro, studio light, droplets, backlit ears. One car on the street. A pink colored car. An umbrella on top of a spoon. A single clock is sitting on a table. Figure 15: Comparison of 10 randomly chosen prompts for different alignment methods using Stable Diffusion v1.5 and ImageReward. 28 Prompt TRS (Ours)RSZODTS*FDOC-Flow A keyboard made of water, the water is made of light, the light is turned off. A red colored dog. Hyper- realistic photo of an abandoned industrial site during a storm. A cat on the left of a dog. A giraffe underneath a microwave. A baby fennec sneezing onto a strawberry, detailed, macro, studio light, droplets, backlit ears. One car on the street. A pink colored car. An umbrella on top of a spoon. A single clock is sitting on a table. Figure 16: Comparison of 10 randomly chosen prompts for different alignment methods using Stable Diffusion v1.5 and HPSv2. 29 Prompt TRS (Ours)RSZODTS*FDOC-Flow A keyboard made of water, the water is made of light, the light is turned off. A red colored dog. Hyper- realistic photo of an abandoned industrial site during a storm. A cat on the left of a dog. A giraffe underneath a microwave. A baby fennec sneezing onto a strawberry, detailed, macro, studio light, droplets, backlit ears. One car on the street. A pink colored car. An umbrella on top of a spoon. A single clock is sitting on a table. Figure 17: Comparison of 10 randomly chosen prompts for different alignment methods using SDXL-Lightning and ImageReward. 30 Prompt TRS (Ours)RSZODTS*FDOC-Flow A keyboard made of water, the water is made of light, the light is turned off. A red colored dog. Hyper- realistic photo of an abandoned industrial site during a storm. A cat on the left of a dog. A giraffe underneath a microwave. A baby fennec sneezing onto a strawberry, detailed, macro, studio light, droplets, backlit ears. One car on the street. A pink colored car. An umbrella on top of a spoon. A single clock is sitting on a table. Figure 18: Comparison of 10 randomly chosen prompts for different alignment methods using SDXL-Lightning and HPSv2. 31