A Learned Proximal Alternating Minimization Algorithm and Its Induced Network for a Class of Two-block Nonconvex and Nonsmooth Optimization

Funding This research was partially supported by NSF grant DMS-2152961 and Simons Foundation Collaboration Grant (Award Number: 584918).

Data Availability The original dataset we used in all experiments are from Multi-modal Brain Tumor Segmentation Challenge 2018.menze2014multimodal

Recent years have witnessed remarkable success of deep learning across various real-world applications. However, a purely data-driven approach may fail to approximate the desired functions, especially when training data are scarce. It is well known that the scarcity of data leads to overfitting and challenges in interpretation. To mitigate these issues, the unrolling/unfolding neural networks (UNNs) have been developed and have shown promising results in solving inverse problems arising from computer vision and medical imaging. The UNNs are multi-phase neural networks, in which each phase mimics one iteration of the optimization algorithms for solving the inverse problems. However, despite their promising performance in practice, most of the existing UNNs only superficially resemble steps of optimization algorithms. Consequently, their outputs do not really yield solutions to any interpretable variational models. This results is lack of theoretical justifications and convergence guarantees for their outputs.

Recently a novel class of unrolling methods known as learned optimization algorithms (LOA) LDA ; ldct ; wanyuMiccai has been developed. The goal of LOA is to tackle the challenges associated with solving a class of inverse problems arising from various image reconstruction and synthesis problems. These LOAs strategically combine the learned variational model with the prior domain knowledge of underlying physical processes to enhance its interpretability. These methods also utilize optimization to guarantee the convergence.

The optimization scheme induces a highly structured deep network, aligning its architecture precisely with the iterative algorithm. Hence, the network inherits the convergence property of the algorithm, and it outputs an approximation of the solution to the variational model. As a result, it is interpretable and parameter efficient. However, the existing LOAs only consider a single block of variables.

Motivated by a wide range of applications in machine learning that involve multi-modalities or multi-domains problems, such as multi-task learning, transfer learning, multi-modal learning and fusion, dual domain image reconstruction, and image synthesis, in this work we extend the idea of LDA to deal with learned multi-block nonconvex and nonsmooth optimization problems. More specifically, we develop a novel convergent learned proximal alternating minimization (LPAM) algorithm to solve the following class of learnable optimization problems:

where each learnable function is a parameterized function. $\Theta=(\theta_{1},\theta_{2},\theta)$ is the set of learned parameters. Each function in (1) is possibly nonconvex and nonsmooth. We shall also study the convergence and iteration complexity of the algorithm, and provide experimental results on the application on the joint multi-contrast MRI reconstruction with significantly under-sampled data.

In this work we propose a novel learned proximal alternating minimization algorithm addressing nonsmooth and nonconvex two-block variational models with provable convergence and iteration complexity. The proposed LPAM algorithm for solving ($M_{1}$) determines the architecture of the deep neural network, hence the design of the LPAM considers not only the convergence and efficiency, but also the ability to assist the training of the network parameters $\Theta$, i.e., reducing the error in minimizing the loss function.

Our main idea for designing this algorithm is as follows: (i) We tackle the nonsmoothness by an appropriate smoothing technique with automatic diminishing smoothing effect; (ii) We modify the PALM scheme by incorporating the residual learning architecture, which has proven to be highly effective in deep network training he2016deep ; and (iii) When some modified PALM iterates fail to satisfy certain conditions, we employ the block coordinate decent (BCD) iterates as a safeguard to ensure convergence. Moreover, we prove that a subsequence generated by the proposed LPAM has accumulation points and all of them are Clarke stationary points of the nonsmooth nonconvex problem.This algorithm can be readily extended to multi-block nonsmooth and nonconvex optimization.

This algorithm significantly broadens its applicability, as it relaxes the strict convergence conditions of the existing algorithms. The PALM or iPALM algorithm assumes the joint term $H(\mathbf{x}_{1},\mathbf{x}_{2})$ to be smooth and possibly nonconvex, and the proximal points associated with $H_{i}(\mathbf{x}_{i})$ ($i=1,2$) are easy to obtain. Their global convergence results are built on KL property of the objective function. These assumptions in general cannot be met by the learnable parameterized functions representing neural networks. The proposed LPAM algorithm will not require these conditions. Each function in (1) could be nonconvex and nonsmooth. Without those restrictions, the LPAM algorithm is flexible and applicable to deep learning problems in various applications. Moreover, the convergence of LPAM can still be established in the sense of sub-sequence convergence to Clarke stationary points (see detail in Section 3). But it cannot have global convergence to a critical point as the PALM or iPALM algorithm.

The convergence result of the BCD-net is based on the assumptions that the paired image refining neural networks are asymptotically nonexpansive, the data fidelity is differentiable, and the MBIR problem is not difficult to solve. The proposed LPAM algorithm does not require these assumptions. The convergence of LPAM is assured by the convergence of the BCD algorithm for smooth optimization and the property of the limit of the gradient of the smoothed objective functions as smoothing factor tends zero. By removing the constraints on the model, the PALM can significantly broaden its range of applications.

The remainder of the paper is organized as follows. In Section 5 we construct the LPAM algorithm motivated by the PALM algorithm and the BCD algorithm, and introduce the three steps that constitute the LPAM algorithm. In Section 6, we prove the convergence of the LPAM algorithm and deduce the iteration complexity of this algorithm. In Section 7, we apply the LPAM algorithm to the joint reconstruction of T1 and T2 MRI images with significantly under-sampled data to show the promising performance of the proposed method. Furthermore, we provide convergence analysis and comparison with several existing methods. All the figures are listed at the end of the article.

Throughout the paper, we use the following notations without further notification.

1. 1.

   $H(\mathbf{x}_{1},\mathbf{x}_{2};\theta)$ denotes the function $H$ of $\mathbf{x}_{1}$ and $\mathbf{x}_{2}$ given a collection of parameters $\theta$. $H_{\epsilon}(\mathbf{x}_{1},\mathbf{x}_{2};\theta)$ denotes the smooth approximation of $H(\mathbf{x}_{1},\mathbf{x}_{2};\theta)$ with a smoothing parameter $\epsilon$.

2. 2.

   For a differentiable function $f(\mathbf{x}_{1},\mathbf{x}_{2})$, we denote the gradient of $f$ with respect to $\mathbf{x}_{i}$ by $\nabla_{i}f(\mathbf{x}_{1},\mathbf{x}_{2})$ $(i=1,2)$. The gradient $f$ on the domain $(\mathbf{x}_{1},\mathbf{x}_{2})\in R^{n}\times R^{m}$ is represented by $\nabla_{1,2}f(\mathbf{x}_{1},\mathbf{x}_{2})$.

3. 3.

   We use $\|\cdot\|$ to denote the standard $L^{2}$ norm in Euclidean space, and write $\langle\cdot,\cdot\rangle_{\mathbb{R}^{n}}$ for the standard inner product on $\mathbb{R}^{n}$.

4. 4.

   $\chi\mathrel{\mathop{\ordinarycolon}}=\mathbb{R}^{n}\times\mathbb{R}^{m}$.

In this section we propose a learned proximal alternating minimization (LPAM) algorithm for solving the problem (1). Since $\theta$ is learned, we will omit it and simply write (1) as (2).

For $H_{1}$,$H_{2}$, $H$ and $\Phi$ we assume

The proposed LPAM algorithm consists of three stages. In the first stage we convert the nonconvex and nonsmooth problem (2) to a nonconvex smooth optimization problem by an appropriate smoothing procedure. The smoothing method is not unique, but we require the smooth approximation $\Phi_{\epsilon}$ of $\Phi\mathrel{\mathop{\ordinarycolon}}=H_{1}(\mathbf{x}_{1})+H_{2}(\mathbf{x}_{2})+H(\mathbf{x}_{1},\mathbf{x}_{2})$, which is defined as

to satisfy following conditions:

• •

  (C1:) $\nabla H_{1,\epsilon}$, $\nabla H_{2,\epsilon}$ and $\nabla_{1,2}H_{\epsilon}$ are Lipschitz continuous in $\mathbb{R}^{n}$, $\mathbb{R}^{m}$ and $\mathbb{R}^{n}\times\mathbb{R}^{m}$, respectively.

• •

  (C2:) $\lim_{\epsilon\to 0^{+},\mathbf{z}_{i}\to\mathbf{x}_{i}}H_{i,\epsilon}(\mathbf{z}_{i})=H_{i}(\mathbf{x}_{i})$, $i=1,2$, for any given $\mathbf{x}_{i}$ in the corresponding domain. Also for any $\mathbf{X}=(\mathbf{x}_{1},\mathbf{x}_{2})\in\mathbb{R}^{n}\times\mathbb{R}^{m}$, we assume that when $\epsilon\to 0^{+}$ and ${\bf Z}=(\mathbf{z}_{1},\mathbf{z}_{2})\to\mathbf{X}$, $\lim_{\epsilon\to 0^{+},{\bf Z}\to\mathbf{X}}H_{\epsilon}({\bf Z})=H(\mathbf{X})$.

• •

  (C3:) There is a continuous and non-negative function $\mathfrak{m}\mathrel{\mathop{\ordinarycolon}}[0,\infty)\to[0,\infty)$ such that $\mathfrak{m}(0)=0$,

  $$\Phi_{\epsilon}(\mathbf{x}_{1},\mathbf{x}_{2})+\mathfrak{m}(\epsilon)\leq\Phi_{\delta}(\mathbf{x}_{1},\mathbf{x}_{2})+\mathfrak{m}(\delta)\,\,\mbox{for all}\,(\mathbf{x}_{1},\mathbf{x}_{2})\in\chi,\,\mbox{and}\,0<\epsilon\leq\delta.$$

• •

  (C4:) For any $\mathbf{X}^{*}=(\mathbf{x}_{1},\mathbf{x}_{2})\in\mathbb{R}^{n}\times\mathbb{R}^{m}$ and any $\mathbf{X}^{k}\to\mathbf{X}^{*}$, $\lim_{\mathbf{X}^{k}\to\mathbf{X}^{*},\epsilon\to 0^{+}}\nabla_{1,2}\Phi_{\epsilon}(\mathbf{X}^{k})\in\partial^{c}\Phi(\mathbf{X}^{*})$, where $\partial^{c}\Phi(\mathbf{X}^{*})$ stands for the Clarke subdifferential of $\Phi$ (see Definition 1).

In the second stage we solve the smoothed nonconvex problem with a fixed smoothing factor $\epsilon$, i.e.

In light of the substantial improvement in practical performance by ResNet he2016deep , we propose a modified PALM that incorporates the architecture of the ResNet for solving (4). With $\epsilon=\epsilon_{k}>0$, PALM bolte2014proximal generates a sequence of iterates $\{\mathbf{x}_{1}^{k+1},\mathbf{x}_{2}^{k+1}\}$ by

and

where $\alpha_{k}$ and $\beta_{k}$ are step sizes. In some situations, the proximal points $\mathbf{u}_{1}^{k+1}$ and $\mathbf{u}_{2}^{k+1}$ are not easy to compute, so we replace $H_{\epsilon_{k}}(\mathbf{u},\mathbf{x}_{2}^{k})$ by its linear approximation at $\mathbf{z}_{1}^{k+1}$ together with a proximal term $\frac{1}{2p_{k}}\|\mathbf{u}-\mathbf{z}_{1}^{k+1}\|^{2}$. Similarly, we replace $H_{\epsilon_{k}}(\mathbf{u}_{1}^{k+1},\mathbf{u})$ by its linear approximation at $\mathbf{z}_{2}^{k+1}$ together with a proximal term $\frac{1}{2q_{k}}\|\mathbf{u}-\mathbf{z}_{2}^{k+1}\|^{2}$. Thus, we have

where we recall that $\nabla_{1}H_{\epsilon_{k}}$ stands for the gradient for the first $n$ variables.

In deep learning approach, the step sizes $\alpha_{k}$, $\tau_{k}$, $\beta_{k}$ and $\gamma_{k}$ could be learnable hyper-parameters. The advantage for taking the scheme (5) and (6) to update $u_{1}$ and $u_{2}$ is that this scheme is consistent with the architecture of residual learning he-1 ; he-2 , which only learns the correction needed for $\mathbf{z}_{i}^{k+1}$ $(i=1,2)$ and avoids gradient vanishing in network training, and thus improves solution quality in practice. However, the convergence of the sequence $\{(\mathbf{u}_{1}^{k+1},\mathbf{u}_{2}^{k+1})\}$ is not guaranteed. Inspired by the proof of convergence in bolte2014proximal , we propose the following:

If $\mathbf{u}_{1}^{k+1}$ and $\mathbf{u}_{2}^{k+1}$ satisfy

and

for some $a>0$, we take $\mathbf{x}_{1}^{k+1}=\mathbf{u}_{1}^{k+1}$, $\mathbf{x}_{2}^{k+1}=\mathbf{u}_{2}^{k+1}$.

If one of (7) and (8) is violated, we compute $(\mathbf{v}_{1}^{k+1},\mathbf{v}_{2}^{k+1})$ by BCD algorithm with a simple line-search strategy for better step sizes as follows: Let $\bar{\alpha},\bar{\beta}\in(0,1)$, we compute

The first order optimality condition leads to

If for some $\delta\in(0,1)$,

we set

otherwise we reduce step sizes $(\bar{\alpha},\bar{\beta})$ by $\rho(\bar{\alpha},\bar{\beta})\rightarrow(\bar{\alpha},\bar{\beta})$ where $0<\rho<1$, and recompute $\mathbf{v}_{1}^{k+1},\mathbf{v}_{2}^{k+1}$ until the condition (12) holds. Then take $\mathbf{x}_{1}^{k+1}=\mathbf{v}_{1}^{k+1}$, $\mathbf{x}_{2}^{k+1}=\mathbf{v}_{2}^{k+1}$.

In the third stage, we check if the $\ell_{2}$-norm of the gradient of the smoothed objective function has been reduced enough, so that we can perform the second stage with a reduced smoothing factor. By gradually reducing the smoothing factor, we obtain a subsequence of the iterates that converges to a Clarke stationary point of the original nonconvex and nonsmooth problem.

The scheme of LPMA is given below (For notational simplicity, we omit all learned network parameters).

We would like to comment here for the LPAM algorithm. Line 2 to line 18 can be viewed as the steps of the inner iteration, which, for a fixed $\varepsilon$, generates a sequence $\{\mathbf{x}_{1}^{k},\mathbf{x}_{2}^{k}\}$ such that $\|\nabla_{1,2}\Phi_{\epsilon}(\mathbf{x}_{1}^{k},\mathbf{x}_{2}^{k})\|\to 0$ as $k\to\infty$ (see the first statement of Lemma 1 in the next section). This guarantees the reduction of $\|\nabla_{1,2}\Phi_{\varepsilon}(\mathbf{x}_{1}^{k},\mathbf{x}_{2}^{k})\|$ for a fixed $\varepsilon$. When this value drops below a prefixed value, line 19 decreases $\varepsilon$, and then we repeat the procedures from line 2 to line 18 to generates a new sequence $\{\mathbf{x}_{1}^{k},\mathbf{x}_{2}^{k}\}$ corresponding to the reduced $\varepsilon$. We continue these steps until the termination condition in the line 20 is met. The sequence formed by the iterates, in which the reduction criterion for $\varepsilon$ is met, has at least one accumulation point and each accumulation point is a Clarke Stationary Point of the original problem. This will be proved in Theorem 1 in Section 6.

Remark: In order to make the LPMA algorithm match the ResNet architecture, if the function $H_{\epsilon}(\mathbf{x}_{1},\mathbf{x}_{2})$ is learnable and $H_{1,\epsilon}(\mathbf{x}_{1})$ and $H_{2,\epsilon}(\mathbf{x}_{2})$ are given, it is better to use the scheme above. If the functions $H_{1,\epsilon}(\mathbf{x}_{1})$ and $H_{2,\epsilon}(\mathbf{x}_{2})$ are learnable and $H_{\epsilon}(\mathbf{x}_{1},\mathbf{x}_{2})$ is given, it is better to change the steps 3-6 to the following:

and keep the other steps in Algorithm 1 unchanged.

Since we deal with a nonconvex and nonsmooth optimization problem, the Clarke subdifferential, which is based on the concept of generalized directional derivatives, is employed here to characterize the optimality of the solutions to (2).

Given $(\mathbf{x}_{1}^{k},\mathbf{x}_{2}^{k})$, in the case that $(\mathbf{u}_{1}^{k+1},\mathbf{u}_{2}^{k+1})$ generated by Algorithm 1 satisfies the conditions (7)-(8), we take $(\mathbf{x}_{1}^{k+1},\mathbf{x}_{2}^{k+1})=(\mathbf{u}_{1}^{k+1},\mathbf{u}_{2}^{k+1})$. From (7)-(8), it is easy to get

If $(\mathbf{u}_{1}^{k+1},\mathbf{u}_{2}^{k+1})$ fails to satisfy the condition (7) or (8), the algorithm computes $(\mathbf{v}_{1}^{k+1},\mathbf{v}_{2}^{k+1})$ by (11) with a line search strategy. Let $\ell_{k}$ be the number of the required line search steps to meet the condition (12), and then we take $(\mathbf{x}_{1}^{k+1},\mathbf{x}_{2}^{k+1})=(\mathbf{v}_{1}^{k+1},\mathbf{v}_{2}^{k+1})$. This gives the following

Since $\nabla_{1,2}\Phi_{\epsilon}$ is $L_{\epsilon}$-Lipschitz continuous, we have

The combination of (19) and (20) yields

Hence the condition (12) is met if

Hence, for any $k=1,2,\ldots$ the maximum line search steps $\ell_{max}$ required for (12) satisfies

so we abuse the notation $\ell_{max}$ by defining it as

This proves the second statement of this lemma. Moreover, the fact that $\ell_{k}\leq\ell_{max}$ for any $k=0,1,2,\ldots$, implies

Next we prove the other two statements. Note that from (19) the condition (12) can be rewritten as:

Now we estimate the last term in (25) in terms of $\nabla_{1,2}\Phi_{\epsilon}(\mathbf{x}_{1}^{k},\mathbf{x}_{2}^{k})$. First by the Lipschitz continuity of $\nabla_{2}\Phi_{\epsilon}$ we have

Clearly this implies

Here if the right hand side is positive, we will choose $\sigma\in(0,1)$ to be determined, and take advantage of the following inequality: $(a-b)^{2}\geq\frac{\sigma}{2}a^{2}-\sigma b^{2}$ together with (19) to obtain

Inserting (6) into (25), we have

If $\bar{\beta}^{2}\rho^{2\ell_{k}}L_{\epsilon}^{2}<\frac{2}{3}$, we just choose $\sigma=\frac{1}{2}$. Otherwise we set $\sigma$ as $\sigma=1/(2\bar{\beta}^{2}\rho^{2\ell_{k}}L_{\epsilon}^{2})$. In either case the coefficient of $\|\nabla_{1}\phi_{\epsilon}(\mathbf{x}_{1}^{k},\mathbf{x}_{2}^{k})\|^{2}$ is less than $-\frac{\delta}{2}\bar{\alpha}^{2}\rho^{2\ell_{k}}$. For the coefficient of $\|\nabla_{2}\phi_{\epsilon}(\mathbf{x}_{1}^{k},\mathbf{x}_{2}^{k})\|^{2}$, in the first case it is less than $-\frac{\delta}{4}\bar{\beta}^{2}\rho^{2\ell_{k}}$. In the second case it is less than $-\frac{\delta}{4L_{\epsilon}^{2}}$. Then it is easy to verify that in both cases, by (24) both coefficients are less than $-\frac{\delta\min(\bar{\alpha},\bar{\beta})^{2}\rho^{2}}{4\max(\bar{\alpha},\bar{\beta})^{2}L_{\epsilon}^{2}}$ Thus we have