The Trainable Alternating Gradient Shrinkage method

Transcription

1 The Trainable Alternating Gradient Shrinkage method Carlos Malanche, supervised by Ha Nguyen April 24, 2018 LIB (prof. Michael Unser), EPFL

2 Quick example 1

3 The reconstruction problem Challenges in image reconstruction Restore blured and noisy pictures with acceptable quality Input signal in the order of megapixels. Not too expensive (neither in memory, neither in time). Hopefully flexible towards different corruption conditions 2

4 The reconstruction problem Nevertheless, solutions for this problem are already here with complex convolutional neural networks, giving an amazing performance. The goal here is to develop a method that mixes theoretical knowledge in inverse problems with machine learning. 3

5 Mathematical model

6 The Augmented Lagrangian Model to solve: y = }{{} Measure L. sys. {}}{ H x }{{} Signal + n }{{} Noise The MAP estimator for this model is written as:, x, n, y R n, H R n n ˆx MAP (y) = argmin{ 1 2 Hs y 2 +σ 2 φ(ls)} s We turn the problem into a two variable problem by defining the innovation as u = Ls 4

7 The Augmented Lagrangian With that new variable, we can write the Augmented Lagrangian of this system: L A (s, u, α) = 1 2 Hs y 2 +σ 2 φ(u) α, Ls u + µ Ls u 2 2 Where α are the Lagrange multipliers and µ is a penalty parameter, to who we ll come back later. 5

8 Proposed solution We can find an iterative solution for each of our three variables. Taking as an advantage the linearity of the operator H, the update rule for s is found by taking a derivative and equating to zero, which yields: s (k+1) = argmin s L A (s, u (k), α (k) ) = (H T H + µl T L) 1 (H T y + µl T (u (k) + α(k) µ )) 6

9 Proposed solution We update the Lagrange multipliers before updating the innovation variable as it will make computations easier after. α (k+1) = α (k) µ(ls (k+1) u (k) ) We notice that the convergence of these multipliers comes from the convergence on the other two variables 7

10 Proposed solution With a bit of algebra over the Augmented Lagrangian, we can see that the update rule for the innovation reduces to solving a proximal u (k+1) = argmin u L A (s (k+1), u, α (k+1) ) = prox σ 2 µ φ(ls(k+1) 1 µ α(k+1) ) We will rename this proximal as a shrinkage operator, denoted by T ( ). 8

11 Proposed solution Putting our three update rules together, we arrive to the Alternating Direction Method of Multipliers s (k+1) = (H T H + µl T L) 1 (H T y + µl T (u (k) + α(k) µ )) α (k+1) = α (k) µ(ls (k+1) u (k) ) u (k+1) = T (Ls (k+1) 1 µ α(k+1) ) We will call a k-stage method that which final estimation is given by ˆx := s (k) 9

12 The learning part

13 Shrinkage Interpolation As the shrinkage function might not even have an analytic expression (it can be difficult to find or very expensive to approximate), we consider that we have a cubic B-Spline interpolation of this function T (v) = i=m i= M c i ψ( v i=m i) = i= M c i ψ i (v); When this function is applied to tensors of rank higher than 0, it is simply applied element-wise. 10

14 Gradient Descent Using a Squared Loss function, we want c to be such that it minimizes it. J(c, µ) = 1 2 ˆx(c, µ, y train) x train 2 Using Gradient Descent c (r+1) = c (r) η c c J(c (r), µ) 11

15 The c explicit expression This gradient can be expressed as c J = z (k) c LA T (ˆx x), with A = H T H + µl T L := c α (k) + µ c u (k) holding the following iterative property and z (k) c Where z (k) c = z (k 1) c (I D (k) + µla T L T (2D (k) I )) + µψ (k) Ψ (k) i,j := D (k) i,j := z (0) c = 0 (k) T (v j ) c i (k) T (v j ) v (k) i = ψ i (v (k) j ) = δ i,j T (v k i ) 12

16 Gradient Descent We think also of a gradient descent approach to estimate the penalty parameter µ. µ (r+1) = µ (r) η µ µ J(c, µ (r) ) This penalty parameter will allow us to test robustness of a model against blur kernels and noise that it was not trained to restore. 13

17 Implementation

18 All done in Matlab The software Matlab was used as it has already been optimized to perform matrix operations (and it also provides easy access to GPU computation when it is possible). I kept object oriented programming in mind to make benchmarking, debugging and general use easier and faster for people that is not familiar with the code. 14

19 Small code diagram 15

20 Particular configuration As a simple first choice, we decided to take as filter L a stack of two whitening filters, the discrete forward derivatives in the canonical directions of the picture. 16

21 Results

22 System specs All the tests that are in this section were carried out in a Dell laptop, with high end components but still a consumer grade laptop. Processor Frequency RAM GPU Intel Core i7-6700hq, 4 cores, 8 threads 2.6GHz base (3.5GHz turbo) 8GB GTX960m, 640 CUDA cores, 1096MHz, 4GB GDDR5 17

23 Training time Evolution of PSNR in a training set of grayscale pictures. 18

24 Two or one shrinkage spline In theory, the advantage of using a shrinkage spline per direction of the gradient, statistically, has no relevance. The graph to the right shows this, as in the BSD68 benchmark, we see just a marginal difference between using two splines or one. 19

25 The shrinkage spline We verify the behavior of the shrinkage training in three different scenarios: Noise-only case Blur-only case Noise and Blur case. For the following tests, µ was fixed to

26 Noise-only case Noise σ = 25 Result obtained with a batch-sgd of 3 images per iteration, 100 iterations. Training time: 66 seconds 21

27 Blur-only case Blur σ = 1.6, kernel Result obtained with a batch-sgd of 3 images per iteration, 100 iterations. Training time: 64 seconds 22

28 Noise and Blur case Noise σ = 25 Blur σ = 1.6, kernel Result obtained with a batch-sgd of 3 images per iteration, 100 iterations. Training time: 65 seconds 23

29 Against gradient based methods σ TAGS CSF 5 pw TV-Opt We compare with the CSF 5 pw and the Anisotropic TV. CSF 5 pw contains 540 parameters to be trained, against 54 of our model. 24

30 The State-of-the-Art σ TAGS CSF 5 pw TV-Opt TNRD DnCNN BM3D DnCNN and TNRD take around 20+ hours of training, against 12 minutes of our model. 25

31 Noise Robustness In the beginning, we discarded the dependence on the factor σ 2 µ 1 (which can, in fact, be f (σ)µ 1 ) in our proximal interpolation. To add more value to our model, we wanted to test if, across different noise levels, keeping this value as a constant would improve the performance of the model running in levels of noise different from the training noise. 26

32 Noise Robustness Similar behavior across noise levels Shrinkages keeping µ constant 27

33 Noise Robustness Train models for odd σ [15, 45] Declare as σ 0, µ 0 the parameters for a chosen base model TAGS 0. Heuristically scale the base model as µ = σσ 0 µ 1 0 and benchmark with noise = σ. Learn only µ in a training set with noise σ, starting with model TAGS 0. 28

34 Noise Robustness Base model with σ = 25, µ = 1.9 µ training time 30 seconds. 29

35 Noise Robustness The SSIM tells a similar story. 30

36 Against Anisotropic TV [2] 31

37 Some Examples

38 The parrot No blur, σ =

39 The parrot TAGS TV 33

40 Slight detail enhancement TAGS TV The original image was corrupted with σ = 45, no blur. 34

41 Slight detail enhancement The original image was corrupted with σ = 45, box blur

42 Slight detail enhancement On the left side, TAGS method. Right side, TV. The denoising was calibrated such that the outside of the spiral was black in both cases. 36

43 Thank you for your attention

44 Bibliography Uwe Schmidt & Stefan Roth, Shrinkage Fields for Effective Image Restoration, The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2014, pp Ulugbek S. Kamilov. A Parallel Proximal Algorithm for Anisotropic Total Variation Minimization, IEEE Transactions on Image Processing ( Volume: 26, Issue: 2, Feb. 2017, pages ) C.J.Schuler, H.C.Burger, S.Harmeling, andb.scholkopf. A Machine Learning Approach for Non-blind Image Deconvolution. CVPR, pages , 2013 Neal Parikh and Stephen Boyd, Proximal Algorithms Foundations and Trends R in Optimization Vol. 1, No. 3 (2013) c 2013 N. 37