PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Yutong He, Jinkyoo Park

arXiv:2610.21125v1cs.LG

TL;DR

Few-step generators are difficult to align when rewards are black-box, while direct fine-tuning and per-sample test-time search have drawbacks. NPO learns an amortized noise-space flow policy around a frozen generator; across image, language, and protein tasks, it offers a better reward-quality trade-off and much lower inference cost than test-time search.

  • Problem

    Few-step generative models need alignment to desired properties, but downstream rewards may be scalar and non-differentiable, leaving reward-gradient-free alignment relatively underexplored.

    Source: §1 · p. 1

  • Method

    NPO freezes the few-step generator and learns an amortized flow policy that transports Gaussian input noise toward the reward-tilted noise distribution.

    Source: §1 · p. 2

  • Results

    Across image, language, and protein tasks, NPO improves the reward-quality trade-off over direct fine-tuning and matches or exceeds test-time noise optimization with 16 to 549× lower inference cost.

    Source: §1 · p. 2

  • Takeaways & Limitations

    NPO provides an effective way to adapt few-step generators while preserving sample quality and fast inference across image, language, and protein generation.

    Source: §6 · p. 9

    The noise policy still requires about 2.2-2.6× the inference cost of the frozen generator.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Few-step generators need alignment to black-box rewards, but direct fine-tuning can damage sample quality and per-sample noise search sacrifices fast inference. NPO instead amortizes noise-space alignment around a frozen generator, achieving favorable reward-quality trade-offs across tasks.

1 INTRODUCTION

  • Direct fine-tuning transfers poorly to few-step generators, either making limited progress or over-optimizing reward at the expense of sample quality.The authors hypothesize that changing generator parameters risks losing knowledge encoded during pretraining.Source: §1 · p. 1
  • Test-time noise search can improve reward without changing the pretrained generator, but must be repeated for every sample and incurs substantial inference cost.This sacrifices the efficiency advantage of few-step generation.Source: §1 · p. 2
  • NPO learns an amortized flow policy that transports Gaussian noise toward a reward-tilted input distribution while keeping the few-step generator frozen.ELBO-based likelihood estimation from sampled noise endpoints improves training sample efficiency relative to trajectory-based estimation.Source: §1 · p. 2
  • Across image, language, and protein tasks, NPO improves the reward-quality trade-off over direct fine-tuning and matches or exceeds test-time noise optimization with 16 to 549× lower inference cost.It also reaches the same reward levels as trajectory-based noise-policy training with fewer reward evaluations.Source: §1 · p. 2

2 RELATED WORK

Prior work reduces generation cost with few-step models, while reward alignment has used gradient-based, RL, variational, and noise-space approaches. NPO combines reusable noise-space optimization with black-box rewards across diverse few-step generators.

Few-step generative models

  • Few-step methods reduce sampling cost, but aligning these generators with downstream rewards remains underexplored.The literature includes consistency models, distillation, and flow-map approaches.Source: §2 · p. 2

Reward alignment of generative models

  • Reward alignment spans reward backpropagation, RL, reward-aware distillation, variational inference, and methods tailored to few-step generators.Some approaches modify pretrained generators and can lead to reward hacking, while others rely on model-specific structures; NPO keeps the generator frozen and applies across diverse few-step generators.Source: §2 · pp. 2–3

Initial-noise optimization

  • Noise-space alignment ranges from per-sample test-time optimization to amortized noise distributions and policies; NPO supports black-box rewards with an expressive flow policy and endpoint forward-noising updates.Test-time methods incur repeated inference-time search, while NPO's endpoint updates are designed to improve sample efficiency.Source: §2 · p. 3

3 PROBLEM SETUP & BACKGROUND

The setup defines alignment as steering a pretrained few-step generator toward high-reward outputs without reward gradients. It then motivates policy optimization and explains why trajectory-based likelihood estimation is difficult for deterministic samplers.

3 PROBLEM SETUP & BACKGROUND

  • Few-step alignment seeks high-reward outputs from a pretrained generator using scalar rewards, without assuming access to reward gradients.The generator maps Gaussian source noise to generated outputs, and its base output distribution is the pushforward of that source distribution.Source: §3 · p. 3, §3.1 · p. 3

3.1 PROBLEM SETUP AND REWARD-TILTED OUTPUT DISTRIBUTION

  • The KL-regularized objective balances expected reward against deviation from the pretrained output distribution, defining an optimal reward-tilted distribution that reweights toward high-reward outputs.Sampling from that distribution is the stated equivalent view of reward alignment.Source: §3.1 · p. 3

3.2 TRAJECTORY-BASED POLICY OPTIMIZATION

  • When reward gradients are unavailable, policy-gradient optimization uses a score-function identity and therefore requires estimating the trainable distribution's endpoint likelihood.Trajectory-based methods estimate likelihood from simulated reverse trajectories and intermediate samples.Source: §3.2 · p. 3
  • Trajectory-based likelihood estimation does not naturally fit deterministic ODE and flow-map samplers, and proposed surrogate methods can be expensive, unstable, or alter the sampling process.NPO's overview instead updates a noise policy using black-box reward and forward-noised points of optimized noise.Source: §3.2 · p. 4, Figure 2 · p. 4

4 NOISE POLICY OPTIMIZATION

NPO aligns a frozen few-step generator by learning a flow policy that transports Gaussian input noise toward the reward-tilted noise distribution. Its sections formalize the target, specify the policy and alignment objective, and describe endpoint-based optimization.

4 NOISE POLICY OPTIMIZATION

  • NPO learns an amortized flow policy over a frozen generator’s input noise to approximate the reward-tilted noise distribution.The method uses an ELBO-based objective to approximate that target in noise space.Source: §4 · p. 4

4.1 NOISE-SPACE REPRESENTATION OF THE TARGET

  • Changing the input-noise distribution is sufficient: its pushforward through the deterministic frozen generator equals the reward-tilted output distribution.This reduces sampling from the target output distribution to sampling from the corresponding target noise distribution.Source: §4.1 · p. 4

4.2 FLOW POLICIES AND FORWARD NOISING

  • A learned velocity field transports Gaussian source noise to policy endpoints, with forward-noising pairs used to train the flow.The procedure samples endpoints, evaluates their rewards, forms advantages, samples forward-noising pairs, and updates the policy.Source: §4.2 · pp. 4–5

4.3 NOISE-SPACE ALIGNMENT OBJECTIVE

  • The noise-space objective maximizes reward while regularizing the learned policy toward the initial noise distribution using a fixed reference velocity field.The reference field has the initial noise distribution as both its source and endpoint distributions.Source: §4.3 · p. 5

4.4 ELBO-BASED POLICY OPTIMIZATION

  • NPO estimates endpoint likelihoods with an ELBO-based estimator on sampled endpoints, avoiding simulated rollout trajectories during policy optimization.Training uses an old policy for endpoint sampling, a group-mean reward advantage, and a squared deviation from the reference field for KL regularization.Source: §4.4 · p. 5

5 EXPERIMENTS

Experiments evaluate NPO across image, language, and protein generation, comparing reward, sample quality, inference cost, and training efficiency. Across these tasks, NPO improves reward-quality trade-offs and reduces the cost of test-time search.

5.1.1 EXPERIMENTAL SETUP

  • The experiments span one- and few-step generators, with image alignment evaluated using aesthetic and JPEG-compressibility rewards and semantic-fidelity and diversity metrics.Image comparisons include direct fine-tuning, test-time noise optimization, and noise-policy learning baselines.Source: §5 · p. 5, §5.1 · p. 6, §5.1.1 · p. 6

5.1.2 RESULTS

  • NPO improves image rewards while better preserving semantic fidelity and diversity than direct fine-tuning, whose four-step reward gains can accompany substantial quality losses.On SANA-Sprint, the comparison is under four-step sampling; NPO also reaches comparable reward to test-time search at approximately 16–205× lower inference cost.Source: §5.1.2 · pp. 6–7, Figure 4 · p. 6

5.2.1 EXPERIMENTAL SETUP

  • Language experiments steer four-step ELF toward politics and science using an entailment-probability reward, while measuring fluency and diversity.The setup uses a DeBERTa-v3-large NLI model for reward and GPT-2 Large perplexity and unigram entropy for quality.Source: §5.2.1 · p. 7

5.2.2 RESULTS

  • NPO achieves higher topic rewards while keeping perplexity and unigram entropy close to the frozen generator, whereas direct fine-tuning degrades generation quality.The experiments use four-step ELF and assess fluency with GPT-2 Large perplexity and diversity with unigram entropy.Source: §5.2.1 · p. 7, §5.2.2 · p. 8, Figure 6 · p. 8
  • On ELF, NPO generates samples 84–549× faster than test-time search and achieves higher reward than noise-policy baselines under the same evaluation budget.The comparisons cover both topic-steering topics.Source: §5.2.2 · p. 8, Figure 7 · p. 8

5.3.1 EXPERIMENTAL SETUP

  • Protein experiments target diverse secondary-structure compositions with added emphasis on β-sheet content and assess plausibility through designability.Pairwise TM-score is also reported, with lower values indicating greater structural diversity.Source: §5.3.1 · pp. 8–9, §5.3 · p. 8

5.3.2 RESULTS

  • NPO substantially improves protein secondary-structure reward while retaining much higher designability than direct fine-tuning, which reduces Flow-GRPO designability to zero.The comparison uses a five-step backbone generator and evaluates designability, composition reward, and structural diversity.Source: §5.3.2 · p. 9, Figure 8 · p. 9, §5.3.1 · p. 9
  • For protein generation, NPO achieves higher target reward while sampling approximately 118–151× faster than test-time optimization and outperforming noise-policy baselines at equal reward-evaluation budgets.Table 2 reports the highest composition score with five evaluations.Source: §5.3.2 · p. 9, Table 2 · p. 9

6 CONCLUSION

NPO aligns few-step generators without reward gradients by freezing the pretrained generator and learning an amortized flow policy over its input noise. Experiments across image, language, and protein generation report preserved sample quality and fast inference, though policy inference still costs more than using the frozen generator.

6 CONCLUSION

  • NPO freezes the pretrained generator and learns an amortized flow policy over its input noise, using ELBO-based likelihood estimation instead of trajectory-based objectives.The paper studies alignment without access to reward gradients.Source: §6 · p. 9
  • Across image, language, and protein generation, NPO adapts few-step generators while preserving sample quality and fast inference.Source: §6 · p. 9
  • The noise policy still requires about 2.2-2.6× the inference cost of the frozen generator.The authors suggest distilling it into a one-step or few-step map to reduce this overhead.Source: §6 · p. 9

AI USE STATEMENT

The appendices state the conditions and distributional identities underlying NPO, then describe implementation choices for its reference flow, regularization, and likelihood estimation. The paper also specifies that LLMs were used only for manuscript polishing and limited code assistance.

AI USE STATEMENT

  • LLMs were used only for minor manuscript polishing and limited code debugging or refactoring, not for research formulation, core methods, or experiment design.Source: AI USE STATEMENT · p. 10

A.1 SOLUTION TO KL-REGULARIZED REWARD MAXIMIZATION

  • Under finite-normalizer and finite-expectation assumptions, KL-regularized reward maximization is minimized by the reward-tilted distribution.The appendix applies this result to both output and noise spaces.Source: §A.1 · p. 15

A.2 PROOF OF THE NOISE-SPACE PUSHFORWARD IDENTITY

  • For a deterministic generator, the reward-tilted noise distribution pushes forward to the reward-tilted output distribution.Source: §A.2 · p. 15

B IMPLEMENTATION DETAILS OF NPO

  • The implementation appendix covers the reference flow, reference regularization, ELBO-based likelihood estimation, and noise-policy architecture and training.Source: B IMPLEMENTATION DETAILS OF NPO · p. 16

B.1 REFERENCE FLOW

  • The reference flow interpolates between Gaussian noise distributions and uses a linear interpolation schedule in the experiments.A variance-preserving interpolation is compared in Appendix G.2.Source: §B.1 · p. 16

B.2 REFERENCE REGULARIZATION

  • Reference regularization approximates the KL divergence as a weighted policy–reference velocity difference, using uniform weighting and an experiment-specific coefficient.Source: §B.2 · p. 16

B.3 ELBO WEIGHTING AND ESTIMATION

  • The ELBO likelihood estimator uses a velocity-regression score with path-KL, constant, or adaptive time weighting; all experiments use adaptive weighting.The adaptive weighting self-normalizes the objective for numerical robustness.Source: §B.3 · p. 16
  • The Monte Carlo expectation for the ELBO estimator uses independently sampled forward-noising pairs per endpoint, with M = 3 in all experiments.Source: §B.3 · p. 16

B.4 NOISE POLICY ARCHITECTURE AND TRAINING

The noise policy uses domain-specific architectures, with settings and optimization details collected in Table 3. Image policies use a conditional 2D U-Net, while language and protein policies use a DiT over input positions.

B.4 NOISE POLICY ARCHITECTURE AND TRAINING

  • Image noise policies use a conditional 2D U-Net conditioned on the generator text encoder's prompt embedding through cross-attention.Source: §B.4 · p. 16
  • Language and protein policies use a DiT whose tokens are positions in the one-dimensional generator input sequence.Flow time enters both architectures through a sinusoidal embedding.Source: §B.4 · p. 17
  • Table 3 lists the noise policy's architecture and optimization settings.Source: Table 3 · p. 17

C IMPLEMENTATION OF THE BASELINES

Baseline implementations share generators, rewards, and evaluation protocols within each domain, with parameter-matched LoRA settings for direct fine-tuning comparisons. The section also specifies noise-policy and per-sample test-time search procedures.

C IMPLEMENTATION OF THE BASELINES

  • Within each domain, methods share the generator, reward, and evaluation protocol; LoRA ranks are chosen to match the noise policy's trainable parameter count.Table 4 reports the LoRA ranks and resulting parameter counts.Source: C IMPLEMENTATION OF THE BASELINES · p. 18, Table 4 · p. 18
  • Noise Flow-GRPO trains a noise policy with a GRPO objective, while Flow-GRPO fine-tunes generator attention layers with LoRA.Their training and evaluation sampling procedures differ, including added noise for Flow-GRPO likelihood estimation during training.Source: C IMPLEMENTATION OF THE BASELINES · p. 18
  • FAV-B estimates reward gradients with antithetic evolution strategies and regresses LoRA parameters onto particles moved by the resulting velocity.Its reward-call cost is K(2N + 1) per condition.Source: C IMPLEMENTATION OF THE BASELINES · p. 18
  • ODS trains toward a reward-weighted noise target using relative trajectory balance and 25-step Euler–Maruyama sampling.Its training samples mix prior, replay-buffer, and current-policy trajectories.Source: C IMPLEMENTATION OF THE BASELINES · p. 18
  • Best-of-N, ReNO, and ZeNO search each sample's initial noise at inference while leaving the generator unchanged.Best-of-N selects the highest-reward noise; ReNO uses reward gradients, and ZeNO uses a zeroth-order estimate.Source: C IMPLEMENTATION OF THE BASELINES · pp. 18–19

D.1 GENERATOR, REWARDS AND EVALUATION METRICS

The image experiments use frozen few-step generators with aesthetic and JPEG-compressibility rewards, while the language experiments steer an unconditional generator toward politics or science. The sections specify their evaluation metrics, selected hyperparameters, and sampling protocols.

D.1 GENERATOR, REWARDS AND EVALUATION METRICS

  • SANA-Sprint and SD-Turbo are evaluated at one and four steps with aesthetic and JPEG-compressibility rewards, using ImageReward and LPIPS to assess quality and diversity.SANA-Sprint generates 1024 × 1024 images; SD-Turbo generates 512×512 images.Source: §D.1 · p. 20

D.2 HYPERPARAMETERS

  • For each image method, the learning rate and coefficient are swept, then a setting with stable target-reward increases across seeds is selected.The swept ranges and selected settings are listed in Tables 5–7.Source: §D.2 · p. 20, Table 5 · p. 20, Table 6 · p. 20

D.3 EVALUATION PROTOCOL

  • Image evaluation uses 45 simple-animal prompts, 16 samples per prompt, and five seeds; target reward, ImageReward, and within-prompt LPIPS are averaged.LPIPS is computed across all sample pairs sharing a prompt.Source: §D.3 · p. 21

E.1 GENERATOR, REWARDS AND EVALUATION METRICS

  • The language experiments use four-step ELF with a deterministic probability-flow ODE sampler and DeBERTa entailment probability as the topic reward; fluency and diversity use GPT-2 perplexity and unigram entropy.The topics are politics and science; sampler times are placed at quantiles of ELF’s logit-normal distribution.Source: §E.1 · p. 22

E.2 HYPERPARAMETERS

  • For each language method, learning rates and coefficients are swept, and the setting with stable target-reward increases across seeds is selected.Tables 8 and 9 report swept ranges and selected settings.Source: §E.2 · p. 22, Table 8 · p. 22, Table 9 · p. 22

E.3 EVALUATION PROTOCOL

The protein experiments evaluate a deterministic five-step backbone generator using secondary-structure diversity as its target reward, alongside designability and structural diversity metrics. The SD-Turbo comparison finds that NPO preserves quality and diversity better than direct fine-tuning at comparable reward.

F.1 GENERATOR, REWARD AND EVALUATION METRICS

  • Protein evaluation uses a five-step generator for 150-residue backbones and rewards diversity of α-helix, β-sheet, and coil composition, assigned with P-SEA.Designability is the fraction with minimum refolded-structure Cα RMSD below 2 ˚A; diversity is mean pairwise TM-score, with lower values indicating higher diversity.Source: §F.1 · p. 23

F.2 HYPERPARAMETERS

  • Protein-method learning rates and coefficients are swept, then settings with stable target-reward increases across seeds are selected.Tables 10 and 11 give swept ranges and selected settings at L = 150.Source: §F.2 · p. 23, Table 10 · p. 23, Table 11 · p. 23

F.3 EVALUATION PROTOCOL

  • Protein evaluation scores 256 backbones for target reward and the first 32 for designability and pairwise TM-score, using five seeds.Means and standard deviations are reported across seeds.Source: §F.3 · p. 23

G.1 COMPARISON WITH DIRECT FINE-TUNING ON SD-TURBO

  • NPO reaches a comparable target reward to Flow-GRPO while keeping LPIPS diversity at or above the frozen generator and ImageReward above Flow-GRPO.Flow-GRPO reaches the highest target reward, but ImageReward and LPIPS diversity decrease substantially, most severely under compressibility; FAV-B barely improves either target reward.Source: §G.1 · p. 24, Figure 9 · p. 24

G.2 ABLATION OF THE NPO OBJECTIVE

The NPO objective is robust to most tested design choices, with path-KL weighting the notable exception. Qualitative examples show changes in image, text, and protein outputs under reward optimization, including more β-sheet content in NPO-generated protein backbones.

G.2 ABLATION OF THE NPO OBJECTIVE

  • Except for path-KL weighting, NPO performs similarly across the tested objective variants, supporting robustness to most design choices.The ablations use SANA-Sprint with one sampling step under the aesthetic reward, changing one choice at a time and using three seeds.Source: §G.2 · p. 24
  • Adaptive and simple ELBO weightings perform similarly, while path-KL barely improves reward; beyond three forward-noising pairs, there is no gain, and regularization and interpolation variants perform similarly.The study therefore uses adaptive weighting, M = 3, and linear interpolation; adding GRPO normalization or ratio clipping does not change results, so the simpler update is retained.Source: §G.2 · pp. 24–26

H.1 IMAGE DOMAIN

H.2 LANGUAGE DOMAIN

  • During topic steering, direct-fine-tuning samples collapse into repetitions of a few tokens as optimization proceeds, matching decreases in perplexity and unigram entropy.Figures 17 and 18 show samples for two initial noises at three training checkpoints.Source: §H.2 · p. 28, Figure 17 · p. 28, Figure 18 · p. 29

H.3 PROTEIN DOMAIN

  • For the same initial noise, NPO-generated protein backbones contain substantially more β-sheet content than pretrained-generator backbones while remaining well-formed.The reward favors diverse secondary-structure compositions with additional emphasis on β-sheet content.Source: §H.3 · p. 28, Figure 19 · p. 29
Loading 2610.21125v1…