Loop Unrolled Shallow Equilibrium Regularizer (LUSER) - A Memory-Efficient Inverse Problem Solver

In inverse problems we aim to reconstruct some underlying signal of interest from potentially corrupted and often ill-posed measurements. Classical optimization-based techniques proceed by optimizing a data consistency metric together with a regularizer. Current state-of-the-art machine learning approaches draw inspiration from such techniques by unrolling the iterative updates for an optimization-based solver and then learning a regularizer from data. This loop unrolling (LU) method has shown tremendous success, but often requires a deep model for the best performance leading to high memory costs during training. Thus, to address the balance between computation cost and network expressiveness, we propose an LU algorithm with shallow equilibrium regularizers (LUSER). These implicit models are as expressive as deeper convolutional networks, but far more memory efficient during training. The proposed method is evaluated on image deblurring, computed tomography (CT), as well as single-coil Magnetic Resonance Imaging (MRI) tasks and shows similar, or even better, performance while requiring up to 8 times less computational resources during training when compared against a more typical LU architecture with feedforward convolutional regularizers.

[1]  K. Kothari,et al.  Conditional Injective Flows for Bayesian Imaging , 2022, IEEE Transactions on Computational Imaging.

[2]  M. Soltanolkotabi,et al.  HUMUS-Net: Hybrid unrolled multi-scale network architecture for accelerated MRI reconstruction , 2022, NeurIPS.

[3]  P. Maass,et al.  LoDoPaB-CT, a benchmark dataset for low-dose computed tomography reconstruction , 2021, Scientific data.

[4]  R. Willett,et al.  Deep Equilibrium Architectures for Inverse Problems in Imaging , 2021, IEEE Transactions on Computational Imaging.

[5]  Vladlen Koltun,et al.  Multiscale Deep Equilibrium Models , 2020, NeurIPS.

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

[7]  Patrick Putzky,et al.  Invert to Learn to Invert , 2019, NeurIPS.

[8]  J. Z. Kolter,et al.  Deep Equilibrium Models , 2019, NeurIPS.

[9]  Laurent El Ghaoui,et al.  Implicit Deep Learning , 2019, SIAM J. Math. Data Sci..

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

[11]  Christopher Ré,et al.  Low-Memory Neural Network Training: A Technical Report , 2019, ArXiv.

[12]  Pascal Vincent,et al.  fastMRI: An Open Dataset and Benchmarks for Accelerated MRI , 2018, ArXiv.

[13]  Yuichi Yoshida,et al.  Spectral Normalization for Generative Adversarial Networks , 2018, ICLR.

[14]  Xiaogang Wang,et al.  Deep Learning Face Attributes in the Wild , 2014, 2015 IEEE International Conference on Computer Vision (ICCV).

[15]  Homer F. Walker,et al.  Anderson Acceleration for Fixed-Point Iterations , 2011, SIAM J. Numer. Anal..

[16]  Christian Oestreicher,et al.  A history of chaos theory , 2007, Dialogues in clinical neuroscience.

[17]  R. Aster Fundamentals of Seismic Wave Propagation , 2005 .

[18]  Samy Wu Fung,et al.  Fixed Point Networks: Implicit Depth Models with Jacobian-Free Backprop , 2021, ArXiv.