Why do we regularise in every iteration
for imaging inverse problems?

The possibility to skip the computation of the proximal operator in some iterations according to a probability $p$, accelerating the algorithms, without affecting convergence, was discussed in[8]. There, the ProxSkip algorithm was introduced to tackle federated learning applications which also rely on computations of expensive proximal operators. ProxSkip introduces a control variable $\bm{h}_{k}$, see Algorithm 4. When the proximal step is not applied, the control variable remains constant. Hence, if at iteration $k$, no proximal step has been applied previously, the accumulated error is passed to $\bm{x}_{k+1}$ without incurring an additional computational cost. If at the iteration $k$ the proximal step is applied, the error is reduced and the control variable will be updated accordingly.

In [8] it was shown that the ProxSkip converges provided that $f$ in (2) is L-smooth and $\mu$-strongly convex, and probability $p$ satisfies

In the case of equality in (4), the algorithm converges (in expectation) at a linear rate with $\gamma=\frac{1}{L}$ and the iteration complexity is $\mathcal{O}(\frac{L}{\mu}\log(\frac{1}{\varepsilon}))$. In addition, the total number of proximal evaluations (in expectation) are only $\mathcal{O}(\frac{1}{\sqrt{p}}\log(\frac{1}{\varepsilon}))$.

Our aim is to showcase for the first time via extended numerical experiments the computational benefits of ProxSkip for a variety of imaging inverse problems, including challenging real-world tomographic applications. In particular, we show that ProxSkip can outperform the accelerated version of its non-skip analogue, namely FISTA. At the same time, we introduce a novel PDHGSkip version of the PDHG, Algorithm 6, which we motivate via numerical experiments. We anticipate that this will spark further research around developing skip-versions of a variety of proximal based algorithms used nowadays.

For all our imaging experiments we consider the following optimisation problem that contains a quadratic distance term as the fidelity term, with the (isotropic) total variation as the regulariser, i.e.,

where $\bm{D}\bm{u}=(\bm{D}_{y},\bm{D}_{x})$ is the finite difference operator under Neumann boundary conditions.

To showcase the algorithmic properties we consider a toy example, with the dual formulation of the classical TV denoising (ROF) which reads

where $\bm{\mathrm{div}}$ is the discrete divergence operator such that $\bm{\mathrm{div}}=-\bm{D}^{T}$. The solutions $\bm{u}^{*}$ and $\bm{q^{*}}$ of the primal and dual (ROF) problems are linked via $\bm{u}^{*}=\bm{b}+\bm{\mathrm{div}}\bm{q^{*}}$. A simple algorithm to solve (6) was introduced in [9], based on a Projected Gradient Descent (ProjGD) iteration $\bm{q_{k+1}}=\mathcal{P}_{C}(\bm{q_{k}}+\gamma\bm{D}(\bm{\mathrm{div}\bm{q_{k}}}+\bm{b}))$ which is globally convergent under a fixed stepsize $\gamma\leq\frac{2}{\|\bm{D}\|^{2}}$ [10], with

and $d$ is the image dimension. This approach became quite popular in the following years for both its simplicity and efficiency, [11], and for avoiding computing smooth approximations of TV. The projection $\mathcal{P}_{C}$ can be identified as the proximal operator of indicator function of the feasibility set $C$. Thus, ProjGD is a special case of PGD and we can apply the ProxSkip Algorithm 4. Note that due to the divergence operator this problem is not strongly convex. In fact, this is the case for the majority of the problems of the type (1), typically due to the non-injectivity of $\bm{A}$. Thus, this example also shows that the strong convexity assumption could potentially be relaxed for imaging inverse problems.

To ensure that any biases from algorithms under evaluation are avoided, the “exact” solution $\bm{u}^{*}$ is calculated using an independent high-precision solver, in particular, the MOSEK solver from the CVXpy library [12, 13], see Figure 1.

Both the ProjGD and ProxSkip algorithms use the stepsize $\gamma=\frac{1}{L}=\frac{1}{8}$, where $L$ is the Lipschitz constant of $\mathcal{F}^{\prime}(\cdot)$. For every iteration, we monitor the $\ell^{2}$ error $\|\bm{u}_{k}-\bm{u}^{*}\|_{2}$ between the iterate $\bm{u}_{k}=\bm{b}+\bm{\mathrm{div}\bm{q_{k}}}$ and estimated exact solution $\bm{u}^{*}$. We use $p=[0.01,0.1,0.3,0.5]$ and 50000 iterations as a stopping criterion.

In Figure 2 (top-right), it can be observed that these two algorithms are almost identical in terms of the $\ell_{2}$ error. Note that ProxSkip and ProjGD coincide only when $p=1$. Indeed, one can detect some discrepancies during first 100 iterations, which quickly dissipate with only a few applications of the projection $\mathcal{P}_{C}$, see bottom row of Figure 2. In Figure 2 (top-left), we plot the $\ell_{2}$ error with respect to CPU time. The shown CPU time is the average over 30 independent runs of each algorithm. We observe a clearly superior performance of ProxSkip, for all values of $p$. This serves as a first demonstration of the advantage of ProxSkip in terms of computational time without affecting the quality of the image. In Section 3, we present a more emphatic computational impact using heavier proximal steps.

In order to be consistent with the convergence theory of ProxSkip where strong convexity is a requirement, one can add a small quadratic term to the objective function. This is a commonly used in imaging applications and allows the use of accelerated versions of first-order methods. For the (6) problem this results in

It is known that the corresponding primal problem of (8) is the standard Huber-TV denoising [7] which involves a quadratic smoothing of the $\|\cdot\|_{2,1}$-norm around an $\varepsilon$-neighbourhood of the origin. Among other effects, this reduces the staircasing artifacts of TV, see last image of Figure 1.

We repeat the same experiment as in the previous section for (8) with 5000 iterations and over 30 independent runs of the algorithm, using the probability  $p=0.04767$, given by (4). This results in on average using only 215 projection steps during 5000 iterations. In addition to ProjGD, we also compare its accelerated version, denoted by AProjGD [14] which is essentially Algorithm 1 with the acceleration step of Algorithm 3. In Figure 3 (left), we observe that AProjGD performs better than both ProjGD and ProxSkip with respect to iteration number. While a similar behaviour is observed with respect to the CPU time, see Figure 3 (right), the average time for AProjGD and ProxSkip to reach a relative error less than $10^{-6}$ is nearly identical. One would expect a similar speed-up when AProjGD is combined with skipping techniques but we leave this for future research.

We now consider more challenging imaging tasks, involving proximal operators that do not admit closed form solutions, and which thus require computationally intensive iterative solvers.

We start with a deblurring problem in which $\bm{A}$ is a convolution operator, see Figure 4.

We solve (5) and compare ISTA, FISTA and ProxSkip with different values of $p$. Here the proximal operator corresponds to a TV denoiser for which we employ AProjGD with a fixed number of iterations as an inner iterative solver, see next section for other feasible options. In the framework of inexact regularisation another option is to terminate the inner solver based on some metric and predefined threshold [15]. As noted therein, the number of required inner iterations typically increases up to $10^{3}$, as the outer algorithm progresses, leading to higher computational costs over time. To explore both the computationally easy and hard cases, we run the inner solver with 10 and 100 iterations, and use a warm-start strategy. Warm-start is a vital assumption for inexact regularisation as it avoids semi-convergence, where the error stagnates and fails to reach high precision solutions. To avoid biases towards proximal-gradient based solutions, the “exact” solution $\bm{u}^{\ast}$ is computed using 200000 iterations of PDHG with diagonal preconditioning [16]. Outer algorithms are terminated if either 3000 iterations are reached or the relative distance error is less than $\varepsilon=10^{-5}$. The reported CPU time is averaged over 10 runs.

In Figure 5, we first observe that solving the inner TV problem with 10 iterations seriously affects the convergence of FISTA. Notably, ProxSkip versions are less affected even though besides the error introduced by the inexact solver, there is also an error from skipping the proximal. By raising the number of inner iterations to 100, and thus increasing the accuracy of the inner solver, we observe that FISTA exhibits an early decay albeit with some oscillations and it terminates after around 500 iterations, reaching an accuracy of $\varepsilon=10^{-5}$, see Figure 5 (bottom-left). On the other hand, ISTA and the ProxSkip require many more iterations to reach the same level of accuracy. However, remarkably, in this regime ProxSkip is significantly faster, in terms of CPU time, than ISTA and even outperforms FISTA when $p=0.05$ and $0.1$, see Figure 5 (bottom-right). In Table 1, we report all the information for the best three algorithms that outperformed FISTA with 100 inner iterations in terms of CPU time.

We note that there are versions of FISTA [17, 18, 19, 20] with improved performance, also avoiding oscillations. However, our purpose here is not an exhaustive comparison but to show that a simple version of ProxSkip outperforms the simplest version of FISTA. We anticipate the future development of more sophisticated skip-based algorithms including ones based on accelerated methods.

Skipping the proximal can also be applied to primal-dual type algorithms. Here, for convex $f,g$ and a bounded linear operator $\mathcal{K}$, the optimisation framework is

In[21], a skip-version of PDHG [7] was proposed, which we denote here by PDHGSkip-1. It allows not only to skip one of the proximal steps but also the forward or backward operations of $\mathcal{K}$, depending on the order of the proximal steps, see Algorithm 5. In step 4, a Bernoulli operator is used, defined as $\mathcal{B}_{p}(\bm{x})=\bm{x}/p$ with probability $p$ and $0$ otherwise. Again, strong convexity of both $f$ and $g$ is required for convergence. However, our imaging experiments – in the absence of strong convexity – revealed a relatively slow performance, even with optimised step sizes $\sigma,\tau$ and $\omega+1=1/p$ according to [21, Theorem 7],

To address the slow convergence, we introduce a modification: PDHGSkip-2, see Algorithm 6. The difference to PDHGSkip-1 is that the adjoint $\mathcal{K}^{T}$ and the proximal operator of $g$ are now separated, see steps 4-7. Note that in both versions, when $p=1$ the control variable $\bm{h}$ vanishes and PDHG is recovered. To illustrate the difference of the two versions in practice, we show in Figure 6 a simple comparison on tomography, see next section for details. Apart from the clear acceleration of PDHGSkip-2 over PDHGSkip-1. We also observe a staircasing pattern for the relative error for PDGHSkip-1, see detailed zoom of the first iterations. This is expected since in most iterations, where the proximal step is skipped, one variable vanishes without contributing to the next iterate. Hence, the update remains unchanged.

In general, $\ell^{2}$-$\mathrm{TV}$ problems can be solved using implicit or explicit formulations of PDHG. In the implicit case, $f$ is the $\|\cdot\|_{2}^{2}$ term, $\mathcal{K}=\bm{A}$ and $g$ is the $\mathrm{TV}$ term. In the explicit case, $f$ is a separable sum of $\|\cdot\|_{2}^{2}$ and $\|\cdot\|_{2,1}$ composed with block operator $\mathcal{K}=[\bm{A},\bm{D}]^{T}$ and $g$ can be a zero function or a non-negativity constraint. Here, every proximal step (Steps 6, 9 in Algorithm 6) has an analytic solution. This significantly reduces the cost per-iteration, but also requires more iterations to reach a desired accuracy [22]. Inexact regularisation is usually preferred and typically reduces the number of iterations. Hence, the number of calls of the forward and backward operations of $\bm{A}$ is reduced, which in certain applications gives a considerable speed-up.

For our final case study, we solve (5) under a non-negativity constraint for $\bm{u}$ for a real-world tomographic reconstruction. Here $\bm{A}$ is the discrete Radon transform and $\bm{b}$ is a noisy sinogram of a real chemical imaging tomography dataset, representing post partial oxidation of methane reaction Ni-Pd/CeO2-ZrO2/Al2O3 catalyst [23, 24], see Figure 7. The initial dataset was acquired for 800 projection angles with 695$\times$695 detector size and 700 vertical slices. For demonstration purposes and to be able to perform multiple runs for computing more representative CPU times, we consider one vertical slice with half the projections and $2\times$ rebinned detector size. Same conditions for the inexact solver and stopping rules are used as in previous section. For these optimisation problems, one can use algorithms that fit both the general frameworks (2) and (9), see [25].

The two algorithms use the same values for step sizes $\sigma$ and $\tau$, satisfying $\sigma\tau\|\mathcal{K}\|^{2}\leq 1$. In the comparisons presented in Figure 8, we observe a similar trend for the proximal-gradient based algorithms as in Section 3.1. When we use 10 iterations for the inner solver, FISTA fails to reach the required accuracy and stagnates, which is not the case for the ProxSkip algorithms even for the smallest $p$. This is accompanied by a significant CPU time speed-up, which further increases when more inner iterations are used. In fact, by increasing the number of inner iterations, the computational gain is evident with around 90% speed-up compared to FISTA.

The best overall performance with respect to CPU time is achieved by PDHGSkip-2. There, distance errors are identical to the case $p=1$, in terms of iterations, except for the extreme case $p=0.1$ which oscillates in the early iterations. For 10 inner iterations and $p=0.3,0.5$ we observe a slight delay towards $\varepsilon=10^{-5}$, due to error accumulation caused by skipping the proximal operator and limited accuracy of the inner solver. This is corrected when we increase the number of inner iterations. Moreover, we see no difference with respect to CPU time, for 10 inner iterations and $p=0.7,1.0$. This is expected since the computational cost to run 10 iterations of AProjGD is relatively low. However, it demonstrates that we can have the same reconstruction using the proximal operator 70% of the time. Finally, the computational gain is more apparent when we increase the number of inner iterations hence the computational cost of the inner solver.

We note that such expensive inner steps are used by open source imaging libraries and are solved with different algorithms and stopping rules. In CIL [26, 27], the AProjGD is used to solve (6) with 100 iterations as default stopping criterion. In PyHST2 [28], AProjGD is used with 200 iterations and a duality gap is evaluated every 3 iterations. In TIGRE [29], the PDHG algorithm with adaptive step sizes is applied to (6), [30], with 50 iterations. In DeepInv ¹¹1https://github.com/deepinv/deepinv PDHG is applied to (5) (with $\bm{A}=\mathbf{Id}$) with 1000 iterations and error distance between two consecutive iterates. Finally, in Tomopy[31], PDHG is applied to (5) (with $\bm{A}=\mathbf{Id}$) and the number of iterations is specified by the user. All these default options are optimised and tested for particular real-world tomography applications like the ones encountered in synchroton facilities for instance. Alternatively, one can avoid inexact solvers for TV denoising [32]. There, the proximal is replaced by a combination of wavelet and scaling transforms and is computed using a componentwise shrinkage operator. Even in this case, we expect an improvement by skipping this operator. Overall, we anticipate a computational gain proportional to the cost savings achieved by omitting specific mathematical operations.