SGD-Net: Efficient Model-Based Deep Learning With Theoretical Guarantees

Deep unfolding networks have recently gained popularity for solving imaging inverse problems. However, the computational and memory complexity of data-consistency layers within traditional deep unfolding networks scales with the number of measurements, limiting their applicability to large-scale imaging inverse problems. We propose SGD-Net as a new methodology for improving the efficiency of deep unfolding through stochastic approximations of the data-consistency layers. Our theoretical analysis shows that SGD-Net can be trained to approximate batch deep unfolding networks to an arbitrary precision. Our simulations on intensity diffraction tomography and sparse-view computed tomography show that SGD-Net can match the performance of the traditional batch network at a fraction of training and testing complexity.

[1]  L. Rudin,et al.  Nonlinear total variation based noise removal algorithms , 1992 .

[2]  Dimitri P. Bertsekas,et al.  On the Douglas—Rachford splitting method and the proximal point algorithm for maximal monotone operators , 1992, Math. Program..

[3]  M. Glas,et al.  Principles of Computerized Tomographic Imaging , 2000 .

[4]  Robert D. Nowak,et al.  Wavelet-based image estimation: an empirical Bayes approach using Jeffrey's noninformative prior , 2001, IEEE Trans. Image Process..

[5]  Mário A. T. Figueiredo,et al.  Wavelet-Based Image Estimation : An Empirical Bayes Approach Using Jeffreys ’ Noninformative Prior , 2001 .

[6]  Robert D. Nowak,et al.  An EM algorithm for wavelet-based image restoration , 2003, IEEE Trans. Image Process..

[7]  I. Daubechies,et al.  An iterative thresholding algorithm for linear inverse problems with a sparsity constraint , 2003, math/0307152.

[8]  Antonin Chambolle,et al.  A l1-Unified Variational Framework for Image Restoration , 2004, ECCV.

[9]  Yurii Nesterov,et al.  Introductory Lectures on Convex Optimization - A Basic Course , 2014, Applied Optimization.

[10]  Eero P. Simoncelli,et al.  Image quality assessment: from error visibility to structural similarity , 2004, IEEE Transactions on Image Processing.

[11]  Michael Elad,et al.  Image Denoising Via Sparse and Redundant Representations Over Learned Dictionaries , 2006, IEEE Transactions on Image Processing.

[12]  H. Robbins A Stochastic Approximation Method , 1951 .

[13]  Alessandro Foi,et al.  Image Denoising by Sparse 3-D Transform-Domain Collaborative Filtering , 2007, IEEE Transactions on Image Processing.

[14]  Léon Bottou,et al.  The Tradeoffs of Large Scale Learning , 2007, NIPS.

[15]  Marc Teboulle,et al.  Fast Gradient-Based Algorithms for Constrained Total Variation Image Denoising and Deblurring Problems , 2009, IEEE Transactions on Image Processing.

[16]  Yonina C. Eldar,et al.  Convex Optimization in Signal Processing and Communications , 2009 .

[17]  Convex Optimization in Signal Processing and Communications , 2010 .

[18]  José M. Bioucas-Dias,et al.  Fast Image Recovery Using Variable Splitting and Constrained Optimization , 2009, IEEE Transactions on Image Processing.

[19]  Michael K. Ng,et al.  Solving Constrained Total-variation Image Restoration and Reconstruction Problems via Alternating Direction Methods , 2010, SIAM J. Sci. Comput..

[20]  Marc Teboulle,et al.  Gradient-based algorithms with applications to signal-recovery problems , 2010, Convex Optimization in Signal Processing and Communications.

[21]  Stephen P. Boyd,et al.  Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers , 2011, Found. Trends Mach. Learn..

[22]  Dimitri P. Bertsekas,et al.  Incremental proximal methods for large scale convex optimization , 2011, Math. Program..

[23]  Mathews Jacob,et al.  A Fast Majorize–Minimize Algorithm for the Recovery of Sparse and Low-Rank Matrices , 2012, IEEE Transactions on Image Processing.

[24]  Karen O. Egiazarian,et al.  BM3D Frames and Variational Image Deblurring , 2011, IEEE Transactions on Image Processing.

[25]  Brendt Wohlberg,et al.  Plug-and-Play priors for model based reconstruction , 2013, 2013 IEEE Global Conference on Signal and Information Processing.

[26]  Jeffrey A. Fessler,et al.  Accelerating Ordered Subsets Image Reconstruction for X-ray CT Using Spatially Nonuniform Optimization Transfer , 2013, IEEE Transactions on Medical Imaging.

[27]  Michael Unser,et al.  Hessian Schatten-Norm Regularization for Linear Inverse Problems , 2012, IEEE Transactions on Image Processing.

[28]  Stefan Roth,et al.  Shrinkage Fields for Effective Image Restoration , 2014, 2014 IEEE Conference on Computer Vision and Pattern Recognition.

[29]  Stephen P. Boyd,et al.  Proximal Algorithms , 2013, Found. Trends Optim..

[30]  Wei Yu,et al.  On learning optimized reaction diffusion processes for effective image restoration , 2015, 2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[31]  Thomas Brox,et al.  U-Net: Convolutional Networks for Biomedical Image Segmentation , 2015, MICCAI.

[32]  Jimmy Ba,et al.  Adam: A Method for Stochastic Optimization , 2014, ICLR.

[33]  Saeed Ghadimi,et al.  Accelerated gradient methods for nonconvex nonlinear and stochastic programming , 2013, Math. Program..

[34]  Jian Sun,et al.  Deep ADMM-Net for Compressive Sensing MRI , 2016, NIPS.

[35]  Demetri Psaltis,et al.  Optical Tomographic Image Reconstruction Based on Beam Propagation and Sparse Regularization , 2016, IEEE Transactions on Computational Imaging.

[36]  C. McCollough TU-FG-207A-04: Overview of the Low Dose CT Grand Challenge. , 2016, Medical physics.

[37]  Michael Elad,et al.  Turning a denoiser into a super-resolver using plug and play priors , 2016, 2016 IEEE International Conference on Image Processing (ICIP).

[38]  Jan Sijbers,et al.  Fast and flexible X-ray tomography using the ASTRA toolbox. , 2016, Optics express.

[39]  Charles A. Bouman,et al.  Plug-and-Play Priors for Bright Field Electron Tomography and Sparse Interpolation , 2015, IEEE Transactions on Computational Imaging.

[40]  Hassan Mansour,et al.  Learning Optimal Nonlinearities for Iterative Thresholding Algorithms , 2015, IEEE Signal Processing Letters.

[41]  Lei Zhang,et al.  Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for Image Denoising , 2016, IEEE Transactions on Image Processing.

[42]  Hassan Mansour,et al.  A Plug-and-Play Priors Approach for Solving Nonlinear Imaging Inverse Problems , 2017, IEEE Signal Processing Letters.

[43]  Stanley H. Chan,et al.  Plug-and-Play ADMM for Image Restoration: Fixed-Point Convergence and Applications , 2016, IEEE Transactions on Computational Imaging.

[44]  Jonas Adler,et al.  Solving ill-posed inverse problems using iterative deep neural networks , 2017, ArXiv.

[45]  Leslie N. Smith,et al.  Cyclical Learning Rates for Training Neural Networks , 2015, 2017 IEEE Winter Conference on Applications of Computer Vision (WACV).

[46]  Prateek Jain,et al.  Non-convex Optimization for Machine Learning , 2017, Found. Trends Mach. Learn..

[47]  Jan Kautz,et al.  Loss Functions for Image Restoration With Neural Networks , 2017, IEEE Transactions on Computational Imaging.

[48]  Jong Chul Ye,et al.  A deep convolutional neural network using directional wavelets for low‐dose X‐ray CT reconstruction , 2016, Medical physics.

[49]  Michael Unser,et al.  Deep Convolutional Neural Network for Inverse Problems in Imaging , 2016, IEEE Transactions on Image Processing.

[50]  Feng Lin,et al.  Low-Dose CT With a Residual Encoder-Decoder Convolutional Neural Network , 2017, IEEE Transactions on Medical Imaging.

[51]  Michael Elad,et al.  The Little Engine That Could: Regularization by Denoising (RED) , 2016, SIAM J. Imaging Sci..

[52]  Jeffrey A. Fessler,et al.  Physics-driven deep training of dictionary-based algorithms for MR image reconstruction , 2017, 2017 51st Asilomar Conference on Signals, Systems, and Computers.

[53]  Michael Unser,et al.  Convolutional Neural Networks for Inverse Problems in Imaging: A Review , 2017, IEEE Signal Processing Magazine.

[54]  Andreas Hauptmann,et al.  Model-Based Learning for Accelerated, Limited-View 3-D Photoacoustic Tomography , 2017, IEEE Transactions on Medical Imaging.

[55]  Ulugbek Kamilov,et al.  Efficient and accurate inversion of multiple scattering with deep learning , 2018, Optics express.

[56]  Bernard Ghanem,et al.  ISTA-Net: Interpretable Optimization-Inspired Deep Network for Image Compressive Sensing , 2017, 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.

[57]  Richard G. Baraniuk,et al.  prDeep: Robust Phase Retrieval with a Flexible Deep Network , 2018, ICML.

[58]  Jonas Adler,et al.  Learned Primal-Dual Reconstruction , 2017, IEEE Transactions on Medical Imaging.

[59]  Laura Waller,et al.  Learning-Based Image Reconstruction via Parallel Proximal Algorithm , 2018, IEEE Signal Processing Letters.

[60]  Aggelos K. Katsaggelos,et al.  Using Deep Neural Networks for Inverse Problems in Imaging: Beyond Analytical Methods , 2018, IEEE Signal Processing Magazine.

[61]  Daniel Rueckert,et al.  A Deep Cascade of Convolutional Neural Networks for Dynamic MR Image Reconstruction , 2017, IEEE Transactions on Medical Imaging.

[62]  Charles A. Bouman,et al.  Plug-and-Play Unplugged: Optimization Free Reconstruction using Consensus Equilibrium , 2017, SIAM J. Imaging Sci..

[63]  Lei Tian,et al.  High-throughput intensity diffraction tomography with a computational microscope. , 2018, Biomedical optics express.

[64]  Kaiming He,et al.  Group Normalization , 2018, ECCV.

[65]  Michael Unser,et al.  CNN-Based Projected Gradient Descent for Consistent CT Image Reconstruction , 2017, IEEE Transactions on Medical Imaging.

[66]  Jorge Nocedal,et al.  Optimization Methods for Large-Scale Machine Learning , 2016, SIAM Rev..

[67]  Jong Chul Ye,et al.  Framing U-Net via Deep Convolutional Framelets: Application to Sparse-View CT , 2017, IEEE Transactions on Medical Imaging.

[68]  Jeffrey A. Fessler,et al.  BCD-Net for Low-dose CT Reconstruction: Acceleration, Convergence, and Generalization , 2019, MICCAI.

[69]  Michael Elad,et al.  DeepRED: Deep Image Prior Powered by RED , 2019, ICCV 2019.

[70]  Ulugbek S. Kamilov,et al.  Online Regularization by Denoising with Applications to Phase Retrieval , 2019, 2019 IEEE/CVF International Conference on Computer Vision Workshop (ICCVW).

[71]  Philip Schniter,et al.  Regularization by Denoising: Clarifications and New Interpretations , 2018, IEEE Transactions on Computational Imaging.

[72]  Brendt Wohlberg,et al.  An Online Plug-and-Play Algorithm for Regularized Image Reconstruction , 2018, IEEE Transactions on Computational Imaging.

[73]  Mathews Jacob,et al.  Dynamic MRI using model‐based deep learning and SToRM priors: MoDL‐SToRM , 2019, Magnetic resonance in medicine.

[74]  Mathews Jacob,et al.  MoDL: Model-Based Deep Learning Architecture for Inverse Problems , 2017, IEEE Transactions on Medical Imaging.

[75]  Xiaohan Chen,et al.  Plug-and-Play Methods Provably Converge with Properly Trained Denoisers , 2019, ICML.

[76]  Michael Rabbat,et al.  fastMRI: A Publicly Available Raw k-Space and DICOM Dataset of Knee Images for Accelerated MR Image Reconstruction Using Machine Learning. , 2020, Radiology. Artificial intelligence.

[77]  Brendt Wohlberg,et al.  Provable Convergence of Plug-and-Play Priors With MMSE Denoisers , 2020, IEEE Signal Processing Letters.

[78]  Ulugbek S. Kamilov,et al.  RARE: Image Reconstruction Using Deep Priors Learned Without Groundtruth , 2019, IEEE Journal of Selected Topics in Signal Processing.

[79]  M. Jacob,et al.  J-MoDL: Joint Model-Based Deep Learning for Optimized Sampling and Reconstruction , 2020, IEEE Journal of Selected Topics in Signal Processing.

[80]  Yu Sun,et al.  Block Coordinate Regularization by Denoising , 2019, IEEE Transactions on Computational Imaging.

[81]  Jiaming Liu,et al.  SIMBA: Scalable Inversion in Optical Tomography Using Deep Denoising Priors , 2019, IEEE Journal of Selected Topics in Signal Processing.

[82]  Charles A. Bouman,et al.  Plug-and-Play Methods for Magnetic Resonance Imaging: Using Denoisers for Image Recovery , 2019, IEEE Signal Processing Magazine.

[83]  Steen Moeller,et al.  Self‐supervised learning of physics‐guided reconstruction neural networks without fully sampled reference data , 2019, Magnetic resonance in medicine.

[84]  Suvrit Sra,et al.  Why Gradient Clipping Accelerates Training: A Theoretical Justification for Adaptivity , 2019, ICLR.

[85]  Zongben Xu,et al.  ADMM-CSNet: A Deep Learning Approach for Image Compressive Sensing , 2020, IEEE Transactions on Pattern Analysis and Machine Intelligence.

[86]  Luc Van Gool,et al.  Deep Unfolding Network for Image Super-Resolution , 2020, 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[87]  Mingyi Hong,et al.  Dense Recurrent Neural Networks for Accelerated MRI: History-Cognizant Unrolling of Optimization Algorithms , 2020, IEEE Journal of Selected Topics in Signal Processing.

[88]  Alexandros G. Dimakis,et al.  Deep Learning Techniques for Inverse Problems in Imaging , 2020, IEEE Journal on Selected Areas in Information Theory.

[89]  Jeffrey A. Fessler,et al.  Momentum-Net: Fast and Convergent Iterative Neural Network for Inverse Problems , 2019, IEEE Transactions on Pattern Analysis and Machine Intelligence.

[90]  Laura Waller,et al.  Memory-Efficient Learning for Large-Scale Computational Imaging , 2019, IEEE Transactions on Computational Imaging.

[91]  Fanhua Shang,et al.  A Novel Learned Primal-Dual Network for Image Compressive Sensing , 2021, IEEE Access.