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

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Schweiger, Niklas, Cremers, Daniel, Ram, Karnik
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912967287635968
author Schweiger, Niklas
Cremers, Daniel
Ram, Karnik
author_facet Schweiger, Niklas
Cremers, Daniel
Ram, Karnik
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2603_14504
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Trust-Region Noise Search for Black-Box Alignment of Diffusion and Flow Models
Schweiger, Niklas
Cremers, Daniel
Ram, Karnik
Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
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.
title Trust-Region Noise Search for Black-Box Alignment of Diffusion and Flow Models
topic Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2603.14504