Inverse problems naturally occur in many applications in computer vision and medical imaging. A successful classical approach relies on the concept of variational regularization EnglHanke ; VariationalRegularization . It combines knowledge about how data is generated in the forward operator with a regularization functional that encodes prior knowledge about the image to be reconstructed.
The success of neural networks in many computer vision tasks has motivated attempts at using deep learning to achieve better performance in solving inverse problemsUnser ; learnedPDHG ; Jo . A major difficulty is the efficient usage of knowledge about the forward operator and noise model in such data driven approaches, avoiding the necessity to relearn the physical model structure.
The framework considered here aims to solve this by using neural networks as part of variational regularization, replacing the typically hand-crafted regularization functional with a neural network. As classical learning methods for regularization functionals do not scale to the high dimensional parameter spaces needed for neural networks, we propose a new training algorithm for regularization functionals. It is based on the ideas in Wasserstein generative adversarial models WassersteinGAN , training the network as a critic to tell apart ground truth images from unregularized reconstructions.
Our contributions are as follows:
We introduce the idea of learning a regularization functional given by a neural network, combining the advantages of the variational formulation for inverse problems with data-driven approaches.
We propose a training algorithm for regularization functionals which scales to high dimensional parameter spaces.
We show desirable theoretical properties of the regularization functionals obtained this way.
We demonstrate the performance of the algorithm for denoising and computed tomography.
2.1 Inverse Problems in Imaging
Let and be reflexive Banach spaces. In a generic inverse problem in imaging, the image is recovered from measurement , where
denotes the linear forward operator and is a random noise term. Typical tasks in computer vision that can be phrased as inverse problems include denoising, where is the identity operator, or inpainting, where
is given by a projection operator onto the complement of the inpainting domain. In medical imaging, common forward operators are the Fourier transform in magnetic resonance imaging (MRI) and the ray transform in computed tomography (CT).
2.2 Deep Learning in Inverse Problems
One approach to solve (1) using deep learning is to directly learn the mapping using a neural network. While this has been observed to work well for denoising and inpainting DeepDenoising , the approach can become infeasible in inverse problems involving forward operator with a more complicated structure fullyLearned2 and when only very limited training data is available. This is typically the case in applications in medical imaging.
Other approaches have been developed to tackle inverse problems with complex forward operators. In Unser an algorithm has been suggested that first applies a pseudo-inverse to the operator , leading to a noisy reconstruction. This result is then denoised using deep learning techniques. Other approaches learnedGradientDescent ; variationalNetworks ; Jo propose applying a neural network iteratively. Learning proximal operators for solving inverse problems is a further direction of research learnedPDHG ; learnedProx .
2.3 Variational regularization
Variational regularization is a well-established model-based method for solving inverse problems. Given a single measurement , the image is recovered by solving
where the data term ensures consistency of the reconstruction with the measurement and the regularization functional allows us to insert prior knowledge onto the solution . The functional is usually hand-crafted, with typical choices including total variation (TV) RudinTV which leads to piecewise constant images and total generalized variation (TGV) TGV , generating piecewise linear images.
3 Learning a regularization functional
In this paper, we design a regularization functional based on training data. We fix a-priori a class of admissible regularization functionals and then learn the choice from data. Existing approaches to learning a regularization functionals are based on the idea that should be chosen such that a solution to the variational problem
But this is computationally feasible only for small sets of admissible functions . In particular, it does not scale to sets parametrized by some high dimensional space .
We hence apply a novel technique for learning the regularization functional that scales to high dimensional parameter spaces. It is based on the idea of learning to discriminate between noisy and ground truth images.
In particular, we consider approaches where the regularization functional is given by a neural network with network parameters . In this setting, the class is given by the functions that can be parametrized by the network architecture of for some choice of parameters . Once is fixed, the inverse problem (1) is solved by
3.1 Regularization functionals as critics
Denote by independent samples from the distribution of ground truth images and by independent samples from the distribution of measurements
. Note that we only use samples from both marginals of the joint distribution
of images and measurement, i.e. we are in the setting of unsupervised learning.
The distribution on measurement space can be mapped to a distribution on image space by applying a -potentially regularized- pseudo-inverse . In Unser it has been shown that such an inverse can in fact be computed efficiently for a large class of forward operators. This in particular includes Fourier and ray transforms occurring in MRI and CT. Let
be the distribution obtained this way. Here, denotes the push-forward of measures, i.e. for . Samples drawn from will be corrupted with noise that both depends on the noise model as well as on the operator .
A good regularization functional is able to tell apart the distributions and - taking high values on typical samples of and low values on typical samples of MartinLearning . It is thus clear that
being small is desirable. With this in mind, we choose the loss functional for learning the regularizer to be
The last term in the loss functional serves to enforce the trained network to be Lipschitz continuous with constant one WGANimproved . The expected value in this term is taken over all lines connecting samples in and .
Training a neural network as a critic was first proposed in the context of generative modeling in GAN . The particular choice of loss functional has been introduced in WassersteinGAN to train a critic that captures the Wasserstein distance between the distributions and . A minimizer to (6) approximates a maximizer to the Kantorovich formulation of optimal transport Villani .
Relaxing the hard Lipschitz constraint in (7) into a penalty term as in (6) was proposed in WGANimproved . Tracking the gradients of for our experiments demonstrates that this way the Lipschitz constraint can in fact be enforced up to a small error.
3.2 Distributional Analysis
Here we analyze the impact of the learned regularization functional on the induced image distribution. More precisely, given a noisy image drawn from , consider the image obtained by performing a step of gradient descent of size over the regularization functional
This yields a distribution of noisy images that have undergone one step of gradient descent. We show that this distribution is closer in Wasserstein distance to the distribution of ground truth images than the noisy image distribution . The regularization functional hence introduces the highly desirable incentive to align the distribution of minimizers of the regularization problem (5) with the distribution of ground truth images.
Henceforth, assume the network has been trained to perfection, i.e. that it is a 1-Lipschitz function which achieves the supremum in (7). Furthermore, assume is almost everywhere differentiable with respect to the measure .
Assume that admits a left and a right derivative at , and that they are equal. Then,
On the other hand, for a.e. one can bound
for any . Hence, in particular the difference quotient is bounded
for any and . By dominated convergence, this allows us to conclude
Under the weak assumptions in (WGANimproved, , Corollary 1), we have for a.e. . This allows to compute the rate of decay of Wasserstein distance explicitly to
Note that the above calculations also show that the particular choice of loss functional is optimal in terms of decay rates of the Wasserstein distance, introducing the strongest incentive to align the distribution of reconstructions with the ground truth distribution amongst all regularization functionals. To make this more precise, consider any other regularization functional with norm-bounded gradients, i.e. .
Denote by the flow associated to . Set . Then
An analogous computation as above shows
3.3 Analysis under data manifold assumption
Here we discuss which form of regularization functional is desirable under the data manifold assumption and show that the loss function (6) in fact gives rise to a regularization functional of this particular form.
Assumption 1 (Weak Data Manifold Assumption).
Assume the measure is supported on the weakly compact set , i.e.
This assumption captures the intuition that real data lies in a curved lower-dimensional subspace of .
If we consider the regularization functional as encoding prior knowledge about the image distribution, it follows that we would like the regularizer to penalize images which are away from
. An extreme way of doing this would be to set the regularization functional as the characteristic function of. However, this choice of functional comes with two major disadvantages: First, solving (5) with methods based on gradient descent becomes impossible when using such a regularization functional. Second, the functional effectively leads to a projection onto the data manifold, possibly causing artifacts due to imperfect knowledge of DataManifoldProjection .
An alternative to consider is the distance function to the data manifold , since such a choice provides meaningful gradients everywhere. This is implicitly done in RED . In Theorem 2, we show that our chosen loss function in fact does give rise to a regularization functional taking the desirable form of the distance function to .
the data manifold projection, where denotes the set of points for which such a projection exists. We assume . This can be guaranteed under weak assumptions on and .
Assume the measures and satisfy
i.e. for every measurable set , we have
We hence assume that the distortions of the true data present in the distribution of pseudo-inverses
are well-behaved enough to recover the distribution of true images from noisy ones by projecting back onto the manifold. Note that this is a much weaker than assuming that any given single image can be recovered by projecting its pseudo-inverse back onto the data manifold. Heuristically, Assumption2 corresponds to a low-noise assumption.
is given by the distance function to the data manifold
First show that is Lipschitz continuous with Lipschitz constant . Let be arbitrary and denote by a minimizer to . Indeed,
where we used the triangle inequality in the last step. This proves Lipschitz continuity by exchanging the roles of and .
Now, we prove that obtains the supremum in 18. Let be any -Lipschitz function. By assumption 2, one can rewrite
As is Lipschitz, this can be bounded via
The distance between and is by definition given by . This allows to conclude via
Remark 2 (Non-uniqueness).
The functional (18) does not necessarily have a unique maximizer. For example, can be changed to an arbitrary -Lipschitz function outside the convex hull of .
Following the well-developed stability theory for classical variational problems EnglHanke
, we derive a stability estimate for the adversarial regularizer algorithm. The key difference to existing theory is that we do not assume the regularization functionalis bounded from below. Instead, this is replaced by a Lipschitz assumption on .
Theorem 3 (Weak Stability in Data Term).
We make Assumption 3. Let be a sequence in with in the norm topology and denote by a sequence of minimizers of the functional
Then has a weakly convergent subsequence and the limit is a minimizer of .
The assumptions and the proof are contained in Appendix A.
5 Computational Results
5.1 Parameter estimation
Applying the algorithm to new data requires choosing a regularization parameter . Making the assumption that the ground truth images are critical points of the variational problem (5), can be estimated efficiently from the noise level, using the fact that the regularization functional has gradients of unit norm. This leads to the formula
where denotes the adjoint and the noise distribution. In all experiments, the regularization parameter has been chosen according to this formula without further tuning.
As a toy problem, we compare the performance of total variation denoising RudinTV , a supervised denoising neural network approach DeepDenoising based on the UNet UNET architecture and our proposed algorithm on images of size cut out of images taken from the BSDS500 dataset BSDS
. The images have been corrupted with Gaussian white noise. We report the average peak signal-to-noise ratio (PSNR) and the structural similarity index (SSIM)SSIM in Table 1.
The results in Figure 1 show that the adversarial regularizer algorithm is able to outperform classical variational methods in all quality measures. It achieves results of comparable visual quality than supervised data-driven algorithms, without relying on supervised training data.
|Total Variation RudinTV||26.3||.836|
|Denoising N.N. DeepDenoising||28.8||.908|
|Adversarial Regularizer (ours)||28.2||.892|
5.3 Computed Tomography
Computer Tomography reconstruction is an application in which the variational approach is very widely used in practice. Here, it serves as a prototype inverse problem with non-trivial forward operator. We compare the performance of total variation CTTV ; RudinTV , post-processing Unser , Regularization by Denoising (RED) RED and our proposed regularizers on the LIDC/IDRI database LUNA of lung scans. The denoising algorithm underlying RED has been chosen to be the denoising neural network previously trained for post-processing. Measurements have been simulated by taking the ray transform, corrupted with Gaussian white noise. With different angles taken for the ray transform, the forward operator is undersampled. The code is available online 111https://github.com/lunz-s/DeepAdverserialRegulariser.
that Post-Processing has been trained with PSNR as target loss function. Again, total variation is outperformed by a large margin in all categories. Our reconstructions are of the same or superior visual quality than the ones obtained with supervised machine learning methods, despite having used unsupervised data only.
We have proposed an algorithm for solving inverse problems, using a neural network as regularization functional. We have introduced a novel training algorithm for regularization functionals and showed that the resulting regularizers have desirable theoretical properties. Unlike other data-based approaches in inverse problems, the proposed algorithm can be trained even if only unsupervised training data is available. This allows to apply the algorithm to situations where -due to a lack of appropriate training data- machine learning methods have not been used yet.
The variational framework enables us to effectively insert knowledge about the forward operator and the noise model into the reconstruction, allowing the algorithm to be trained on little training data. It also comes with the advantages of a well-developed stability theory and the possibility of adapting the algorithms to different noise levels by changing the regularization parameter , without having to retrain the model from scratch.
The computational results demonstrate the potential of the algorithm, producing reconstructions of the same or even superior visual quality as the ones obtained with supervised approaches on the LIDC dataset, despite the fact that only unsupervised data has been used for training. Classical methods like total variation are outperformed by a large margin.
Our approach is particularly well-suited for applications in medical imaging, where usually very few training samples are available and ground truth images to a particular measurement are hard to obtain, making supervised algorithms impossible train.
The algorithm admits some extensions and modifications.
Local Regularizers. The regularizer is restricted to act on small patches of pixels only, giving the value of the regularization functional by averaging over all patches. This allows to harvest many training samples from a single image, making the algorithm trainable on even less training data. Local Adversarial Regularizers can be implemented by choosing a neural network architecture consisting of convolutional layers followed by a global average pooling.
Recursive Training. When applying the regularization functional, the variational problem has to be solved. In this process, the regularization functional is confronted with partially reconstructed images, which are neither ground truth images nor exhibit the typical noise distribution the regularization functional has been trained on. By adding these images to the samples the regularization functional is trained on, the neural network is enabled to learn from its own outputs. First implementations show that this can lead to an additional boost in performance, but that the choice of which images to add is very delicate.
We thank Sam Power, Robert Tovey, Matthew Thorpe, Jonas Adler, Erich Kobler, Jo Schlemper, Christoph Kehle and Moritz Scham for helpful discussions and advice.
The authors acknowledge the National Cancer Institute and the Foundation for the National Institutes of Health, and their critical role in the creation of the free publicly available LIDC/IDRI Database used in this study. The work by Sebastian Lunz was supported by the EPSRC grant EP/L016516/1 for the University of Cambridge Centre for Doctoral Training, the Cambridge Centre for Analysis and by the Cantab Capital Institute for the Mathematics of Information. The work by Ozan Öktem was supported by the Swedish Foundation for Strategic Research grant AM13-0049. Carola-Bibiane Schönlieb acknowledges support from the Leverhulme Trust project on ‘Breaking the non-convexity barrier’, EPSRC grant Nr. EP/M00483X/1, the EPSRC Centre Nr. EP/N014588/1, the RISE projects CHiPS and NoMADS, the Cantab Capital Institute for the Mathematics of Information and the Alan Turing Institute.
-  Jonas Adler and Ozan Öktem. Solving ill-posed inverse problems using iterative deep neural networks. Inverse Problems, 33(12), 2017.
-  Jonas Adler and Ozan Öktem. Learned primal-dual reconstruction. IEEE Transactions on Medical Imaging, 37(6):1322–1332, 2018.
-  Pablo Arbelaez, Michael Maire, Charless Fowlkes, and Jitendra Malik. Contour detection and hierarchical image segmentation. IEEE Trans. Pattern Anal. Mach. Intell., 33(5).
-  Maria Argyrou, Dimitris Maintas, Charalampos Tsoumpas, and Efstathios Stiliaris. Tomographic image reconstruction based on artificial neural network (ANN) techniques. In Nuclear Science Symposium and Medical Imaging Conference (NSS/MIC). IEEE, 2012.
-  Martín Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. International Conference on Machine Learning, ICML, 2017.
-  Samuel Armato, Geoffrey McLennan, Luc Bidaut, Michael McNitt-Gray, Charles Meyer, Anthony Reeves, Binsheng Zhao, Denise Aberle, Claudia Henschke, Eric Hoffman, et al. The lung image database consortium (LIDC) and image database resource initiative (IDRI): a completed reference database of lung nodules on ct scans. Medical physics, 38(2), 2011.
-  Martin Benning, Guy Gilboa, Joana Sarah Grah, and Carola-Bibiane Schönlieb. Learning filter functions in regularisers by minimising quotients. In International Conference on Scale Space and Variational Methods in Computer Vision. Springer, 2017.
-  Ashish Bora, Ajil Jalal, Eric Price, and Alexandros Dimakis. Compressed sensing using generative models. arXiv preprint arXiv:1703.03208, 2017.
-  Luca Calatroni, Chung Cao, Juan Carlos De Los Reyes, Carola-Bibiane Schönlieb, and Tuomo Valkonen. Bilevel approaches for learning of variational imaging models. RADON book series, 8, 2012.
-  Antonin Chambolle and Thomas Pock. An introduction to continuous optimization for imaging. Acta Numerica, 25, 2016.
-  Heinz Werner Engl, Martin Hanke, and Andreas Neubauer. Regularization of inverse problems, volume 375. Springer Science & Business Media, 1996.
-  Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural information processing systems (NIPS), 2014.
-  Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. Improved training of wasserstein GANs. Advances in Neural Information Processing Systems (NIPS), 2017.
-  Kerstin Hammernik, Teresa Klatzer, Erich Kobler, Michael Recht, Daniel Sodickson, Thomas Pock, and Florian Knoll. Learning a variational network for reconstruction of accelerated MRI data. Magnetic resonance in medicine, 79(6), 2018.
Kyong Hwan Jin, Michael McCann, Emmanuel Froustey, and Michael Unser.
Deep convolutional neural network for inverse problems in imaging.IEEE Transactions on Image Processing, 26(9), 2017.
-  Florian Knoll, Kristian Bredies, Thomas Pock, and Rudolf Stollberger. Second order total generalized variation (TGV) for MRI. Magnetic Resonance in Medicine, 65(2), 2011.
-  Karl Kunisch and Thomas Pock. A bilevel optimization approach for parameter learning iniational models. SIAM Journal on Imaging Sciences, 6(2), 2013.
-  Rowan Leary, Zineb Saghi, Paul Midgley, and Daniel Holland. Compressed sensing electron tomography. Ultramicroscopy, 131, 2013.
-  Tim Meinhardt, Michael Moeller, Caner Hazirbas, and Daniel Cremers. Learning proximal operators: Using denoising networks for regularizing inverse imaging problems. In International Conference on Computer Vision (ICCV), 2017.
-  Paul Milgrom and Ilya Segal. Envelope theorems for arbitrary choice sets. Econometrica, 70(2), 2002.
-  Yaniv Romano, Michael Elad, and Peyman Milanfar. The little engine that could: Regularization by denoising (RED). SIAM Journal on Imaging Sciences, 10(4), 2017.
-  Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical Image Computing and Computer-Assisted Intervention. Springer, 2015.
-  Leonid I Rudin, Stanley Osher, and Emad Fatemi. Nonlinear total variation based noise removal algorithms. Physica D: nonlinear phenomena, 60(1-4), 1992.
-  Otmar Scherzer, Markus Grasmair, Harald Grossauer, Markus Haltmeier, and Frank Lenzen. Variational methods in imaging. Springer, 2009.
-  Jo Schlemper, Jose Caballero, Joseph V Hajnal, Anthony Price, and Daniel Rueckert. A deep cascade of convolutional neural networks for mr image reconstruction. In International Conference on Information Processing in Medical Imaging. Springer, 2017.
-  Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
-  Zhou Wang, Alan Bovik, Hamid Sheikh, and Eero Simoncelli. Image quality assessment: from error visibility to structural similarity. IEEE transactions on image processing, 13(4), 2004.
-  Junyuan Xie, Linli Xu, and Enhong Chen. Image denoising and inpainting with deep neural networks. In Advances in Neural Information Processing Systems (NIPS), 2012.
A Stability Theory
All of the following three conditions hold true.
is lower semi-continuous with respect to the weak topology on and Lipschitz with respect to the metric induced by the norm.
is continuous or equivalently weak-to-weak continuous
One of the following two conditions hold true
These assumptions are standard in the classical stability theory for inverse problems , with the difference that we assume to be Lipschitz instead of being bounded from below.
Next we show that in the particular setting where is the distance function to the manifold , the weak continuity assumption on is always satisfied.
The map is weakly lower semi-continuous.
Let be a sequence in with weakly. Pick any subsequence of , for convenience still denoted by . Denote by elements in such that
As all , we can extract a weakly convergent subsequence, denoted by , such that weakly for some . By lower semi-continuity of the norm, estimate
Lemma 2 (Coercivity).
Let . Then under assumptions 3,
as , uniformly in all with .
Assume first . WLOG assume . Then
as , uniformly in with . The last inequality uses the assumption that is Lipschitz.
In the case as the statement follows immediately. ∎
Theorem 4 (Existence of Minimizer).
Under assumptions 3, there exists a minimizer of
Let be a sequence in such that
as . Then by Lemma 2, is bounded in norm allowing to extract a weakly convergent subsequence . As the norm is weakly lower-semi continuous and so is by assumption, we obtain
thus proving that is indeed a minimizer. ∎
Theorem 5 (Weak Stability in Data Term).
Assume 3. Let be a sequence in with in the norm topology and denote by a sequence of minimizers of the functional
Then has a weakly convergent subsequence and the limit is a minimizer of .
By Lemma 2, the sequence is bounded and hence contains a weakly convergent subsequence . Note that the map is continuous with respect to the norm topology. On the other hand, as is by assumption weak to weak continuous and both the norm and are weakly lower semi-continuous, the overall loss is weakly lower semi-continuous, so
Together we obtain
proving that the limit point is indeed a minimizer of . ∎
B Implementation details
) activations and two final dense layer for all experiments with the adversarial regularizer algorithm. The network was optimized with RMSProp. We solved the variational problem using gradient descent with fixed step size. The regularization parameter was chosen according to the heuristic given in paper.
The comparison experiments with Post-Processing and the Denoising Neural Network used a UNet style architecture, with four down-sampling (strided convolution, stride ) and four up-sampling (transposed convolution) convolutional layers with skip-connections after every down-sampling step to the corresponding up-sampled layer of the same image resolution. Again, leaky ReLU activations were used. The network was optimized using Adam. As training loss we used the distance to the ground truth, with no further regularization terms on the network parameters.
In the experiments with total variation, the regularization parameter was chosen using line search, picking the parameter that leads to the best PSNR value. The minimization problem was solved using primal-dual hybrid gradient descent (PDHG) .
C Further Computational Results
From left to right: Ground truth, Noisy Image, TV, Denoising Neural Network, Adversarial Reg.
From left to right: Ground truth, FBP, TV, Post-Processing, Adversarial Reg.
From left to right: Ground truth, FBP, TV, Post-Processing, Adversarial Reg.
Below the Sinogram used for reconstruction of the images.