Noise-Space Reward Alignment for Few-Step Generative Models

Woocheol Shin; Taeyoung Yun; Jiwon Sung; Hyeongyu Kang; Jeongmin Park; Jaewoo Lee; Kyuil Sim; Dohyun Kim

cs.LGOct 15, 2026

Noise-Space Reward Alignment for Few-Step Generative Models

Authors: Woocheol Shin, Taeyoung Yun, Jiwon Sung, Hyeongyu Kang, Jeongmin Park, Jaewoo Lee, Kyuil Sim, Dohyun Kim,

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 $q_{\psi}$ that maps Gaussian noise $z_{1}$ to optimized noise $z_{0}$ for a frozen generator $g_{\theta}$ . The policy is updated with $\mathcal{L}_{\mathrm{NPO}}$ , computed from the black-box reward of $g_{\theta}(z_{0})$ and $M$ forward-noised points $z_{t}$ of each $z_{0}$ , given by forward-noising pairs $(t,\epsilon)$ .

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	$1.0\times$
BoN-32	6.56	32	2.51	$44.2\times$
ReNO	6.52	13	2.42	$42.6\times$
ZeNO	6.50	482	30.78	$542.8\times$
NPO (Ours)	6.52	1	0.15	$2.6\times$

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.

Panel of Figure 7 — 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)	$-0.84$	5	0.07	$1.0\times$
BoN-256	$-0.28$	1,280	17.71	$259.6\times$
ZeNO	$-0.72$	2,410	22.72	$333.1\times$
NPO (Ours)	0.74	5	0.15	$2.2\times$

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	$(48,96,192)$	$256$	$256$
Layers	$2$ per resolution	depth $6$ , $4$ heads	depth $6$ , $4$ heads
Weight decay	$10^{-4}$	$10^{-4}$	$0$
Shared across domains

Appendix

Table 3: Architecture and optimization settings of the noise policy.

Generator	LoRA rank $r$	LoRA parameters	Noise-policy parameters
SANA-Sprint	$28$	$14.45$ M	$14.43$ M
SD-Turbo	$68$	$14.11$ M	$14.16$ M
ELF	$137$	$7.58$ M	$7.58$ M
Protein	$106$	$7.33$ M	$7.32$ M

Appendix

Table 4: LoRA rank of the direct baselines and the resulting number of trainable parameters.

Method	Coefficient
SANA-Sprint , learning rate $\in\{3\text{e-}5,1\text{e-}4,3\text{e-}4,1\text{e-}3,3\text{e-}3\}$
NPO (ours)	$\lambda_{\mathrm{ref}}\in\{0,3\text{e-}5,1\text{e-}4,3\text{e-}4,1\text{e-}3\}$
Noise Flow-GRPO	$\beta\in\{0,4\text{e-}3,1\text{e-}2,4\text{e-}2,1\text{e-}1\}$
ODS	$\beta\in\{50,100,500,1000,5000\}$
Flow-GRPO	$\beta\in\{0,1\text{e-}3,1\text{e-}2,1\text{e-}1,1\}$
FAV-B	$\beta\in\{1\text{e-}3,1\text{e-}2,1\text{e-}1,1\}$

Appendix

Table 5: Swept ranges in the image domain.

\lambda_{\mathrm{ref}}

is the reference coefficient of NPO.

\beta

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)	$3\text{e-}5$	$\lambda_{\mathrm{ref}}=1\text{e-}4$	$3\text{e-}5$	$\lambda_{\mathrm{ref}}=3\text{e-}4$
Noise Flow-GRPO	$3\text{e-}5$	$\beta=1\text{e-}2$	$3\text{e-}5$	$\beta=1\text{e-}1$
ODS	$3\text{e-}5$	$\beta=500$	$3\text{e-}5$	$\beta=500$
Flow-GRPO	$3\text{e-}4$	$\beta=1\text{e-}3$	$1\text{e-}3$	$\beta=1\text{e-}3$
FAV-B	$3\text{e-}4$	$\beta=1$	$1\text{e-}3$	$\beta=1$

Appendix

Table 6: Selected settings on SANA-Sprint.

	Aesthetic		Compressibility
Method	Learning rate	Coefficient	Learning rate	Coefficient
NPO (ours)	$3\text{e-}5$	$\lambda_{\mathrm{ref}}=3\text{e-}4$	$1\text{e-}5$	$\lambda_{\mathrm{ref}}=1\text{e-}4$
Noise Flow-GRPO	$3\text{e-}5$	$\beta=4\text{e-}2$	$1\text{e-}5$	$\beta=1\text{e-}1$
ODS	$3\text{e-}5$	$\beta=500$	$1\text{e-}5$	$\beta=500$
Flow-GRPO (four steps)	$3\text{e-}4$	$\beta=1\text{e-}2$	$1\text{e-}4$	$\beta=1\text{e-}3$
Flow-GRPO (one step)	$3\text{e-}4$	$\beta=1\text{e-}2$	$1\text{e-}4$	$\beta=1\text{e-}3$

Appendix

Table 7: Selected settings on SD-Turbo.

Method	Coefficient
All methods , learning rate $\in\{1\text{e-}4,3\text{e-}4,1\text{e-}3,3\text{e-}3\}$
NPO (ours)	$\lambda_{\mathrm{ref}}\in\{0,3\text{e-}5,1\text{e-}4,3\text{e-}4,1\text{e-}3\}$
Noise Flow-GRPO	$\beta\in\{0,4\text{e-}3,1\text{e-}2,4\text{e-}2,1\text{e-}1\}$
ODS	$\beta\in\{50,100,500,1000,5000\}$
Flow-GRPO	$\beta\in\{0,1\text{e-}3,1\text{e-}2,1\text{e-}1,1\}$
FAV-B	$\beta\in\{1\text{e-}3,1\text{e-}2,1\text{e-}1,1\}$

Appendix

Table 8: Swept ranges on the topic-steering task.

\lambda_{\mathrm{ref}}

is the reference coefficient of NPO.

\beta

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)	$1\text{e-}4$	$\lambda_{\mathrm{ref}}=3\text{e-}4$	$1\text{e-}4$	$\lambda_{\mathrm{ref}}=3\text{e-}4$
Noise Flow-GRPO	$1\text{e-}4$	$\beta=1\text{e-}2$	$1\text{e-}4$	$\beta=1\text{e-}2$
ODS	$1\text{e-}4$	$\beta=50$	$1\text{e-}4$	$\beta=50$
Flow-GRPO	$1\text{e-}3$	$\beta=1\text{e-}3$	$1\text{e-}3$	$\beta=1\text{e-}2$
FAV-B	$1\text{e-}3$	$\beta=1$	$1\text{e-}3$	$\beta=1$

Appendix

Table 9: Selected settings on the topic-steering task.

Method	Coefficient
All methods , learning rate $\in\{1\text{e-}4,3\text{e-}4,1\text{e-}3,3\text{e-}3\}$
NPO (ours)	$\lambda_{\mathrm{ref}}\in\{3\text{e-}5,1\text{e-}4,3\text{e-}4,1\text{e-}3\}$
Noise Flow-GRPO	$\beta\in\{1\text{e-}2,4\text{e-}2,1\text{e-}1,4\text{e-}1\}$
ODS	$\beta\in\{100,300,500,1000\}$
Flow-GRPO	$\beta\in\{4\text{e-}4,1\text{e-}3,4\text{e-}3,1\text{e-}2\}$
FAV-B	$\beta\in\{1\text{e-}4,3\text{e-}4,1\text{e-}3,3\text{e-}3,1\text{e-}2\}$

Appendix

Table 10: Swept ranges in the biological domain.

\lambda_{\mathrm{ref}}

is the reference coefficient of NPO.

\beta

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)	$3\text{e-}4$	$\lambda_{\mathrm{ref}}=1\text{e-}4$
Noise Flow-GRPO	$3\text{e-}4$	$\beta=1\text{e-}2$
ODS	$3\text{e-}4$	$\beta=500$
Flow-GRPO	$1\text{e-}4$	$\beta=4\text{e-}4$
FAV-B	$1\text{e-}4$	$\beta=3\text{e-}4$

Appendix

Table 11: Selected settings at

L=150

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 …

Appendix

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 …

Appendix

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. $\alpha$ -helices are drawn in blue , $\beta$ -sheets in red , and coils in green .

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

May 12, 2026

Gradient-Free Noise Optimization for Reward Alignment in Generative Models

Existing reward alignment methods for diffusion and flow models rely on multi-step stochastic trajectories, making them difficult to extend to deterministic generators. A natural alternative is noise-space optimization, but existing approaches require backpropagation through the generator and reward pipeline, limiting applicability to differentiable settings. To address this, here we present ZeNO (Zeroth-order Noise Optimization), a gradient-free framework that formulates noise optimization as a path-integral control problem, estimable from zeroth-order reward evaluations. When instantiated with an Ornstein--Uhlenbeck reference process and a one-step receding horizon, the repeated update connects to Langevin dynamics targeting a reward-tilted distribution. ZeNO enables effective inference-time scaling and demonstrates strong performance across diverse generators and reward functions, including a protein structure generation task where backpropagation is infeasible.

Jeongsol Kim, Hongeun Kim, Jian Wang +1

KAIST AI · Snap Inc.

Diffusion Model Alignment Stochastic Optimal Control Zeroth-Order Optimization Generative Modeling Derivative-Free Optimization Gradient-Free Optimization

Ask PDF

Sep 28, 2026

FestDPO: Few-step Generator Alignment with Direct Preference Optimization

Few-step generative models can generate high-fidelity samples within a few function evaluations. Despite this efficiency, generated samples may not exhibit desirable properties. When these properties are difficult to encode as an explicit reward function, direct preference optimization (DPO) can align generative models using pairwise preference feedback without training a separate reward model. However, extending DPO to few-step generative models is challenging because few-step generative models are generally implicit, making the likelihood evaluation required by DPO intractable. To address this challenge, we introduce Few-step DPO (FestDPO), an extension of DPO for few-step generative models that leverages nonparametric likelihood estimation from empirical samples. By exploiting the fast sampling capabilities of few-step generative models, our approach makes sample-based approximation of DPO loss computationally feasible. Furthermore, the sample-based formulation makes FestDPO agnostic to the model family and sampling procedure. Our toy experiment demonstrates that FestDPO matches the reward-tilted target distribution across four few-step generators. For real-world tasks, we evaluate FestDPO in two domains: text-to-image generation and protein backbone generation. In text-to-image generation, FestDPO outperforms preference optimization baselines in both win rates against the base models and human evaluation scores. In protein backbone generation, it achieves a higher

β

-sheet fraction and better structural designability than the baselines.

Jaewoo Lee, Kyuil Sim, Hyeongyu Kang +3

KAIST · MongooseAI · Omelet

Direct Preference Optimization T2I Generation

Ask PDF

May 26, 2026

Aligning Few-Step Generative Models by Amortizing Sample-based Variational Inference

Aligning a few-step generative model is challenging, since existing alignment frameworks typically rely on restrictive assumptions: a tractable likelihood, a specific ODE/SDE solver, or a particular model family. We introduce FAV, Few-step Generative Models Alignment via Sample-based Variational Inference, a general alignment framework that requires only sample access to the generator and the reference distribution. We cast alignment as sampling from a reward-tilted distribution anchored to a reference distribution. We leverage Stein Variational Gradient Descent as a sample-based variational inference scheme and amortize its particle updates into the generator parameters via fixed-point regression. We evaluate FAV on two domains: robotics manipulation and image generator alignment. On generative policy alignment for robotic manipulation, FAV outperforms prevailing policy extraction baselines across 56 offline and 30 offline-to-online RL tasks. For image generator alignment, FAV fine-tunes diverse few-step backbones, including GAN, drifting model, consistency models, and flow maps, scaling from ImageNet-

256

to 1024

^2

text-to-image synthesis. Code is available at https://github.com/Jaewoopudding/FAV.

Jaewoo Lee, Hyeongyu Kang, Dohyun Kim +9

1KAIST · 2MongooseAI · 3Mila – Quebec AI Institute +3

Stein Variational Gradient Descent

Ask PDF

Abstract

RL for Generative Models

Explore similar work

Gradient-Free Noise Optimization for Reward Alignment in Generative Models

FestDPO: Few-step Generator Alignment with Direct Preference Optimization

Aligning Few-Step Generative Models by Amortizing Sample-based Variational Inference