cs.LGOct 15, 2026
SaveNoise-Space Reward Alignment for Few-Step Generative Models
Organizations: KAIST · Stanford University · MongooseAI · Carnegie Mellon University ·
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.
Figures & tables
Figure 1: Motivation Figure. (a) Reward learning curve of NPO, direct fine-tuning, and trajectory-based noise policy. (b) Generated samples over training for NPO and direct fine-tuning. (c) Comparison with test-time scaling methods. As shown in the figure, NPO achieves a better reward-quality trade-off than baselines while maintaining an inference cost comparable to the pretrained model.
Figure 2: Overview of NPO. NPO trains a noise policy that maps Gaussian noise to optimized noise for a frozen generator . The policy is updated with , computed from the black-box reward of and forward-noised points of each , given by forward-noising pairs .
Figure 3: Reward learning curve of NPO and direct fine-tuning baselines on 1-step generation. NPO efficiently learns to generate high-rewarding samples compared to the baselines.
Figure 4: Reward-quality trade-off of NPO and direct fine-tuning baselines on SANA-Sprint with 4-step sampling. NPO achieves a better trade-off than the baselines.
| Method | Aesthetic | NFE | Time (s) | Rel. cost |
| SD-Turbo (Base) | 6.12 | 1 | 0.06 | |
| BoN-32 | 6.56 | 32 | 2.51 | |
| ReNO | 6.52 | 13 | 2.42 | |
| ZeNO | 6.50 | 482 | 30.78 | |
| NPO (Ours) | 6.52 | 1 | 0.15 |
Table 1: Comparison of NPO with test-time search methods on SD-Turbo. NPO achieves comparable aesthetic scores with a single generator evaluation.
Figure 5: Reward learning curve of NPO and noise-policy learning baselines on few-step generation. NPO efficiently learns to generate high-rewarding samples compared to the baselines.
Figure 6: Reward-quality trade-off of NPO against direct fine-tuning on the topic-steering task. NPO maintains perplexity and entropy comparable to those of the base model while achieving higher reward than the baselines.
Figure 7: Results on the topic-steering task with ELF. (a) NPO achieves comparable reward to test-time search methods with only four generator evaluations. (b) NPO learns to generate high-reward samples more efficiently than noise-policy learning baselines.
Figure 8: Comparison of NPO against baselines on biological structure generation. NPO achieves a better reward-quality trade-off and training efficiency than the baselines.
| Method | SS Composition | NFE | Time (s) | Rel. cost |
| SiD-Protein (Base) | 5 | 0.07 | ||
| BoN-256 | 1,280 | 17.71 | ||
| ZeNO | 2,410 | 22.72 | ||
| NPO (Ours) | 0.74 | 5 | 0.15 |
Table 2: Comparison of NPO with test-time search on biological structure generation tasks. NPO achieves the highest composition score with five evaluations.
Appendix figures & tables20 assets
Supplementary material from the paper’s appendix.
Appendix
| Domain-specific | |||
| Image | Language | Protein | |
| Backbone | U-Net | DiT | DiT |
| Channel widths / hidden size | |||
| Layers | per resolution | depth , heads | depth , heads |
| Weight decay | |||
| Shared across domains | |||
Table 3: Architecture and optimization settings of the noise policy.
| Generator | LoRA rank | LoRA parameters | Noise-policy parameters |
| SANA-Sprint | M | M | |
| SD-Turbo | M | M | |
| ELF | M | M | |
| Protein | M | M |
Table 4: LoRA rank of the direct baselines and the resulting number of trainable parameters.
| Method | Coefficient |
| SANA-Sprint , learning rate | |
| NPO (ours) | |
| Noise Flow-GRPO | |
| ODS | |
| Flow-GRPO | |
| FAV-B | |
Table 5: Swept ranges in the image domain. is the reference coefficient of NPO. is the KL coefficient of Noise Flow-GRPO and Flow-GRPO, the reward exponent of ODS, and the reward-gradient scale of FAV-B.
| Aesthetic | Compressibility | |||
| Method | Learning rate | Coefficient | Learning rate | Coefficient |
| NPO (ours) | ||||
| Noise Flow-GRPO | ||||
| ODS | ||||
| Flow-GRPO | ||||
| FAV-B | ||||
Table 6: Selected settings on SANA-Sprint.
| Aesthetic | Compressibility | |||
| Method | Learning rate | Coefficient | Learning rate | Coefficient |
| NPO (ours) | ||||
| Noise Flow-GRPO | ||||
| ODS | ||||
| Flow-GRPO (four steps) | ||||
| Flow-GRPO (one step) | ||||
Table 7: Selected settings on SD-Turbo.
| Method | Coefficient |
| All methods , learning rate | |
| NPO (ours) | |
| Noise Flow-GRPO | |
| ODS | |
| Flow-GRPO | |
| FAV-B | |
Table 8: Swept ranges on the topic-steering task. is the reference coefficient of NPO. is the KL coefficient of Noise Flow-GRPO and Flow-GRPO, the reward exponent of ODS, and the reward-gradient scale of FAV-B.
| Politics | Science | |||
| Method | Learning rate | Coefficient | Learning rate | Coefficient |
| NPO (ours) | ||||
| Noise Flow-GRPO | ||||
| ODS | ||||
| Flow-GRPO | ||||
| FAV-B | ||||
Table 9: Selected settings on the topic-steering task.
| Method | Coefficient |
| All methods , learning rate | |
| NPO (ours) | |
| Noise Flow-GRPO | |
| ODS | |
| Flow-GRPO | |
| FAV-B | |
Table 10: Swept ranges in the biological domain. is the reference coefficient of NPO. is the KL coefficient of Noise Flow-GRPO and Flow-GRPO, the reward exponent of ODS, and the reward-gradient scale of FAV-B.
| Method | Learning rate | Coefficient |
| NPO (ours) | ||
| Noise Flow-GRPO | ||
| ODS | ||
| Flow-GRPO | ||
| FAV-B |
Table 11: Selected settings at .
Figure 9: Reward-quality trade-off of NPO and direct fine-tuning baselines on SD-Turbo with four sampling steps. (a, b) Aesthetic score against ImageReward and LPIPS diversity. (c, d) Compressibility against ImageReward and LPIPS diversity.
Figure 10: Ablation of the ELBO weighting.
Figure 11: Ablation of the number of forward-noising pairs.
Figure 12: Sensitivity to the regularization strength.
Figure 13: Ablation of the interpolation.
Figure 14: Ablation of the policy-gradient objective.
Figure 15: Qualitative comparison under the aesthetic reward.
Figure 16: Qualitative comparison under the compressibility reward.
| Base | 1 | As everyone has seen in the NFL, Tyry Johnson didn’t appear, but he’s one of the best runers in the NFL. We’ll get to see from this. 22 DeJ McCrney,WR, WR, WR, Hardt : He’s been a guy. |
| 2 | If you’re one of the most adventured fans on the game, we’ll look at you a little to let you to see what Global Wring is up for. World Wring is one as one of the most engagingd games on the console that plays absolutely well about world order. | |
| NPO | 1 | As everything happens in a pond, modern scientists can’t exactly understand how energy is pumped by any thing that is moving through the air. But they might be able to do something from a piece of the lake into an atmosphere or through a piece inside … |
| 2 | This week, the nine month since Elson launched, we’ve talked with ToE and developers and producer Peter Chaleon to see what ORon is looking for. It’s one of the big features that we’ve seen into the Nintendo- game slate. | |
| Flow-GRPO | 1 | In an interview released late Sunday, House President Donald Trump didn’t say if he would place a tax on same-sex couples. His comments would be made public on July 5 at a seven-day music festival in Wilmington, Va. |
| 2 | ”We’re not in a final situation, and we’ven’t known what could happen next,” Kailke said. While House Republican Republicans are voting against the seven-point vote, the Senate’s Financeding Committee and the House Republican National Committee didn’ … | |
| FAV-B | 1 | Other people who live in the world who have’t known or not never known knew other who have often two good friends who liked both met each other other others who who talked to other on Twitter and after each other often many times and talked with othe … |
Figure 17: Generated text samples on the science topic. Red characters indicate a collapsed sample with no content left to quote.
| Base | 1 | As everyone has seen in the NFL, Tyry Johnson didn’t appear, but he’s one of the best runers in the NFL. We’ll get to see from this. 22 DeJ McCrney,WR, WR, WR, Hardt : He’s been a guy. |
| 2 | If you’re one of the most adventured fans on the game, we’ll look at you a little to let you to see what Global Wring is up for. World Wring is one as one of the most engagingd games on the console that plays absolutely well about world order. | |
| NPO | 1 | As it appeared earlier this month, House President Donald Trump didn’t express the reason he acted the right way about gun change in the country. The issue will be on display on TV, at a six-day music event in Detroit, Texas where he lost the 2016 el … |
| 2 | This week, the year-Julsith anniversary of Microsoft gathered with U&C President and CEO Mike Melbeck to discuss what OpRon is working on its three-year project, the company’s chief of development officer and vice president of computer design. | |
| Flow-GRPO | 1 | STTE: GETE: I don’t don’t believe it. He’s a little bit odd. he’s tricky. I’m attempting to get him get a bit cranky. He doesn’t dance. he doesn’t. He’s very anxious. He think he’s more delicate and he’s so intelligent at something. |
| 2 | Tese: I’m perseted [I’m really so intrigued]. What does that perspective Peter how would you get ask? Tese: I get. I’m one of the two. He’s a very sophisticated player. I he’s competent. he’s a very basic rate. He’s very remarkably sophisticated. | |
| FAV-B | 1 | ”All the happens in the batting is very well kept and working extremely and he’s a great guy that is efficient at the speed and I’ll be able to play and at great level in the defensive and is able to quickly as he handles the punch and that’s the how … |
Figure 18: Generated text samples on the politics topic. Red characters indicate a collapsed sample with no content left to quote.
Figure 19: Qualitative comparison on protein structure generation. Top: backbones generated by the pretrained generator. Bottom: backbones generated by NPO from the same initial noise. -helices are drawn in blue , -sheets in red , and coils in green .
More in this topic
RL for Generative Models
Reinforcement learning for generative models uses reward feedback to steer generated outputs toward goals, preferences, and constraints.
See all 154 papers