Stochastic Optimisation Framework using the Core Imaging Library and Synergistic Image Reconstruction Framework for PET Reconstruction

Abstract

We introduce a stochastic framework into the open–source Core Imaging Library (CIL) which enables easy development of stochastic algorithms. Five such algorithms from the literature are developed, Stochastic Gradient Descent, Stochastic Average Gradient (-Amélioré), (Loopless) Stochastic Variance Reduced Gradient. We showcase the functionality of the framework with a comparative study against a deterministic algorithm on a simulated 2D PET dataset, with the use of the open-source Synergistic Image Reconstruction Framework. We observe that stochastic optimisation methods can converge in fewer passes of the data than a standard deterministic algorithm.

I Introduction

Iterative reconstruction methods have been applied with great success for solving challenging optimisation problems, such as total variation (TV) regularisation. Since iterative methods are computationally demanding due to the increasingly large data sizes, a range of stochastic optimisation algorithms have been proposed in the literature to reduce the computational effort.

In this work, we extend the optimisation functionality of the Core Imaging Library (CIL) [jorgensen2021core, papoutsellis2021core] with a stochastic framework that enables develo** a range of stochastic algorithms found in the literature: Stochastic Gradient Descent (SGD), Stochastic Average Gradient (SAG) [schmidt2017minimizing], SAG- Amélioré (SAGA) [defazio2014saga], Stochastic Variance Reduced Gradient (SVRG) [johnson2013accelerating] and Loopless SVRG (LSVRG) [kovalev2020don]. These add to the Stochastic Primal-Dual Hybrid-Gradient (SPDHG) currently available in CIL, see [papoutsellis2021core] and references therein. We demonstrate the use of this framework on a Positron Emission Tomography (PET) application, thanks to the combined use of the Synergistic Image Reconstruction Framework (SIRF) [ovtchinnikov2020sirf].

The developed framework allows for an easy comparison between stochastic gradient estimators. In this summary, we observe that stochastic optimisation methods can converge in fewer passes of the data than a deterministic benchmark.

Figure 1: Simulated 2D ¹⁸F-FDG PET thorax dataset, reference solution $x^{*}$ and error plots for FISTA, Prox-SGD, Prox-LSVRG and SPDHG at 100 data passes.

II Stochastic Framework

We consider optimisation problems of the form

$$x^{*}\!=\!\operatorname*{arg\,min}_{x\in\mathbb{X}}\{F(x):=f(x)+g(x)\}\equiv\sum_{i=1}^{n}f_{i}(x)+g(x),$$ (1)

where $\mathbb{X}$ is a finite dimensional space. Functions $f_{i},f,g:\mathbb{X}\rightarrow\mathbb{R}$ for $i\in\{1,...,n\}$, are proper and convex, where $f_{i}$ are $L$-smooth and represent the fitness to the data. Regulariser $g$ has a proximal operator which either has a closed-form representation or can be efficiently solved up to some precision.

Proximal gradient descent [beck2009fast] (also known as ISTA, or forward-backward splitting) is a classical deterministic algorithm to solve (1) by the iterations

$$x_{k+1}=\text{prox}_{\gamma_{k}g}\left(x_{k}-\gamma_{k}\nabla f(x_{k})\right),\quad k=0,1,2,\dots$$ (2)

for step size $\gamma_{k}$ and initial guess $x_{0}$. When $g\equiv 0$ this reduces to gradient descent. Instead of computing the full gradient $\nabla f(x_{k})$ in each iteration, stochastic optimisation algorithms employ an estimator $\tilde{\nabla}f(x_{k})$, typically using the information of only one randomly selected function $f_{i}$.

The stochastic framework in CIL consists of four components that can be combined in a plug-and-play fashion: i) functions providing stochastic estimators for the gradient of $f$, ii) sampling methods which take in a set of probabilities $p_{i}$ for choosing each of the functions $f_{i}$, iii) a partitioner to split up the data, defining the $f_{i}$’s, and, iv) algorithms to solve (1).

Thus far we have implemented 5 functions that provide stochastic estimators for the gradient of $f$. When these are used in combination with GD or ISTA algorithms of CIL, they correspond to SGD (Prox-SGD), SAG (Prox-SAG) [schmidt2017minimizing], SAGA (Prox-SAGA) [defazio2014saga], SVRG (Prox-SVRG) [johnson2013accelerating] and LSVRG (Prox-LSVRG) [kovalev2020don]. Due to our flexible design, the above stochastic estimators can also be combined with Nesterov-type accelerated algorithms, e.g., FISTA [beck2009fast], see [Driggs2020] for more details.

III Methodology and Results

For the numerical study, we use a simulated 2D ¹⁸F-FDG PET dataset from SIRF . Simulated Poisson noise is applied to the acquisition data. The data is partitioned into 32 subsets with equidistant projection views. Kullback-Leibler data fitting term is used for $f_{i}$s and $f$. TV with a non-negativity constraint is used for the regulariser, $g=\alpha\text{TV}$, where $\alpha=0.1$. Results are shown in figure 1.

The optimal reconstruction $x^{*}$ was obtained using 500 data passes of SPDHG. All the algorithms are warm-started with one data pass of Prox-SGD. Functions $f_{i}$ are selected randomly with replacement. Algorithmic parameters such as step size, update frequency for Prox-SVRG, and probability for Prox-LSVRG were optimised using a parameter search. In figures 2 and 3, we compare the different stochastic algorithms in CIL, and their performance with respect to “data passes”, i.e. how many times the algorithm has processed all the acquisition data in expectation. All the proposed stochastic algorithms have a faster convergence rate to the optimal solution than the deterministic FISTA.

Figure 2: Distance from the optimal solution $x^{*}$ with respect to “data passes”.

Figure 3: Distance from the optimal objective value $F(x^{*})$ with respect to “data passes”.

IV Discussion and Future Work

This contribution describes an open-source framework that enables investigating a large variety of optimisation algorithms in many different contexts, including CT, MR, PET and SPECT image reconstruction. Presented results for PET show that stochastic algorithms in CIL can converge in fewer data passes than a deterministic counterpart. However, comparing the stochastic algorithms fairly will require a thorough investigation of algorithmic parameters, such as step size regimes. We leave this comparison for future work.

In addition, work is in progress to further empirically validate the stochastic framework by applying it to real PET data and to expand the versatility of the stochastic framework by extending its applicability to a wider array of stochastic algorithms, diverse imaging modalities, and integrating additional methodologies such as acceleration and pre-conditioning.