Few-step generative models improved for desired outputs with noise-space rewards
Noise-Space Reward Alignment for Few-Step Generative Models
Machine Learning
Summary
Generating high-quality content quickly is hard when you want the output to have certain qualities, especially if you can only evaluate it as a black box. The authors found that existing methods to fine-tune or control these quick generators either push them too far or lose their speed advantage. They developed a method called Noise Policy Optimization that learns how to adjust the starting noise input to the generator, guiding it to produce outputs that score better on the desired property without changing the generator itself. This approach keeps fast generation while improving quality based on rewards, and works well for images, text, and protein structures.
What this means in practice
- •For image generation teams: Generate images with specific desired features using few evaluation steps while maintaining quality and efficiency.
- •For protein design engineers: Speed up generation of protein structures that meet functional criteria using reward-guided noise adjustments.
Authors
Woocheol Shin, Taeyoung Yun, Jiwon Sung, Hyeongyu Kang, Jeongmin Park, Jaewoo Lee, Kyuil Sim, Dohyun Kim, Yutong He, Jinkyoo Park
Abstract
Few-step generative models produce high-quality samples in one or a few evaluations. However, steering them toward desired properties remains difficult, especially when only black-box reward evaluations are available. While one could apply approaches developed for iterative diffusion and flow-based models, such as direct RL fine-tuning or test-time scaling, we find that neither transfers well to the few-step regime: direct fine-tuning suffers from severe reward over-optimization, and test-time scaling sacrifices the efficiency advantage of few-step generation. We propose Noise Policy Optimization (NPO), which explicitly targets the alignment problem for few-step generative models. The key idea is to learn an amortized flow policy over the initial noise space of the generator. The trained policy transports Gaussian noise to a reward-tilted noise distribution, whose pushforward distribution through the frozen generator is exactly the reward-tilted output distribution. By keeping the generator frozen, we can effectively utilize the knowledge encoded in the pretrained model, yielding a better reward-quality trade-off than direct fine-tuning. The amortized policy moves the search cost into training, so inference adds only a lightweight noise sampling step on top of the pretrained model. To improve sample efficiency during training, we use ELBO-based likelihood estimation that requires only sampled endpoints rather than simulated trajectories. Across image, language, and protein structure generation, NPO consistently achieves a better reward-quality trade-off than direct fine-tuning, and matches test-time search methods with 16 to 549 times lower inference cost.