Unrolling: A principled method to develop deep neural networks

1. Unrolling:A principled method to develop deep neural networks Chris Metzler, Ali Mousavi, Richard Baraniuk Rice University

2. Solving Imaging Inverse Problems Traditional Methods (e.g., ADMM) Deep Neural Networks (DNNs) Blackbox methods which learn complex functions What happens between layers is an open research question Priors are learned with training data Demonstrate state-of-the-art performance on a variety of imaging tasks • Have well-understood behavior • E.g., iterations are refining an estimate • Convergence guarantees • Interpretable priors • Are easy to apply

3. Solving Imaging Inverse Problems Traditional Methods (e.g., ADMM) Deep Neural Networks (DNNs) Blackbox methods which learn complex functions What happens between layers is an open research question Priors are learned with training data Demonstrate state-of-the-art performance on a variety of imaging tasks • Have well-understood behavior • E.g., iterations are refining an estimate • Convergence guarantees • Interpretable priors • Are easy to apply Is there a way to combine the strengths of both?

4. Describe unrolling; a process to turn an iterative algorithm into a deep neural network • Can use training data • Is interpretable • Maintains convergence and performance guarantees • Apply unrolling to the Denoising-based AMP algorithm to solve the compressive imaging problem This talk

5. Unrolling an algorithm

6. Example: Iterative Shrinkage/Thresholding

7. The unrolling process [Gregor and Lecun 2010, Kamilov and Mansour 2015, Borgerding and Schniter 2016]

8. The unrolling process [Gregor and Lecun 2010, Kamilov and Mansour 2015, Borgerding and Schniter 2016]

9. has free parameters. A and AH can also be treated as parameters to learn • Feed training data, i.e., (x, y) pairs, through the network, calculate errors, backpropagate Training the Network Error Gradients [Gregor and Lecun 2010, Kamilov and Mansour 2015, Borgerding and Schniter 2016]

10. Performance [Borgerding and Schniter 2016]

11. Unrolling Denoising-based Approximate Message Passing

12. The Compressive Imaging Problem Target (x) Sensor Array (y)

13. The Compressive Imaging Problem Target (x) Sensor Array (y)

14. The Compressive Imaging Problem Target (x) Sensor Array (y)

15. The Compressive Imaging Problem Target (x) Sensor Array (y)

16. The Compressive Imaging Problem Target (x) Sensor Array (y)

17. The Compressive Imaging Problem Target (x) Sensor Array (y) • Every measurement has a cost; $, time, bandwidth, etc. • Compressed sensing enables us to recover x with fewer measurements

18. Compressive Imaging Applications Higher Resolution Seismic Imaging Higher Resolution Synthetic Aperture Radar Faster Medical Imaging Low Cost High Speed Imaging Low Cost Infrared/Hyperspectral Imaging Veeraraghavan et al. 2011 Baraniuk 2007

19. Compressive Imaging Problem Set of all natural images

20. Traditional Methods vs DNNs Traditional Methods (D-AMP) DNNs (Ali’s talk) Black box Why does it work? When will it fail? Less accurate Much faster Learns priors with training data • Well understood behavior • Recovery guarantees • Noise sensitivity analysis • More accurate • Slower • Uses denoisers to impose priors Learned D-AMP gets the strengths of both

21. Denoising-based AMP M. et al. 2016

22. Neural Network (NN) Recovery Methods Mousavi and Baraniuk. 2017

23. Unroll D-AMP to form a NN Learned D-AMP • Efficiently train a 200 layer network • Demonstrate proposed algorithm is fast, flexible, and effective • >10x faster than D-AMP • Handles arbitrary right-rotationally-invariant matrices • State-of-the-art recovery accuracy Our contribution

24. Unroll D-AMP

25. Unroll D-AMP

26. Unroll D-AMP Need a denoiser that can be trained

27. Deep convolution networks are now state-of-the-art image denoisers [Zhang et al. 2016] CNN-based Denoiser

28. Place a DNN-based denoiser unrolled D-AMP to form Learned D-AMP • We need to train this huge network Learned D-AMP

29. The Challenge: Network is 200 layers deep and has over 6 million free parameters • Solution: [1] proved that for D-AMP layer-by-layer training is optimal Training [1] M. et al. "From denoising to compressed sensing." IEEE Transactions on Information Theory 62.9 (2016): 5117-5144.

30. 400 training images were broken into ~300000 50x50 patches • Patches were randomly flipped and/or rotated during training • 10 identical networks were trained to remove additive white Gaussian noise at 10 different noise levels • Stopped training when validation error stopped improving; generally 10 to 20 epochs • Trained on 3584 core Titan X GPU for between 3 and 5 hours Training Details

31. High-Res 10% Gaussian Matrix BM3D-AMP (199 sec) Learned DVAMP (62 sec)

32. Performance: Gaussian Measurements 1 dB Better than BM3D-AMP

33. Runtime: Gaussian Measurements >10x Faster than BM3D-AMP

34. Computational Complexity Computational complexity is dominated by matrix multiplies

35. Performance: Coded Diffraction Measurements 2 dB Better than BM3D-AMP

36. Runtime: Coded Diffraction Measurements >17x Faster than BM3D-AMP (at 128x128)

37. High-Res 5% Coded Diffraction Matrix TVAL3 (6.85 sec) Learned DVAMP (1.22 sec)

38. High-Res 5% Coded Diffraction Matrix BM3D-AMP (75.04 sec) Learned DVAMP (1.22 sec)

39. Unrolling turns an iterative algorithm into a deep neural net • Illustrated here with D-AMP  Learned D-AMP • Learned-DAMP is fast, flexible, and effective • >10x faster than D-AMP • Handles arbitrary right-rotationally-invariant matrices • State-of-the-art recovery accuracy Summary

40. NSF GRFP Acknowledgments