Stochastic Three-Operator Splitting Algorithms for Nonconvex and Nonsmooth Optimization Arising from FLASH Radiotherapy

Radiation therapy (RT) has long been a cornerstone of cancer treatment. With the continuous advancement of medical technology, there are various radiation therapy techniques available, such as IMPT and PBS [54, 18, 27]. Although these techniques have been widely adopted in the medical field, complete eradication of cancerous tissue still relies on the dose, which is limited by the risk of severe radiation-induced side effects. A new technique called flash radiotherapy (FLASH-RT) has emerged. FLASH-RT delivers ultra-high dose rates, several orders of magnitude higher than conventional therapy [34]. Clinical data [34, 52] indicates that FLASH-RT can achieve better control with fewer side effects. This is primarily because flash radiotherapy irradiates tissues at ultra-high dose rates ($\geq$ 40 Gy/s), potentially reducing the toxicity of normal tissues while retaining tumor-killing effects. Proton irradiation can achieve this ultra-high dose rate, due to the unique depth-dose characteristics of proton beams, allowing the delivery of ultra-high dose rates to deep tissues [37, 8]. Additionally, clinical data [48, 34] has demonstrated that FLASH proton radiotherapy, due to its superior immune response capabilities, has become an important treatment modality for certain acute and advanced cancers.

While FLASH proton radiotherapy may become the primary radiotherapy method for certain tumors, current research on FLASH proton radiotherapy planning has mainly focused on optimizing dose distribution. However, each spot has constraints not only on dose distribution but also on dose rate distribution, highlighting the importance of exploring a new treatment optimization method called FLASH proton radiotherapy with dose and dose rate optimization[19, 55, 14]. Mathematically, it can be formulated as the following constraint nonconvex optimization:

where $x$ is the spot weight to be optimized, $b$ is the dose objective, $A$ represents the forward system matrix, $Q$ is a block matrix composed of the forward matrix that only maps $x$ to the dose in the ROI, $u$ is a vector associated with DVH constraints, and $\alpha$ represents the planning monitor unit. Here $H$ exhibits a large-scale summation structure, representing the sum of planning objectives. See section 4 for more details. Currently, most of the research on solving (1.1) employs the nonstandard alternating direction method of multipliers (ADMM) algorithm to solve the convex relaxation problem. A classical technique for solving constraint optimization is to transform it into a multiple summation problem via indicator functions. Here, we employ this technique to equivalently transform (1.1) into the following large-scale nonconvex optimization:

where $C=\big{\{}x\in\mathbb{R}^{n}|x_{i}\in\{0\}\bigcup[\alpha,+\infty),\,\forall\,i=1,2,\dots,n\big{\}}=\mathbb{R}_{+}^{n}/(0,\alpha)^{n}$, the set $D=\big{\{}x\in\mathbb{R}^{n}|Qx\geq u\big{\}}$. This transformation allows us to apply classical splitting algorithms to solve (1.2). The three-operator splitting algorithm [6] is an efficient algorithm for solving this kind of composite optimization. However, in our case, $H$ has a large-scale summation structure, and computing $\nabla H$ can be computationally expensive, reducing the efficiency of the method. This prompts us to propose a novel stochastic three-operator splitting algorithm and motivates the research presented below.

Driven by the FLASH proton radiotherapy problem (1.2) and to broaden the applicability of our research, we consider the following general optimization problem:

where $F,G$ and $H$ are proper lower semi-continuous functions. We do not assume convexity for $F$, $G$, or $H$ throughout this paper. Corresponding to the FLASH proton radiotherapy problem (1.2), we know $H=\left\|Ax-b\right\|^{2}$ with a large-scale summation structure, while $F$ and $G$ correspond to indicator functions of two constraint sets. Problem (1.3) not only corresponds to the FLASH proton radiotherapy problem but also captures many other applications, particularly in signal/image processing and computer vision, that involve multiple regularization terms. For instance, deep learning-based models in image processing often include both sparsity regularization terms and neural network-based regularization terms [31]. Non-negative low-rank matrix decompositions have both low-rank and non-negativity constraints [29, 40, 33]. In compressive sensing, a highly significant type of regularization known as L1-L2 sparse regularization allows the compressive sensing problem to be formulated as a three-term composite nonconvex optimization problem [57, 43]. In robust principal component analysis (RPCA), the optimization problem includes constraints not only on the data term but also on the low-rank factor and sparse factor, which can be formulated as a three-term composite minimization [17]. Operator splitting is an important method for solving composite optimization problems of type (1.3). Classic algorithms such as Forward-Backward Splitting (FBS) and Douglas-Rachford Splitting (DRS) can be employed when (1.3) $H=0$. In [15], Attouch, Bolte, and Svaiter established the convergence of FBS for nonconvex and nonsmooth optimization under the assumption that $\nabla F$ is Lipschitz continuous. In [12], Li and Pong applied the DRS algorithm to nonconvex feasibility problems and established its convergence under similar assumptions.

The above splitting algorithms are not suitable for (1.3) because they require computing the proximal operator of $G+H$ in each iteration. Computing this operator is much more expensive or even infeasible compared to computing the proximal operators of $G$ and $H$ separately. To address this limitation, Davis and Yin [6] proposed a three-operator splitting algorithm for solving (1.3). In the three-operator splitting algorithm, each sub-problem involves solving the proximal operator of a single function, either $F$ or $G$. When the proximal operators of $F$ and $G$ are easy to solve, the three-operator splitting algorithm becomes a simple and effective method. Davis and Yin [6] established the convergence of a three-operator splitting algorithm for convex optimization. Liu and Yin in [56] proved the convergence of the three-operator splitting method when $F$ is weakly convex, and $G$ and $H$ are twice continuously differentiable with bounded Hessians. In [9] Bian and Zhang established the global convergence of a three-operator splitting algorithm for solving nonconvex optimization problems and provided local convergence rates.

However, when solving large-scale problems, the TOS algorithm may suffer from computational inefficiency due to the time-consuming computation of $\nabla H$, thereby reducing the efficiency of the TOS method. To overcome this challenge, we propose a stochastic version of the TOS algorithm by replacing $\nabla H$ with its stochastic approximation. This will lead to a new stochastic three-operator splitting (STOS) algorithm that allows for more efficient computations. We present STOS in Algorithm 1 below.

Note that the STOS algorithm is also an extension of these classical algorithms. When $F=0$, Algorithm 1 becomes the stochastic Forward-Backward Splitting (FBS) algorithm. When $H=0$, the Algorithm 1 corresponds to the classical Douglas-Rachford Splitting (DRS) algorithm. When $F=0$ and $G=0$, this algorithm becomes the stochastic gradient descent algorithm. Here, ${\widetilde{\nabla}}H$ can be either unbiased gradient estimators (such as SGD, etc.) or variance-reduced gradient estimators (such as SAGA, SVRG, SARAH, etc.).

Stochastic splitting algorithms for solving large-scale nonconvex optimization problems have received extensive interest in the last two decades.

When $F=G=0$ in (1.3), the classical stochastic gradient descent algorithm [22] can solve this kind of problem. In [46], Ghadimi and Lan proposed the randomized stochastic gradient (RSG) method to solve (1.3) with $F=G=0$ and presented the first non-asymptotic convergence analysis of SGD. They demonstrated that under the assumptions of $H$ being smooth, ${\widetilde{\nabla}}H$ being an unbiased and bounded variance estimator, the RSG method achieves a solution $x_{*}$ satisfying $\mathbb{E}[\|\nabla H(x_{*})\|^{2}]\leq\varepsilon$ with an iterative complexity $\mathcal{O}(1/\varepsilon^{2})$. Further, the stochastic gradient descent algorithm has also been combined with certain variance-reduced gradient estimators. For instance, Shamir [42] proposed nonconvex SVRG and considered the problem of computing several leading eigenvectors. Allen-zhu and Hazan [59] also proposed SVRG, showing its linear convergence rate. Later, Reddi et al. [49] introduced the nonconvex SAGA method and demonstrated its linear convergence to the global optimal solution for a class of nonconvex optimization.

When $F=0$ in (1.3), this problem can be solved by the classical proximal gradient descent (PGD) algorithm. Ghadimi, Lan, and Zhang [47] introduced a randomized stochastic projected gradient (RSPG) algorithm to solve (1.3) with a constraint set $X$ when $F=0$. It is worth noting that when the constraint set $X=\mathbb{R}^{n}$, the RSPG algorithm becomes the stochastic PGD algorithm. Further, the PGD algorithm combined with variance-reduced gradient estimators has also received a lot of attention. Reddi et al. [50] integrate the proximal gradient descent with variance-reduced estimators SAGA and SVRG, referred to as PROXSAGA and PROXSVRG, respectively. They showed that under the assumptions of $G$ being nonsmooth convex and $H$ being smooth nonconvex, achieving an $\varepsilon$-accurate solution requires calling the number of proximal oracles to be $\mathcal{O}(1/\varepsilon)$. Pham et al. [41] combined proximal gradient descent with the SARAH stochastic estimator, referred to as ProxSARAH. Under the assumptions of $H$ being smooth nonconvex and $G$ nonconvex nonsmooth, they demonstrated that ProxSARAH finds an $\varepsilon$-stationary point with complexity $\mathcal{O}(N+N^{1/2}\varepsilon^{-2})$ for $H$ having finite-sum structure.

The study of (1.3) with three-term composite structures is limited in the nonconvex stochastic setting. Metel and Takeda [38] proposed the stochastic proximal gradient method for solving (1.3) and showed that the iteration complexity is $\mathcal{O}(1/\varepsilon^{3})$ when the distance to the subdifferential mapping of loss is less than $\varepsilon$. Yurtsever, Mangalick, and Sra [3] proposed a stochastic version of a three-operator splitting algorithm to solve (1.3) by combining unbiased and bounded-variance stochastic gradient estimators, and showed that finding an $\varepsilon$-stationary point requires an iteration complexity of $\mathcal{O}(1/\varepsilon^{3})$ based on a variational inequality. Driggs et al. [7] proposed a stochastic proximal alternating linearization algorithm when $H$ is a cross-term in (1.3), which combines a class of variance-reduced stochastic gradient estimators. The algorithm finds a $\varepsilon$-stationary point with a complexity of $\mathcal{O}(1/\varepsilon)$. Further, for a class of large-scale three-block composite optimization problems in deep learning, Bian, Liu, and Zhang [11] propose a new stochastic three-block splitting algorithm for (1.3) with a cross-term.

In this paper, we propose a stochastic three-operator splitting (STOS) algorithm to solve the nonconvex and nonsmooth problem (1.3). The main contributions of our paper can be summarized as follows:

• 1.

  We propose a stochastic three-operator splitting (STOS) algorithm for solving (1.3), which combines two types of stochastic gradient estimators: unbiased estimators and variance-reduced estimators. Compared to the stochastic TOS in [3] which only considers unbiased gradient estimator, our STOS (Algorithm 1) combines both unbiased gradient estimators and variance-reduced gradient estimators. These include the most commonly used gradient estimators. Hence, our algorithm can select the appropriate gradient estimators for different applications.
  

• 2.

  When the unbiased gradient estimator is incorporated, we establish the convergence and obtain an $\mathcal{O}(1/\varepsilon)$ convergence rate for the STOS algorithm under the variance-bounded assumption. When a variance-reduced gradient estimator is incorporated, we construct a stability function for the STOS algorithm. Furthermore, if the objective function is semi-algebraic, we obtain the convergence and convergence rate of STOS. Compared to the stochastic TOS in [3], our theoretical analysis does not require the convexity of $F$ and $G$.
  

• 3.

  In the numerical experiments, we apply the STOS algorithm to solve the FLASH proton radiotherapy with dose and dose rate optimization in radiation therapy. In all current research work, nonstandard ADMM algorithms are used to solve the convex relaxation of this problem. In this paper, we directly solve this nonconvex problem using the stochastic three-operator splitting algorithm. To the best of our knowledge, this is the first time a stochastic algorithm has been adopted to solve the FLASH proton radiotherapy optimization. All test results demonstrate that STOS significantly outperforms the nonstandard ADMM algorithm.
  

The rest of this paper is organized as follows. We first present some notation and preliminaries in Section 2. The principal theoretical analysis of our proposed algorithm can be found in Section 3, followed by numerical results in Section 4. In Section 5, we give some concluding remarks.

A class of gradient estimators that are often used in stochastic algorithms are unbiased gradient estimators, such as the most classical stochastic gradient descent (SGD) estimator [21]. In this paper, we also analyze STOS (Algorithm 1) combined unbiased gradient estimators. We define an unbiased gradient estimator.

We provide the specific form of SGD, there exist other sampling methods that can produce an unbiased stochastic gradient, such as sampling with replacement and independent sampling without replacement. For more details, we refer to readers to [2].

Variance-reduced gradient estimators have been proposed and widely adopted in stochastic algorithms, such as SVRG (stochastic variance reduced gradient) [44], SAG (stochastic average gradient) [39], and SDCA (stochastic dual coordinate ascent) [51]. In [7], a unified definition is provided for a class of variance-reduced gradient estimators. In this paper, we analyze the combination of STOS (Algorithm 1) with such variance-reduced gradient estimators taken from[7]. For convenience, we recall the definition of variance-reduced gradient estimators below.

The variance-reduced gradient estimator defined in Definition 2.5 is widely used in many algorithms, such as SAGA [1], SARAH [35], SAG [39], and SVRG [30, 44]. It has been verified in [7, 10] that these four estimators are variance-reduced gradient estimators of this type. Below, we provide the definitions for these four estimators.

The Kurdyka-Łojasiewicz (KL) property plays a crucial role in the convergence analysis of nonconvex optimization problems. We first recall the definition of KL property, which is particularly applicable to semi-algebraic functions. For a more detailed account, we refer readers to [16, 15, 24, 23, 25] and the references therein. Let $\varepsilon_{1}$ and $\varepsilon_{2}$ be two real numbers satisfying $-\inf<\varepsilon_{1}<\varepsilon_{2}<+\inf$. We define the set $[\varepsilon_{1}<F<\varepsilon_{2}]\overset{\mathrm{def}}{=}\{x\in\mathbb{R}^{m_{1}}:\varepsilon_{1}<F(x)<\varepsilon_{2}\}$.

Roughly speaking, the KL property allows for sharp optimization up to reparameterization via $\varphi$, which is known as a desingularizing function for $F$. A typical class of KL functions are semi-algebraic functions, such as the $\ell_{0}$ pseudo-norm and the rank function. More specifically, we can find a detailed result on the KL property in [16, Section 4.3]. Further arguments can be found in [24, Section 2] and [23, Corollary 16]. Semi-algebraic functions are easier to identify and encompass a wide range of possibly nonconvex functions that arise in applications, as shown in [16, 15, 24, 23, 25]. To clarify, we revisit the definition of semi-algebraic functions.

In this subsection, we consider the stochastic nonconvex problem (1.3) with $H$ being the expectation of a function of a random variable, i.e., $H(x)=\mathbb{E}_{\xi}H(x,\xi)$, where $\xi$ is a random variable with distribution $\mathcal{P}$:

As a special case of (3.2), if $\xi$ is a uniformly random vector defined on a finite support set $\Omega:=\{\xi_{1},\xi_{2},\cdots,\xi_{N}\}$, then (3.2) simplifies to the following finite-sum minimization problem:

where $H_{i}(x):=H(x,\xi_{i})$ for $i=1,\cdots,N$. Here, the stochastic gradient estimator ${\widetilde{\nabla}}H$ in Algorithm 1 adopts the following SGD estimator:

where $J_{t}$ is a set of $b$ $i.i.d.$ samples from distribution $\mathcal{P}$. Our convergence analysis relies on the following assumptions.

In the following, we establish the convergence analysis of STOS (Algorithm 1) with unbiased estimators. We show that STOS convergences to an $\varepsilon$-stationary point and obtain its convergence rate. The proof of Theorem 3.5 is given in Appendix A.2.

We have provided the convergence guarantee for STOS (Algorithm 1) when combined with unbiased gradient estimators. However, some variance-reduced gradient estimators, such as SARAH, can make the algorithm more effective for certain large-scale image processing problems [10]. To allow STOS to use a wider variety of stochastic gradient estimators, we further provide the convergence guarantee for STOS combined with the variance-reduced gradient estimator (Definition 2.5) and obtain its convergence rate. The convergence analysis is based on the Kurdyka-Łojasiewicz inequality. For this purpose, it is crucial to construct an energy function that decreases along the sequence generated by Algorithm 1. Thus, we define

for $C_{1}>0$. Denote $\Phi_{t}=\Phi(x_{t},y_{t},z_{t},y_{t-1})$. Next, we give the energy function $\Psi_{t}$ associated with STOS (Algorithm 1) as follows:

The decreasing property of $\Psi_{t}$ in expectation is given in the following theorem.

The proof of Theorem 3.7 is provided in Appendix A.3. Next, we will establish the global convergence in expectation for the whole sequence generated by Algorithm 1. See Appendix A.3 for details of the proof.

Furthermore, we derive the convergence rates for the sequence $\{(x_{t},y_{t},z_{t})\}_{t\geq 0}$ generated by Algorithm 1 based on the KL exponent. The calculation of the KL exponent is detailed in [13]. We omit the proof of convergence rate, which is similar to the proof of Theorem 3.7 in [10].