Year: 2019

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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORKS

heading: 2.1 STEIN RISK UNBIASED ESTIMATION FOR PREDICTION ERROR
text: 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.

heading: 2.2 PIECEWISE LINEAR REPRESENTATION BY E-D CNN WITH RELUS
text: 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.

heading: 3 MAIN CONTRIBUTION
text: 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.

heading: 3.1 DIVERGENCE SIMPLIFICATION IN SURE
text: 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.

heading: 3.2 BAGGING ESTIMATOR
text: 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) 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 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.

heading: 4 IMPLEMENTATION DETAILS
text: The final loss function for our bagging estimator is given by
` (Θ,Ξ) := 1
P P∑ i=1 ‖x(i) − K∑ k=1 wk(Ξ)FΘ(x, Lk)‖2, (13)
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.
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 wk, our attention network consists of two fully connected layer. The input dimension of the attention network is R1×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 R1×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.

heading: 5 EXPERIMENTAL RESULTS
text: Experiments were conducted for various inverse problems such as accelerated magnetic resonance imaging (MRI), energy-dispersive X-ray spectroscopy (EDX) (Solé et al., 2007), and image superresolution in Appendix.

heading: 5.1 ACCELERATED MRI
text: 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.

heading: 5.2 ENERGY-DISPERSIVE X-RAY SPECTROSCOPY IMAGING
text: As for another experiment, we use the energy-dispersive X-ray spectroscopy (EDX) (Solé et al., 2007) data set which were measured by scanning transmission electron microscopy. EDX (Solé 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.
We used 28 cases from the EDX dataset. The specimen are composed of quantum-dots, where core and shell consist of Cadmium (Cd), Selenium (Se), Zinc (Zn), and Sulfur (S). The regression network was trained using image patches, whose sizes was 128 × 128. For stable training, the intensity of data was normalized between [0, 1]. The initial learning rate was set to 10−3, and it was divided in
half at every 50 epochs until it reaches around 10−5. For the bootstrap subsampling, the number of random subsampling mask wasK = 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.

heading: 6 CONCLUSION
text: 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.

heading: A APPENDIX: SUPER-RESOLUTION APPLICATION
text: 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×105 iterations for total 106 iterations. For the training of our bagging network, the number of random subsampling maskK 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).
Thanks to the bootstrapping and aggregation using attention network, the data distribution can be fully exploited to restore the high resolution components, which results in the properly reconstructed details of the image as shown in Fig. 6. As described in Table 4, our method can also improve the quantitative performance of the super-resolution task.

Review summaries:
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.

The authors consider an encoder decoder setup for linear deblurring problem and propose efficient boosting estimators. Specifically, they use the Stein's unbiased risk estimator for the problem when the noise is gaussian. In the case when the encoder and decoder is represented by a convolutional neural network with RELU activations, they show how they can exploit the recent theoretical results that show the kernel type results to make their procedure efficient. They then propose using a set of models (boosting) and prove that the boosted loss function lower bounds the "nonboosted" loss function.

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.