Source-linked AI summary
Tuning-free Plug-and-Play Proximal Algorithm for Inverse Imaging Problems
Kaixuan Wei, Angelica Aviles-Rivero, Jingwei Liang, Ying Fu, Carola-Bibiane Schönlieb, Hua Huang
TL;DR
PnP performance depends on internal parameters that must be tuned across different images and imaging conditions. The paper uses a learned reinforcement-learning policy to select these parameters, reporting strong results across inverse imaging problems, including performance comparable to oracle settings.
Problem
PnP performance is highly sensitive to penalty, denoising-strength, and terminal-time choices, whose optimal settings vary by image and imaging conditions.
Source: §1 · p. 2Method
A policy network uses mixed model-free and model-based reinforcement learning to select PnP parameters, including denoising strength, penalty, and terminal time.
Source: §1 · p. 2, §3.3 · p. 5Results
Takeaways & Limitations
The policy customizes parameters for different images and often provides faster practical convergence and better empirical performance than handcrafted criteria.
Source: §1 · p. 2, §5 · p. 8Oracle parameter settings require ground truth and are generally impractical in real testing scenarios.
Source: §4.2 · p. 6
Abstract
from arXiv · showhide
Plug-and-play (PnP) is a non-convex framework that combines ADMM or other proximal algorithms with advanced denoiser priors. Recently, PnP has achieved great empirical success, especially with the integration of deep learning-based denoisers. However, a key problem of PnP based approaches is that they require manual parameter tweaking. It is necessary to obtain high-quality results across the high discrepancy in terms of imaging conditions and varying scene content. In this work, we present a tuning-free PnP proximal algorithm, which can automatically determine the internal parameters including the penalty parameter, the denoising strength and the terminal time. A key part of our approach is to develop a policy network for automatic search of parameters, which can be effectively learned via mixed model-free and model-based deep reinforcement learning. We demonstrate, through numerical and visual experiments, that the learned policy can customize different parameters for different states, and often more efficient and effective than existing handcrafted criteria. Moreover, we discuss the practical considerations of the plugged denoisers, which together with our learned policy yield state-of-the-art results. This is prevalent on both linear and nonlinear exemplary inverse imaging problems, and in particular, we show promising results on Compressed Sensing MRI and phase retrieval.
1. Introduction
PnP combines proximal optimization with denoisers, but its performance depends on parameters that vary across images and imaging conditions. The paper casts parameter selection as sequential decision-making and reports that its learned policy performs well across inverse imaging problems.
1. Introduction
- PnP replaces proximal operators with off-the-shelf denoisers, combining image-denoising priors with optimization methods for inverse imaging.The denoiser implicitly defines the regularizer.Source: §1 · pp. 1–2
- PnP performance is sensitive to penalty, denoising-strength, and terminal-time choices, whose optima vary with the image and imaging conditions.Manual tuning takes repeated trials, while semi-automated criteria can lead to suboptimal local minima.Source: §1 · p. 2
- A policy selects internal parameters sequentially as a reinforcement-learning problem, seeking faster convergence and better restoration accuracy.The policy maps observations to actions and is trained to maximize cumulative reward.Source: §1 · p. 2
- The approach is reported to outperform state-of-the-art techniques on compressed sensing MRI and phase retrieval, with learned settings comparable to oracle-tuned parameters.The oracle settings are tuned using inaccessible ground truth.Source: §1 · p. 2
2. Related Work
Prior PnP research spans proximal algorithms, imaging applications, denoiser priors, and convergence analysis. This work focuses on automatic parameter selection, distinguishing its approach from automated regularizer-specific tuning, deep unrolling, and other image-recovery uses of reinforcement learning.
2. Related Work
- PnP research covers proximal algorithms, imaging applications, denoiser priors, and convergence analysis.Examples include BM3D and deep-learning denoisers, alongside methods such as ADMM and primal-dual algorithms.Source: §2 · p. 3
- This work targets PnP parameter selection, where poor choices can severely degrade results.Prior automated selectors described here are restricted to specified regularizer classes and do not apply to sophisticated non-convex or learned PnP priors.Source: §2 · p. 3
- Unlike deep unrolling, the PnP approach adaptively selects stop time and penalty parameters for varying input states while using an off-the-shelf denoiser.Deep-unrolling methods treat parameters as trainable and require a fixed iteration count for end-to-end training.Source: §2 · p. 3
- Where earlier image-recovery studies used reinforcement learning for restoration tools, endpoints, filters, or network paths, this work uses mixed model-free and model-based deep RL to select PnP parameters.The cited prior examples include Yu et al. (2018), Zhang et al. (2019b), Furuta et al. (2019), and Yu et al. (2019).Source: §2 · p. 3
3. Tuning-free PnP Proximal Algorithm
The method casts PnP parameter selection as sequential decision-making and learns a policy to choose parameters and termination based on the current optimization state. Its environment models PnP-ADMM, while mixed model-free and model-based reinforcement learning trains the two policy components.
3. Tuning-free PnP Proximal Algorithm
- The approach targets PnP parameter selection and is described as applicable to other proximal methods sharing PnP-ADMM’s fixed points.The section divides the method into automated parameter selection, an environment model, and policy learning.Source: §3 · pp. 3–4
3.1. RL Formulation for Automated Parameter Selection
- The MDP treats optimization states as states, terminal time and denoising/penalty values as actions, and optimization steps as state transitions.The policy seeks to maximize expected discounted return; its terminal-time sub-policy makes a discrete stop/continue decision, while the other selects continuous parameters.Source: §3.1 · p. 4
3.2. Environment Model
- The environment uses PnP-ADMM dynamics and a reward that balances PSNR gain against a penalty for continuing rather than terminating.The denoiser is a residual U-Net with a tunable noise-level input; the data-fidelity proximal step uses a fast solution when feasible, otherwise one gradient-descent step.Source: §3.2 · pp. 4–5
3.3. RL-based policy learning
- Model-free reinforcement learning trains the discrete termination sub-policy, while model-based learning optimizes continuous denoising-strength and penalty choices through environment dynamics.The shared-feature policy uses a likelihood estimator for the discrete decision and backpropagation through the reward, value network, and transition function for the continuous decision.Source: §3.3 · p. 5
- Table 2 compares policies for CS-MRI on seven medical images across x2/x4/x8 acceleration factors at noise level 15, reporting PSNR and iteration counts.An asterisk marks each policy’s best PSNR across iterations under optimal early stopping.Source: Table 2 · p. 6
4. Experiments
The experiments evaluate the method on CS-MRI and phase retrieval, representing linear and nonlinear inverse imaging problems, respectively. They examine implementation choices, denoiser priors, learned policies, and comparisons with competing methods.
4. Experiments
- The experiments focus on CS-MRI and phase retrieval as representative linear and nonlinear inverse imaging tasks.CS-MRI tests x2/x4/x8 acceleration and noise levels 5/10/15 with one policy network; phase retrieval likewise uses one policy across settings.Source: §4 · p. 6, §4.1 · p. 6
4.2. Compressed sensing MRI
- For CS-MRI, the forward model uses partially sampled Fourier measurements with additive white Gaussian noise.The data-fidelity term is a squared-error objective with a proximal operator specified in prior work.Source: §4.2 · p. 6
- Residual UNet and MemNet outperform DnCNN in denoising and CS-MRI, and residual UNet is chosen because it is significantly faster than MemNet.The denoiser comparison exhaustively searches denoising-strength and penalty parameters to isolate prior effects.Source: §4.2 · p. 6
- Poor fixed parameter choices yield poor results and make early stopping important for recovery, while early stopping can reduce computation when parameters are properly assigned.The policy comparisons include handcrafted, fixed, fixed-optimal, oracle, and learned policies; the learned policy chooses its terminal time automatically.Source: §4.2 · p. 6
- The learned policy yields fast practical convergence and excellent performance, sometimes outperforming the ground-truth-tuned oracle in the ×2 case.The mixed model-free and model-based policy performs better than the model-free policy, which produces suboptimal parameters and fails to learn early stopping behavior.Source: §4.2 · p. 7
- Across Medical7 and 50 MICCAI 2013 medical images, the method significantly outperforms IRCNN, particularly in the difficult ×8 case; at ×2, it reaches desirable performance in 5 iterations versus IRCNN’s 30 and BM3D-MRI’s 70.Source: §4.2 · p. 7
4.3. Phase retrieval
- Phase retrieval tests use four coded diffraction pattern measurements with phase-only spatial light modulator patterns and varying Poisson noise.Each measurement operator combines a Fourier transform with a diagonal modulation matrix.Source: §4.3 · pp. 7–8, Table 4 · p. 7
- For phase retrieval, the method achieves state-of-the-art performance against three classic and three PnP algorithms, producing cleaner and clearer reconstructions.The comparison uses twelve images and noisy intensity-only coded diffraction pattern measurements.Source: §4.3 · p. 8, Figure 4 · p. 8
5. Conclusion
The proposed method introduces reinforcement learning into PnP to select parameters for different images. Experiments show that its learned policy often produces effective settings and can approach the performance of oracle parameters tuned using inaccessible ground truth.
5. Conclusion
- The learned policy customizes parameters for different images and often approaches the performance of oracle settings tuned using inaccessible ground truth.The paper presents this as the main strength of its tuning-free PnP algorithm for inverse imaging problems.Source: §5 · p. 8
Related papers
This paper cites
- ADAM: A METHOD FOR STOCHASTIC OPTIMIZATION2014
- U-Net: Convolutional Networks for Biomedical Image Segmentation2015
- Playing Atari with Deep Reinforcement Learning2013
- Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for Image Denoising2016
- MemNet: A Persistent Memory Network for Image Restoration2017
- The Little Engine that Could: Regularization by Denoising (RED)2016