Technical Summary: Gradient Preconditioning for Efficient and Reliable Reward-Guided Generation

Problem Statement

The paper addresses two critical limitations in test-time latent optimization for one-step generative models (e.g., diffusion distillations like FLUX, SDXL-Turbo). While optimizing the latent vector at inference time allows for reward maximization without retraining, existing approaches suffer from:

1. Reward Hacking: As optimization progresses, the latent vector drifts away from the standard white Gaussian noise prior. This causes the model to generate unrealistic artifacts or "hacks" the reward function, degrading sample quality.
2. Inefficiency: Current methods typically require multiple gradient steps per sample to converge, often taking minutes per image, which hinders practical deployment.

Existing solutions rely on regularization (soft penalties) to encourage Gaussian properties (e.g., $\ell_2$ norm, kurtosis, or spectral penalties). However, the authors argue that regularization is inherently a "soft" mechanism that offers no guarantee of maintaining noise-like characteristics, often fails to prevent reward hacking when optimizers find shortcuts, and introduces hyperparameters that require tuning.

Methodology

The authors propose Gradient Preconditioning, a method that replaces soft regularization with a hard, geometric constraint on the optimization direction.

1. The Core Concept

Instead of penalizing deviations from white Gaussian noise, the method projects the reward gradient onto a carefully designed white Gaussian noise feasible set ($\mathcal{G}$) before applying the update.

• Mechanism: At each optimization step, the reward gradient $\nabla_x J$ is projected onto $\mathcal{G}$. The resulting vector becomes the update direction.
• Effect: This reshapes every update into a "noise-aligned" direction. It drives faster reward ascent within the noise-compatible subspace while preventing the latent vector from drifting into regions that cause reward hacking.
• Advantage: Unlike regularization, this requires no additional hyperparameters (like a weighting coefficient $\lambda$) and does not require the latent vector itself to lie strictly within the set, only the update direction.

2. Design of the Feasible Set ($\mathcal{G}$)

The feasible set is constructed to tightly characterize white Gaussian noise in the spectral domain while admitting a closed-form projection.

• Compact Spectral Domain: The authors define a bijective mapping $\mathcal{F}$ that transforms the real-valued latent vector $x \in \mathbb{R}^N$ into a complex-valued vector $y \in \mathbb{C}^{N/2}$, removing the Hermitian redundancy of the Discrete Fourier Transform (DFT) coefficients.
• Blockwise Constraints: The set $\mathcal{G}_C$ in the compact spectral domain is defined by enforcing that every block of size $B$ satisfies specific $\ell_1$ and $\ell_2$ norm constraints matching the theoretical expectations of a complex Gaussian distribution $\mathcal{CN}(0, I)$.
  • $\ell_1$ constraint: $\|y^{(p)}\|_1 = \sqrt{\frac{\pi}{2}} B$
  • $\ell_2$ constraint: $\|y^{(p)}\|_2^2 = B$
• Projection Algorithm: The projection onto the intersection of $\ell_1$ and $\ell_2$ spheres is solved in closed form. The algorithm involves sorting magnitudes within each block and solving a quadratic equation to find the optimal Lagrange multipliers.
• Complexity: The projection is implemented via Fast Fourier Transform (FFT) and block-wise sorting, resulting in $O(N \log N)$ complexity. In experiments with the FLUX model, this adds only 0.04% to the runtime.

Key Contributions

1. Gradient Preconditioning Framework: A novel optimization strategy that projects gradients onto a noise-feasible set, effectively preventing reward hacking without soft penalties.
2. Tight Noise Characterization: A spectral-domain feasible set defined by blockwise $\ell_1$ and $\ell_2$ constraints that captures both marginal Gaussian statistics and spatial uncorrelatedness (white noise properties) more tightly than prior marginal or spectral regularizers.
3. Efficient Closed-Form Projection: The derivation of an $O(N \log N)$ projection algorithm that is computationally negligible compared to the generative model forward pass.
4. Hyperparameter-Free Operation: The method eliminates the need for tuning regularization coefficients ($\lambda$), which are a source of instability in prior works.

Experimental Results

The method was evaluated on one-step text-to-image models (FLUX, SDXL-Turbo, SANA-Sprint, SD-Turbo) using four reward models: Aesthetic Score, PickScore, HPSv2, and ImageReward.

• Performance vs. State-of-the-Art: On FLUX with Aesthetic Score optimization, the proposed method achieved a comparable reward level to the state-of-the-art regularization-based method (MPGR) using only 30% of the wall-clock time (60 iterations vs. 200 iterations).
• Reward vs. Quality: The method consistently reached higher target rewards (e.g., Aesthetic Score ~8.9 vs. ~7.1 for baselines) while preserving or improving held-out human-preference metrics (PickScore, HPSv2).
• Stability: Qualitative results showed that the method reliably produces realistic images aligned with text prompts, whereas regularization-based methods often yielded inconsistent artifacts or failed to maintain prompt fidelity.
• Diversity: Analysis using Inception Score (IS) and Vendi Score confirmed that the method does not cause mode collapse; sample diversity remained within the variance of unoptimized baselines.
• Ablation: Removing the latent projection (keeping only gradient projection) still yielded significant improvements over baselines, but the full method (gradient + latent projection) provided the best stability and held-out scores.

Significance and Claims

The paper claims that by resolving both the efficiency bottleneck and the reward hacking problem, gradient preconditioning makes test-time latent optimization substantially more practical for real-world applications.

The authors emphasize that their approach offers a principled alternative to soft regularization. By enforcing a hard geometric constraint on the update direction, they guarantee that the optimization trajectory remains aligned with the white Gaussian noise prior. This ensures that the generative model's capacity is exploited for reward maximization without compromising the statistical properties required for realistic generation. The negligible computational overhead and lack of hyperparameters make the method a drop-in replacement for existing gradient-based optimization pipelines in one-step generative models.