Learning convex regularizers satisfying the variational source condition for inverse problems
Abstract
Variational regularization has remained one of the most successful approaches for reconstruction in imaging inverse problems for several decades.
With the emergence and remarkable success of deep learning in recent years, a considerable amount of research has gone into data-driven modeling of the regularizer in the variational setting.
Our work extends a recently proposed method, referred to as adversarial convex regularization (ACR), that seeks to learn a data-driven convex regularizer via adversarial training in an attempt to combine the power of data with the classical convex regularization theory.
Speciﬁcally, we leverage the variational source condition (SC) during training to enforce that the ground-truth images minimize the variational loss corresponding to the learned convex regularizer.
This is achieved by adding an appropriate penalty term to the ACR training objective.
The resulting regularizer (abbreviated as ACR-SC) performs on par with standard ACR, but unlike ACR, comes with a quantitative convergence rate estimate.
Introduction
Linear inverse problems seek to recover an unknown parameter x ∗ ∈ X from its noisy measurement y δ
= Ax ∗ + e ∈ Y , where A : X → Y is a bounded linear operator between the Hilbert spaces X and Y , and e denotes measurement noise with (cid:107) e (cid:107) Y ≤ δ .
The clean measurement corresponds to δ =
and is denoted by y 0 .
Inverse problems are typically ill-posed, in the sense that A is either non-invertible or poorly conditioned, leading to noise ampliﬁcation in the solution obtained via naïve inversion.
Variational methods [4, 13], wherein one seeks to trade-off data-ﬁdelity with a regularizer, has traditionally been the most popular and successful approach for computing a stable solution to ill-posed inverse problems and are rooted in a rigorous function-analytic foundation.
With the advent of deep learning, considerable attention has been devoted to leveraging the availability of data for solving inverse problems [3].
A particularly notable deep learning-based endeavor has sought to model the regularizer via over-parametrized deep neural networks and learn it in a datadriven manner, instead of using hand-crafted functionals [9, 10, 12, 8].
Such methods inherit the theoretical guarantees (such as stability and convergence) offered by the variational framework while effectively utilizing the power of data, provided that the regularizer fulﬁlls certain conditions (such as convexity [12]).
This work essentially builds upon the approach introduced in [12].
By incorporating an additional penalty term during training, our approach encourages the ground-truth images to satisfy the so-called variational source condition (SC), which leads to precise convergence rate estimates.
Variational source condition
In [5], convergence rate estimates for variational reconstruction with convex regularizers were derived under the source condition, which provides an additional regularity condition on the solution of ill-posed inverse problems (see [4] for a detailed discussion).
Here, we brieﬂy recall the key results to make the exposition self-contained.
Consider the variational regularization approach that minimizes an energy functional, given by
[Equation]
Here, { ψ θ } θ ∈ Θ is a convex regularizer parametrized by an input-convex neural network (ICNN) [2].
Deﬁnition 1 (Source condition) .
The ψ θ -minimizing solution is deﬁned as
[Equation]
The variational source condition is satisﬁed if there exists some ˜ w ∈ Y such that A ∗ ˜ w ∈
∂ ψ θ (˜ x ) .
One can show that the set of ˜ x satisfying the SC is the same as the set of solutions to the variational problem (1), see [5, Proposition 1].
This serves as the main motivation behind our approach, in that we encourage the ground-truth images during training to satisfy the SC and thereby be the solution to the variational problem corresponding to the learned regularizer.
Let D ψ θ ( x 1 , x 2 ) :
= (cid:8) ψ θ ( x 1 ) − ψ θ ( x 2 ) − (cid:104) u, x 1 − x 2 (cid:105) (cid:12) (cid:12) u ∈ ∂ψ θ ( x 2 ) (cid:9) be the Bregman distance corresponding to ψ θ .
Then, the following convergence rate can be established (see [5]):
Theorem 1.
Let (cid:107) y δ − y
(cid:107) ≤ δ , ˜ x be a ψ θ -minimizing solution as deﬁned in (2) , and suppose that
the
SC holds.
Then, for each minimizer x λ of (1) , there exists d ∈ D ψ θ ( x λ , ˜ x ) such that d ≤ λ (cid:107) [Equation] 2 2 +
[Equation]
2 λ holds.
Therefore, choosing λ ∝ δ leads to an O ( δ ) convergence rate.
Consequently, SC enables one to derive quantitative convergence rate estimates in terms of the
Bregman distance induced by the learned convex regularizer.
In contrast, [12] shows convergence, without any quantitative rate estimate, while having to incorporate an additional Tikhonov term.
2.1 Learning ACR with the source condition (ACR-SC)
If the forward operator A is invertible, SC dictates that (cid:96) sc ( x ; θ ) = (cid:107) ( A ∗ ) − 1 ∇ x ψ ( x ; θ ) (cid:107) < ∞ must be satisﬁed whenever x is a solution to (1).
Therefore, the smaller the quantity (cid:96) sc ( x ; θ ) is, the more suitable x would be as a solution to (1).
A natural way to encourage this is to penalize
L sc ( θ ) = 1 n (cid:80) ni
=1 (cid:96) sc
( x i ; θ ) , where x i ∼ P r are the ground-truth images.
For non-invertible A , one can replace the inverse with the Moore-Penrose pseudo-inverse.
Denoting by P n the distribution of the undesirable images in the adversarial training framework, the overall training objective becomes
[Equation]
Here, x i ∼ P r and z i ∼ P n are the ground-truth and noisy/undesirable images in the training dataset, respectively.
The gradient penalty is given by L gp ( θ )
= 1 n (cid:80) n i =1 ( (cid:107)∇ ψ θ ( (cid:15)x i + (1 − (cid:15) )
z i ) (cid:107) 2 − 1) 2 , where (cid:15) ∼ uniform [0 , 1] , which encourages the regularizer to be 1-Lipschitz (cf.
Notably, the training loss in (3) can be computed without direct supervision (i.e., without pairs of noisy and ground-truth images), thus offering more ﬂexibility.
For solving the variational problem (1) with the learned convex ψ θ , one can employ a simple iterative sub-gradient algorithm for minimization.
Notably, an advantage of the enforced SC is that Bregman iteration techniques are available for (1) and well understood (cf.
[4] for details).
We show some preliminary promising results here (see Figure 1) and will further pursue this approach in future research.
Numerical results
We ﬁrst show a proof-of-concept of the proposed ACR-SC method considering a denoising problem on MNIST and subsequently compare it with a number of state-of-the-art supervised and unsupervised data-driven methods for the prototypical inverse problem of computed tomography (CT) reconstruction from sparse-view projection data.
In both experiments, the reconstruction quality is evaluated in terms of the peak signal-to-noise ratio (PSNR) and the structural similarity index (SSIM) with respect to the target ground-truth.
[Table]
For the denoising experiment on MNIST, the clean digits are corrupted with an additive white
Gaussian noise with σ = 0 .
The regularizer is trained over 10 epochs with a batch-size of 64 (with λ sc = 2 . 0 and λ gp = 10 . 0 ), by minimizing the training loss (3).
Subsequently, the learned regularizer is plugged into the variational framework (1).
We employ both gradient-descent (with λ = 5 ,iterations with a step-size of 0 . 01 ) and Bregman iterations [4] (with λ = 25 ) for solving (1).
The numerical examples shown in Fig.
1 demonstrate that ACR-SC performs a reasonable denoising and signiﬁcantly improves the PSNR and SSIM over the noisy input.
As shown in Fig.
1, Bregman iterations do a noticeably better job of circumventing loss of contrast in the reconstruction.
For Sparse-view CT reconstruction, We adopt the same experimental setting considered in [12] to ensure a fair comparison, and we brieﬂy recall it here for the sake of completeness.
For training the models, we use the publicly available data for the Mayo-clinic low-dose CT challenge [11].
All models are trained on 9 patients (2250 2D slices) and tested on the remaining one patient (128 slices).
Parallel-beam projection data are simulated in ODL [1] with 200 uniformly spaced angles andlines per angle, followed by additive Gaussian noise contamination (with σ = 2 . 0 ).
For comparison, we consider the classical ﬁltered back-projection (FBP) and total variation (TV) regularization as two representative model-based approaches.
Among data-driven approaches, we compare with two supervised methods, namely (i) U-Net-based post-processing of FBP [7] and (ii) learned primal-dual (LPD); and two unsupervised approaches, namely (i) the adversarial regularization (AR) method introduced in [10] and (ii) its convex variant ACR [12].
Together with PSNR and SSIM, the reconstruction times are also reported to facilitate an easier assessment of the quality-vs.-time trade-off for different approaches (c.f. Table 1).
The architecture of the ICNN that models the regularizer in ACR-SC is taken to be identical to the one considered in [12].
However, unlike [12], our regularizer does not have the sparsifying ﬁlter-bank and the squared- (cid:96) 2 terms.
We choose λ sc = 2 . 0 and λ gp = 5 . 0 , and the regularizer is trained over 10 epochs with a batch-size of four.
Adam optimizer with η, β 1 , β 2 = 10 − 5 , 0 . 90 , 0 . 99 is used for training.
The variational problem with the resulting regularizer is solved via gradient-descent ( λ = 0 .
05 , step-size 0.8, 400 iterations).
The numerical results reported in Table 1 and Fig.
[Figure]
indicate that ACR-SC is only marginally inferior to ACR in terms of the quality metrics, while it still outperforms classical hand-crafted regularizers such as total-variation (TV).
The slight deterioration as compared to ACR can be attributed to the additional regularity constraint imposed by the SC.
Conclusions
We addressed the problem of learning data-adaptive convex regularizers for inverse problems and proposed an approach to augment the existing adversarial learning framework with a suitable penalty that enforces the variational source condition.
Incorporating the source condition penalty leads to only minor degradation in the numerical performance, but it simultaneously offers a theoretical grounding for Bregman iterations and paves the way for deriving convergence rate estimates for the resulting variational reconstruction problem as the noise-level approaches zero.