Gradient-Free Noise Optimization for Reward Alignment in Generative Models

Technical Summary: ZeNO (Zeroth-order Noise Optimization)

Problem Statement
Existing reward alignment methods for generative models (e.g., diffusion and flow models) typically rely on optimizing multi-step stochastic trajectories. While effective for these settings, such approaches are difficult to extend to deterministic or one-step generators (e.g., consistency models, distilled diffusion models), which do not naturally expose the same trajectory structure. A natural alternative is to optimize the input noise directly. However, standard gradient-based noise optimization requires backpropagation through both the generator and the reward function. This limits applicability to differentiable pipelines and can be unreliable when gradients reflect spurious local sensitivities rather than meaningful semantic improvements. Furthermore, many practical reward functions (e.g., protein designability, compositional text-image alignment) are non-differentiable or black-box, rendering gradient-based methods infeasible.

Methodology: ZeNO
The authors propose ZeNO (Zeroth-order Noise Optimization), a gradient-free framework that formulates noise optimization as a path-integral control (PIC) problem. The core idea is to treat the evolution of the input noise as a controllable stochastic process, optimizing it using only zeroth-order reward evaluations.

1. Path-Integral Control Formulation:
   The method models noise optimization as a controlled Stochastic Differential Equation (SDE). The objective is to steer the noise trajectory toward regions that yield high-reward terminal states while regularizing the control effort. The optimal control $u^*(z, t)$ is derived in a path-integral form:
   $$u^*(z, t) = \frac{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda) d w_t]}{\mathbb{E}_{z:t \to T} [\exp(r(G_\theta(z_T))/\lambda)]}$$
   Crucially, this estimator depends only on terminal-cost evaluations (rewards) and does not require differentiation through the generator $G_\theta$ or the reward function $r$.

2. Reference Dynamics (Ornstein–Uhlenbeck):
   To ensure the updated noise remains compatible with pre-trained generators, ZeNO adopts the Ornstein–Uhlenbeck (OU) process as the reference dynamics. Unlike variance-exploding SDEs, the OU process preserves the Gaussian stationary distribution $N(0, I)$, which matches the prior of most pre-trained generators. This choice connects the noise update to Langevin dynamics implicitly targeting a reward-tilted distribution $p^*(z) \propto p(z) \exp(r(G_\theta(z))/\lambda)$.

3. Linearized Estimator and Variance Reduction:
   The standard path-integral estimator suffers from high variance due to exponential weighting. The authors introduce a linearized approximation (Proposition 1) by subtracting the mean reward $\bar{r}$ as a baseline and linearizing the exponential term under the assumption of small reward variance:
   $$u^*(z_t, t) \approx \mathbb{E}[(r(G_\theta(z_T)) - \bar{r}) d w_t] / \lambda$$
   This estimator coincides with the score function estimator used in Natural Evolution Strategies and empirically outperforms exponential variants.

4. One-Step Receding-Horizon:
   To avoid the computational cost of long-horizon rollouts, ZeNO employs a one-step receding-horizon approximation ($H=1$). At each iteration, $N$ perturbed noise candidates are sampled from the reference dynamics, evaluated by the reward model, and the linearized control estimate is used to update the current noise state. This allows for inference-time scaling by increasing the number of particles ($N$) or iterations ($M$) without modifying the generator.

Key Contributions

• Gradient-Free Alignment: ZeNO enables reward alignment for deterministic and one-step generators without requiring backpropagation through the generator or reward model.
• Black-Box Compatibility: The method is applicable to non-differentiable and black-box reward functions, a setting where existing gradient-based noise optimization methods (e.g., ORIGEN, ReNO) fail.
• Theoretical Connection: The framework establishes a theoretical link between zeroth-order noise optimization and Langevin sampling targeting a reward-tilted distribution, providing a principled justification for the update rule.
• Inference-Time Scaling: ZeNO introduces a flexible mechanism to trade off computational budget (via $N$ and $M$) for improved reward performance.

Experimental Results
The authors evaluate ZeNO across diverse generators (SDXL-Turbo, DMD2, LCM, DFGAN) and reward functions (Aesthetic score, PickScore, GenEval, and protein designability).

• Image Generation: ZeNO consistently improves reward scores across all tested generators and reward types. It achieves competitive or superior performance compared to gradient-based baselines (ORIGEN, $\Psi$-Sampler) and Best-of-N sampling. Notably, on the LCM generator, gradient-based methods degraded below the uncontrolled baseline, whereas ZeNO maintained improvement, highlighting the robustness of zeroth-order methods against local gradient sensitivities.
• Non-Differentiable Rewards: In the GenEval task (object counting and positioning), where rewards are non-differentiable, ZeNO outperformed Best-of-N on counting tasks, demonstrating that iterative noise optimization extracts more value from the same computational budget than independent sampling.
• Protein Structure Generation: Applied to a one-step Riemannian MeanFlow (RMF) protein backbone generator, ZeNO significantly improved designability (measured by self-consistency RMSD). It achieved designability levels comparable to a 5-step RMF baseline using only a single flow evaluation, demonstrating its ability to compensate for reduced sampling steps via noise optimization.
• Diversity: Analysis using the Vendi Score indicates that ZeNO improves reward while preserving sample diversity, avoiding the collapse to a single mode often seen in gradient-based approaches.

Significance and Claims
The paper claims that ZeNO provides a practical and theoretically grounded solution for aligning the growing class of efficient, deterministic, and one-step generative models with reward functions. Its primary significance lies in its ability to operate in settings where gradients are unavailable or unreliable, such as with black-box reward models or complex biological design tasks. By formulating noise optimization as a path-integral control problem, ZeNO offers a scalable, generator-agnostic framework for inference-time alignment that does not require retraining or modifying the underlying generative model. The authors note a limitation: the method requires multiple reward evaluations per update step, which can be computationally demanding if the reward model itself is expensive.