Enhancing Diffusion-based Restoration Models via Difficulty-Adaptive Reinforcement Learning with IQA Reward

TL;DR

This paper tackles fidelity and hallucination issues in diffusion-based image restoration by introducing an IQA-based reinforcement learning framework. It leverages Multimodal Large Language Model–based IQA rewards to drive distribution-level alignment, focusing RL on hard, ground-truth distant samples before progressively reintroducing supervised fine-tuning for fine-grained alignment. A difficulty-aware weighting scheme selects which samples to prioritize, uses stepwise rewards across diffusion steps, and constrains updates with a KL term to stay close to the pretrained baseline. Across diverse restoration tasks and datasets, the approach yields consistent improvements over strong baselines and SFT-only methods, and is designed to be plug-and-play for existing diffusion restorers, providing a practical path to higher fidelity and realism in restored images.

Abstract

Reinforcement Learning (RL) has recently been incorporated into diffusion models, e.g., tasks such as text-to-image. However, directly applying existing RL methods to diffusion-based image restoration models is suboptimal, as the objective of restoration fundamentally differs from that of pure generation: it places greater emphasis on fidelity. In this paper, we investigate how to effectively integrate RL into diffusion-based restoration models. First, through extensive experiments with various reward functions, we find that an effective reward can be derived from an Image Quality Assessment (IQA) model, instead of intuitive ground-truth-based supervision, which has already been optimized during the Supervised Fine-Tuning (SFT) stage prior to RL. Moreover, our strategy focuses on using RL for challenging samples that are significantly distant from the ground truth, and our RL approach is innovatively implemented using MLLM-based IQA models to align distributions with high-quality images initially. As the samples approach the ground truth's distribution, RL is adaptively combined with SFT for more fine-grained alignment. This dynamic process is facilitated through an automatic weighting strategy that adjusts based on the relative difficulty of the training samples. Our strategy is plug-and-play that can be seamlessly applied to diffusion-based restoration models, boosting its performance across various restoration tasks. Extensive experiments across multiple benchmarks demonstrate the effectiveness of our proposed RL framework.

Figures (7)

• Figure 1: IQA scores in this figure are computed using the MLLM-based IQA model DeQA-Score you2025teaching on the widely-used real-world low-light image enhancement benchmark LOL-real yang2021sparse. The original IQA score of the diffusion-based method (DiffBIR lin2024diffbir) via SFT shows a significant gap compared to the IQA score of the ground truth. Thus, IQA can serve to distinguish the distribution between output images and ground truths. After our RL training, it aligns more closely with the ground-truth distribution.
• Figure 2: Visual comparisons using reconstruction error and IQA as the reward function. The reward guided by IQA leads to better visual performance, producing results that are closer to the ground truth, whereas the reconstruction-error-based reward fails.
• Figure 3: The schematic diagram of our RL algorithm (the stars $\bigstar$ of various colors in (a) and (b) represent the change of sample's latent state during training). (a) Although diffusion-based restoration models via SFT can generate samples with satisfactory perception, there still exist challenging cases far away from ground truths. (b) Our proposed RL process is dynamic: guided by IQA-based rewards, focuses on these underperforming samples, first exploring potential optimization directions and gradually aligning with the target distribution. Once the results approach the target distribution, SFT is combined to ensure fine-grained reference-based alignment. (c) Results optimized via RL not only demonstrate improved fidelity but also retain generative advantages of diffusion-based models. The radar map shows the quantitative comparisons between the current SOTA diffusion-based method (e.g., DiffBIR) via SFT and its variants with our RL strategy across multiple datasets, including LOL-real (low-light image enhancement, LLIE), Rain100H (image deraining, I.D.), GoPro (image motion deblurring, I.M.D.), and DPDD (image defocus deblurring, I.D.D.). More results are provided in the experimental section.
• Figure 4: The illustration of our RL process, especially including the policy modeling and difficulty weights computation. The computation of RL loss can be viewed in Fig. \ref{['fig:pipeline2']}.
• Figure 5: The illustration of our RL process that compute the loss on multiple time steps with reward normalization and KL constraint.