Accelerating Multi-Block Constrained Optimization Through Learning to Optimize

Optimization is a fundamental and essential process in machine learning and data science [4]. It involves fine-tuning algorithms, models, or systems to make them perform at their best on complex tasks. The goal of optimization is to find the best solution from a set of possible choices (i.e., feasible set), often with the aim of minimizing or maximizing a particular function (i.e., loss function). In the context of machine learning, optimization often refers to the process of adjusting a model’s parameters to minimize the error or loss function. This is typically achieved through various classical algorithmic frameworks such as gradient descent (GD) [8], stochastic gradient descent (SGD) [50], and more advanced methods like Adam [31] and RMSprop [27], to mention just a few. Machine learning approaches can play significant roles in designing efficient optimization algorithms and greatly improving the performance of well-designed algorithms. This field of research is known as Learning to Optimize (L2O) which attracts growing attention [14]. In particular, given a class of optimization problems, machine learning approaches can be used to effectively predict the performance of different hyperparameters and configurations of an optimization algorithm, thereby guiding the search for optimal configurations and even automating the process of tuning [29, 20].

Learn to optimize (L2O). L2O is a modern and effective approach towards designing optimization algorithms that reach a new level of efficiency. A comprehensive survey of L2O can be found in [14]. Our work is closely related to different subjects of L2O, including the configuration and hyperparameter learning [29, 20], plug-and-play approaches [56], and algorithm unrolling [43]. Particularly, a series of works targeting unrolling algorithmic frameworks for the Lasso model [55] (which as well as its dual problem are special cases of ($P(\xi)$)) has attracted increasing attention [24, 44, 47, 62, 1, 17, 25, 39]. The fundamental algorithmic framework in these works is the Iterative Shrinkage Thresholding Algorithm (ISTA) [3]. Despite being characterized by hyperparameter learning, our approach can also be viewed as the application of the algorithm unrolling to the MPALM. In this way, our work continues this line of research and provides an alternative L2O approach for efficiently solving challenging optimization problems from practical applications, including Lasso-type problems.

ADMM-type methods. Incorporating machine learning techniques with ADMM-type methods has drawn growing attention in recent years. The two-block ADMM was employed as the fundamental algorithmic framework in image denoising, in which the involved proximal operators are replaced with learned operators [56, 6, 10, 53]. Great empirical success and comprehensive theoretical guarantees have been achieved in these works, making the plug-and-play ADMM approach one of the most popular and effective algorithms for image science. Just like unrolling the ISTA, unrolling the two-block ADMM and its linearized variant (also known as the primal dual hybrid gradient method [9]) has also attracted much attention; see e.g., [54, 49, 2, 17, 15, 57] and references therein. The unrolled two-block ADMM-type methods have also been demonstrated to have excellent empirical performance. There are other machine learning techniques that are helpful for improving the practical performance of the two-block ADMM. For instance, [60] successfully applied a reinforcement learning approach for selecting the hyperparameters in the two-block ADMM for distributed optimal power flow. While fruitful results have been established for L2O with two-block ADMM-type methods, results on combining L2O with multi-block ADMM-type methods, including MPALM, remain limited. We demonstrate in the present paper that the MPALM-based L2O approach is also effective when compared with existing state-of-the-art L2O approaches.

Optimal transport (OT). Recent years have seen a blossoming of interest in developing efficient solution methods for OT, which have numerous important applications, due partly to the essential metric property of its optimal solution [18, 48]. To solve OT problems efficiently at scale, the most popular first-order method is perhaps the Sinkhorn’s algorithm [18]. Sinkhorn’s algorithm is a simple iterative method for finding optimal solutions of entropy-regularized OT problems. Thus, it only provides approximate solutions to the original OT problems. To get high quality approximations, one needs to choose a small entropy regularization parameter, which causes the convergence of the Sinkhorn’s algorithm to be extremely slow. Moreover, a small entropy regularization parameter may also result in numerical issues, making the Sinkhorn’s algorithm unstable. Hence, developing effective algorithmic frameworks for solving OT problems remains an active research direction [23, 41, 32, 40, 16, 28, 59, 58]. However, works focusing on L2O approaches for OT remain limited, to the best of our knowledge. Our work provides a feasible approach for solving OT problems reliably via combining L2O with MPALM.

The following notation will be used throughout this paper. Let $\mathbb{E}$ be a finite dimensional Euclidean space, the standard inner product is denoted as $\left\langle\cdot,\cdot\right\rangle$ and the associated induced norm is denoted as $\left\lVert\cdot\right\rVert$. Let $S:\mathbb{E}\to\mathbb{E}$ be a self-adjoint positive semidefinite linear operator, the weighted-norm associated with $S$ is denoted as $\|\cdot\|_{S}$, i.e., for any $x\in\mathbb{E}$, $\left\lVert x\right\rVert_{S}=\sqrt{\left\langle x,Sx\right\rangle}$. For a differential function $f:\mathbb{E}\to\mathbb{R}$, its gradient is denoted as $\nabla f$. If $f:\mathbb{E}\to\mathbb{R}\cup\{\pm\infty\}$ is an extended-valued convex function, then the effective domain of $f$ is denoted as $\mathrm{dom}(f):=\{x\;:\;f(x)<\infty\}$. Let $x\in\mathrm{dom}(f)$, then the subgradient of the convex function $f$ at $x$ is denoted as $\partial f(x):=\{v\;:\;f(x^{\prime})\geq f(x)+\left\langle v,x^{\prime}-x\right\rangle,\;\forall x^{\prime}\}$.

In this section, we provide a detailed introduction of the main algorithmic framework proposed in [12], namely the majorized proximal augmented Lagrangian method (MPALM), and explain how to cleverly choose the proximal term for the augmented Lagrangian function such that the resulting proximal ALM subproblem can be solved efficiently and analytically. The key idea of designing the proximal term is motivated by the Gauss-Seidel iterative method for solving symmetric positive definite linear systems.

Since the problem ($P(\xi)$) is a constrained convex optimization problem, we see that the first-order optimality conditions (also called the Karush-Kuhn-Tucker conditions) [51] for problem ($P(\xi)$) are given by

For any point $(x^{*},y^{*})\in\mathbb{X}\times\mathbb{Y}$ that satisfies the above first-order optimality conditions, one can show that $y^{*}$ is a solution to problem ($P(\xi)$) and $x^{*}$ is the corresponding Lagrange multiplier (also known as the dual optimal solution). For the rest of this paper, we assume that the following assumption holds.

Assumption 1 is commonly employed in the literature, as it presented one of the weakest conditions ensuring the solvability of the problem ($P(\xi)$). Assuming the solvability of the first-order optimality conditions (1), one may apply the classic augmented Lagrangian method for solving ($P(\xi)$). To this end, given a penalty parameter $\sigma>0$, we define the augmented Lagrangian function as

Then given an initial point $x^{0}\in\mathbb{X}$ and an increasing sequence of penalty parameter $\{\sigma_{k}\}$, the classical augmented Lagrangian iteratively performs the following updating scheme:

where $k\geq 0$ denotes the iteration counter. Though admitting excellent convergence properties [52], the classical augmented Lagrangian method faces certain challenges: (1) $y^{k+1}$ with high accuracy can be difficult to obtain since one needs to solve a (possibly nonsmooth) convex optimization problem. (2) The potential separable structure in $y$ is not explicitly explored. We next show how to address these challenges via replacing the augmented Lagrangian function $\mathcal{L}_{\sigma}$ in the classical augmented Lagrangian method with a convex function (as the sum of a convex quadratic function and $g$) such that the resulting subproblems in updating $y$ are much easier to solve via fully exploring the potential separable structure in $y$. To this end, we need to make the following mild assumption with respect to the function $f_{\xi}$.

For later usage, we shall partition $\Sigma$ into $p\times p$ sub-blocks as $\Sigma=[\Sigma_{ij}]_{1\leq i,j\leq p}$,

where $\Sigma_{ij}:\mathbb{Y}_{j}\to\mathbb{Y}_{i},\;1\leq i,j\leq p$, are linear mappings. We can see that Assumption 2 is indeed quite mild. For example, if $f_{\xi}$ has a Lipschitz continuous gradient for every $\xi\in\Xi$ with a uniformly bounded Lipschitz constant, then Assumption 2 holds. In the latter case, $\Sigma$ can be chosen as a diagonal matrix. Here, we allow non-diagonal $\Sigma$ in order to obtain a more accurate majorization of $f_{\xi}$ when defining the majorized augmented Lagrangian function ¹¹1For example, consider $f_{\xi}(y)=\frac{1}{2}\left\langle y,\mathcal{Q}y\right\rangle+\left\langle b_{\xi},y\right\rangle$ where $\mathcal{Q}\succeq 0$ does not depend on $\xi$. We see that $f_{\xi}$ satisfies Assumption 2. In this case, we can set $\Sigma=\lambda_{\textrm{max}}(\mathcal{Q})I$ or $\Sigma=\mathcal{Q}$. Obviously, the latter choice leads to an exact majorization of $f_{\xi}$..

Then, for a given $\xi\sim\mathcal{P}$ and the augmented Lagrangian penalty parameter $\sigma>0$, we can define the majorized augmented Lagrangian function (with slightly abuse of notation $\mathcal{L}$)

Using the majorized augmented Lagrangian function, the main algorithmic framework, namely the majorized proximal augmented Lagrangian method (ALM), for solving problem ($P(\xi)$) is presented in Algorithm 1.

Here, we assume that the subproblem in Line 2 of Algorithm 1 can be solved exactly. We shall see later that the two applications considered in this paper satisfy this assumption. However, we have to emphasize that solving the subproblem inexactly is more realistic and important; see e.g., [12]. The convergence properties of Algorithm 1 are summarized as follows. Note that under additional conditions, including a certain error-bound condition, Algorithm 1 can further be shown to have a local Q-linear convergence rate [12]. Since these additional conditions are generally not easy to verify, we decide not to present the result here for simplicity. However, such a linear convergence is empirically observed based on numerical experience.

Given the above convergence properties presented in Theorem 1, the next essential task is to choose the linear mapping $\mathcal{S}:\mathbb{Y}\to\mathbb{Y}$ such that the subproblems can be solved efficiently via exploring the separable structure associated with $y$. [36] provides an elegant approach for choosing $\mathcal{S}$ based on the idea of solving symmetric positive definite linear systems by Gauss-Seidel iterative method. Here, we shall briefly describe how to achieve this goal. To this end, we consider the mapping $\mathcal{S}$ which can be decomposed as the sum of two self-adjoint mappings, i.e., $\mathcal{S}:=\widetilde{\mathcal{S}}+\widehat{\mathcal{S}}$ where $\widetilde{\mathcal{S}}:=\mathrm{Diag}(\widetilde{\mathcal{S}}_{11},\;\dots,\widetilde{\mathcal{S}}_{pp})$ is block-diagonal with $\widetilde{\mathcal{S}}_{ii}:\mathbb{Y}_{i}\to\mathbb{Y}_{i}$, for $i=1,\dots,p$, and $\widehat{\mathcal{S}}$ to be determined shortly. Note that the objective function for the ALM subproblem in Algorithm 1 can be simplified as

where ”const” means a quantity that does not depend on $y$. For later usage, we write $\mathcal{Q}:=\sigma\mathcal{A}\mathcal{A}^{*}y+\widetilde{\mathcal{S}}+\Sigma:=\mathcal{U}+\mathcal{D}+\mathcal{U}^{*}$,

Under Assumption 3, $\mathcal{D}$ is positive definite and hence nonsingular. Using the above decomposition, we propose to choose $\widehat{\mathcal{S}}:=\mathcal{U}\mathcal{D}^{-1}\mathcal{U}^{*}$, which is called the SGS-operator for $\mathcal{Q}$, denoted by $\mathrm{SGS}(\mathcal{Q})$. With this particular choice of $\widehat{\mathcal{S}}$, we can show in the following theorem that one cycle of the block symmetric Gauss-Seidel update exactly solves the ALM subproblem in Line 2 of Algorithm 1.

Applying Theorem 2, one immediately obtains a symmetric Gauss-Seidel-based majorized proximal ALM (see Algorithm 2). Moreover, one can verify that that Assumption 3 ensures the convergence of the algorithm, as stated in Theorem 1. Consequently, we obtain an algorithm that takes a similar form of the multi-block ADMM with provable convergence. This is an appealing feature for practical applications since it allows one updating one block of the decision variables while keeping the remaining blocks fixed; see, e.g., the applications considered in Section 4. One can also observe that for $i=2,\dots p$, solving the associated subproblems involves only solving linear systems, which can be done nearly exactly via elementary linear algebra routines. For $i=1$, the computation of $y_{1}$ is the proximal mapping of $g$. It is well-known that many important functions $g$ admit analytical proximal mapping or can be approximate efficiently and accurately [46]. Thus, the associated ALM subproblem can also be solved in low computational costs.

Though Algorithm 1 has attractive convergence properties, its practical performance is sensitive to the choice of the penalty parameter $\sigma$, based on our numerical experience. Practitioners typically adjust $\sigma$ dynamically in order to obtain better numerical performance [33]. From a high level point of view, the following algorithm with adaptive penalty parameters, i.e., Algorithm 2, is commonly adopted in practice.

However, the criterion guiding the adjustment of the penalty parameters can be highly heuristic, which depends on the problems being solved and often requires advanced domain knowledge. To alleviate this difficulty, we propose to apply the data-driven supervised learning approach for learning the penalty parameters $\{\sigma_{j}\}$ such that the resulted algorithm performs empirically well when the random variable $\xi$ is sampled from a fixed distribution $\mathcal{P}$. In particular, for a fixed initial point $(x^{0},y^{0})\in\mathbb{X}\times\mathbb{Y}$ and a fixed step size $\tau\in(0,2)$, we see that the output of Algorithm 2 depends only on the random variable $\xi\in\Xi$ and the penalty parameters $\{\sigma_{j}\}$. To emphasize this dependency, we shall denote the output of Algorithm 2 as $x^{K}(\xi,\{\sigma_{j}\})$, with slightly abuse of notation. Similarly, we denote $x^{*}(\xi)$ to be the optimal solution of problem ($P(\xi)$) depending on $\xi\in\Xi$. In order to find a good strategy of selecting the values of the penalty parameters $\{\sigma_{j}\}$, we consider solving the following empirical risk minimization problem:

It is clear that the number of elements in the set $\{\sigma_{j}\}$ is at most $J:=\left\lfloor\frac{K}{K_{0}}\right\rfloor+1$. Thus, the (ERM) is usually a small-scale optimization problem, which can be solved efficiently by existing algorithms, such as SGD and ADAM. We note here that these commonly used optimizers rely on the back-propagation to compute the stochastic gradient estimators of the objective functions in (ERM). This implicitly requires that the computations in Algorithm 2 do not break the computational tree in order to keep track of the gradient information. If the back-propagation fails in practice, stochastic gradient based optimizers are no longer applicable. In this case, we may rely on grid search to find good penalty parameters, though it can be costly.

For the rest of this section, we consider applying the proposed hyperparameter learning approach for solving two class of important problems that attract growing attention in recent years. This first problem is the Lasso problem, which plays an important role in compressed sensing and sparse coding. And the second problem is the discrete optimal transport problem, which has many applications in modern machine learning. Solving them efficiently has been and will remain an active research direction in the literature.