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

2025-11-03

AI Summary

Reinforcement learning (RL) is difficult to integrate effectively into diffusion-based image restoration models, as these models prioritize fidelity over generative diversity. Method: This paper proposes DiffRL, a novel training framework that synergistically combines RL with diffusion models. Its core innovations are: (1) a multimodal large language model (MLLM)-driven image quality assessment (IQA) module serving as a differentiable reward proxy; and (2) a difficulty-adaptive weighting mechanism that dynamically balances supervised fine-tuning and RL optimization, enabling coarse-to-fine progressive training. The framework is plug-and-play and requires no architectural modification to the underlying diffusion model. Contribution/Results: DiffRL achieves significant performance gains across multiple image restoration tasks—including deblurring, denoising, and super-resolution—particularly demonstrating enhanced robustness on complex or low-quality samples. Extensive experiments validate its effectiveness and strong generalization capability.

Technology Category

Computer Vision: Diffusion Models for VisionMachine Learning: Large Multimodal Models (LMMs)Search and Optimization: Learning to Search

Application Category

Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for rankingUser Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendationGraph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

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.

Problem

• Optimizing reinforcement learning for diffusion-based image restoration models
• Developing difficulty-adaptive training using IQA rewards for challenging samples
• Enhancing fidelity in image restoration through dynamic RL-SFT integration

Innovation

• Difficulty-adaptive RL weighting for challenging samples
• MLLM-based IQA reward replacing ground-truth supervision
• Plug-and-play RL integration with diffusion restoration models