Trust-Region Noise Search for Black-Box Alignment of Diffusion and Flow Models

2026-03-15

AI Summary

Existing inference-time alignment methods are often constrained by differentiability requirements, dependence on specific model architectures, or poor computational efficiency. This work proposes Trust-Region Search (TRS), a noise-based optimization algorithm that treats both the generative model and the reward model as black boxes, optimizing only the initial noise vector to strike an effective balance between global exploration and local exploitation. TRS requires no access to internal model structures, is compatible with diverse generative frameworks—including diffusion models and flow models—and supports arbitrary reward functions, all while requiring minimal hyperparameter tuning. Empirical results across text-to-image generation, molecular design, and protein engineering demonstrate that TRS consistently outperforms existing baselines and sophisticated trajectory optimization approaches, highlighting its remarkable versatility and efficiency.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Constraint Satisfaction and Optimization: Search
• Machine Learning: Optimization

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

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.

Problem

• black-box alignment
• diffusion models
• flow models
• noise optimization
• reward modeling

Innovation

• trust-region
• black-box optimization
• noise search
• diffusion models
• inference-time alignment