BOOSTING ENCODER-DECODER CNN FOR INVERSE PROBLEMS

1 INTRODUCTION . Suppose that the measurement x is corrupted with additive noises : x = µ+w , w ∼ N ( 0 , σ2I ) , ( 1 ) where µ ∈ Rn denotes the unknown mean vector , and σ2 is the variance . Consider a deep neural network model F ( x ) : = FΘ ( x ) with the weight Θ , which is trained with the data x producing an estimate µ̂ = FΘ ( x ) . Then , our goal is to estimate the prediction error : Err ( µ̂ ) = Ex , x∗‖x∗ − µ̂‖2 = Ex , x∗‖x∗ − F ( x ) ‖2 , ( 2 ) which quantifies how well µ̂ can predict a test data x∗ , independently drawn from the same distribution of x ( Efron , 2004 ; Tibshirani & Rosset , 2018 ) . The problem of estimating the prediction error is closely related to the generalizability of neural network ( Anthony & Bartlett , 2009 ) . Moreover , this problem is tightly linked to the classical approaches for model order selection in statistics literature ( Stoica & Selen , 2004 ) . One of the most investigated statistical theories to address this question is so-called covariance penalties approaches such as Mallow ’ s Cp ( Mallows , 1973 ) , Akaike ’ s information criterion ( AIC ) ( Akaike , 1974 ) , Stein ’ s unbiased risk estimate ( SURE ) ( Donoho & Johnstone , 1995 ) , etc . This paper is particularly interested in estimating the prediction error of encoder-decoder convolutional neural networks ( E-D CNNs ) such as U-Net ( Ronneberger et al. , 2015 ; Han & Ye , 2018 ; Ye & Sung , 2019 ) . E-D CNNs have been extensively used for various inverse problems ( Ye et al. , 2018 ) . Recent theoretical results showed that , thanks to the ReLU nonlinearities , the input space is partitioned into non-overlapping regions so that input images in each region share the same linear representation , but not across different partitions ( Ye & Sung , 2019 ) . One of the most important contributions of this paper is that this property can be exploited to derive a simplified form of prediction error estimator that can be used for neural network training . In particular , we provide a close form expression for the divergence term in the SURE estimator , which suggests that the divergence term can be neglected if a proper batch normalization is used . Furthermore , the simplified form of the SURE estimate for the prediction error leads to a bootstrap and aggregation scheme to prevent the CNN from converging to a trivial solution so that it can improve the neural network performance . The applicability of the new method is demonstrated using various inverse problems such as accelerated MRI , energy-dispersive X-ray spectroscopic imaging , deconvolution , and etc , which clearly show that our method can significantly improve the image quality compared to the existing approaches . 2 RELATED WORKS . 2.1 STEIN RISK UNBIASED ESTIMATION FOR PREDICTION ERROR . The Stein Risk Unbiased Estimator ( SURE ) for the prediction error can be represented by ( Donoho & Johnstone , 1995 ) : ŜURE ( x ) : = ‖x− F ( x ) ‖2 + 2σ2div { F ( x ) } , ( 3 ) where div ( · ) denotes the divergence . Recently , the authors in ( Soltanayev & Chun , 2018 ) employed the SURE for CNN-based image denoising without references by minimizing the following loss : ` SURE ( Θ ) : = 1 P P∑ i=1 ‖x ( i ) − FΘ ( x ( i ) ) ‖2 + 2σ2div { FΘ ( x ( i ) ) } . ( 4 ) where { x ( i ) } Pi=1 denotes the training samples . Although the application of SURE for unsupervised denoising is an important advance in theory , there are several practical limitations . In particular , due to the difficulty of calculating the divergence term , the authors relied on MonteCarlo SURE ( Ramani et al. , 2008 ) which calculates the divergence term using MonteCarlo simulation . This introduces additional hyperparameters , on which the final results critically depend . 2.2 PIECEWISE LINEAR REPRESENTATION BY E-D CNN WITH RELUS . Recently , the authors in ( Ye & Sung , 2019 ) showed that the output of a CNN composed of κ layer of encoder and decoder without skipped connection can be represented by the following basis-like representation : y : = F ( x ) = ∑ i 〈bi ( x ) , x〉b̃i ( x ) = B̃ ( x ) B ( x ) > x , ( 5 ) where bi ( x ) and b̃i ( x ) denote the i-th column of the following frame basis and its dual : B ( x ) = E1Σ1 ( x ) E2 · · ·Σκ−1 ( x ) Eκ , ( 6 ) B̃ ( x ) = D1Σ̃1 ( x ) D2 · · · Σ̃κ−1 ( x ) Dκ , ( 7 ) where Σl ( x ) and Σ̃l ( x ) denote the diagonal matrix with 0 and 1 values that are determined by the ReLU output in the previous convolution steps . El andDl refer to the encoder and decoder matrices , respectively , which are determined by pooling ( resp . unpooling ) operation and convolution filters ( Ye & Sung , 2019 ) . For the case of encoder-decoder CNN with the skipped connection , similar basis expression can be obtained Ye & Sung ( 2019 ) . Note that the expression ( 5 ) is nonlinear due to the dependency on the input signal x . Moreover , as the ReLU nonlinearity is applied after the convolution operation , the on-and-off activation pattern of each ReLU determines a binary partition of the feature space across the hyperplane that is determined by the convolution . Accordingly , one of the most important observation in Ye & Sung ( 2019 ) is that the input space X is partitioned into multiple non-overlapping regions so that input images for each region share the same linear representation , but overall representation is still non-linear . This implies that two different input images may be automatically switched to two distinct linear representations that are different from each other . 3 MAIN CONTRIBUTION . In ( 3 ) and ( 4 ) , the calculation of the divergence term is not trivial for general neural networks . This is why the authors in ( Soltanayev & Chun , 2018 ) employed the MonteCarlo SURE . Thanks to the basis-like expression in ( 5 ) , here we show that there exists a simple explicit form of the divergence term for the case of E-D CNNs . Then , the batch normalization is shown to make the divergence term trivial . 3.1 DIVERGENCE SIMPLIFICATION IN SURE . In Proposition 6 of ( Ye & Sung , 2019 ) , the authors show that ∂F ( x ) /∂x = B̃ ( x ) B ( x ) > for the case of E-D CNN with ReLUs . Accordingly , we have div { F ( x ) } = Tr ( ∂F ( x ) ∂x ) = Tr ( B̃ ( x ) B ( x ) > ) = Tr ( B ( x ) > B̃ ( x ) ) = ∑ i 〈bi ( x ) , b̃i ( x ) 〉 , ( 8 ) where Tr ( A ) denotes the trace ofA . This close form representation ( 8 ) leads to a further simplification of divergence term by exploiting the property of the batch normalization . Recall that batch normalization has been extensively used to make the training stable ( Ioffe & Szegedy , 2015 ; Hoffer et al. , 2018 ; Cho & Lee , 2017 ; Miyato et al. , 2018 ; Ulyanov et al. , 2016 ) . It has been consistently shown that the batch normalization is closely related to the norm of the Jacobian matrix ∂F ( x ) /∂x , which is equal to B̃ ( x ) B ( x ) > . For example , in their original paper ( Ioffe & Szegedy , 2015 ) , the authors conjectured that “ Batch Normalization may lead the layer Jacobians to have singular values close to 1 , which is known to be beneficial for training ” . By extending the idea in ( Ioffe & Szegedy , 2015 ) to multiple layers , the batch normalization can be understood as to make the covariance of the network output and input similar . For example , for the uncorrelated input with Cov [ x ( i ) ] = σ2I , the batch normalization may result in Cov [ F ( x ( i ) ) ] ' σ2I . Furthermore , for sufficiently smaller σ , we have σ2I ' Cov [ F ( x ( i ) ) ] = Cov [ B̃ ( x ( i ) ) B ( x ( i ) ) > x ( i ) x > ( i ) B ( x ( i ) ) B̃ ( x ( i ) ) > ] = B̃ ( x ( i ) ) B ( x ( i ) ) > Cov [ x ( i ) x > ( i ) ] B ( x ( i ) ) B̃ ( x ( i ) ) > = σ2B̃ ( x ( i ) ) B ( x ( i ) ) > B ( x ( i ) ) B̃ ( x ( i ) ) > , because the corresponding piecewise linear representation does not change for the small perturbation of the input ( Ye & Sung , 2019 ) . Therefore , we have B̃ ( x ( i ) ) B ( x ( i ) ) > ' I , since B̃ ( x ( i ) ) B ( x ( i ) ) > is a square matrix . This suggests that divEL { F ( x ( i ) ) } = ∑ i 〈bi ( x ( i ) ) , b̃i ( x ( i ) ) 〉 = Tr ( B̃ ( x ( i ) ) B ( x ( i ) ) > ) ' n , ( 9 ) where n is the dimension of the input signal . As the resulting divergence term is just a constant , the contribution of the divergence term is considered trivial and can be neglected . 3.2 BAGGING ESTIMATOR . Another important drawback of SURE-based denoising network ( 4 ) is that it is difficult to prevent the network from learning a trivial identity mapping . More specifically , if FΘ ( x ) = x , the cost function in ( 4 ) becomes zero . One way to avoid this trivial solution is to guarantee that the divergence term at the optimal network parameter should be negative . However , given that the divergence term comes from the degree of the freedom ( Efron , 2004 ) and the amount of excess optimism in estimating the prediction error ( Tibshirani & Rosset , 2018 ) , enforcing negative value may be unnatural . Moreover , with the divergence simplification in ( 9 ) , the cost function ( 4 ) can be simplified as ` SURE−ED ( Θ ) : = 1 P P∑ i=1 ‖x ( i ) − FΘ ( x ( i ) ) ‖2 + 2σ2n ( 10 ) so that it is much apparent that a trivial solution for ( 10 ) is FΘ ( x ) = x . To prevent the neural network from converging to this trivial solution , here we propose a bootstrap aggregation ( bagging ) scheme ( Breiman , 1996 ) . Bagging is a classical machine learning technique which uses bootstrap sampling and aggregation of the results to reduce the variance and to improve the accuracy of the base learner . The rationale for bagging is that it may be easier to train several simple weak learners and combine them into a more complex learner than to learn a single strong learner . More specifically , we use the following bagging estimator : µ̃ ( x ) = K∑ k=1 wkFΘ ( Lkx ) , ( 11 ) and the corresponding prediction error estimate : Err ( µ̃ ) : = Ex , x∗‖x∗ − µ̃‖2 = Ex , x∗‖x∗ − K∑ k=1 wkFΘ ( Lkx ) ‖2 , where { Lk } Kk=1 denotes the diagonal matrix whose diagonal elements are either 0 or 1 depending on the bootstrap subsampling patterns , and { wk } Kk=1 is the corresponding weights . The corresponding empirical prediction error estimate is then given by ` ( Θ , { wk } ) : = 1 P P∑ i=1 ‖x ( i ) − K∑ k=1 wkFΘ ( Lkx ( i ) ) ‖2 , ( 12 ) Note that x ( i ) 6= Lkx ( i ) due to the subsampling patterns so that the optimal neural network FΘ that minimizes ( 12 ) can not be the trivial identity mapping . Furthermore , the following proposition shows that the prediction error for the bagging estimator ( 11 ) is not bigger than the average of the individual error estimate . Proposition 1 . Let wk ≥ 0 and ∑ k wk = 1 . Then , we have 1 P P∑ i=1 ∑ k wk‖x ( i ) − F ( Lkx ( i ) ) ‖2 ≥ 1 P P∑ i=1 ∥∥∥x ( i ) −∑wkF ( Lkx ( i ) ) ∥∥∥2 . Proof . For any x ∈ { x ( i ) } Pi=1 , we have∑ k wk‖x− F ( Lkx ) ‖2 = x > x− 2x > ∑ k wkF ( Lkx ) + ∑ k wk‖F ( Lkx ) ‖2 ≥ x > x− 2x > ∑ k wkF ( Lkx ) + ‖ ∑ wkF ( Lkx ) ‖2 = ∥∥∥x−∑wkF ( Lkx ) ∥∥∥2 , where we use the Jensen ’ s inequality for the inequality . By summing up for all training data { x ( i ) } Pi=1 , we conclude the proof . The gap increases more when the F ( Lkx ) provides diverse output for each realization of the index Lk ( Breiman , 1996 ) . In fact , the authors in ( Ye & Sung , 2019 ) showed that the input space results in non-overlapping partitions with different linear representations so that by changing the subsampling pattern Lk , the distinct representation may be selected , which can make the corresponding bagging estimate more accurate .

Summary: The authors address the problem of efficiently employing the SURE estimator as a network training regularizer. They show that for CNN autoencoders this can be efficiently computed. Their other contribution is a bagging/boosting technique which is proved to avoid trivial solutions. The proposed architecture, motivated by the theoretical statements, is shown to outperform classic and 2019 state of the art image reconstruction algorithms in MRI and EDX.