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 develops a learning based approach for end-to-end reconstruction of operators for ill-posed problems. In particular it examines the problem of image reconstruction. The main contribution is to combine traditional variational approach with an optimal transport based regularizer.