End-to-end reconstruction meets data-driven regularization for inverse problems

1 Introduction . Inverse problems are ubiquitous in imaging applications , wherein one seeks to recover an unknown model parameter x ∈ X from its incomplete and noisy measurement , given by yδ = A ( x ) + e ∈ Y . Here , the forward operator A : X→ Y models the measurement process in the absence of noise , and e , with ‖e‖2 ≤ δ , denotes the measurement noise . For example , in computed tomography ( CT ) , the forward operator computes line integrals of x over a predetermined set of lines in R3 and the goal is to reconstruct x from its projections along these lines . Without any further information about x , inverse problems are typically ill-posed , meaning that there could be several possible reconstructions that are consistent with the measured data , even without any noise . The variational framework circumvents ill-posedness by encoding prior knowledge about x via a regularization functionalR : X→ R. In the variational setting , one solves min x∈X LY ( yδ , A ( x ) ) + λR ( x ) , ( 1 ) where LY : Y× Y→ R+ measures data-fidelity andR penalizes undesirable or unlikely solutions . The penalty λ > 0 balances the regularization strength with the fidelity of the reconstruction . The variational problem ( 1 ) is said to be well-posed if it has a unique solution varying continuously in yδ . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . The success of deep learning in recent years has led to a surge of data-driven approaches for solving inverse problems [ 5 ] , especially in imaging applications . These methods come broadly in two flavors : ( i ) end-to-end trained models that aim to directly map the measurement or a model-based noisy reconstruction to the corresponding true parameter [ 11 , 2 ] and ( ii ) learned regularization methods that seek to find a data-adaptive regularizer instead of handcrafting it [ 15 , 17 , 13 ] . Techniques in both categories have their relative advantages and demerits . Specifically , end-to-end approaches offer fast reconstruction of astounding quality , but lack in terms of theoretical guarantees and need supervised data ( i.e. , pairs of input and target images ) for training . On the contrary , learned regularization methods inherit the provable well-posedness properties of the variational setting and can be trained on unpaired training data , however the reconstruction entails solving a high-dimensional optimization problem , which is often slow and computationally demanding . Our work derives ideas from learned optimization and adversarial machine learning , and makes an attempt to combine the best features of both aforementioned paradigms . In particular , the proposed method offers the flexibility of training on unpaired samples , produces fast reconstructions comparable to end-to-end supervised methods in quality , while enjoying the well-posedness and stability guarantees of the learned regularization framework . We first provide an overview of the recent literature on data-driven solutions for inverse problems before detailing our contributions . 1.1 Related works . End-to-end learned methods for inverse problems either map the measurement directly to the clean image [ 40 , 22 ] , or learn to eliminate artifacts from a reconstruction produced by a model-based technique [ 11 ] . Such approaches are data-intensive and may generalize poorly when trained on limited data . Iterative unrolling [ 20 , 38 , 1 , 19 , 12 ] , with its origin in the seminal work by Gregor and LeCun on data-driven sparse coding [ 10 ] , employs reconstruction networks that are inspired by optimization-based approaches and hence are interpretable . The unrolling paradigm enables one to encode the knowledge about the acquisition physics into the network architecture [ 2 ] , thereby achieving data-efficiency . Nevertheless , end-to-end trained methods are supervised , and it is often challenging to obtain a large ensemble of paired data , especially in medical imaging applications . Data-driven regularization methods aim to learn a regularizer in the variational setting instead of handcrafting it . Some notable approaches in this paradigm include adversarial regularization ( AR ) [ 17 ] and its convex counterpart [ 21 ] , network Tikhonov ( NETT ) [ 15 ] , total deep variation ( TDV ) [ 13 ] , etc. , wherein one explicitly parametrizes the regularization functional using a neural network . The regularization by denoising ( RED ) and more general Plug-and-play ( PnP ) approaches aim to solve inverse problems by using a pre-trained denoiser inside an algorithm for minimizing the variational objective [ 24 , 25 , 39 , 36 , 26 , 16 , 7 , 31 ] . Such methods achieve remarkable performances for various inverse problems when equipped with sophisticated pre-trained denoisers [ 39 , 36 , 26 , 16 , 7 ] , and are essentially equivalent to data-driven regularization , subject to additional constraints on the denoiser [ 23 ] . The deep image prior technique [ 33 ] is training-free , but it regularizes the solution by restricting it to be in the range of a generator and can thus be interpreted as a deep learning-based regularization scheme . It is relatively easier to analyze learned regularization schemes using classical functional analysis [ 28 ] , but they tend to perform worse than end-to-end supervised methods . Moreover , these methods require solving a high-dimensional , potentially non-convex , variational optimization , leading to slow reconstruction with no provably convergent optimization algorithm in general . Data-driven methods for inverse problems can be classified into different categories based on the training protocol . Supervised approaches need data in the form of pairs of noisy measurements ( or some model-based reconstructions ) and target images ( or ground-truths ) [ 2 , 40 , 11 ] , while fully unsupervised methods require access to only the noisy measurements for training [ 29 , 32 , 8 , 37 ] . Our approach lies in between these two extremes , as it requires samples from the marginal distributions of the measurements and the ground-truth images , but not from their joint distribution . Recently , such unpaired training approaches have been developed in [ 30 , 14 ] based on the optimal transport theory . The training objective in [ 14 ] only seeks to match the ground-truth distribution via a Wasserstein generative adversarial network ( GAN ) without any data-consistency , whereas [ 30 ] shows that the cycle-GAN loss between the image and the measurement spaces is an upper bound on the optimal transport loss with a penalized least squares objective as the transportation cost . Unlike [ 30 ] and [ 14 ] , our approach minimizes a fidelity term in the data space alone and not in the image space ( together with an optimal transport-based loss ) , and allows for a more straightforward theoretical analysis using the standard variational framework . 1.2 Specific contributions . Our work seeks to combine iterative unrolling with data-adaptive regularization via an adversarial learning framework , and hence is referred to as unrolled adversarial regularization ( UAR ) . The proposed approach learns a data-adaptive regularizer parametrized by a deep neural network , along with an iteratively unrolled reconstruction network that minimizes the corresponding expected variational loss in an adversarial setting . Unlike AR [ 17 ] where the undesirable images are taken as the pseudo-inverse reconstruction and kept fixed throughout the training , we update them with the output of the unrolled reconstruction network in each training step , and , in turn , use them to further improve the regularizer . Thanks to the Kantorovich-Rubinstein ( KR ) duality [ 4 ] , the alternating learning strategy of the reconstruction and the regularizer networks is equivalent to minimizing the expected data-fidelity over the distribution of the measurement , penalized by the Wasserstein-1 distance between the distribution of the reconstruction and the ground-truth . Once trained , the iteratively unrolled operator produces a fast end-to-end reconstruction . We show that this efficient reconstruction can be improved further by a refinement step that involves minimizing the variational loss with the corresponding regularizer , starting from this initial estimate . The refinement step not only produces reconstructions that are superior to state-of-the-art unsupervised methods and are competitive with supervised methods , but also facilitates well-posedness and stability analyses akin to classical variational approaches [ 28 ] . Our theoretical results on the learned unrolled operator and the regularizer are corroborated by strong experimental evidence for the CT inverse problem , and the illustrative inpainting and denoising examples ( see Sec . B in the supplementary document ) . The unpaired training framework for UAR is particularly useful for medical imaging applications , where it is in general challenging to collect a large corpus of paired noisy input and reference images . 2 The proposed unrolled adversarial regularization ( UAR ) approach . In this section , we give a brief mathematical background on optimal transport , followed by a detailed description of the UAR framework , including the training protocol and the network architectures . 2.1 Background on Optimal transport . Optimal transport theory [ 9 , 34 ] has recently gained prominence in the context of measuring the distance between two probability measures . In particular , given two probability measures π1 and π2 on Rn , the Wasserstein-1 distance between them is defined as W1 ( π1 , π2 ) : = inf µ∈Π ( π1 , π2 ) ∫ ‖x1 − x2‖2 dµ ( x1 , x2 ) , ( 2 ) where Π ( π1 , π2 ) denotes all transport plans having π1 and π2 as marginals . The Wasserstein distance has proven to be suitable for deep learning tasks , when the data is assumed to be concentrated on low-dimensional manifolds in Rn . It has been shown that in such cases , the Wasserstein distance provides a usable gradient during training [ 4 ] , as opposed to other popular divergence metrics . By the KR duality , the Wasserstein-1 distance can be computed equivalently by solving a maximization problem over the space of 1-Lipschitz functions ( denoted by L1 ) as W1 ( π1 , π2 ) = sup R∈L1 ∫ R ( x1 ) dπ1 ( x1 ) − ∫ R ( x2 ) dπ2 ( x2 ) , ( 3 ) provided that π1 and π2 have compact support [ 27 ] . Finally , we recall the definition of push-forward of probability measures , which is used extensively in our theoretical exposition . Given a probability measure π on A and a map T : A → B , the push-forward of π by T ( denoted as T # π ) is defined as a probability measure on B such that T # π ( B ) = π ( T−1 ( B ) ) , for all measurable B ⊂ B . 2.2 Training strategy and network parametrization for UAR . The principal idea behind UAR is to learn an unrolled deep network Gφ : Y→ X for reconstruction , together with a regularization functionalRθ : X→ R parametrized by another convolutional neural network ( CNN ) . The role of Rθ is to discern ground-truth images from images produced by Gφ , while Gφ learns to minimize the variational loss withRθ as the regularizer . As the images produced Algorithm 1 Learning unrolled adversarial regularization ( UAR ) . 1 . Input : Training data-set { xi } Ni=1 ∼ πx and { yj } N j=1 ∼ πyδ , initial reconstruction network parameter φ and regularizer parameter θ , batch-size nb = 1 , penalty λ = 0.1 , gradient penalty λgp = 10.0 , Adam optimizer parameters ( β1 , β2 ) = ( 0.50 , 0.99 ) . 2 . Learn a baseline regularizer : • for training steps k = 1 , 2 , · · · , 10 Nnb , do : ( i.e. , # epochs = 10 , # mini-batches = N nb ) – Sample xj ∼ πx , yj ∼ πyδ , and j ∼ uniform [ 0 , 1 ] ; for 1 ≤ j ≤ nb . Compute uj = A† ( yj ) and x ( ) j = jxj + ( 1− j ) uj . – θ ← Adamη , β1 , β2 ( θ , ∇ J̃1 ( θ ) ) , where η = 10−4 , and J̃1 ( θ ) = 1 nb nb∑ j=1 [ Rθ ( xj ) −Rθ ( uj ) + λgp ( ∥∥∥∇Rθ ( x ( ) j ) ∥∥∥ 2 − 1 ) 2 ] . 3 . Learn a baseline reconstruction operator : ( with 5 epochs , Nnb batches per epoch ) • for training steps k = 1 , 2 , · · · , 5 Nnb , do : – Sample yj ∼ πyδ , and compute J̃2 ( φ ) = 1nb ∑nb j=1 ∥∥yj −A ( Gφ ( yj ) ) ∥∥22 +λRθ ( Gφ ( yj ) ) . – φ← Adamη , β1 , β2 ( φ , ∇ J̃2 ( φ ) ) , with step-size η = 10−4 . 4 . Jointly trainRθ and Gφ adversarially : ( over 25 epochs , with Nnb batches in each epoch ) • for k = 1 , 2 , · · · , 25 Nnb , do : – Sample xj , yj , and j ∼ uniform [ 0 , 1 ] ; for 1 ≤ j ≤ nb . Compute uj = Gφ ( yj ) and x ( ) j = jxj + ( 1− j ) uj . – θ ← Adamη , β1 , β2 ( θ , ∇ J̃1 ( θ ) ) , where J̃1 ( θ ) is as in Step 2 , with η = 2× 10−5 . – Update φ← Adamη , β1 , β2 ( φ , ∇ J̃2 ( φ ) ) twice , with J̃2 ( φ ) as in Step 3 , and η = 2× 10−5 . 5 . Output : The trained networks Gφ andRθ . by Gφ get better , Rθ faces a progressively harder task of telling them apart from the ground-truth images , thus leading to an improved regularizer . On the other hand , as the regularizer improves , the quality of reconstructions obtained using Gφ improves as a consequence . Thus , Gφ andRθ help each other improve as the training progresses via an alternating update scheme .

Summary: This paper proposes a deep learning based reconstruction for ill-posed inverse problems that combines ideas from end-to-end supervised training with learned regularization in a variational setting. The method adversarially trains the weights of a deep neural network (used as a prior for regularization) while also learning an end-to-end, unrolled deep network. The unrolled network takes the form of a learned primal-dual algorithm, and the regularizer aims to discern between ground-truth images and images generated by the unrolled network. After training, the end-to-end network can be used for reconstruction alone, or as an initialization for a separate reconstruction that strictly uses the learned regularizer in a variational setting. The authors claim that this combines the best of both worlds, as an end-to-end model can be used for fast reconstruction, while the variational model can be used for refinement (while also being more amenable to analysis). The authors provide theoretical and empirical evidence that their approach is successful and outperforms other state-of-the-art methods.