Optimizing Visual Generative Models via Distribution-wise Rewards
Ruihang Li 1 2 3 Mengde Xu 3 Shuyang Gu 3 Leigang Qu 4 Fuli Feng 1 Han Hu 3 Wenjie Wang 1
arXiv:2607.02291v1 [cs.LG] 2 Jul 2026
Abstract
Chang et al., 2026; Crowson et al., 2024; Wang et al., 2024) and training strategies (Yu et al., 2024b; Huang et al., 2024; Hang et al., 2024). In the post-training stage, reinforcement learning (RL) with sample-wise reward models (Fan et al., 2023; Wu et al., 2023b; Kirstain et al., 2023; Xu et al., 2023; Wang et al., 2025) is employed to align model outputs with human preferences. Nevertheless, reinforcement fine-tuning driven by sample-wise rewards is prone to reward hacking (Weng, 2024; Amodei et al., 2016; Everitt et al., 2017; Gao et al., 2023; Wen et al., 2024; Liu et al., 2025; Li et al., 2025a), often introducing visual artifacts and diminishing the diversity of generated images (Ku et al., 2024; Xue et al., 2025; Miao et al., 2024; Liu et al., 2025), as shown in Figure 1. In contrast, distribution-wise metrics quantify diversity and mode coverage, penalizing generators that miss modes or exhibit low diversity (Borji, 2022; Ku et al., 2024; Cai et al., 2025). Early studies also confirmed their consistency with human judgment and their sensitivity to subtle shifts in the real distribution (Heusel et al., 2017; Borji, 2022), indicating greater robustness compared to sample-wise metrics. Moreover, alignment with a reference distribution that captures holistic, high-level attributes such as visual quality and aesthetics beyond the reach of sample-wise metrics, opens new avenues for improvement.
Conventional reinforcement learning strategies for visual generation typically employ samplewise 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 subsetreplace 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.
In this work, we propose a RL approach based on distribution-wise rewards to improve coverage of the realworld data distribution, achieving both high visual fidelity in samples and broad generation diversity. Quantifying the discrepancy between two distributions is a well-studied problem, with established metrics like KL divergence (Joyce, 2011), MMD (Gretton et al., 2006) and Wasserstein distance (Villani, 2009). In the field of image generation, Fréchet Inception Distance (FID) (Heusel et al., 2017; Jayasumana et al., 2024; Chong & Forsyth, 2020) is a widely used metric for assessing the degree of fit between the learned and real image distribution (Karras et al., 2022; 2024; Chang et al., 2026; Crowson et al., 2024; Wang et al., 2024; Yu et al., 2024b; Huang et al., 2024; Hang et al., 2024). Compared to sample-wise metrics like CLIP Score (Hessel et al., 2021) and HPS (Wu et al., 2023b;a), distributionbased metrics provide a better evaluation of how well the generative model covers the real distribution and can identify incorrect fits (Heusel et al., 2017; Gretton et al., 2006; Villani, 2009). As a widely used metric in image generation,
1. Introduction Visual generative models are designed to approximate the complex probability distribution of real-world images and videos. Existing studies have advanced this objective by improving network architectures (Karras et al., 2022; 2024; 1 University of Science and Technology of China 2 Shanghai Innovation Institute 3 Hunyuan Frontier Lab, Tencent 4 National University of Singapore. Correspondence to: Wenjie Wang <[email protected]>, Leigang Qu <[email protected]>.
Proceedings of the 43 rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026. Copyright 2026 by the author(s).
1
Optimizing Visual Generative Models via Distribution-wise Rewards
Without RL
Sample-wise RL
Distribution-wise RL (Ours)
Figure 1. Visualization of class-conditional image generation using varied initial noises. The baseline model (without RL, first row, FID 8.30) frequently produces visual artifacts, such as incorrect text rendering, spurious elements, distortion, and vignetting. Applying a sample-wise RL reward 1 leads to severe reward hacking (second row, FID 34.26), causing a collapse in sample diversity and introducing artifacts like bizarre rainbow patterns. In contrast, our distribution-wise reward (third row, FID 5.77) significantly mitigates these defects, enhancing overall generation quality and better aligning the learned distribution with the real data.
FID has been validated to correlate well with human perception of visual quality, while also providing a balanced assessment of both fidelity and diversity (Heusel et al., 2017; Salimans et al., 2016; Barratt & Sharma, 2018). Given these advantages, we choose FID as the distribution-wise metric to measure the generative model’s fitting capability and use it as the reward signal for reinforcement fine-tuning.
In this work, we propose distribution-wise reward for RL training. Specifically, we use a novel subset-replace strategy to obtain dense distribution-wise reward signals at a low compute cost. First, we generate a reference set of images and compute its FID against the target distribution as a starting point. During rollouts, a small subset of this reference set is replaced by newly generated samples, and the FID of the updated set is used as a dense reward signal. While this signal can be used to directly fine-tune the entire model, and indeed shows promise on models like SiT (Ma et al., 2024), such an approach still requires an SDE-based training formulation (Fan et al., 2023; Liu et al., 2025; Xue et al., 2025; He et al., 2025), inheriting the train-inference inconsistency issue. Inspired by EDM2 (Karras et al., 2024), we explore a more effective optimization strategy: applying our reward signal to search for optimal post-hoc model merging coefficients, instead of fine-tuning all parameters directly. This paradigm decouples the RL optimization from the denoising process, thereby eliminating the potential train-inference gap caused by SDE.
Training with distribution-wise rewards remains underexplored. Existing RL approaches for image generation (Black et al., 2023; Fan et al., 2023; Xue et al., 2025; Liu et al., 2025; Li et al., 2025a) generally treat the denoising process as a Markov Decision Process (MDP) in a stochastic environment (Fan et al., 2023; Liu et al., 2025), employing sample-wise reward models (Fan et al., 2023; Wu et al., 2023b; Kirstain et al., 2023; Xu et al., 2023; Wang et al., 2025) to obtain reward signals for each denoising trajectory, and utilize Group Relative Policy Optimization (GRPO) (Shao et al., 2024; Guo et al., 2025) to optimize the entire state–action sequence. However, directly optimizing with distribution-wise rewards requires computing statistical metrics on a huge set of images (e.g., 50K samples for FID), incurring significant computational cost. Besides, such distribution-wise metrics can’t provide reward signals for each individual denoising trajectories that is necessary for RL training. Moreover, we observed that performance improvements from RL fine-tuning in a SDE-based stochastic environment (Fan et al., 2023; Liu et al., 2025; He et al., 2025; Wang & Yu, 2025) for exploration do not fully translate to the faster, ODE-based deterministic sampling used during inference process. This discrepancy highlights a significant train-inference inconsistency and motivates the search for alternative optimization methods that avoids the performance gap between SDE-based training and ODEbased inference.
Specifically, the subset-replace strategy first computes a base FID on a class-balanced reference set of moderatelysized generated images. During the rollout phase, a small subset (0.01× of the reference set) of images in the reference set are randomly replaced with newly generated samples of the same corresponding classes. The FID of this partially updated set (replaced FID) is then computed, and its negative value serves as the reward signal for the related subset of images. Experiments on SiT (Ma et al., 2024) demonstrate that our method significantly reduces the FID from 8.30 to 5.77, and the FDDINOv2 from 230.39 to 164.88. For post-hoc model merging coefficient optimization, our strategy improves the FID-50K from 3.74 to 3.52 on the EDM2 (Karras et al., 2024) model, highlighting its power as a lightweight,
1
We use ImageReward (Xu et al., 2023) as the sample-wise reward model, formatting prompts as “an image of {class name}” to adapt to the class-conditional generation setting, and
train both the sample-wise and distribution-wise RL models until the reward saturates.
2
Optimizing Visual Generative Models via Distribution-wise Rewards ② Replace Old Subset with New Images
① Reference Set Generation
③ Reward Calculation & Normalization
Checkpoints Pool
learnable embedding
𝒛 EMANet
…
Replacing
periodically
...
... 𝑅'
Diffusion Policy Network
M!
𝐴"%
𝐴" &
𝐰 = 𝑤! , 𝑤" , … , 𝑤#! 𝑤$ ∼ 𝒩(𝑤 ! $ , 𝜎$ )
Gaussian sampling multiple times (!)
M#$%&$
())
M#$%&$
... 𝐴" '
𝐰 != 𝑤 !! , 𝑤 !" , … , 𝑤 ! #! 𝝈 = 𝜎! , 𝜎" , … , 𝜎#!
M"!
…
(*)
M#$%&$
(+)
…
M#$%&$
…
R(+)
Subset-replace Strategy
R(!)
R())
R(*)
Policy Gradient Update
Subset
generate
𝑅&
Batch-level Advantage Normalization
𝑅%
Reference Set 𝒢
Batch-level Advantage Normalization Policy Gradient Update
(2) Post-hoc Model Merging with RL
(1) Subset-replace Strategy
Figure 2. Illustration of our proposed RL framework with distribution-wise rewards. (1) Subset-replace Strategy: Initially, a reference set is generated using the diffusion policy. During rollout, a random subset is replaced with newly generated samples in the same classes. The distribution-wise metric of the resulting set acts as a reward, which is then normalized into an advantage signal to update the model via policy gradient. The reference set is regenerated periodically. (2) Post-hoc Model Merging with RL: The distribution-wise reward signal can guide a lightweight policy to learn the optimal weights for merging a pool of model checkpoints. This efficiently creates an improved model, while allowing the rollout process to utilize ODE-based inference.
plug-and-play module for enhancing pretrained models.
by applying policy gradients to the score function (Song et al., 2020), enabling preference-aligned image generation (Black et al., 2023; Fan et al., 2023; Fan & Lee, 2023; Lee et al., 2023). Offline Direct Preference Optimization was later introduced for text-to-image tasks (Wallace et al., 2024), though distributional shift in pairwise data motivated online methods with step-aware preference models (Yuan et al., 2024; Liang et al., 2025). More recently, GRPO-based approaches (Tong et al., 2025; Liu et al., 2025; Xue et al., 2025) have advanced RL-enhanced generation with samplewise reward models, with (Liu et al., 2025; Xue et al., 2025) extending GRPO to flow matching via ODE–SDE reformulation. (Liu et al., 2025; Xue et al., 2025; Li et al., 2025a) found that reward hacking occurs in the RL process. In this work, we explore the potential to mitigate this issue with distribution-wise rewards. (He et al., 2025; Li et al., 2025a) further employ hybrid SDE–ODE to rollout denoising trajectories to accelerate training. (Wang & Yu, 2025) points out the SDE formulation in common RL practices is injecting greater level of noise than the original ODE, leading to a train-inference inconsistency. In this paper, we applies RL to optimize post-hoc model merging coefficients, eliminating the need for SDE-based rollouts and resolving the train-inference inconsistency of SDE-based RL.
Our contributions are summarized as follows: 1. We analyze the limitations of reinforcement learning with sample-wise reward functions, showing that they are susceptible to reward hacking, which degrades distributional fidelity and introduces artifacts while reducing diversity. 2. We propose a RL framework with distribution-wise rewards by the subset-replace strategy. This provides a robust alternative to conventional sample-wise rewards, which are vulnerable to reward hacking. Through extensive experiments, we derive an effective and optimal training recipe that reduces the FID-50K of SiT from 8.30 to 5.77 and the FDDINOv2 score from 230.39 to 164.88 without requiring additional training data or architectural modifications. 3. To resolve the train-inference inconsistency of SDEbased RL, we propose a post-hoc optimization of model merging coefficients with distribution-wise reward signals using ODE-based denoising procedure. This training paradigm improves EDM2’s FID-50K score from 3.74 to 3.52, validating a more consistent and effective approach to model refinement.
Distribution-wise Metrics. Distribution-wise metrics are widely used in training and evaluating neural networks. KL Divergence (Joyce, 2011), which is often included as a regularization term in RL (Fan et al., 2023; Liu et al., 2025; He et al., 2025; Shao et al., 2024; Guo et al., 2025), mea-
2. Related Work Reinforcement Learning in Image Generation. Early works adapted reinforcement learning to diffusion models 3
Optimizing Visual Generative Models via Distribution-wise Rewards
sures the difference between distributions but can be unstable when one distribution assigns zero probability to regions where the other has non-zero probability. Maximum Mean Discrepancy (MMD) (Gretton et al., 2006) compares distributions by their means in a Reproducing Kernel Hilbert Space. While non-parametric and robust, MMD can struggle with high-dimensional data and is sensitive to outliers (Lerasle et al., 2019). Frechet Inception Distance (FID) (Heusel et al., 2017), on the other hand, has become the preferred metric to evaluate image generation models (Karras et al., 2022; 2024; Chang et al., 2026; Crowson et al., 2024; Wang et al., 2024; Yu et al., 2024b; Huang et al., 2024; Hang et al., 2024). By comparing feature distributions of real and generated data using a pre-trained Inception network (Szegedy et al., 2016; Heusel et al., 2017), FID reflects how well a generative model fits the real image distribution with lower computational cost and greater statistical robustness. In this work, we introduce a tractable online formulation of the FID, allowing it to be effectively used as a direct distribution-wise reward signal to guide RL in image generation.
linear interpolations between the two samples are defined as xt = (1 − t)x0 + tx1 ,
t ∈ [0, 1].
(1)
A time-dependent velocity field vθ (xt , t) is then learned by minimizing the flow-matching objective (Lipman et al., 2022), given by LFM (θ) = Et, x0 , x1 ∥ v − vθ (xt , t)∥22 ,
v = x 1 − x0 . (2)
Denoising as a MDP. (Black et al., 2023; Liu et al., 2025) cast the iterative denoising procedure in flow matching models as a Markov Decision Process (MDP) (S, A, ρ0 , P, R), where R is the reward of this denoising trajectory. Given a class label c ∈ C, at step t, the state is written as st ≜ (c, t, xt ), the action corresponds to the model’s prediction at ≜ xt−1 , and the policy is defined by π(at | st ) ≜ pθ (xt−1 | xt , c). The transition is deterministic, i.e., P (st+1 | st , at ) ≜ (δc , δt−1 , δxt−1 ), and the initial distribution is specified as ρ0 (s0 ) ≜ (p(c), δT , N (0, I)), where δy denotes the Dirac delta distribution centered at y. 3.2. Subset-Replace Strategy
Model Merging. Model averaging (Izmailov et al., 2018; Polyak & Juditsky, 1992; Tarvainen & Valpola, 2017; Yaz et al., 2018) has become an widely-adopted techniques in the pre-training of state-of-the-art image synthesis models (Balaji et al., 2022; Dhariwal & Nichol, 2021; Ho et al., 2022; Karras et al., 2019; Nichol & Dhariwal, 2021; Peebles & Xie, 2023; Ma et al., 2024; Karras et al., 2022). In the domain of large language models, several studies have similarly explored the use of model averaging during both pre-training (Li et al., 2025b; 2022; Sanyal et al., 2023; Liu et al., 2024; Yang et al., 2023; Dubey et al., 2024; Tian et al., 2025) and post-training (Ilharco et al., 2022; Yu et al., 2024a; Zhou et al., 2024) to improve overall performance and enhance training stability. However, existing approaches such as exponential moving average (EMA) (Morales-Brotons et al., 2024) perform model merging during training, which makes tuning their hyperparameters computationally expensive. (Karras et al., 2024) addresses this limitation by introducing a post-hoc EMA strategy, where the optimal averaging profile is determined through grid search after training. Building on this idea, we propose to optimize the model merging hyperparameters with reinforcement learning, guided by reward signals rather than exhaustive search.
Existing RL approaches in diffusion models generally formulate the denoising process as a MDP in a stochastic environment (Fan et al., 2023; Liu et al., 2025; Xue et al., 2025; Li et al., 2025a), where a sample-wise reward (Xu et al., 2023; Wang et al., 2025; Wu et al., 2023b; Kirstain et al., 2023) is used as the optimization signal for each denoising trajectory. Directly replacing this with a distribution-wise reward is infeasible: computing such reward typically requires a very large number of trajectories (about 50k images and their denoising trajectories for FID), and assigning the same scalar reward to all trajectories leads to overly sparse feedback, providing little guidance for optimization. To address these limitations, we propose a subset-replace strategy for computing distribution-wise rewards, as demonstrated in Figure 2. Specifically, we first construct a classbalanced moderately-sized reference set G of N generated images with the initial pretrained model. During rollout, a small subset of n images g ⊆ G is randomly replaced with newly generated samples g ′ of the same classes. We then compute the FID of the partially updated set (G \ g) ∪ g ′ , denoted as replaced FID, whose negative value is used as the reward signal for the associated n denoising trajectories, as shown in Equation 4. To mitigate discrepancies between the reference set and the current model distribution, the reference set is periodically regenerated using the latest model during training. Compared with directly using FID-50K as the reward signal, this strategy substantially reduces computational cost while yielding denser and more informative rewards for model optimization.
3. Method 3.1. Preliminaries Flow Matching. Let x0 ∼ X0 be drawn from the real data distribution and x1 ∼ X1 from a noise distribution. Following the rectified flow framework (Liu et al., 2022), 4
Optimizing Visual Generative Models via Distribution-wise Rewards Table 1. FID (Heusel et al., 2017; Salimans et al., 2016; Barratt & Sharma, 2018) and FDDINOv2 (Stein et al., 2023; Karras et al., 2024) results on ImageNet 256×256. Our results demonstrate that finetuning pretrained visual generative models with a distribution-wise reward function is highly effective. This approach significantly enhances the visual quality of generated images within a minimal number of training steps while preserving generative diversity. We validate that the proposed subset-replace strategy provides a robust distribution-wise reward signal for both Rejection Sampling (RS) and Policy Gradient Reinforcement Training (RL). Applying our method to a pretrained SiT model reduces the FID-50K score from 8.30 to 6.98 (RS) and 5.77 (RL), validating its efficacy in enhancing perceptual quality. The FDDINOv2 metric exhibits a congruent reduction, confirming the generalizability of this improvement across different feature extractors and demonstrating that our approach is not overfitting to a single metric’s feature space.
Model
Training Steps
FID ↓
FDDINOv2 ↓
ADM ADM-U LDM-8 LDM-4 DiT-XL/2 DiT-XL/2
1.98M 1.98M 4.8M 178K 400K 7M
10.94 7.49 15.51 10.56 19.50 9.60
-
SiT-XL/2 SiT-XL/2 + Ours (RS) + Ours (RL)
400K 7M + 120 + 450
17.20 8.30 6.98 5.77
230.39 183.75 164.88
the same class distribution, and calculate the reward R as: R(gi′ ) = −FID[(G \ gi ) ∪ gi′ , G],
(4)
where G denotes the ground-truth image set of the same size as G. Then, the advantage of i-th subset is calculated by: Âi =
R(gi′ ) − mean({R(gi′ )}B i=1 ) . ) std({R(gi′ )}B i=1
(5)
Considering the complete denoising trajectory i,j i,j (xi,j T , xT −1 , . . . , x0 ) of the j-th image in the ith subset, the resulting image subset is given by i,2 i,n gi′ = {xi,1 0 , x0 , . . . , x0 }. Reinforcement fine-tuning then optimizes the policy model θ by maximizing the following objective as Liu et al. (2025): JFlow-RL (θ) = Ec∼C,{xi }G i=1 ∼πθ
old
(·|c) f (r, Â, θ, ε, β),
(6)
where πθold is the initial pretrained policy, and f (r, Â, θ, ε, β) B
=
We apply the subset-replace strategy to obtain distributionwise reward signals, and perform direct reinforcement finetuning of diffusion models based on them. Following (Fan et al., 2023; Liu et al., 2025), we learn a policy πθ that maximizes the expected cumulative reward, typically formulated as:
n
T −1
1 X1X1 X B i=1 n j=1 T t=0 min rti,j (θ) Âi , clip rti,j (θ), 1 − ε, 1 + ε Âi ! − β DKL (πθ || πref ) ,
rti,j (θ) =
max E(s0 ,a0 ,...,sT ,aT )∼πθ θ " T # X R(st , at ) − β DKL πθ (· | st ) || πref (· | st ) .
i,j pθ (xi,j t−1 | xt , c) i,j pθold (xi,j t−1 | xt , c)
.
3.3. Post-hoc Model Merging with Distribution-wise Reward
t=0
(3) where the KL-divergence DKL from a reference policy πref , scaled by β, serves as a regularization penalty. We adopt a lightweight variant (Shao et al., 2024; Hu, 2025) of traditional policy gradient methods (Schulman et al., 2015; 2017), which estimates the advantage without requiring a value function. Our early experiments presented in Section 4.3 found that batch-level normalization outperforms group-level normalization under our setting, as also observed in (Hu, 2025; Xie et al., 2025).
While directly applying our distribution-wise reward signal for fine-tuning with subset-replace strategy is a straightforward approach, our experiments in Section 4.3 expose an issue of train-inference inconsistency. Specifically, while existing RL methods on diffusion models (Fan et al., 2023; Xue et al., 2025; Liu et al., 2025) rely on SDEs to provide the stochasticity necessary for the RL process, we observe that the performance gains from this stochastic environment fail to transfer robustly to the ODE-based deterministic samplers (Karras et al., 2022; 2024; Ma et al., 2024) used for standard inference. To bridge this gap, we introduce a posthoc optimization strategy inspired by EDM2 (Karras et al., 2024). Our method uses RL with distribution-wise rewards to find optimal model merging coefficients, thereby eliminating the dependence on complex SDE solvers (Fan et al., 2023; Liu et al., 2025; Xue et al., 2025) during RL training.
To formalize the above process, let the reference set G consist of N generated images. At each iteration, a subset g of n randomly selected images is replaced. Considering rollouts with batch size B, the replaced subset is denoted by B {gi }B i=1 , with the corresponding class labels {ci }i=1 . We B ′ B substitute {gi }i=1 with a new subset {gi }i=1 that preserves 5
Optimizing Visual Generative Models via Distribution-wise Rewards Figure 3. Ablation studies on hyperparameters in RL with subset-replace strategy. (a) Reference set size. The relationship between set size and FID-50K is non-monotonic. While performance generally improves as the size increases from 2,500 to 10,000, the 7,500-sample set exhibits significant degradation, performing worse than even smaller sets. (b) Number of images to replace. We evaluate replacing 50, 100, and 200 images in the subset-replace strategy. A smaller replacement size of 50 images yields the best FID-5K performance after 100 training steps. (c) Impact of rollout sample selection strategies. Selecting the global top 25% of samples is optimal. Per-process selection methods are inferior, and retaining low-quality samples hinders training.
Reference Set Size 2500 7500 5000 10000
8.25
FID-5K
FID-50K
8.50
#Replace Images 50 100 200
20
8.00
18
20
FID-5K
8.75
16
7.75
14
7.50 0
50
100
Training Steps
150
18 All Global Top 25% Local Top 25% Local Top 50% Local T+B 25%
16 14
0
50
100
150
Training Steps
(a)
200
0
50
(b)
Model merging is a widely used technique in deep learning, and early works in large language models (Li et al., 2025b; Yu et al., 2024a; Zhou et al., 2024) and visual generation models (Balaji et al., 2022; Dhariwal & Nichol, 2021; Ho et al., 2022; Karras et al., 2019; Nichol & Dhariwal, 2021; Peebles & Xie, 2023; Ma et al., 2024; Karras et al., 2022) has demonstrated its effectiveness in stabilizing training and improving model performance. The most common approach is Exponential Moving Average (EMA) (Morales-Brotons et al., 2024), which maintains a separate EMA copy of the model and updates it throughout training. However, this requires fixing the merging hyperparameters in advance, often resulting in suboptimal choices. (Karras et al., 2024) shows that by carefully designing the averaging formulation of model replicas during training, it is possible to approximate the EMA version after training. This allows the merging hyperparameters to be adjusted retrospectively based on downstream performance metrics.
100
Training Steps
150
(c)
wi ∼ N (w̄i (z; πθema ), σi (z; πθema ))
(8)
and their corresponding probabilities pwi are computed as: (wi − w̄i )2 1 pwi = p . (9) exp − 2σi2 2πσi2 We regard the coefficients involved in constructing the merged model Mmerge as a vector w = (w1 , w2 , . . . , wNc ). The reward corresponding to each w is computed using the subset-replace strategy. During rollouts, we generate a batch of B such coefficient vectors {w(j) }B j=1 , with the (j)
corresponding merged models denoted as {Mmerge }B j=1 . For (j)
each model Mmerge , we first construct a reference set Gj , s from which Ns subsets {gk }N k=1 are selected. For each subset gk , we replace it with Nr newly generated sets of images ′ r {gk,p }N p=1 , obtaining a reward collection
To formulate it, let Nc sequential checkpoints along the c training trajectory be denoted as {Mi }N i=1 , where Mi represents the parameters of the i-th model. These checkpoints are then merged into a single final model Mmerge , where each checkpoint is assigned a weighting coefficient wi . The merged model is computed as: Nc X Mmerge = w i Mi (7)
(j)
s , Nr {Rk,p }N k=1,p=1 .
Finally, the overall reward for coefficient vector w(j) is defined as the simple average: N
R(j) =
i=1
We optimize the model merging coefficients wi using RL. To introduce the stochasticity and related probabilities required for the RL procedure, we employ a simple MLP policy network πθema (EMANet) to generate the mean w̄i and standard deviation σi of each coefficient from a learnable input embedding z. The final values wi are then sampled from a Gaussian distribution
N
s X r 1 X (j) Rk,p . Ns Nr p=1
(10)
k=1
We compute the advantages at the batch level (Hu, 2025) across B reward values and use them to update parameters θema of the policy model. Since the stochasticity in the RL process originates from the coefficient vectors w generated by πθema , it is unnecessary to introduce additional 6
Optimizing Visual Generative Models via Distribution-wise Rewards Table 2. FID results on ImageNet 512×512. The EDM2 baseline results are achieved through post-hoc model merging, with coefficients optimized via extensive grid search, as detailed in (Karras et al., 2024). Results show that using RL to obtain better model merging coefficients is an effective method to boost the performance of pretrained models.
randomness in the diffusion denoising process. Therefore, we employ efficient ODE sampling (Karras et al., 2022; 2024) throughout the image generation process.
4. Experiments 4.1. Reinforcement Fine-tuning with Distribution-wise Reward We use ImageNet (Deng et al., 2009) in 256×256 resolution as our main dataset following (Ma et al., 2025), and perform full parameter reinforcement fine-tuning on SiT (Ma et al., 2024). To lower the training cost, we adopt the denoising reduction technique introduced in (Liu et al., 2025): the number of denoising steps is set to 50 during training and 250 steps during evaluation, following the optimal inference settings in (Ma et al., 2024). We first validated the feasibility of the subset-replace strategy as well as the distribution-wise reward signal under the rejection sampling fine-tuning (RS) setting, and then applied it to the standard RL setting. During RS training, we only use the samples with the highest distribution-wise reward values. Table 1 summarizes FID-50K results of our methods as well as several earlier pretrained models on the ImageNet dataset, following the widely-used evaluation protocol (Karras et al., 2024; Peebles & Xie, 2023; Ma et al., 2024). To demonstrate that our method’s efficacy generalizes across different feature representations, we also report FDDINOv2 (Stein et al., 2023; Karras et al., 2024) scores. This metric computes the Fréchet Distance of DINOv2 (Oquab et al., 2023) features on 50K ImageNet images, for which we adopt the same evaluation setting from (Karras et al., 2024).
Model
FID ↓
ADM (Dhariwal & Nichol, 2021) ADM-U DiT-XL/2 (Peebles & Xie, 2023)
23.24 9.96 12.03
EDM2-XS (Karras et al., 2024) + RL-EMA EDM2-S + RL-EMA
3.74 3.52 2.57 2.52
any additional curated training data or architectural modifications. Further incorporating RL, SiT-XL achieves an FID-50K of 5.77 and an FDDINOv2 of 164.88 with a small amount of additional training, substantially improving the ability to model image distribution. 4.2. Post-hoc Model Merging with Distribution-wise Reward Following prior settings (Karras et al., 2024), we perform experiments on ImageNet (Deng et al., 2009) (512×512) with models of various sizes to demonstrate the generality of our method. The results are presented in Table 2. We set Nc = 8 to compose the final model Mmerge . Starting from latest official checkpoints (Karras et al., 2024), we select checkpoints for every 192 × 220 training images, resulting in a checkpoint pool of Nc = 8 checkpoints. A simple 3-layer MLP is employed as the policy network to obtain the model merging coefficients w, with the sampling standard deviation fixed to 1.
For batch-level advantage normalization, we compute the mean and standard deviation across all processes. In the RL practice, we found that optimization becomes challenging when training on the entire set of rollout samples. To mitigate this, we retain only the top 25% of samples ranked by advantage for parameter update, and further perform detailed ablation experiments in Section 4.3. We adopt an on-policy RL setting in which each rollout sample is used only once for updating the model. Besides, we parallelize reference set generation by distributing tasks across processes and synchronizing the full set to all workers. To balance efficiency and quality, we refresh the reference set with the current model every 10 steps. We performed experiments on 16 NVIDIA Hopper GPUs, and the experiment that yielded the best FID-50K score completed in approximately 20 hours.
As shown in Table 2, by optimizing several parameters (Nc = 8 in our setting), our method reduces FID from 3.74 to 3.52 on EDM2-XS and from 2.57 to 2.52 on EDM2-S. These results demonstrate that reinforcement learning can effectively optimize model-merging coefficients, yielding further improvements to pretrained models without resorting to complex SDE solvers or training techniques such as denoising reduction (Liu et al., 2025), which has been observed to cause model collapse issues at certain denoising steps.
Experimental results in Table 1 demonstrate that a simple subset-replace strategy provides an effective distributionwise reward signal for model optimization. Under the simple RS setting, SiT-XL reduces the FID-50K from 8.30 to 6.98 and FDDINOv2 from 230.39 to 183.75, without requiring
4.3. Ablation Study We systematically evaluate the influence of key hyperparameters and components in our subset-replace strategy, following the experimental protocol in Section 4.1. 7
Optimizing Visual Generative Models via Distribution-wise Rewards Figure 4. Analysis of key design choices for our RL training pipeline. (a) Batch-level advantage normalization for advantages outperforms group-level constantly, yielding faster convergence regardless of whether all or only the top 25% of rollout samples are used for training. (b) The performance gap from training-inference inconsistency. A model trained with SDE-based rollouts shows a steadily improving FID score when evaluated with an SDE solver while its performance stagnates when using an ODE solver at the same 250 denoising steps. (c) RL training after Rejection Sampling fine-tuning (RS) provided no performance gain, likely due to overfitting from the RS phase. We therefore adopted a pure RL approach.