Reinforced Diffusion: Learning to Push the Limits of Anisotropic Diffusion for Image Denoising

TL;DR

This work treats image denoising as a learnable diffusion process by casting diffusion as a Markov decision process and solving it with a multi-agent deep reinforcement learning framework. Each pixel hosts an agent that selects from a small, discrete 9-action set (8 neighbor-average actions plus a do-nothing option), enabling a content-adaptive, edge-preserving diffusion over multiple iterations, with states given by the current image $u^{(t)}$ and rewards based on reductions in reconstruction error. The approach uses an asynchronous actor-critic (A3C) style training with a shared fully convolutional policy, achieving strong performance against traditional diffusion-based methods and competitive results with CNN-based denoisers; data augmentation further boosts gains. Overall, the study demonstrates that a learnable diffusion model, optimized via DRL, can push the limits of anisotropic diffusion for denoising and adapt to diverse noise types, with potential extensions to self-supervised learning and other image restoration tasks.

Abstract

Image denoising is an important problem in low-level vision and serves as a critical module for many image recovery tasks. Anisotropic diffusion is a wide family of image denoising approaches with promising performance. However, traditional anisotropic diffusion approaches use explicit diffusion operators which are not well adapted to complex image structures. As a result, their performance is limited compared to recent learning-based approaches. In this work, we describe a trainable anisotropic diffusion framework based on reinforcement learning. By modeling the denoising process as a series of naive diffusion actions with order learned by deep Q-learning, we propose an effective diffusion-based image denoiser. The diffusion actions selected by deep Q-learning at different iterations indeed composite a stochastic anisotropic diffusion process with strong adaptivity to different image structures, which enjoys improvement over the traditional ones. The proposed denoiser is applied to removing three types of often-seen noise. The experiments show that it outperforms existing diffusion-based methods and competes with the representative deep CNN-based methods.

Figures (4)

• Figure 1: Sub-actions. The red triangle denotes the position corresponding to the agent and the black circle denotes the position of neighbor pixel.
• Figure 2: The network architecture of the deep reinforcement framework.
• Figure 3: Denoising process of the proposed method and the action map at each time step.
• Figure 4: Visualization of the composited average operations.

Theorems & Definitions (1)

• Remark 1