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

Abstract:
We propose a new approach for learning end-to-end reconstruction operators based on unpaired training data for ill-posed inverse problems. The proposed method combines the classical variational framework with iterative unrolling and essentially seeks to minimize a weighted combination of the expected distortion in the measurement space and the Wasserstein-1 distance between the distributions of the reconstruction and the ground-truth. More specifically, the regularizer in the variational setting is parametrized by a deep neural network and learned simultaneously with the unrolled reconstruction operator. The variational problem is then initialized with the output of the reconstruction network and solved iteratively till convergence. Notably, it takes significantly fewer iterations to converge as compared to variational methods, thanks to the excellent initialization obtained via the unrolled operator. The resulting approach combines the computational efficiency of end-to-end unrolled reconstruction with the well-posedness and noise-stability guarantees of the variational setting. Moreover, we demonstrate with the example of image reconstruction in X-ray computed tomography (CT) that our approach outperforms state-of-the-art unsupervised methods and that it outperforms or is at least on par with state-of-the-art supervised data-driven reconstruction approaches.

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

heading: 1.1 Related works
text: 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.

heading: 1.2 Specific contributions
text: 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.

heading: 2 The proposed unrolled adversarial regularization (UAR) approach
text: 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.

heading: 2.1 Background on Optimal transport
text: 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.

heading: 2.2 Training strategy and network parametrization for UAR
text: 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.

heading: 2.2.1 Adversarial training
text: Let us denote by πx the ground-truth distribution and by πyδ the distribution of the noisy measurement. The UAR approach seeks to learn Gφ andRθ simultaneously starting from an appropriate initialization. At the kth iteration of training, the parameters φ of the reconstruction network are updated as
φk ∈ arg min φ J (1) k (φ), where J (1) k (φ) := Eπyδ [∥∥A(Gφ(yδ))− yδ∥∥22 + λRθk(Gφ(yδ))] , (4) for a fixed regularizer parameter θk. Subsequently, the regularizer parameters are updated as
θk+1 ∈ arg max θ:Rθ∈L1 J (2) k (θ), where J (2) k (θ) := Eπyδ
[ Rθ ( Gφk(yδ) )] − Eπx [Rθ (x)] . (5)
The learning protocol for UAR only requires unpaired samples of x and yδ , since the loss functionals J
(1) k (φ) and J (2) k (θ) can be computed based solely on the marginals πx and πyδ . The alternating update algorithm in (4) and (5) essentially seeks to solve the min-max variational problem given by
min φ max θ:Rθ∈L1 Eπ yδ ∥∥A(Gφ(yδ))− yδ∥∥22 + λ(Eπyδ [Rθ (Gφ(yδ))]− Eπx [Rθ (x)]) . (6) Thanks to KR duality in (3) and the definition of push-forward, (6) can be reformulated as
min φ Eπ yδ ∥∥A(Gφ(yδ))− yδ∥∥22 + λW1 ((Gφ)#πyδ , πx) . (7)
We refer the reader to Section 3 for a mathematically rigorous statement of this equivalence as well as for a well-posedness theory of the problem in (7). Note that the equivalence of the alternating minimization procedure and the variational problem in (7) holds only if the regularizer is fully optimized in every iteration. Nevertheless, in practice, the reconstruction and regularizer networks are not fully optimized in every iteration. Instead, one refines the parameters by performing one (or a few) Adam updates on the corresponding loss functionals. Notably, if W1 ( (Gφk)#πyδ , πx ) = 0, i.e., the parameters of G are such that the reconstructed images match the ground-truth in distribution, the loss functional J (2)k (θ) and its gradient vanish, leading to no further update of θ. Thus, both networks stop updating when the outputs of Gφ are indistinguishable from the ground-truth images. The concrete training steps are listed in Algorithm 11.

heading: 2.2.2 Variational regularization as a refinement step
text: The unrolled operator Gφ∗ trained by solving the min-max problem in (6) provides reasonably good reconstruction when evaluated on X-ray CT, and already outperforms state-of-the-art unsupervised methods (c.f. Section 4). We demonstrate that the regularizer Rθ∗ obtained together with Gφ∗ by solving (6) can be used in the variational framework to further improve the quality of the end-to-end reconstruction Gφ∗(yδ) for a given yδ ∈ Y. Specifically, we solve the variational problem
min x∈X ‖A(x)− yδ‖2 + λ′
( Rθ∗(x) + σ‖x‖22 ) , (8)
where λ′, σ ≥ 0, by applying gradient descent, initialized with Gφ∗(yδ). The additional Tikhonov term in (8) ensures coercivity of the overall regularizer, making it amenable to the standard wellposedness analysis [28]. Practically, it improves the stability of the gradient descent optimizer for (8). Nevertheless, one essentially gets the same reconstruction with σ = 0 subject to early stopping (100 iterations). Notably, the fidelity term in (8) is the `2 distance, instead of the squared-`2 fidelity. We have empirically observed that this choice of the fidelity term improves the quality of the reconstruction, possibly due to the higher gradient of the objective in the initial solution Gφ∗(yδ). Since the end-to-end reconstruction gives an excellent initial point, it takes significantly fewer iterations for gradient-descent to recover the optimal solution to (8), and therefore UAR retains its edge in reconstruction time over fully variational approaches with learned regularizers (e.g., AR [17] or its convex version [21]).

heading: 2.2.3 Architectures of the iteratively unrolled generator and the deep regularizer
text: The objective of Gφ is to approximate the minimizer of the variational loss withRθ as the regularizer. Therefore, an iterative unrolling strategy akin to [2] is adopted for parametrizing Gφ. Iterative unrolling aims to approximate the variational minimizer via a primal-dual-style algorithm [6], with the proximal operators in the image and measurement spaces replaced with trainable CNNs. Although the variational loss in our case is non-convex, this parametrization for Gφ is chosen because of its expressive power over a generic network. Initialized with x(0) = A† yδ and h(0) = 0, Gφ produces a reconstruction x(L) by iteratively applying the CNNs Λ
φ (`) p and Γ φ (`) d in X and Y, respectively:
h(`+1) = Γ φ (`) d
( h(`), σ(`)A(x`),yδ ) , andx(`+1) = Λ
φ (`) p
( x(`), τ (`)A∗(h`+1) ) , 0 ≤ ` ≤ L−1.
The step-size parameters σ(`) and τ (`) are also made learnable and initialized as σ(`) = τ (`) = 0.01 for each layer `. All learnable parameters in the generator are denoted by the shorthand notation φ for brevity. The number of layers L is typically much smaller (we take L = 20) than the number of iterations needed by an iterative primal-dual scheme to converge, thus expediting the reconstruction by two orders of magnitude (as compared to variational methods) once trained.
The regularizerRθ is taken as a deep CNN with six convolutional layers, followed by one averagepooling and two dense layers in the end.

heading: 3 Theoretical results
text: The theoretical properties of UAR are stated in this section and their proofs are provided in the supplementary document. Throughout this section, we assume that X = Rn and Y = Rk, and
1Codes at https://github.com/Subhadip-1/unrolling_meets_data_driven_regularization.
A1. πx is compactly supported and πyδ is supported on a compact set K ⊂ Rk for every δ ≥ 0.
We then consider the following problem:
inf φ sup R∈L1
J1 ( Gφ,R|λ, πyδ ) := Eπ
yδ ∥∥yδ −AGφ(yδ)∥∥22 + λ (Eπyδ [R(Gφ(yδ))]− Eπx [R(x)]) . (9)
Problem (9) is identical to the min-max variational problem defined in (6), with the only difference that the maximization in R is performed over the space of all 1-Lipschitz functions. Basically, we consider the theoretical limiting case where the neural networks Rθ are expressive enough to approximate all functions in L1 with arbitrary accuracy. We make the following assumptions on Gφ:
A2. Gφ is parametrized over a finite dimensional compact set K, i.e. φ ∈ K. A3. Gφn → Gφ pointwise whenever φn → φ. A4. supφ∈K ‖Gφ‖∞ <∞.
Assumptions A2-A4 are satisfied, for instance, when Gφ is parametrized by a neural network whose weights are kept bounded during training. These assumptions apply to all results in this section.

heading: 3.1 Well-posedness of the adversarial loss and noise stability
text: Here, we prove well-posedness and stability to noise for the optimal reconstructions. As a consequence of the KR duality, (9) can be equivalently expressed as
inf φ J2 ( Gφ|λ, πyδ ) := Eπ yδ ∥∥yδ −AGφ(yδ)∥∥22 + λW1(πx, (Gφ)#πyδ) . (10) In the next theorem, we state this equivalence, showing the existence of an optimal Gφ andR for (9). Theorem 1. Problems (9) and (10) admit an optimal solution and
inf φ sup R∈L1
J1 ( Gφ,R|λ, πyδ ) = inf φ J2 ( Gφ|λ, πyδ ) . (11)
Moreover, if (Gφ∗ ,R∗) is optimal for (9), then Gφ∗ is optimal for (10). Conversely, if Gφ∗ is optimal for (10), then (Gφ∗ ,R∗) is optimal for (9), for allR∗ ∈ arg maxR∈L1 Eπyδ [ R(Gφ∗(yδ)) ] −Eπx [R(x)].
Next, we study the stability of the optimal reconstruction Gφ∗ to noise. We consider Gφn , where
φn ∈ arg inf φ
J2 ( Gφ|λ, πyδn ) , (12)
and show that Gφn → Gφ∗ as δn → δ, thus establishing noise-stability of the unrolled reconstruction. Theorem 2 (Stability to noise). Suppose, for given a sequence of noise levels δn → δ ∈ [0,∞), it holds that πyδn → πyδ in total variation. Then, with φn as in (12), Gφn → Gφ∗ up to sub-sequences.

heading: 3.2 Effect of λ on the end-to-end reconstruction
text: In order to analyze the effect of the parameter λ in (10) on the resulting reconstruction Gφ∗ , it is convenient to introduce the following two sets:
ΦL := { φ : Eπ
yδ ∥∥yδ −AGφ(yδ)∥∥22 = 0} and ΦW := {φ : (Gφ)#πyδ = πx} . We assume that both ΦL and ΦW are non-empty, which is tantamount to asking that the parametrization of the end-to-end reconstruction operator is expressive enough to approximate a right inverse of A (ΦL 6= ∅) and a transport map from πyδ to πx (ΦW 6= ∅), and therefore is not very restrictive (keeping in view the excellent approximation power of unrolled deep architectures). Proposition 1. Let Gφ∗ be a minimizer for (10). Then, it holds that
• Eπy ∥∥yδ −AGφ∗(yδ)∥∥22 ≤ λW1(πx, (Gφ)#πyδ), for every φ ∈ ΦL.
• W1(πx, (Gφ∗)#πyδ) ≤ 1
λ Eπ yδ
∥∥yδ −AGφ(yδ)∥∥22 , for every φ ∈ ΦW.
The previous proposition shows in a quantitative way that for small λ, the optimal Gφ∗ has less expected distortion in the measurement space as the quantity Eπ
yδ ∥∥yδ −AGφ∗(yδ)∥∥22 is small. On the other hand, if λ is large, then the optimal Gφ∗ maps πyδ closer to the ground-truth distribution πx; as the quantity W(πx, (Gφ∗)#πyδ) is small. Therefore, the regularization is stronger in this case. We extend this analysis by studying the behavior of the unrolled reconstruction as λ converges to 0 and to +∞. Consider a sequence of parameters λn > 0 and the minimizer of the objective in (10) with parameter λn:
φ′n ∈ arg inf φ
J2 ( Gφ|λn, πyδ ) . (13)
Theorem 3. Let λn → 0. Then, there exists φ∗1 ∈ arg min φ∈ΦL W1(πx, (Gφ)#πyδ) such that Gφ′n → Gφ∗1 up to sub-sequences, and lim
n→∞ 1 λn inf φ J2 ( Gφ|λn, πyδ ) = W1(πx, (Gφ∗1 )#πyδ).
Theorem 4. Let λn → +∞. Then, there exists φ∗2 ∈ arg min φ∈ΦW Eπ yδ ∥∥yδ −AGφ(yδ)∥∥22 such that Gφ′n → Gφ∗2 up to sub-sequences, and limn→∞ infφ J2 ( Gφ|λn, πyδ ) = Eπ yδ
∥∥yδ −AGφ∗2 (yδ)∥∥22. Theorems 3 and 4 characterize the optimal end-to-end reconstruction Gφ∗ as λ → 0 and λ → ∞, respectively. Specifically, if λ→ 0, Gφ∗ minimizes the Wasserstein distance between reconstruction and ground-truth among all the reconstruction operators that achieve zero expected data-distortion. In particular, Gφ∗ is close to the right inverse of A that minimizes the Wasserstein distance. Therefore, when λ is very small, we expect to obtain a reconstruction that is close to the unregularized solution in quality. If λ → ∞ on the other hand, the operator Gφ∗ is close to a transport map between πx and πyδ , i.e., (Gφ∗)#πyδ = πx, which minimizes the expected data-distortion. Therefore the reconstruction produces realistic images, but they are not consistent with the measurement. These theoretical observations are corroborated by the numerical results (c.f. Section 4, Fig. 2). One has to thus select a λ that optimally trades-off data-distortion with the Wasserstein distance to achieve the best reconstruction performance.

heading: 3.3 End-to-end reconstruction vis-à-vis the variational solution
text: The goal of this section is two-fold. Firstly, we theoretically justify the fact that the end-to-end reconstruction performs well, despite minimizing the expected loss over the distribution πyδ . Secondly, we analyze the role of the regularizer in the variational setting in refining the end-to-end reconstruction.
It is important to remark that the the end-to-end reconstruction is trained with respect to the average variational loss computed using samples from πyδ and πx. Therefore, the end-to-end reconstruction cannot learn a point-wise correspondence between measurement and model parameter, but only a distributional correspondence. Despite that, the end-to-end reconstruction achieves excellent performance for a given measurement yδ. A justification of this phenomenon is given formally by the next proposition.
Proposition 2. Let (Gφ∗ ,R∗) be an optimal pair for (9) such thatR∗ ≥ 0 almost everywhere under (Gφ)# πyδ . Define M1 := Eπyδ ∥∥yδ −AGφ∗(yδ)∥∥22 and M2 := W1(πx, (Gφ∗)#πyδ). Then, the following two upper bounds hold for every η > 0:
• Pπ yδ
{ yδ : ∥∥yδ −AGφ∗(yδ)∥∥22 ≥ η} ≤ M1η . • Suppose,R∗(x) = 0 for πx-almost every x. Then, P(Gφ∗ )#πyδ { x : R∗(x) ≥ η } ≤ M2η .
Proposition 2 provides an estimate in probability of the sets {yδ : ∥∥yδ −AGφ∗(yδ)∥∥22 ≥ η} and
{x : R∗(x) ≥ η}. In particular, if M1 is small, then ∥∥yδ −AGφ∗(yδ)∥∥22 is small in probability. If instead M2 is small, thenR∗(x) is small in probability on the support of (Gφ∗)#πyδ , implying that samples Gφ∗(yδ) are difficult to distinguish from the ground-truth. We remark that the assumption R∗(x) = 0 can be justified using a data manifold assumption as in Section 3.3. of [17]. We now analyze the role of the regularizer R∗ in the optimization of the variational problem (8) that refines the end-to-end reconstruction Gφ∗ . We rely on a similar distributional analysis as the one performed in [17]. For η > 0, consider the transformation by a gradient-descent step onR∗ given by gη(x) = x− η∇R∗(x). Using the shorthand πG∗ := (Gφ∗)#πyδ , and by denoting the distribution of gη(x) as πη := (gη)#πG∗ for x ∼ πG∗ , we have the following theorem. Theorem 5 ([17]). Suppose that η →W1(πη, πx) is differentiable at η = 0. Then, the derivative at η = 0 satisfies ddηW1(πη, πx) ∣∣∣ η=0 = −EπG∗ ‖∇R∗(x)‖22.
This theorem states that a gradient-descent step performed on R∗ at x = Gφ∗(yδ) decreases the Wasserstein distance with respect to the ground-truth distribution πx. Therefore, if the gradientdescent step to solve the variational problem (8) is initialized with the reconstruction Gφ∗(yδ), the next iterate gets pushed closer to the ground-truth distribution πx. We stress that this property holds because of the chosen initialization point, due to the relation between R∗ and Gφ∗ . For a different initialization, this property may not hold.

heading: 4 Numerical results
text: On the application front, we consider the prototypical inverse problem of CT reconstruction from noisy sparse-view projections. The abdominal CT scans for 10 patients, made publicly available by the Mayo-Clinic for the low-dose CT grand challenge [18], were used in our numerical experiments. Specifically, 2250 2D slices of size 512 × 512 corresponding to 9 patients were used to train the models, while 128 slices from the remaining one patient were used for evaluation. The projections were simulated in ODL [3] using a parallel-beam geometry with 200 uniformly spaced angular positions of the source, with 400 lines per angle. Subsequently, Gaussian noise with standard deviation σe = 2.0 was added to the projection data to simulate noisy sinograms. Notably, the CT projection data are typically corrupted by Poisson noise, but the data-likelihood for Poisson noise is not the `22 distance and is therefore not amenable to the theoretical analysis presented in Sec. 3. Hence, we simulate the noisy projection data by adding Gaussian noise to demonstrate a proof-of-concept for the UAR method. Extending the theoretical results to deal with Poisson noise in the data remains to be undertaken as a future work. Additionally, since the UAR approach seeks to
minimize a variational loss, which includes a data-fidelity term arising from the noise distribution, one needs to systematically investigate the robustness of UAR to the change in noise statistics to make it fully applicable to CT reconstruction from real low-dose projection data.
The proposed UAR method is compared with two classical model-based approaches for CT, namely filtered back-projection (FBP) and total variation (TV). The LPD method [2] and U-net-based postprocessing [11] of FBP are chosen as two supervised approaches for comparison. The AR approach [17] and its convex variant [21], referred to as adversarial convex regularizer (ACR), are taken as the competing unpaired training approaches. For LPD and AR, we develop a PyTorch-based implementation based on their publicly available TensorFlow codes2 3, while for ACR, we use the publicly available PyTorch implementation4.
The unrolled network Gφ has 20 layers, with 5×5 filters in both primal and dual spaces to increase the overall receptive field for sparse-view measurements. The hyper-parameters involved in training the
2LPD: https://github.com/adler-j/learned_primal_dual. 3AR: https://github.com/lunz-s/DeepAdverserialRegulariser. 4ACR: https://github.com/Subhadip-1/data_driven_convex_regularization.
UAR are specified in Algorithm 1. We found that training a baseline regularizer and a corresponding baseline reconstruction operator stabilized the training process. Training UAR took approximately three hours per epoch on an NVIDIA Quadro RTX 6000 GPU (with 24 GB of memory).
The average performance on the test images in terms of PSNR and SSIM [35] indicates that UAR (with λ = 0.1 and no refinement) outperforms AR and ACR by approximately 0.5 dB and 2.8 dB, respectively (see Table 1). We would like to emphasize that this gain was found to be consistent across all test images and not just realized on average. With the refinement step, UAR surpasses AR by almost 1 dB and slightly surpasses the U-net-based post-processing. The end-to-end UAR reconstruction is a couple of orders of magnitude faster than AR, while the reduction in reconstruction time is by a factor of four with the refinement. Since the refinement step entails running a few gradient-descent iterations on a high-dimensional variational loss, it makes UAR slower than supervised end-to-end methods. However, thanks to a superior initialization provided by the generator, one requires significantly fewer gradient-descent iterations to refine as compared to a fully variational scheme such as AR. The reconstructions of a representative test image using the competing methods are shown in Fig. 1 for a visual comparison. The effect of λ on the reconstruction of UAR is demonstrated in Fig. 2, which confirms the theoretical results in Section 3.2. The refinement step also visibly improves the reconstruction quality of the end-to-end operator, as shown in Fig. 3.
Some additional numerical examples for the task of CT reconstruction, and some illustrative examples of the performance of UAR on two other important imaging inverse problems, namely inpainting and denoising, are provided in Sec. B of the supplementary document.

heading: 5 Conclusions and outlook
text: To the best of our knowledge, this work makes the first attempt to blend end-to-end reconstruction with data-driven regularization via an adversarial learning framework. The proposed UAR approach retains the fast reconstruction of the former together with provable guarantees of the latter. We rigorously analyze the proposed framework in terms of well-posedness, noise-stability, and the effect of the regularization penalty, and establish a link between the trained reconstruction operator and the corresponding variational objective. We show strong numerical evidence of the efficacy of the UAR approach for CT reconstruction, wherein it achieves the same performance as supervised data-driven post-processing and outperforms competing unsupervised techniques. Our work paves the way to better understand the role of adversarially learned regularizers in solving ill-posed inverse problems, although several important aspects need further investigation. Since the learned regularizer is non-convex, the performance of gradient-descent on the variational objective greatly depends on initialization. This problem is partly addressed by the unrolled reconstruction operator that efficiently computes a better initial point for gradient descent. However, the precise relationship between the end-to-end reconstruction and the variational minimizer for a given measurement vector remains elusive. Moreover, the quality of the reconstruction relies on the expressive power of neural networks and thus suffers from the curse of dimensionality. We believe that addressing such limitations will be important to better understand adversarial regularization methods.

heading: 6 Acknowledgments
text: MC acknowledges support from the Royal Society (Newton International Fellowship NIF\R1\192048 Minimal partitions as a robustness boost for neural network classifiers). CBS acknowledges sup-
port from the Philip Leverhulme Prize, the Royal Society Wolfson Fellowship, the EPSRC grants EP/S026045/1 and EP/T003553/1, EP/N014588/1, EP/T017961/1, the Wellcome Innovator Award RG98755, the Leverhulme Trust project Unveiling the invisible, the European Union Horizon 2020 research and innovation programme under the Marie Skodowska-Curie grant agreement No. 777826 NoMADS, the Cantab Capital Institute for the Mathematics of Information, and the Alan Turing Institute.

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

This paper proposes a method for adversarial regularization, to be used in image reconstruction. This method relies on a joint training of the regularizing part of the system and the reconstruction. Several properties regarding convergence are proven, and the method is shown to outperform similar regularization methods for this task.

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. 

This paper proposes a framework for training an adversarial network, where the discriminator is a regularization functional, for image reconstruction.  A similar work mentioned in the paper is [14]. The key difference between [14] and the proposed method, UAR, is that the former is based on the variational inference while the latter replaces it with an unfolding network, making UAR a pure end-to-end training deep learning model. Since UAR incorporates the forward model in the loss function, it is self-supervised requiring only $y$ to impose supervision.  The author also addresses the theoretical properties of UAR, presenting interesting results on the optimality, influence of $\lambda$, and justification of end-to-end learning & regularization.

This paper proposes a framework, named UAR, for jointly training an iterative deep unfolding neural network (DUNN) and an adversarial discriminator based on Wasserstein distance for solving image inverse problems, inspired by adversarial regularization (AR) method (Lunz’2018). Unlike original AR method that takes pseudo-inverse reconstruction as input images, the proposed method assimilates the reconstruction from DUNN that unrolls the iterative primal-dual algorithm in each training step. The authors also provide theoretical analysis about when such adversarial training with unrolling networks would be successful. The theoretical results are well established by with thorough proofs. Finally, the performance of the proposed method is evaluated for solving CT inverse problem with satisfactory results compared to both supervised and unsupervised deep learning methods such as U-Net, LPD, and AR.