Title: IMPROVING ENCODER-DECODER CNN FOR INVERSE PROBLEMS

Abstract: Encoder-decoder convolutional neural networks (CNN) have been extensively used for various inverse problems. However, their prediction error for unseen test data is difficult to estimate a priori, since the neural networks are trained using only selected data and their architectures are largely considered blackboxes. This poses a fundamental challenge in improving the performance of neural networks. Recently, it was shown that Stein's unbiased risk estimator (SURE) can be used as an unbiased estimator of the prediction error for denoising problems. However, the computation of the divergence term in SURE is difficult to implement in a neural network framework, and the condition to avoid trivial identity mapping is not well defined. In this paper, inspired by the finding that an encoderdecoder CNN can be expressed as a piecewise linear representation, we provide a close form expression of the unbiased estimator for the prediction error. The close form representation leads to a novel bootstrap and aggregation scheme to prevent a neural network from converging to an identity mapping so that it can enhance the performance. Experimental results show that the proposed algorithm provides consistent improvement in various inverse problems.

1. INTRODUCTION:

Suppose that the measurement x is corrupted with additive noises:

EQUATION

where µ ∈ R n 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:

EQUATION

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) :

EQUATION

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:

EQUATION

)

where {x (i) } P i=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:

EQUATION

where b i (x) and bi (x) denote the i-th column of the following frame basis and its dual:

EQUATION

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. E l and D l 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

EQUATION

where Tr(A) denotes the trace of A.

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) ] = σ 2 I, the batch normalization may result in Cov[F (x (i) )] σ 2 I. Furthermore, for sufficiently smaller σ, we have

σ 2 I 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) ) = σ 2 B(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

EQUATION

)

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

SU RE-ED (Θ) := 1 P P i=1 x (i) -F Θ (x (i) ) 2 + 2σ 2 n (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:

EQUATION

and the corresponding prediction error estimate:

Err( μ) := E x,x * x * -µ 2 = E x,x * x * - K k=1 w k F Θ (L k x) 2 ,

where {L k } K k=1 denotes the diagonal matrix whose diagonal elements are either 0 or 1 depending on the bootstrap subsampling patterns, and {w k } K k=1 is the corresponding weights. The corresponding empirical prediction error estimate is then given by

EQUATION

)

Note that x (i) = L k x (i) due to the subsampling patterns so that the optimal neural network F Θ that minimizes (12) cannot 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 w k ≥ 0 and k w k = 1. Then, we have

1 P P i=1 k w k x (i) -F (L k x (i) ) 2 ≥ 1 P P i=1 x (i) - w k F (L k x (i) ) 2 .

Proof. For any x ∈ {x (i) } P i=1 , we have

k w k x -F (L k x) 2 = x x -2x k w k F (L k x) + k w k F (L k x) 2 ≥ x x -2x k w k F (L k x) + w k F (L k x) 2 = x - w k F (L k x) 2 ,

where we use the Jensen's inequality for the inequality. By summing up for all training data {x (i) } P i=1 , we conclude the proof.

The gap increases more when the F (L k x) provides diverse output for each realization of the index L k (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 L k , the distinct representation may be selected, which can make the corresponding bagging estimate more accurate.

4. IMPLEMENTATION DETAILS:

The final loss function for our bagging estimator is given by

EQUATION

)

Figure 1 : Overall network architecture of the proposed method. The acquired data is split into multiple subset using bootstrap subsampling. The data is then processed using a regression network followed by an aggregation module using an attention network.

where we also parameterize the weighting using a neural network with parameter Ξ. Although the standard way of aggregation in bagging estimator is a simple average of the overall results of the regression network, this may not be the best method when the number of bootstrap subsampling is limited. Instead, we propose a weighted average scheme whose weight is calculated by data attention module so that it efficiently combines all data by adaptively incorporating output from various bootstrap sub-sampling patterns. A schematic diagram of the proposed method is illustrated in Fig. 1 , which consists of three building blocks: bootstrap subsampling, a regression network, and an attention network to calculate the weight parameters.

Figure 2 : Architecture for the regression network.

In particular, U-net (Ronneberger et al., 2015) in Fig. 2 was used as our regression network. The network is composed of four stages with convolution, batch normalization, ReLU, and skip connection with concatenation. Each stage is composed of three 3 × 3 convolution layers followed by batch normalization and ReLU, except for the last layer, which is 1 × 1 convolution layer. The number of convolutional filters increases from 64 in the first stage to 1024 in the final stage. To estimate the weight w k , our attention network consists of two fully connected layer. The input dimension of the attention network is R 1×1×K followed by the average pooling of the concatenated output of regression network. The number of hidden node is 64, and the final dimension of the output is R 1×1×K for aggregation.

The overall network was trained using Adam optimization with the momentum β 1 = 0.9 and β 2 = 0.999. The proposed network was implemented in Python using TensorFlow library and trained using an NVidia GeForce GTX 1080-Ti graphics processing unit.

5. EXPERIMENTAL RESULTS:

Experiments were conducted for various inverse problems such as accelerated magnetic resonance imaging (MRI), energy-dispersive X-ray spectroscopy (EDX) (Solé et al., 2007) , and image superresolution in Appendix.

5.1. ACCELERATED MRI:

In accelerated MRI, the goal is to recover high quality MR images from sparsely sampled k-space data to reduce the acquisition time. This problem has been extensively studied using compressed sensing (Lustig et al., 2007) , but recently deep learning approaches have been the main research interest due to its excellent performance at significantly reduced run time complexity (Hammernik et al., 2018; Han & Ye, 2019) . A standard method for neural network training for accelerated MRI is based on the supervised learning, where the MR images from fully sampled k-space data is used as reference and subsampled and zero-filled k-space data is used as the input for the neural network. Specifically, we use the R = 13.45 subsampled k-space data as neural network input and the goal is to obtain high resolution images from the sub-sampled k-space data. As a baseline algorithm, the state-of-the-art deep learning algorithm called k-space deep learning (Han & Ye, 2019) is employed, which interpolates the missing k-data using a CNN.

In the proposed method, the accelerated k-space data is further subsampled using bootstrap sampling with the subsampling ratio of 91 %, and the number of bootstrap sample set K was set to 10. The same k-space deep learning network (Han & Ye, 2019) was used for our regression network using bootstrap samples. We trained the baseline algorithm and our method under the same conditions, except that the proposed network was trained by minimize the loss ( 13). The initial learning rate was set to 10 -2 , and it was divided by half at every 50 epochs until it reaches to around 10 -4 . The training data was generated from Human Connectome Project (HCP) MR dataset (https://db.humanconnectome.org). Among the 34 subject data sets, 28 subject data sets were used for training and validation. The other subject data sets were used for test. We also provide reconstruction results by GRAPPA (Griswold et al., 2002) , which is a standard k-space interpolation method in MRI. Thanks to the bagging, the proposed method provided nearly perfect reconstruction results compared to other algorithms as shown in Fig. 3 . Moreover, its average peak signal-to-noise ratio (PSNR) and structural similarity (SSIM) index for entire test dataset significantly outperform existing methods as shown in Table 1 . To demonstrate the optimality of the proposed architecture, we performed extensive ablation study.

For the ablation study, we used 13 subject data sets from HCP MR dataset, from which seven subject data sets were used for training and validation and the remaining dataset were used for test. The ablation study was performed by excluding some structures from the network and applying the same training processes. First, our theory says that the performance improvement of our method increases with the number of bootstrap subsamples K. To verify this, we performed comparative studies with different K values such as 1, 2, 4 and 10. As shown in Table 2 , the proposed network provided better reconstruction results as the number of bootstrap sample set K increased. This clearly confirmed the importance of the bootstrap subsampling. Second, our derivation suggests that the batch normalization plays a key role to simplify the divergence term. As shown in Fig. 4 and Table 3 , the lack of batch normalization produces a significantly degraded image with an average PSNR drop from 40.047 dB to 25.290 dB. This verified the important role of batch normalization. We also evaluated the importance of the attention network to calculate the weights. With standard bagging, the final output is a simple average of the entire network output for each bootstrap subsamples without the attention network. As shown in Fig. 4 and Table 3 , the simple average produced blurry results than by the attention network and the PSNR and SSIM values were worse. This confirmed that the important of attention network for the weight calculation. 

5.2. ENERGY-DISPERSIVE X-RAY SPECTROSCOPY IMAGING:

As for another experiment, we use the energy-dispersive X-ray spectroscopy (EDX) (Solé et al., 2007) data set which were measured by scanning transmission electron microscopy. EDX (Solé et al., 2007) is an analytical technique used for spectroscopic characterization of a sample. It relies on the emission of characteristic X-rays from a specimen that is generated by a high-energy beam of electrons. Specifically, the incident beam may eject an electron from the inner shell of atoms. Then, an electron from an outer shell then fills the hole, and the energy difference may be released in the form of an X-ray. As the energies of the X-rays are characteristic of the energy differences that are determined by the atomic structure of the emitting element, EDX is widely used for nano-scale quantitative and qualitative elemental composition analysis by measuring X-ray radiation from the interaction with high energy electron and the material (McDowell et al., 2012) .

However, in EDX, the specimens can be quickly damaged by the high energy electrons, so the acquisition time should be reduced to its minimum. This usually results in very noisy and even incomplete images as shown in Fig. 5 (a). Therefore, the goal of this experiment was to remove the noise and reconstruct high resolution EDX images. The main technical difficulty is that there are no noiseless reference data. Due to the lots of missing photons, a widely used approach for EDX analysis is a kernel regression method using an average kernel as shown in Fig. 5(b) . Unfortunately, this often results in severe blurring. Existing unsupervised learning methods such as SURE network (Soltanayev & Chun, 2018) and Noise2Noise (Lehtinen et al., 2018) cannot be used for this purpose, since there are no specific noise models for the EDX and the noiseless clean data is not available for Noise2Noise training. In fact, this difficulty of EDX denoising was our original motivation for this work. For our network training and inference, we use bootstrap subsampled images from the measurement image in Fig. 5 (a). In addition, we used the measurement image in Fig. 5 (a) as the label for training.

half at every 50 epochs until it reaches around 10 -5 . For the bootstrap subsampling, the number of random subsampling mask was K = 30. The regression network was optimized to minimize the loss (13) with respect to Θ first, after which the attention network was trained to properly aggregate the entire interpolated output. As shown in Fig. 5 (d), our method successfully produced high resolution images which are significantly better than existing kernel regression methods. Thanks to the bagging procedure, the regression network is believed to learn the measurement statistics so that the network not only removes the noise but also filled in the unmeasured data. We also compared the performance between the standard bagging and the proposed method. As shown in Fig. 5 (c) and (d), the weighted averaging from the attention module provides better performance over the the standard bagging that uses the simple average. This confirmed the efficacy of the proposed method. 

6. CONCLUSION:

In this paper, we proposed a novel bagging scheme of encoder-decoder CNNs for various inverse problems. The algorithm was derived by the observation that the existing SURE-based unsupervised denoising networks have several practical limitations due to the divergence term and potential to converge to a trivial solution. Inspired by the recent discovery of basis-like representation of the encoder-decoder CNN, we provide a simple approximation of the divergence term. This also led to a novel bagging scheme to prevent from converging to the trivial identity mapping. In the implementation, multiple input data were generated by bootstrap subsampling, after which final result are obtained by aggregating the entire output of network using an attention network. Experimental results with accelerated MRI and EDX experiments showed that the proposed method provides consistent improvement for various inverse problems. A APPENDIX: SUPER-RESOLUTION APPLICATION Single image super-resolution (SR) is a task to estimate a high-resolution (HR) image from its lowresolution (LR) image. To demonstrate the effectiveness of the proposed method, the proposed bagging scheme was also applied to these problems. Deep Back-Projection Network (DBPN) (Haris et al., 2018 ) is a recent state-of-the-art method for this applications, so we employed DBPN as our baseline model to improve its performance by the proposed method. The baseline network is trained using DIV2K (Agustsson & Timofte, 2017) with totally 800 training images on the ×2 and ×4 (in both horizontal and vertical directions) super-resolution tasks. The network consists of three modules: initial feature extraction, back-projection stages, and reconstruction. The initial LR feature maps were extracted using the convolution layer of the initial feature extraction module. By applying the sequence of projection units which alters between configuration of LR and HR feature maps, the extensive HR features maps can be obtained. The final HR image was reconstructed using the overall HR feature maps from the projection units. Training conditions of the baseline algorithm and our method were followed as described in (Haris et al., 2018) . Specifically, we used patch processing whose size is 32×32 for LR image. The initial learning rate was set to 10 -4 , and it was divided by 10 at every 5 × 10 5 iterations for total 10 6 iterations. For the training of our bagging network, the number of random subsampling mask K was set to 32 and 8 for ×2 and ×4 task, respectively. The subsampling ratio for bootstrap sampling was set to 80%. In addition, the entire networks were trained simultaneously to minimize the loss (13).

:



Figures and Tables:

Figure 3: Reconstruction results using accelerated MR data at the acceleration factor of R = 13.45. The PSNR and SSIM index values for each images are written at the corner.

Figure 4: Ablation study for the accelerated MR experiments at the acceleration ratio of R = 13.45. The PSNR and SSIM index values for each images are written at the corner. w/o BN: reconstruction results from an ablated network without batch normalization, Standard bagging: reconstruction results from an ablated network without an attention module.

Figure 5: EDX experimental results. Green and blue particles refer to Zn and Cd, respectively. (a) Input data, and the reconstruction results using (b) the existing kernel regression method, and the proposed methods (c) with simple average (standard bagging) and (d) the attention based weighted average.

Figure 6: Comparisons between DBPN and the proposed method for ×4 super-resolution task.

Quantitative comparison using PSNR and SSIM index.

Table content: <table><tr><td/><td colspan="3">GRAPPA Baseline (Han &amp; Ye, 2019) Ours</td></tr><tr><td>PSNR (dB)</td><td>34.52</td><td>38.67</td><td>42.72</td></tr><tr><td>SSIM</td><td>0.72</td><td>0.89</td><td>0.91</td></tr></table>

Quantitative comparison using PSNR and SSIM index with respect to number of bootstrap subsample sets.

Table content: <table><tr><td>K</td><td>1</td><td>2</td><td>4</td><td>10</td></tr><tr><td colspan="5">PSNR (dB) 38.68 39.81 40.05 40.27</td></tr><tr><td>SSIM</td><td>0.89</td><td>0.90</td><td>0.91</td><td>0.93</td></tr></table>

Quantitative comparison in terms of PSNR and SSIM index in our ablation study.

Table content: <table><tr><td/><td colspan="3">without BN Standard bagging Ours</td></tr><tr><td>PSNR (dB)</td><td>25.29</td><td>35.91</td><td>40.05</td></tr><tr><td>SSIM</td><td>0.49</td><td>0.867</td><td>0.914</td></tr></table>

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computerassisted intervention, pp. 234-241. Springer, 2015. Tibshirani and Saharon Rosset. Excess optimism: How biased is the apparent error of an estimator tuned by SURE? Journal of the American Statistical Association, pp. 1-16, 2018. Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. Instance normalization: The missing ingredient for fast stylization. arXiv preprint arXiv:1607.08022, 2016. Jong Chul Ye and Woon Kyoung Sung. Understanding geometry of encoder-decoder CNNs. Proceedings of the 2019 International Conference on International Conference on Machine Learning (ICML). also available as arXiv preprint arXiv:1901.07647, 2019. Jong Chul Ye, Yoseob Han, and Eunju Cha. Deep convolutional framelets: A general deep learning framework for inverse problems. SIAM Journal on Imaging Sciences, 11(2):991-1048, 2018.

Table content: <table><tr><td>VA Solé, E Papillon, M Cotte, Ph Walter, and J Susini. A multiplatform code for the analysis of</td></tr><tr><td>energy-dispersive X-ray fluorescence spectra. Spectrochimica Acta Part B: Atomic Spectroscopy,</td></tr><tr><td>62(1):63-68, 2007.</td></tr><tr><td>Shakarim Soltanayev and Se Young Chun. Training deep learning based denoisers without ground</td></tr><tr><td>truth data. In Advances in Neural Information Processing Systems, pp. 3257-3267, 2018.</td></tr><tr><td>Petre Stoica and Yngve Selen. Model-order selection: a review of information criterion rules. IEEE</td></tr><tr><td>Signal Processing Magazine, 21(4):36-47, 2004.</td></tr><tr><td>Ryan J</td></tr></table>

Peer review (iclr2020):
This paper proposed a piecewise linear close form expression for the Stein’s unbiased risk estimator and use this formulation to construct a new Encoder-decoder convolutional neural network. The author claimed that this closely related to bagging. Improved experimental results on two inverse problems are presented. Overall, the experiment results are encouraging but the paper need clarification on a few points. 1. In the model description part, the intuition behind the attention modules is never mentioned. It will be nice to explain the intuition and possibly attached the derivation of the loss function the attention modules. 2. the author seems misunderstand the difference between boosting and bagging. The way described in the paper is bagging and in order to do boosting, a sequential type of network structure probably need to be proposed. 3. How will be model performance compared with a simple bagging for the baseline compared in the experiment part?