e-mail comments/suggestions
Some notes on ADMM based on Boyd’s resources 1, assuming some familiarity with duality.
Let us consider the following problem:
\[\min_x \ f(x) \ \st A x = b\]Again, the Lagrangian is:
\[\LL(x, \lambda) = f(x) - \lambda^T\block{Ax - b}\]with dual function:
\[d(\lambda) = \min_x \LL(x, \lambda)\]Assuming that \(f\) is smooth and that we can compute the dual function, its gradient is trivial to compute:
\[\nabla d(\lambda) = Ax^\star(\lambda) - b\]where \(x^\star(\lambda) = \argmin{x}\ \LL(x, \lambda)\). Indeed, by definition of the dual function, we have:
\[d(\lambda) = \LL\block{x^\star(\lambda), \lambda}\]But since \(x^\star(\lambda)\) minimizes \(\LL\block{\cdot, \lambda}\), the derivative along \(x\) at the optimum is zero:
\[\ddd{\LL}{x}\block{x^\star(\lambda), \lambda} = 0\]so that only \(\ddd{\LL}{\lambda}\block{x^\star(\lambda), \lambda} = \block{Ax^\star - b}^T\) appears in the differential of \(d(\lambda)\). We now have a smooth2 function and its gradient, we can simply use gradient ascent with step sizes \(\alpha_k\) to solve the dual problem of maximizing \(d(\lambda)\):
Note that when \(f\) is not smooth, the above procedure provides a subgradient at each iteration and the method is known as dual subgradient ascent.
As usual, there are conditions to be met for the simple gradient ascent to converge. Unfortunately, it is not trivial to get Lipschitz constants for the dual function, therefore it might be difficult to converge robustly in practice.
In order to improve convergence, the Method of Multipliers replaces the initial problem with the following, equivalent one:
\[\min_x \ f(x) + \rho \norm{Ax - b}^2 \ \st A x = b\]whose Lagrangian is called the Augmented Lagrangian. Clearly the added penalty is zero on the feasible set, so the two problems are equivalent. In a sense, doing so regularizes the function \(f\) outside the feasible set by driving solutions towards the feasible set, over which the regularization vanishes.
Applying dual ascent to the regularized problem yields the following optimality conditions (dual feasibility) for the optimization problem at each iteration:
\[\nabla f\block{x_k} + \underbrace{\rho A^T\block{Ax_k - b} - A^T\lambda_k}_{-A^T\block{\lambda_k - \rho\block{Ax_k - b}}} = 0\]This suggests that picking step size \(\alpha_k = \rho\) such that \(\lambda_{k+1} = \lambda_k - \rho\block{Ax_k - b}\) will have the following benefits:
Therefore we might expect the nice properties of implicit integration to somehow ensure convergence. The iteration becomes:
More rigorously, convergence can be shown by considering the dual error \(V_k = \norm{\lambda_{k+1} - \lambda^\star}^2\) where \(\lambda^\star\) is the solution:
\[\begin{aligned} V_{k+1} - V_k &= \norm{\lambda_{k+1} - \lambda_k + \lambda_k - \lambda^\star}^2 - \norm{\lambda_k - \lambda^\star}^2\\ &= \norm{\lambda_{k+1} - \lambda_k}^2 + \norm{\lambda_k - \lambda^\star}^2 + 2\block{\lambda_{k+1} - \lambda_k}^T\block{\lambda_k - \lambda^\star} - \norm{\lambda_k - \lambda^\star}^2 \\ &= \norm{\lambda_{k+1} - \lambda_k}^2 + 2\block{\lambda_{k+1} - \lambda_k}^T\block{\lambda_k - \lambda^\star} \\ &= \norm{\lambda_{k+1} - \lambda_k}^2 + 2\block{\lambda_{k+1} - \lambda_k}^T\block{\lambda_k - \lambda_{k+1} + \lambda_{k+1} - \lambda^\star} \\ &= - \norm{\lambda_{k+1} - \lambda_k}^2 + 2 \block{\lambda_{k+1} - \lambda_k}^T\block{\lambda_{k+1} - \lambda^\star} \\ \end{aligned}\]The second term can be rewritten in terms of \(x\):
\[\begin{aligned} \block{\lambda_{k+1} - \lambda_k}^T\block{\lambda_{k+1} - \lambda^\star} &= -\rho\block{A x_k - b}^T \block{\lambda_{k+1} - \lambda^\star} \\ &= -\rho\block{x_k - x^\star}^TA^T\block{\lambda_{k+1} - \lambda^\star} \\ &= -\rho\block{x_k - x^\star}^T\block{\nabla f\block{x_k} - \nabla f\block{x^\star}} \\ &\leq 0 \end{aligned}\]due to \(\nabla f\) being monotone since \(f\) is convex. This means the method strictly converges while \(x_k\) is not primal feasible, and terminates with a result that is both primal and dual feasible. One can also show that the Method of Multipliers corresponds to the proximal point algorithm applied to the dual function (dual proximal point method), which is more direct in the non-smooth case.
The Method of Multipliers is also known as the Augmented Lagrangian Method (ALM).
At every iteration, \(x\) is solved as:
\[x_k = \argmin{x}\ f(x) - \lambda_k^T\block{Ax - b} + \frac{\rho}{2}\norm{Ax - b}^2\]The above can be simplified by completing the square:
\[\frac{\rho}2\norm{Ax - b - \frac{\lambda_k}{\rho}}^2 = \frac{\rho}{2}\norm{Ax - b}^2 - \lambda_k^T\block{Ax - b} + \frac{1}{2\rho}\norm{\lambda_k}^2\]Therefore, we can introduce \(u_k = \frac{\lambda_k}{\rho}\) and the ALM simplifies to the following:
which is a bit more convenient in practice.
So far, so good: we solved the problem of choosing step sizes \(\alpha_k\) and still get convergence. One practical issue is that the penalty term \(\norm{Ax - b}^2\) introduces coupling between variables that may not exist in function \(f\): while dual ascent could optimize a separable function \(f(x) = g(y) + h(z)\) well, separately (possibly using dedicated, optimized solvers), this is no longer possible in general with the Method of Multipliers.
The Alternating Direction Method of Multipliers (ADMM) improves the situation by working around the coupling introduced by constraint matrix. Let us introduce some notation first: we consider the following problem of minimizing a separable function under affine constraints:
\[\min_{x, z}\ f(x) + g(z)\ \st\ Ax + Bz = c\]Instead of minimizing jointly over both \(x, z\) like the Method of Multipliers would, the minimization is now split into two subproblems: minimizing along \(x\) alone first with \(z\) constant (the \(x\)-update), then along \(z\) alone with \(x\) constant (the \(z\)-update), in a Gauss-Seidel-like fashion:
In other words, this corresponds to alternating two Method of Multiplier solves in the \(x, z\) directions with varying constraint values, hence the name. Crucially, the matrices \(A, B\) remain constant, which enables preprocessing so that \(x, z\)-updates are as efficient as possible. Note that the role played by \(x,z\) is not symmetric. Dual feasibility for \(x_k\) gives:
\[\nabla f\block{x_k} - A^T \lambda_k + \rho A^T\block{Ax_k + Bz_k - c} = 0\]while dual feasibility for \(z_{k+1}\) gives:
\[\underbrace{\nabla g\block{z_{k+1}} - B^T \lambda_k + \rho B^T\block{Ax_k + Bz_{k+1} - c}}_{\nabla g\block{z_{k+1}} - B^T \lambda_{k+1}} = 0\]Therefore, \(z_{k+1}, \lambda_{k+1}\) is automatically dual-feasible for the original problem:
\[\nabla g\block{z_{k+1}} - B^T \lambda_{k+1} = 0\]while \(x_k, \lambda_{k+1}\) is not:
\[\begin{aligned} \nabla f\block{x_k} &- A^T \lambda_k + \rho A^T\block{Ax_k + Bz_k - c} \\ = \nabla f\block{x_k} &- A^T \block{\lambda_k - \rho \block{Ax_k + Bz_k - c}} \\ = \nabla f\block{x_k} &- A^T\block{\lambda_{k+1} - \rho Bz_{k+1} + \rho Bz_k} \\ = \nabla f\block{x_k} &- A^T\lambda_{k+1} - \rho A^TB\block{z_{k+1} - z_k} \\ \end{aligned}\]This suggests that convergence checks should not only consider the primal residual \(\norm{Ax_k + Bz_{k+1} - c}\) (constraints), but also \(\rho\norm{A^TB\block{z_{k+1} - z_k}}\) (stationarity for \(f\)), as described below.
For proving convergence, we consider the following energy:
\[W_k = \frac{1}{\rho} {\underbrace{\norm{\lambda_k - \lambda^\star}}_{V_k}^{}}^2 + \rho {\underbrace{\norm{B\block{z_k - z^\star}}}_{U_k}^{}}^2\]As above, we obtain:
\[V_{k+1} - V_k = -\frac{1}{\rho}\norm{\lambda_{k+1} - \lambda_k}^2 + \frac{2}{\rho} \block{\lambda_{k+1} - \lambda_k}^T\block{\lambda_{k+1} - \lambda^\star}\]and a similar computation gives:
\[U_{k+1} - U_k = -\rho\norm{B\block{z_{k+1} - z_k}}^2 + 2\rho\block{z_{k+1} - z_k}^TB^TB\block{z_{k+1} - z^\star}\]This time, \(\frac{\lambda_{k+1} - \lambda_k}{\rho}\) expands as:
\[\begin{aligned} \frac{\lambda_{k+1} - \lambda_k}{\rho} &= -\block{Ax_k + Bz_{k+1} - c} \\ &= -\block{Ax_k + Bz_{k+1} - \block{Ax^\star + B z^\star}} \\ &= -A\block{x_k - x^\star} - B\block{z_{k+1} - z^\star} \end{aligned}\]and we obtain:
\[\begin{aligned} &\frac{1}{\rho} \block{\lambda_{k+1} - \lambda_k}^T\block{\lambda_{k+1} - \lambda^\star} \\ &=-\block{x_k - x^\star}^TA^T\block{\lambda_{k+1} - \lambda^\star} - \block{z_{k+1} - z^\star}^TB^T\block{\lambda_{k+1} - \lambda^\star} \\ &= -\block{x_k - x^\star}^T\block{\nabla f\block{x_k} - \rho A^TB\block{z_{k+1} - z_k} -\nabla f\block{x^\star}} \\ &\phantom{=\,\,} -\block{z_{k+1} - z^\star}^T\block{\nabla g\block{z_{k+1}} - \nabla g\block{z^\star}}\\ &= -\underbrace{\block{x_k - x^\star}^T\block{\nabla f\block{x_k} -\nabla f\block{x^\star}}}_{\geq 0}\\ &\phantom{=\,\,} -\underbrace{\block{z_{k+1} - z^\star}^T\block{\nabla g\block{z_{k+1}} - \nabla g\block{z^\star}}}_{\geq 0}\\ &\phantom{=\,\,} +\rho \block{x_k - x^\star}^TA^TB\block{z_{k+1} - z_k} \\ \end{aligned}\]Therefore, we’re left with analyzing the sign of
\[\block{B\block{z_{k+1} - z^\star} + A\block{x_k - x^\star}}^TB\block{z_{k+1} - z_k}\]or, equivalently:
\[\begin{aligned} -\block{\lambda_{k+1} - \lambda_k}^TB\block{z_{k+1} - z_k} &= -\block{\nabla g\block{z_{k+1}} - \nabla g\block{z_k}}^T\block{z_{k+1} - z_k}\\ &\leq 0 \end{aligned}\]which entails strict convergence as long as either the primal residual \(\norm{Ax_k + Bz_{k+1} - c}\) or the dual residual \(\rho\norm{B\block{z_{k+1} - z_k}}\) is non-zero. A zero dual residual implies dual feasibility for \(x_k, \lambda_{k+1}\) and since dual feasibility for \(z_{x+1}, \lambda_{k+1}\) always holds, the proof is complete.
As before, introducing \(u_k = \frac{\lambda_k}{\rho}\) yields slightly more convenient equations in practice:
In particular, for the consensus problem \(A=I, B=-I, c = 0\) this gives:
\[\begin{aligned} x_k &= \argmin{x}\ f(x) + \frac{\rho}{2}\norm{x - \block{z_k + u_k}}^2 \\ z_{k+1} &= \argmin{z}\ g(z) \, + \frac{\rho}{2}\norm{z - \block{x_k - u_k}}^2 \\ \end{aligned}\]In some situations one might want to soften constraints in a controllable way, for instance when modelling materials with large-but-finite stiffness in physically-based simulations. A simple way to achieve this is to apply dual regularization to our constrained optimization problem:
\[\min_x\ f(x)\ \st Ax = b\]with Lagrangian
\[\LL(x, \lambda) = f(x) - \lambda^T(Ax - b)\]Dual regularization consists in penalizing Lagrange multipliers as follows:
\[\LL_C(x, \lambda) = f(x) - \lambda^T(Ax - b) - \frac{1}{2}\lambda^T C \lambda\]for some suitable positive semidefinite matrix \(C\) generally called compliance (inverse stiffness) in the context of mechanics. The primal function thus becomes:
\[p(x) = \max_\lambda \LL_C(x, \lambda) = \LL_C\block{x, \lambda^\star(x)}\]where \(\lambda^\star(x) = \argmin{\lambda}\ \LL_C\block{x, \lambda}\) satisfies
\[C \lambda^\star(x) = -\block{Ax - b}\]Therefore, the primal function can be rewritten as:
\[p(x) = f(x) + \frac{1}{2} \block{Ax - b}^T \inv{C} \block{Ax - b}\]which gets rid of constraints entirely, and validates the “constraint softening” interpretation.
Why bother, then? When \(C \to 0\), the primal problem as written above becomes severely ill-conditioned since \(\inv{C} \to \infty\). On the contrary, applying dual ascent to the regularized Lagrangian only ever involves matrix \(C\) and remains perfectly stable: the regularized dual function
\[d(\lambda) = \min_x \LL_C(x, \lambda)\]is minimized by the same \(x^\star(\lambda)\) as the original problem:
\[\begin{aligned} x^\star(\lambda) &= \argmin{x}\ f(x) -\lambda^T(Ax - b) - \frac{1}{2} \lambda^TC\lambda \\ &= \argmin{x}\ f(x) -\lambda^T(Ax - b) \\ &= \argmin{x}\ \LL(x, \lambda) \end{aligned}\]but the gradient of the dual function becomes:
\[\nabla d(\lambda) = -\block{Ax^\star(\lambda) - b} - C \lambda\]Since the \(x\)-update remains unchanged:
\[x_k = x^\star\block{\lambda_k} = \argmin{x}\ \LL(x, \lambda)\]the only modification to dual ascent is to replace the dual update with:
\[\lambda_{k+1} = \lambda_k - \alpha_k \block{\block{A x_k - b} + C \lambda_k}\]Unfortunately, introducing dual regularization as above undermines the nice properties of the method of multipliers. Since the \(x\)-update for the regularized Lagrangian is unchanged, dual feasibility still gives:
\[\nabla f\block{x_k} + \underbrace{\rho A^T\block{Ax_k - b} - A^T\lambda_k}_{-A^T\block{\lambda_k - \rho\block{Ax_k - b}}} = 0\]and the dual update corresponding to implicit integration for the non-augmented problem is still
\[\lambda_{k+1} = \lambda_k - \rho\block{A x_k - b}\]which correspond to the unregularized problem. In other words, either we keep implicit integration but loose regularization, or keep regularization but loose implicit integration. Therefore, we need some way of expressing dual feasibility in terms of \(C\) and deduce the dual update rule from it to recover implicit integration of the regularized problem. Luckily, the regularized problem can also be expressed as another constrained problem without regularization:
\[\min_{x, w}\ f(x) + \frac{1}{2}w^TCw\ \st\ Ax - b + Cw = 0\]to which one can apply the vanilla method of multipliers. The \(x,w\)-update becomes:
\[\block{x_k, w_k} = \argmin{x, w}\ f(x) + \frac{1}{2}w^TCw - \lambda^T\block{Ax - b + Cw} + \frac{\rho}{2}\norm{Ax - b + Cw}^2\]with optimality conditions:
\[\begin{aligned} \nabla f(x) - A^T\lambda + \rho A^T\block{Ax - b} + \rho A^TCw = 0\\ Cw - C^T\lambda + \rho C^TCw + \rho C^T\block{Ax - b} = 0\\ \end{aligned}\]Eliminating \(w\) gives (assuming \(C\) is invertible):
\[\begin{aligned} \block{I + \rho C} w &= \lambda - \rho \block{Ax - b}\\ Cw &= C\block{I + \rho C}^{-1}\block{\lambda - \rho \block{Ax - b}} \\ \end{aligned}\]and we’re left with:
\[\nabla f(x) - A^T\block{\block{I - \rho C\block{I + \rho C}^{-1}}\block{\lambda - \rho \block{Ax - b}}}= 0\]Since
\[I - \rho C\block{I + \rho C}^{-1} = \block{I + \rho C - \rho C}\block{I + \rho C}^{-1} = \block{I + \rho C}^{-1}\]Dual feasibility reduces to:
\[\nabla f(x) - A^T\block{I + \rho C}^{-1}\block{\lambda - \rho \block{Ax - b}} = 0\]and we recover implicit integration by choosing the following dual update:
\[\lambda_{k+1} = \block{I + \rho C}^{-1}\block{\lambda_k - \rho \block{Ax_k - b}}\]Assuming strict convexity for \(f\) ↩
e-mail comments/suggestions