PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimizing Visual Generative Models via Distribution-wise Rewards

Ruihang Li, Mengde Xu, Shuyang Gu, Leigang Qu, Fuli Feng, Han Hu, Wenjie Wang

arXiv:2607.02291v1cs.LGcs.CV

TL;DR

Sample-wise rewards in visual-generation reinforcement learning are vulnerable to reward hacking, artifacts, and reduced diversity. This paper introduces distribution-wise rewards with subset replacement and post-hoc model merging, improving FID-50K from 8.30 to 5.77 for SiT and from 3.74 to 3.52 for EDM2.

  • Problem

    Sample-wise rewards in visual-generation reinforcement learning are susceptible to reward hacking that introduces artifacts and reduces generated-image diversity.

  • Method

    The framework uses distribution-wise rewards computed through subset replacement and optimizes post-hoc model-merging coefficients to address SDE-based training and ODE-based inference inconsistency.

  • Results

    FID-50K improves across SiT and EDM2, decreasing from 8.30 to 5.77 for SiT and from 3.74 to 3.52 for EDM2-XS.

  • Takeaways & Limitations

    The method enhances distributional fidelity and perceptual quality while preserving sample diversity across the evaluated generative-model settings.

  • Takeaways & Limitations

    Experiments focus on class-conditional ImageNet generation, and extending subset replacement to open-vocabulary text-to-image settings requires further exploration.

Abstract

from arXiv · show

Conventional reinforcement learning strategies for visual generation typically employ sample-wise reward functions, yet this practice frequently results in reward hacking that degrades image diversity and introduces visual anomalies. To address these limitations, we present a novel framework that finetunes generative models using distribution-wise rewards, ensuring better alignment with real-world data distributions. Unlike rewards that evaluate samples individually, distribution-wise reward accounts for the data distribution of the samples, mitigating the mode collapse problem that occurs when all samples optimize towards the same direction independently. To overcome the prohibitive computational cost of estimating these rewards, we introduce a subset-replace strategy that efficiently provides reward signals by updating only a small subset of a generated reference set. Additionally, we apply RL to optimize post-hoc model merging coefficients, potentially mitigating the train-inference inconsistency caused by introducing stochastic differential equation (SDE) in regular RL practices. Extensive experiments show our approach significantly improves FID-50K across various base models, from 8.30 to 5.77 for SiT and from 3.74 to 3.52 for EDM2. Qualitative evaluation also confirms that our method enhances perceptual quality while preserving sample diversity.

1. Introduction

The paper identifies reward hacking in sample-wise reinforcement learning as a cause of artifacts, reduced diversity, and degraded distributional fidelity. It proposes distribution-wise rewards with a subset-replace strategy and post-hoc model-merging optimization to improve visual generation.

  • Proposed approach: Distribution-wise rewards aim to improve coverage of the real-world data distribution while achieving high visual fidelity and broad generation diversity.The framework uses FID as its distribution-wise metric and reinforcement fine-tuning reward signal.
  • Proposed approach: The subset-replace strategy generates a reference set, replaces a small subset with rollout samples, and uses the updated set’s FID as a dense, low-cost reward.The reference set is class-balanced, and the replaced subset is 0.01× of the reference set.
  • Motivation: Sample-wise reinforcement fine-tuning is susceptible to reward hacking, degrading distributional fidelity, introducing artifacts, and reducing sample diversity.The paper analyzes these limitations as a motivation for distribution-wise rewards.
  • Results: 8.30 to 5.77: the framework reduces SiT’s FID-50K without additional training data or architectural modifications.The same training recipe reduces the FDDINOv2 score from 230.39 to 164.88.
  • Results: 3.74 to 3.52: post-hoc optimization of EDM2 model-merging coefficients improves FID-50K while addressing train-inference inconsistency in SDE-based RL.The optimization uses distribution-wise reward signals with an ODE-based denoising procedure.

2. Related Work

Prior work applies reinforcement learning to image generation for preference alignment, while distribution-wise metrics provide alternative ways to compare generated and target distributions. Model averaging has also become widely adopted in image synthesis pre-training, with related exploration in large language models.

  • Reinforcement Learning in Image Generation: Reinforcement learning was adapted to diffusion models through policy gradients on the score function, enabling preference-aligned image generation.Later work introduced offline Direct Preference Optimization for text-to-image tasks and online methods with step-aware preference models to address pairwise-data distribution shift.
  • Reinforcement Learning in Image Generation: Offline Direct Preference Optimization was extended to text-to-image tasks, while distributional shift in pairwise data motivated online, step-aware preference models.The passage also notes more recent GRPO-based approaches, though the supplied text is truncated.
  • Distribution-wise Metrics: KL Divergence measures differences between distributions but can be unstable when one distribution assigns zero probability where the other has non-zero probability.Maximum Mean Discrepancy compares distribution means in a Reproducing Kernel Hilbert Space, but can struggle with high-dimensional data and outliers.
  • Model Merging: Model averaging has become widely adopted in the pre-training of state-of-the-art image synthesis models, with similar exploration in large language models.The supplied passage lists numerous studies on model averaging before truncating the discussion of language-model work.

3. Method

The method casts iterative denoising as an MDP and introduces subset-replace distribution-wise rewards to make reinforcement fine-tuning computationally feasible. It further uses RL-optimized model merging coefficients to address train-inference inconsistency from stochastic denoising while enabling efficient ODE sampling.

  • Denoising as a MDP: Denoising is formulated as an MDP whose state contains the class, timestep, and current sample, with model predictions as actions and deterministic transitions.The policy is π(at | st) ≜ pθ(xt−1 | xt, c).
  • Distribution-wise Reward: Direct distribution-wise reward computation is infeasible at scale, motivating subset replacement as an efficient approximation for reinforcement fine-tuning.Existing approaches generally use sample-wise rewards for individual denoising trajectories, whereas distribution-wise rewards require many samples.
  • Distribution-wise Reward: Subset-replace rewards construct a class-balanced reference set of N generated images, replace n same-class images with new samples, and use the negative replaced FID as the signal.The partially updated set is (G \ g) ∪ g′, whose FID measures the distribution-wise effect of the replacement.
  • Reinforcement Fine-tuning: The policy maximizes expected cumulative reward with KL regularization, using a lightweight value-function-free advantage estimator and batch-level normalization.The KL-divergence from a reference policy is scaled by β, and experiments found batch-level normalization superior to group-level normalization in this setting.
  • RL-based Model Merging: Because coefficient sampling supplies RL stochasticity, the method avoids additional denoising randomness and uses efficient ODE sampling to reduce train-inference inconsistency.The inconsistency arises because prior diffusion RL methods rely on SDE stochasticity that does not robustly transfer to inference.
  • RL-based Model Merging: RL optimizes coefficients for merging sequential checkpoints into one model, sampling coefficient vectors with an MLP policy and evaluating them through the subset-replace reward.The merged model assigns weighting coefficients wi to Nc checkpoints, while EMANet generates their means and standard deviations.

4. Experiments

Experiments show that subset-replace rewards and reinforcement learning improve visual generation quality while preserving an efficient training setup. Ablations identify effective normalization, replacement, sample-selection, and model-merging choices, while exposing train-inference inconsistency from SDE rollouts.

  • Subset-replace results: FID-50K reaches 5.77 and FDDINOv2 reaches 164.88 after adding reinforcement learning to SiT-XL with limited additional training.The result substantially improves the model’s ability to represent the image distribution.
  • Model merging: FID decreases from 3.74 to 3.52 on EDM2-XS and from 2.57 to 2.52 on EDM2-S by reinforcement-learning model-merging coefficients.The method improves pretrained models without complex SDE solvers or denoising reduction.
  • Ablations: Batch-level advantage normalization converges faster than group-level normalization with either all rollout samples or the top 25%.The final experiments therefore use batch-level normalization.
  • Ablations: A replacement subset of 50 images achieves optimal generation quality with the lowest computational overhead, using a reference set of 5,000.Reference-set sizes show a non-monotonic performance trend because small sets lack representativeness and large sets add noise.
  • Ablations: Global top 25% sample selection performs best, while lower-quality samples slow convergence and local process rankings are inferior.Pure RL also outperforms RS-then-RL, whose additional training does not improve performance after likely overfitting during reject sampling.

5. Conclusion

The framework replaces sample-wise rewards with distribution-wise rewards enabled by an efficient subset-replace strategy, addressing reward hacking and reduced diversity in visual-generation RL. It improves reported metrics through direct fine-tuning and post-hoc model-merging optimization while resolving SDE-base train-inference inconsistencies.

  • 5. Conclusion: The proposed framework uses distribution-wise rewards and an efficient subset-replace strategy to address reward hacking and reduced diversity from sample-wise rewards.The subset-replace strategy enables the distribution-wise reward signals.
  • 5. Conclusion: 8.30 to 5.77: direct fine-tuning substantially improves SiT’s FID-50K score.The reported FID-50K score changes from 8.30 to 5.77.
  • 5. Conclusion: 230.39 to 164.88: direct fine-tuning improves SiT’s FDDINOv2 score.The reported FDDINOv2 score changes from 230.39 to 164.88.
  • 5. Conclusion: 3.74 to 3.52 and 2.57 to 2.52: post-hoc model-merging optimization reduces EDM2-XS and EDM2-S FID, respectively.The reductions are reported for EDM2-XS and EDM2-S, respectively.
  • 5. Conclusion: Post-hoc model-merging optimization resolves train-inference inconsistencies in SDE-base.This result is reported alongside the FID reductions for EDM2-XS and EDM2-S.

Impact Statement · A. Hyperparameter Details

The paper targets improved alignment and fidelity in visual generative models through distribution-wise reinforcement learning, while specifying Adam-based fine-tuning and policy-gradient hyperparameters. Its impact emphasis is preserving generation diversity and mode coverage while mitigating reward hacking.

  • Impact Statement: The work aims to improve the alignment and fidelity of visual generative models.
  • Impact Statement: Distribution-wise reinforcement learning is presented as a way to mitigate reward hacking associated with sample-wise reinforcement learning.
  • Impact Statement: The approach addresses the tendency of sample-wise reinforcement learning to diminish generation diversity.
  • Impact Statement: Maintaining high diversity and mode coverage is identified as important for the societal impact of generative models.
  • A. Hyperparameter Details: The model uses Adam with β1 = 0.9, β2 = 0.999, no weight decay, and a constant learning rate of 1 × 10−5.
  • A. Hyperparameter Details: Policy-gradient rollouts use global batch size 128, KL-divergence regularization scaler β = 0, and one policy-network update per rollout step.

A.1. More Ablation Studies

The ablations identify 10 training steps as the best reference-set refresh interval, balancing reward representativeness against computational overhead. Continued RL training can instead cause performance deterioration through adaptation bias toward the training denoising schedule rather than general overfitting.

  • Reference Set Refresh Interval: Large refresh intervals make the reference set stale and reduce reward representativeness, whereas small intervals add unnecessary computational overhead.The subset-replace strategy periodically regenerates the reference set using the current model after a fixed number of training steps.
  • Reference Set Refresh Interval: 10 training steps achieves the best FID-5K score among refresh intervals of 5, 10, and 20 steps.This interval provides a balance between reward representativeness and computational overhead.
  • Adaptation bias toward the training denoising schedule: After reaching optimal performance, continued RL training gradually deteriorates performance because of adaptation bias toward the training denoising schedule.The experiments suggest this effect is not caused by general overfitting and occurs under the denoising reduction paradigm.

B. Cross-Metric Evaluation

Cross-metric evaluation shows that distribution-wise reward training improves performance beyond FID and Inception-v3, supporting genuine distributional improvement. The evaluation uses independent metrics, including FDDINOv2 with DINOv2 features, on the fine-tuned SiT model.

  • Cross-metric evaluation: All evaluated metrics improve for the fine-tuned SiT model at 450 training steps, indicating genuine distributional improvement beyond FID.The evaluation uses a comprehensive set of independent metrics on the same fine-tuned model.
  • Cross-metric evaluation: FDDINOv2 uses DINOv2 features rather than Inception-v3, providing evidence that improvements are not specific to the Inception-v3 feature space.KID and MMD use Inception-v3 features, while FDDINOv2 uses DINOv2 features.

C. Reward Variance Analysis

The subset-replace reward remains stable despite replacing only a small subset of the reference set, with negligible replacement-position noise and low variance across tested subset sizes. Best-of-N selection, ratio clipping, and advantage normalization further bound variance impact, yielding stable convergence without destructive policy updates.

  • Overall reward stability: 4.67% reward CV across 450 training steps and 0.14% intra-step FID CV show that replacement-position noise is negligible relative to sample-quality differences.The analysis evaluates reward stability over the full 450-step training period.
  • Variance across replacement sizes: Intra-step FID CV remains very low across replacement subset sizes with the reference set fixed at 5,000 images.The results indicate stable reward computation across all tested configurations.
  • Mechanisms bounding variance impact: Best-of-N selection, ratio clipping with ε = 0.0001, and batch-wide advantage normalization prevent noisy reward estimates from destabilizing policy optimization.These mechanisms filter low-quality samples, limit large single-step policy updates, and standardize rewards across the batch.
  • Mechanisms bounding variance impact: Zero destructive policy updates and stable, monotonic convergence were observed throughout training.These observations support the stability of the variance-control mechanisms during optimization.

D. Computational Cost Analysis · E. Limitations

The subset-replace strategy adds limited computational overhead, with its unique FID matrix computation accounting for 8.0% of step time and using a substantially smaller reward model. Current limitations concern generalization beyond class-conditional ImageNet and the need for adaptive reference-set scheduling.

  • D. Computational Cost Analysis: The method’s per-step computational cost was profiled on 8× L40S GPUs to quantify distribution-wise reward overhead.
  • D. Computational Cost Analysis: 8.0% of total step time comes from FID matrix computation, the only component unique to the distribution-wise reward approach.
  • D. Computational Cost Analysis: 71.5% of total step time is policy training, while 10.3% is rollout generation; both costs are shared with sample-wise RL methods.
  • D. Computational Cost Analysis: The Inception-v3 reward model has 24M parameters, making it 12.7× smaller than the typical CLIP ViT-L model with 304M parameters.
  • E. Limitations: Experiments currently focus on class-conditional ImageNet generation, leaving extension of subset-replace to open-vocabulary text-to-image settings for further exploration.
  • E. Limitations: Open-vocabulary extension requires determining how to construct representative reference sets without predefined class labels.
  • E. Limitations: Reference-set regeneration interval and size are tuned through ablation, while adaptive scheduling based on training dynamics remains future work.

F. Qualitative Results

The qualitative evaluation visualizes uncurated image-generation samples from pretrained and distribution-wise reward RL-finetuned SiT-XL/2 models across five ImageNet classes.

  • Qualitative comparison: The study compares image-generation results from pretrained SiT-XL/2 and distribution-wise reward RL-finetuned models.These visualizations are presented in Figures 6–10.
  • Qualitative comparison: The visualized classes include airliner (404), balloon (417), giant panda (388), lion (291), and zebra (340).Each figure shows uncurated samples for one class label.
Loading 2607.02291v1…