Gradient Preconditioning for Efficient and Reliable Reward-Guided Generation

We propose a gradient preconditioning method that makes reward-guided generation with one-step generative models both efficient and reliable. Test-time noise optimization can unlock substantially better reward-guided generations from pretrained generative models, but it is prone to reward hacking that degrades quality and is often too slow for practical use. We precondition reward gradients by projecting them onto a carefully designed white Gaussian noise feasible set, a compact spectral set with blockwise norm constraints that tightly captures the statistics and spatial uncorrelatedness of white Gaussian noise. This preconditioning reshapes each gradient update into a noise-aligned direction, driving faster and more effective reward ascent while preventing reward hacking. The projection is closed-form and matches the $\mathcal{O}(N \log N)$ complexity of FFT, adding negligible overhead in practice. In experiments on FLUX with four reward models, our approach reaches a comparable Aesthetic Score using only 30\% of the wall-clock time required by the state-of-the-art regularization-based method.

Modern AI image generators work by starting from random noise and transforming it into a realistic image. Once trained, these models can be steered toward better outputs by running an optimization process at generation time, pushing the result to score higher on metrics like aesthetic quality or text alignment. However, this optimization often breaks the images: it exploits weaknesses in the quality metric rather than genuinely improving the output, a problem known as reward hacking. It is also slow, sometimes taking minutes per image. We address both problems by constraining how each optimization step is taken. We ensure that every update stays within the space of valid random noise that the model was designed to handle. This simple constraint can be computed very efficiently, making the optimization both faster and more reliable, and producing high-quality, realistic images that genuinely score better on human preference metrics.