Reward Sharpness-Aware Fine-Tuning for Diffusion Models

Kwanyoung Kim; Byeongsu Sim

cs.LGMar 22, 2026

Reward Sharpness-Aware Fine-Tuning for Diffusion Models

Authors: Kwanyoung Kim, Byeongsu Sim

Organizations: Department of AI Convergence, GIST · Samsung Research

Abstract

Reinforcement learning from human feedback (RLHF) has proven effective in aligning large language models with human preferences, inspiring the development of reward-centric diffusion reinforcement learning (RDRL) to achieve similar alignment and controllability. While diffusion models can generate high-quality outputs, RDRL remains susceptible to reward hacking, where the reward score increases without corresponding improvements in perceptual quality. We demonstrate that this vulnerability arises from the non-robustness of reward model gradients, particularly when the reward landscape with respect to the input image is sharp. To mitigate this issue, we introduce methods that exploit gradients from a robustified reward model without requiring its retraining. Specifically, we employ gradients from a flattened reward model, obtained through parameter perturbations of the diffusion model and perturbations of its generated samples. Empirically, each method independently alleviates reward hacking and improves robustness, while their joint use amplifies these benefits. Our resulting framework, RSA-FT (Reward Sharpness-Aware Fine-Tuning), is simple, broadly compatible, and consistently enhances the reliability of RDRL.

Figures & tables

Figure 1: Qualitative comparison across different diffusion backbones and RDRL frameworks. (Top): SD1.5 [ 38 ] results on Draft-LV [ 6 ] and AlignProp [ 35 ] (Middle–Bottom): Larger backbones (SDXL [ 34 ] and SD3 [ 9 ] ) on ReFL [ 52 ] and DRTune [ 50 ] . Each panel compares the vanilla model, the baseline RDRL method, and the same method combined with RSA-FT (Ours). RSA-FT is compatible with diverse reward-centric diffusion reinforcement learning frameworks and backbones, effectively mitigating reward hacking and producing clear improvements in visual quality and text–prompt alignment.

Figure 2: Illustration of reward hacking in RDRL (Draft-LV). The original reward model raises the HPS v2.1 score but degrades other metrics and visual quality, whereas our flattened model improves all metrics with consistent visuals.

Figure 3: Geometry of reward fine-tuning and our proposed method. Reward models are inherently sharp and prone to adversarial perturbations. Flattening these reward landscapes alleviates their sensitivity and reduces the occurrence of adversarial gradients. (a) Prior methods directly maximize rewards along adversarial gradients from sharp reward surfaces, which often leads to reward hacking. (b) Our method instead leverages gradients from flattened reward models, mitigating hacking by flattening both the image and parameter spaces.

Figure 4: Negative correlation between reward sharpness and human preference. Higher sharpness in the reward model correlates with lower preference quality (Pearson $r_{\text{corr}}=-0.802$ for PickScore, $r_{\text{corr}}=-0.669$ for ImageReward), supporting the hypothesized negative relationship.

Dataset	Method	HPSV2.1 $\uparrow$	PickScore $\uparrow$	ImageReward $\uparrow$
Drawbench	Vanilla	24.02	21.02	-0.147
	Draft-LV	25.59	20.96	-0.062
	+ Ours	26.67 (+1.08)	21.12 (+0.16)	0.035 (+0.09)
	Alignprop	25.12	20.98	-0.033
	+ Ours	29.59 (+4.47)	21.51 (+0.53)	0.268 (+0.301)
	ReFL	31.08	21.57	0.536

Table 1: Quantitative results of various RDRL methods on SD 1.5 (512

\times

512) with our proposed method. Bold text indicates the best performance for each metric.

Dataset	Method	HPSV2.1 $\uparrow$	PickScore $\uparrow$	ImageReward $\uparrow$
Drawbench	Vanilla	27.41	22.31	0.619
	ReFL	29.45	22.41	0.705
	+ Ours	30.31 (+0.86)	22.60 (+0.19)	0.719 (+0.014)
	DRTune	30.04	22.49	0.804
	+ Ours	31.58 (+1.54)	22.61 (+0.12)	0.944 (+0.140)
HPD	Vanilla	28.92	22.56	0.921

Table 2: Quantitative results of various RDRL on SDXL (1024

\times

1024). Bold text indicates the best performance for each metric.

Dataset	Method	HPSV2.1 $\uparrow$	PickScore $\uparrow$	ImageReward $\uparrow$
Drawbench	Vanilla	28.77	22.45	0.914
	ReFL	29.45	22.38	0.921
	+ Ours	29.81 (+0.86)	22.51 (+0.13)	0.953 (+0.014)
	DRTune	29.48	22.41	0.914
	+ Ours	29.89 (+1.54)	22.61 (+0.20)	0.979 (+0.140)
HPD	Vanilla	30.48	22.40	1.123

Table 3: Quantitative results of various RDRL on SD3 (1024

\times

1024). Bold text indicates the best performance for each metric.

Panel of Figure 5 — Figure 5: Human preference study results. The dashed line indicates the 50% mark; crossing it demonstrates a strict preference over the baseline.

Figure 6: Qualitative comparison with and without our method. Each image is generated using the same text prompt and random seed across all methods. Models equipped with our method (RSA-FT) produce images with more accurate text–image alignment and higher visual quality compared to their baselines.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Conceptual illustration of the RDRL framework.

Backbone model	SD 1.5	SDXL	SD 3	Flux.1-dev
Architecture	U-Net	U-Net	MMDiT	MMDiT
Inference Steps ( $N$ )	50	50	28	28
CFG Scale	7.5	5.0	7.0	3.5
Resolution	$512\times 512$	$1024\times 1024$	$1024\times 1024$	$1024\times 1024$

Appendix

Table 4: Hyperparameters and configurations for the backbone text-to-image models used in our experiments.

Configuration	SD 1.5	SDXL / SD 3 / Flux.1-dev
Hardware Resources	2 $\times$ NVIDIA H100	4 $\times$ NVIDIA H100
Training Duration	2500 iterations	2000 iterations
Global Batch Size	32	4
LoRA Rank	4	64
Learning Rate	$2\times 10^{-5}$	$1\times 10^{-5}$

Appendix

Table 5: Detailed training hyperparameters and hardware configurations for the different backbone models.

Method	AT	SAM	HPSV2.1 $\uparrow$	PickScore $\uparrow$	ImageReward $\uparrow$
Draft-LV	✗	✗	26.87	20.67	0.126
(a)	✓	✗	27.30	20.82	0.182
(b)	✗	✓	27.83	20.78	0.157
Ours	✓	✓	28.28	20.91	0.191
Alignprop	✗	✗	25.12	20.35	0.032
(a)	✓	✗	30.98	21.22	0.453

Appendix

Table 6: Quantitative results of various RDRL methods with backbone model SD 1.5. Bold text indicates the best performance for each metric across the different methods.

$\rho$	HPSv2.1	ImageReward	PickScore
0.1	31.20	1.171	22.245
0.01	31.50	1.196	22.402
0.001	31.29	1.185	22.391

Appendix

Table 7: Ablation study on the hyper‑parameter

\rho

$\rho_{w}$	HPSv2.1	ImageReward	PickScore
0.1	31.65	1.182	22.39
0.01	32.01	1.207	22.48
0.001	31.56	1.189	22.40

Appendix

Table 8: Ablation study on the hyper‑parameter

\rho

Reward	Method	HPSV2.1 $\uparrow$	PickScore $\uparrow$	ImageReward $\uparrow$
	ReFL	33.75	24.30	1.197
	+ Ours	34.21	24.33	1.207
	DRTune	34.16	24.01	1.217
HPS + Pick	+ Ours	34.74	24.14	1.259
	ReFL	33.99	24.25	1.198
	+ Ours	34.39	24.30	1.226

Appendix

Table 9: Quantitative comparison on Multi-reward scenarios.

Method	FID $\downarrow$	KID $\downarrow$
ReFL	76.36	0.0513
+ Ours	48.70	0.0006
DRTune	88.82	0.0074
+ Ours	92.92	0.0099

Appendix

Table 10: Divergence from Vanilla sampling.

Method	HPSv2.1 $\uparrow$	PickScore $\uparrow$	ImageReward $\uparrow$
Vanilla	29.76	22.51	0.906
ReFL	30.16	22.57	0.913
+ Ours	30.32	22.64	0.932
DRTune	29.99	22.59	0.926
+ Ours	30.25	22.67	0.961

Appendix

Table 11: uantitative results of various RDRL on Flux.1-dev (1024 × 1024). Bold text indicates the best performance for each metric.

Figure 8: Qualitative comparison on the SD1.5 backbone with and without our method. All images are generated using the same text prompt and random seed across methods. Models trained with our RSA-FT approach produce images with more reliable text–image alignment and noticeably improved visual quality. In addition, reward-induced artifacts commonly observed in RDRL baselines are reduced.

Figure 9: Qualitative comparison on the SDXL backbone with and without our method. Each image is generated using the same text prompt and random seed across all methods. Models equipped with our method (RSA-FT) produce images with more accurate text–image alignment and higher visual quality compared to their baselines. Moreover, RSA-FT suppresses the unwanted structural artifacts seen in the baselines.

Figure 10: Qualitative comparison on the SD3 backbone with and without our method. Each image is generated using the same text prompt and random seed across all methods. Models equipped with our method (RSA-FT) produce images with more accurate text–image alignment and higher visual quality compared to their baselines. Moreover, RSA-FT preserves text rendering quality, introducing fewer distortions than the baseline methods.

Diffusion Model Alignment

Diffusion model alignment steers generated content and internal representations toward prompts, human preferences, safety goals, and task constraints.

See all 185 papers

Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF
Jul 8, 2026Eric Zhu, Abhinav Shrivastava, Soumik MukhopadhyayDiffusion Model Alignment RL for Diffusion Models
Carnegie Mellon University · University of Maryland, College Park
Reinforcement Learning-driven Diffusion Models: Language, Vision, and Beyond
Oct 15, 2026Zhiyuan Ma, Jiabao Wei, Tianyu Wang +21RL for Diffusion Models RL for Generative Models
Huazhong University of Science and Technology · Beijing Institute of Technology · Tsinghua University +9
Diffusion Reinforcement Learning via Centered Reward Distillation
Mar 14, 2026Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière +1RL for Diffusion Models Diffusion Model Alignment
LIX, École Polytechnique, CNRS, IPP · Inria at Univ. Grenoble Alpes, CNRS, LJK

Jul 8, 2026J/K move · Enter open · S save

Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF

Eric Zhu, Abhinav Shrivastava, Soumik Mukhopadhyay

Carnegie Mellon University · University of Maryland, College Park

Diffusion Model Alignment RL for Diffusion Models Sample-Efficient RL

Open paper AskPDF arXiv

Reinforcement learning from human feedback (RLHF) has emerged as a powerful paradigm for aligning generative models with human preferences. However, applying RLHF to diffusion models remains highly feedback inefficient, as existing approaches typically require large amounts of human or reward model evaluations. This limitation reduces the practicality of diffusion RLHF in realworld settings where feedback is the primary bottleneck. In this paper, we propose two complementary strategies that substantially improve the feedback efficiency of diffusion RLHF while preserving generalization to unseen prompts. Our key observation is that reward information in diffusion trajectories is unevenly distributed: not all denoising timesteps or trajectories contribute equally to learning from a reward signal. By emphasizing informative timesteps and trajectories during optimization, we obtain more effective gradient updates. First, we introduce a per-timestep weighting scheme that reweights denoising steps during policy optimization. We theoretically connect this weighting to the optimal convergence properties of proximal policy optimization (PPO) and approximate the resulting weighting trend empirically. Second, we introduce a replay mechanism that prioritizes informative trajectories, enabling the model to reuse past samples instead of repeatedly querying new rewards. Together, these strategies significantly improve the feedback efficiency of diffusion RLHF. Under identical hyperparameter settings, our approach achieves up to a 6

\times

improvement in sample efficiency compared to widely used diffusion RLHF baselines.

Abstract

Diffusion Model Alignment

Explore similar work

Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF

Reinforcement Learning-driven Diffusion Models: Language, Vision, and Beyond

Diffusion Reinforcement Learning via Centered Reward Distillation