I MPROVING E NCODER -D ECODER CNN FOR I NVERSE P ROBLEMS
Anonymous authors
Paper under double-blind review
A BSTRACT
Encoder-decoder convolutional neural networks (CNN) have been extensively used for various inverse problems.
However, their prediction error for unseen test data is difﬁcult 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 difﬁcult to implement in a neural network framework, and the condition to avoid trivial identity mapping is not well deﬁned.
In this paper, inspired by the ﬁnding 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.
I NTRODUCTION
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 (cid:98) µ = F Θ ( x ) .
Then, our goal is to estimate the prediction error :
[Equation]
(2) which quantiﬁes how well (cid:98) µ 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 simpliﬁed 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 simpliﬁed 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 signiﬁcantly improve the image quality compared to the existing approaches.
R ELATED WORKS 2. S TEIN R ISK U NBIASED E STIMATION FOR P REDICTION E RROR
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 ) } 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 difﬁculty 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 ﬁnal results critically depend.
2. P IECEWISE L INEAR R EPRESENTATION BY E-D CNN WITH R E LU S
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 ˜ b i ( 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 ﬁlters (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.
M AIN C ONTRIBUTION
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. D IVERGENCE SIMPLIFICATION IN SURE
In Proposition 6 of (Ye & Sung, 2019), the authors show that ∂ F ( x ) /∂
x = ˜ B ( x ) B ( x ) (cid:62) 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 simpliﬁcation 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 ) (cid:62) .
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 beneﬁcial 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 ) )]
(cid:39) σ 2 I .
Furthermore, for sufﬁciently smaller σ , we have
[Equation]
because the corresponding piecewise linear representation does not change for the small perturbation of the input (Ye & Sung, 2019).
Therefore, we have
[Equation]
since ˜ B ( x ( i ) ) B ( x ( i ) ) (cid:62) is a square matrix.
This suggests that
[Equation]
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. B AGGING E STIMATOR
Another important drawback of SURE-based denoising network (4) is that it is difﬁcult to prevent the network from learning a trivial identity mapping.
More speciﬁcally, 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 simpliﬁcation in (9), the cost function (4) can be simpliﬁed as
[Equation]
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 speciﬁcally, we use the following bagging estimator:
[Equation]
and the corresponding prediction error estimate:
[Equation]
where { L k } Kk =1 denotes the diagonal matrix whose diagonal elements are either 0 or 1 depending on the bootstrap subsampling patterns, and { w k } Kk =1 is the corresponding weights.
The corresponding empirical prediction error estimate is then given by
[Equation]
Note that x ( i ) (cid:54) =
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 ≥
and (cid:80) k
w k
= 1 .
Then, we have
[Equation]
Proof.
For any x ∈ { x ( i ) } Pi =1 , we have
[Equation]
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 ( 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.
I MPLEMENTATION D ETAILS
The ﬁnal loss function for our bagging estimator is given by (cid:96)
[Equation]
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 efﬁciently 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 ﬁlters increases from 64 in the ﬁrst stage to 1024 in the ﬁnal 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 ﬁnal 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 .
The proposed network was implemented in Python using TensorFlow library and trained using an NVidia GeForce GTX 1080-Ti graphics processing unit.
E XPERIMENTAL R ESULTS
Experiments were conducted for various inverse problems such as accelerated magnetic resonance imaging (MRI), energy-dispersive X-ray spectroscopy (EDX) (Sol´e et al., 2007), and image superresolution in Appendix.
5. A CCELERATED 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 signiﬁcantly 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-ﬁlled k -space data is used as the input for the neural network.
Speciﬁcally, 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.
Moreover, its average peak signal-to-noise ratio (PSNR) and structural similarity (SSIM) index for entire test dataset signiﬁcantly outperform existing methods as shown in Table 1.
[Table]
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 conﬁrmed 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 signiﬁcantly degraded image with an average PSNR drop from 40.047 dB to 25.290 dB.
This veriﬁed the important role of batch normalization.
[Table]
We also evaluated the importance of the attention network to calculate the weights.
With standard bagging, the ﬁnal 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 conﬁrmed that the important of attention network for the weight calculation.
[Table]
5. E NERGY - DISPERSIVE X- RAY SPECTROSCOPY I MAGING
As for another experiment, we use the energy-dispersive X-ray spectroscopy (EDX) (Sol ´ e et al., 2007) data set which were measured by scanning transmission electron microscopy.
EDX (Sol ´ e 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.
Speciﬁcally, the incident beam may eject an electron from the inner shell of atoms.
Then, an electron from an outer shell then ﬁlls 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 difﬁculty 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 speciﬁc noise models for the EDX and the noiseless clean data is not available for Noise2Noise training.
In fact, this difﬁculty 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 was K = 30 .
The regression network was optimized to minimize the loss (13) with respect to Θ ﬁrst, 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 signiﬁcantly 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 ﬁlled 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 conﬁrmed the efﬁcacy of the proposed method.
C ONCLUSION
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 ﬁnal 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 A PPENDIX : S UPER - 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 ×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 conﬁguration of LR and HR feature maps, the extensive HR features maps can be obtained.
The ﬁnal 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).
Speciﬁcally, 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 10iterations.
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).
Thanks to the bootstrapping and aggregation using attention network
[Table]
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.
As described in Table 4, our method can also improve the quantitative performance of the super-resolution task
[Table]
[Table]