Duality and ADMM

Composite Objective Function

We have introduced proximal gradient descent to solve optimization problems of the form $\underset{x}{min}f(x)+g(x)$, where $f$ is smooth but $g$ is not differentiable. A key step in proximal gradient descent is solving the proximal operator:

$${\mathrm{p}\mathrm{r}\mathrm{o}\mathrm{x}}_g(x)=\arg \underset{z\in {\mathbb{R}}^d}{min}\{\frac{1}{2}\Vert z-x{\Vert }_2^2+g(z)\}$$

For example, when $g(x)=\lambda \Vert x{\Vert }_1$, the proximal operator is soft-thresholding. For the fused Lasso problem:

$$\underset{\beta }{min}\Vert Y-X\beta {\Vert }_2^2+\lambda \Vert D\beta {\Vert }_1$$

where $D$ is some matrix representing the differential map, applying proximal gradient descent requires solving:

$$\arg \underset{z\in {\mathbb{R}}^d}{min}\{\frac{1}{2}\Vert z-x{\Vert }_2^2+\lambda \Vert Dz{\Vert }_1\}$$

Unlike the ${\ell }_1$-norm, this problem lacks a closed-form solution, making proximal gradient descent inefficient for fused Lasso.

We discuss solving composite objective functions. Generally, we are interested in the following composite optimization problem:

$$\underset{x,y}{min}f(x)+g(y),\text{ s.t. }Ax+By=c$$

for some convex $f$ and $g$. The fused Lasso can be reduced to this form by letting $f(\beta )=\Vert Y-X\beta {\Vert }_2^2$, $g(y)=\lambda \Vert y{\Vert }_1$, so:

$$\underset{\beta }{min}\Vert Y-X\beta {\Vert }_2^2+\lambda \Vert D\beta {\Vert }_1⟺\underset{\beta ,y}{min}f(\beta )+g(y),\text{ s.t. }D\beta -y=0$$

Another example is basis pursuit:

$$\underset{\beta }{min}\Vert \beta {\Vert }_1\text{ s.t. }X\beta -Y=0$$

Duality

To solve the composite optimization problem, we start by reviewing duality in optimization. We define the Lagrange multiplier function of the problem as:

$$L(x,y,\lambda )=f(x)+g(y)+{\lambda }^{\mathrm{\top }}(Ax+By-c)$$

This allows us to convert the primal problem to the dual problem:

$$\underset{\lambda }{max}h(\lambda ),\text{ where }h(\lambda )=\underset{x,y}{min}L(x,y,\lambda )$$

Duality is crucial in optimization. The following theorem implies we can solve the primal problem via the dual problem.

Theorem (Strong Duality): Given convex functions $f,g$, the primal problem has the solution:

$$(x^{\ast },y^{\ast })=\arg \underset{x,y}{min}f(x)+g(y),\text{ s.t. }Ax+By=c$$

Consider the dual problem ${\lambda }^{\ast }=\arg \underset{\lambda }{max}h(\lambda )$, where $h(\lambda )=\underset{x,y}{min}L(x,y,\lambda )$. We can solve $(x^{\ast },y^{\ast })$ via:

$$(x^{\ast },y^{\ast })=\arg \underset{x,y}{min}L(x,y,{\lambda }^{\ast })$$

As $L(x,y,{\lambda }^{\ast })$ is decomposable, we further have:

$$x^{\ast }=\arg \underset{x}{min}f(x)+{\lambda }^{\ast \mathrm{\top }}Ax\text{ and }{y}^{\ast }=\arg \underset{y}{min}g(y)+{\lambda }^{\ast \mathrm{\top }}By$$

The strong duality property is useful for converting a linear constraint problem to an unconstrained dual problem. The following example shows how to find the dual problem of Lasso.

Example: Duality of Lasso

As Lasso is an unconstrained problem, we first convert it to the standard form:

$$\underset{\beta }{min}\frac{1}{2}\Vert Y-X\beta {\Vert }_2^2+\lambda \Vert \beta {\Vert }_1⟺\underset{\beta ,z}{min}\frac{1}{2}\Vert z{\Vert }_2^2+\lambda \Vert \beta {\Vert }_1,\text{ s.t. }Y-X\beta =z$$

The Lagrange function is:

$$L(\beta ,z,\mu )=\frac{1}{2}\Vert z{\Vert }_2^2+\lambda \Vert \beta {\Vert }_1+{\mu }^{\mathrm{\top }}(Y-X\beta -z)$$

Notice that $L(\beta ,z,\mu )$ is decomposable:

$$\underset{\beta ,z}{min}L(\beta ,z,\mu )=\underset{\beta }{min}\{\lambda \Vert \beta {\Vert }_1-{\mu }^{\mathrm{\top }}X\beta \}+\underset{z}{min}\{\frac{1}{2}\Vert z{\Vert }_2^2-{\mu }^{\mathrm{\top }}z\}+{\mu }^{\mathrm{\top }}Y$$

The second problem is quadratic, and we have:

$$z^{\ast }=\arg \underset{z}{min}\{\frac{1}{2}\Vert z{\Vert }_2^2-{\mu }^{\mathrm{\top }}z\}⟹z^{\ast }-\mu =0$$

Applying the first-order optimality condition to the first problem, we get:

$${\beta }^{\ast }=\arg \underset{\beta }{min}\{\lambda \Vert \beta {\Vert }_1-{\mu }^{\mathrm{\top }}X\beta \}⟹0=-X^{\mathrm{\top }}\mu +\lambda g,\text{ for some }g\in \partial \Vert {\beta }^{\ast }{\Vert }_1$$

As $\Vert g{\Vert }_{\mathrm{\infty }}\le 1$, the equality $0=-X^{\mathrm{\top }}\mu +\lambda g$ has a finite solution only if $\lambda \ge \Vert X^{\mathrm{\top }}\mu {\Vert }_{\mathrm{\infty }}$. Moreover, if $\lambda \ge \Vert X^{\mathrm{\top }}\mu {\Vert }_{\mathrm{\infty }}$, multiplying ${\beta }^{\ast }$ on both sides of $0=-X^{\mathrm{\top }}\mu +\lambda g$ gives:

$$0=-{\beta }^{\ast \mathrm{\top }}{X}^{\mathrm{\top }}\mu +\lambda \Vert {\beta }^{\ast }{\Vert }_1$$

Therefore, we have:

$$\underset{\beta }{min}\{\lambda \Vert \beta {\Vert }_1-{\mu }^{\mathrm{\top }}X\beta \}=\{\begin{array}{ll}0,& \text{if }\lambda \ge \Vert X^{\mathrm{\top }}\mu {\Vert }_{\mathrm{\infty }};\\ -\mathrm{\infty },& \text{if }\lambda <\Vert X^{\mathrm{\top }}\mu {\Vert }_{\mathrm{\infty }}.\end{array}$$

Combining the results, the dual problem becomes:

$$\underset{\lambda }{max}\underset{\beta ,z}{min}L(\beta ,z,\mu )=\underset{\lambda }{max}\{-\frac{1}{2}\Vert \mu {\Vert }_2^2+{\mu }^{\mathrm{\top }}Y\}\text{ s.t. }\lambda \ge \Vert X^{\mathrm{\top }}\mu {\Vert }_{\mathrm{\infty }}$$

Therefore, the dual problem of Lasso is:

$$\underset{\mu }{max}L({\beta }^{\ast },z^{\ast },\mu )=\underset{\mu }{max}-\frac{1}{2}\Vert \mu {\Vert }_2^2+{\mu }^{\mathrm{\top }}Y\text{ s.t. }\Vert X^{\mathrm{\top }}\mu {\Vert }_{\mathrm{\infty }}\le \lambda$$

This is equivalent to:

$$\text{**Duality of Lasso**}:\underset{\mu }{min}\Vert Y-\mu {\Vert }_2^2,\text{ s.t. }\Vert X^{\mathrm{\top }}\mu {\Vert }_{\mathrm{\infty }}\le \lambda$$

By the duality, the dual variable ${\mu }^{\ast }=z^{\ast }=Y-X{\beta }^{\ast }$ is the residual. This gives a new geometric insight into Lasso, illustrating that Lasso is a projection, similar to OLS.

Alternating Direction Method of Multipliers (ADMM)

Consider an equivalent form of the problem by adding a quadratic term:

$$\left\{\begin{array}{rl}& \underset{x,y}{min}f(x)+g(y)\\ & \text{ s.t. }Ax+By=c,\end{array}\right\}⟺\left\{\begin{array}{rl}& \underset{x,y}{min}f(x)+g(y)+\frac{\rho }{2}\Vert Ax+By-c{\Vert }_2^2\\ & \text{ s.t. }Ax+By=c.\end{array}\right\}$$

We introduce the augmented Lagrangian:

$$L_{\rho }(x,y,\lambda )=f(x)+g(y)+{\lambda }^{\mathrm{\top }}(Ax+By-c)+\frac{\rho }{2}\Vert Ax+By-c{\Vert }_2^2$$

To solve the dual problem $\underset{\lambda }{max}\underset{x,y}{min}{L}_{\rho }(x,y,\lambda )$, we propose updating the primal and dual variables alternately:

Primal step:

$$\{\begin{array}{l}{x}_{t+1}=\arg \underset{x}{min}{L}_{\rho }(x,y_t,{\lambda }_t);\\ y_{t+1}=\arg \underset{y}{min}{L}_{\rho }(x_{t+1},y,{\lambda }_t);\end{array}$$

Dual step:

$${\lambda }_{t+1}=\arg \underset{\lambda }{max}{L}_{\rho }(x_{t+1},y_{t+1},\lambda )-\frac{1}{2\rho }\Vert \lambda -{\lambda }_t{\Vert }_2^2$$

We add a proximal term in the dual step because the augmented Lagrangian $L_{\rho }(x_{t+1},y_{t+1},\lambda )$ is linear with respect to $\lambda$, making $\underset{\lambda }{max}{L}_{\rho }(x_{t+1},y_{t+1},\lambda )=\mathrm{\infty }$. The proximal term prevents ${\lambda }_{t+1}$ from deviating too much from ${\lambda }_t$.

Plugging the augmented Lagrangian into the primal and dual steps, we have the following algorithm.

Alternating Direction Method of Multipliers (ADMM)

$$\begin{array}{rl}{x}_{t+1}& =\arg \underset{x}{min}f(x)+\frac{\rho }{2}\Vert Ax+B{y}_t-c+{\lambda }_t/\rho {\Vert }_2^2\\ y_{t+1}& =\arg \underset{y}{min}g(x)+\frac{\rho }{2}\Vert A{x}_{t+1}+By-c+{\lambda }_t/\rho {\Vert }_2^2\\ {\lambda }_{t+1}& ={\lambda }_t+\rho (A{x}_{t+1}+B{y}_{t+1}-c)\end{array}$$

If $f$ and $g$ are closed convex functions, the ADMM algorithm converges.

ADMM not first-order method

ADMM is not a first-order method. The sub-problem in the primal step use the functions $f$ and $g$ instead of just their gradients like the proximal gradient descent. So ADMM may converge faster than first-order methods in terms of the number of iterations, but the computational cost per iteration might be higher.

Example: Fused Lasso

Applying ADMM to the fused Lasso:

$$\underset{\beta }{min}\frac{1}{2}\Vert Y-X\beta {\Vert }_2^2+\lambda \Vert D\beta {\Vert }_1⟺\underset{\beta ,z}{min}\frac{1}{2}\Vert Y-X\beta {\Vert }_2^2+\lambda \Vert z{\Vert }_1\text{ s.t. }D\beta -z=0$$

The updating rule is:

$$\begin{array}{rl}{\beta }_{t+1}& =\arg \underset{\beta }{min}\frac{1}{2}\Vert Y-X\beta {\Vert }_2^2+\frac{\rho }{2}\Vert D\beta -z_t+{\lambda }_t/\rho {\Vert }_2^2\\ z_{t+1}& =\arg \underset{z}{min}\lambda \Vert z{\Vert }_1+\frac{\rho }{2}\Vert D{\beta }_{t+1}-z+{\lambda }_t/\rho {\Vert }_2^2\\ {\lambda }_{t+1}& ={\lambda }_t+\rho (D{\beta }_{t+1}-z_{t+1})\end{array}$$

This yields the algorithm:

$$\begin{array}{rl}{\beta }_{t+1}& =(X^{\mathrm{\top }}X+\rho D^{\mathrm{\top }}D{)}^{-1}(X^{\mathrm{\top }}Y+\rho D^{\mathrm{\top }}{z}_t-D^{\mathrm{\top }}{\lambda }_t)\\ z_{t+1}& =\text{SoftThreshold}(D{\beta }_{t+1}+{\lambda }_t/\rho ,\lambda /\rho )\\ {\lambda }_{t+1}& ={\lambda }_t+\rho (D{\beta }_{t+1}-z_{t+1})\end{array}$$

When $D=I$, the algorithm reduces to ADMM for Lasso. Unlike FISTA, ADMM involves matrix inversion, which may be time-consuming when $d$ is large.

Here is the implementation of ADMM for fused Lasso in PyTorch. Notice we can implement the tips in Linear Algebra to pre-compute the Cholesky decomposition of the matrix $(X^{\mathrm{\top }}X+\rho D^{\mathrm{\top }}D)$ to make it more efficient.

import torch

# Soft-thresholding function
def soft_thresholding(y, threshold):
    return torch.sign(y) * torch.clamp(torch.abs(y) - threshold, min=0)

# Define problem parameters
torch.manual_seed(42)
n, d = 10, 5  # Number of samples (n) and features (d)
X = torch.randn(n, d)  # Design matrix
Y = torch.randn(n)  # Response vector

# Define difference matrix D (for fused lasso)
D = torch.eye(d) - torch.eye(d, k=1)  # First-order difference matrix

# Initialize variables
beta = torch.zeros(d)  # Coefficients
z = torch.zeros(d - 1)  # Auxiliary variable
lam = torch.zeros(d - 1)  # Lagrange multiplier

# ADMM parameters
lambda_ = 0.1  # Regularization parameter
rho = 1.0  # ADMM penalty parameter
num_iters = 100  # Number of iterations

# Compute Cholesky decomposition of (X^T X + rho D^T D)
XtX = X.T @ X
DtD = D.T @ D
A = XtX + rho * DtD  # Regularized system matrix

# Cholesky decomposition (more stable than direct inversion)
L = torch.linalg.cholesky(A)  # Compute Cholesky factor L

for t in range(num_iters):
    # Solve for beta using Cholesky decomposition
    b_rhs = X.T @ Y + rho * D.T @ z - D.T @ lam  # Right-hand side
    y = torch.linalg.solve_triangular(L, b_rhs, upper=False)  # Forward substitution
    beta = torch.linalg.solve_triangular(L.T, y, upper=True)  # Backward substitution

    # Update z using soft-thresholding
    z = soft_thresholding(D @ beta + lam / rho, lambda_ / rho)

    # Update lambda (dual variable)
    lam += rho * (D @ beta - z)

When $d$ is large, we can further improve the efficiency by using the Woodbury matrix identity to solve the linear system $$ (X^TX + ρD^TD) = (ρD^TD) - (ρD^TD)X^T(I_n + X(ρD^TD)X^T)X(ρD^TD) $$ The computation will be much more efficient, especially when for the Lasso case $D=I$.

Example: Graphical Lasso

Given i.i.d. samples $X_1,\dots ,X_n\sim N(0,\mathrm{\Sigma })$, we estimate the precision matrix $\mathrm{\Theta }={\mathrm{\Sigma }}^{-1}$ via graphical Lasso:

$$\begin{array}{rl}& \underset{\mathrm{\Theta }}{min}-\log det\mathrm{\Theta }+\text{tr}({\mathrm{\Theta }}^{\mathrm{\top }}\hat{\mathrm{\Sigma }})+\lambda \Vert \mathrm{\Theta }{\Vert }_{1,1}\\ \Updownarrow \\ & \underset{\mathrm{\Theta },\mathrm{\Psi }}{min}-\log det\mathrm{\Theta }+\text{tr}({\mathrm{\Theta }}^{\mathrm{\top }}\hat{\mathrm{\Sigma }})+\lambda \Vert \mathrm{\Psi }{\Vert }_{1,1}\text{ s.t. }\mathrm{\Theta }=\mathrm{\Psi }\end{array}$$

Applying ADMM, we get:

$$\begin{array}{rl}{\mathrm{\Theta }}_{t+1}& =\arg \underset{\mathrm{\Theta }}{min}-\log det\mathrm{\Theta }+\text{tr}({\mathrm{\Theta }}^{\mathrm{\top }}\hat{\mathrm{\Sigma }})+\frac{\rho }{2}\Vert \mathrm{\Theta }-{\mathrm{\Psi }}_t+{\mathrm{\Lambda }}_t/\rho {\Vert }_{\mathrm{F}}^2\\ {\mathrm{\Psi }}_{t+1}& =\arg \underset{\mathrm{\Psi }}{min}\lambda \Vert \mathrm{\Psi }{\Vert }_{1,1}+\frac{\rho }{2}\Vert {\mathrm{\Theta }}_{t+1}-\mathrm{\Psi }+{\mathrm{\Lambda }}_t/\rho {\Vert }_{\mathrm{F}}^2\\ {\mathrm{\Lambda }}_{t+1}& ={\mathrm{\Lambda }}_t+\rho ({\mathrm{\Theta }}_{t+1}-{\mathrm{\Psi }}_{t+1})\end{array}$$

where the Frobenius norm $\Vert A{\Vert }_{\mathrm{F}}^2=\sum _{jk}{A}_{jk}^2$. This simplifies to:

$$\begin{array}{rl}{\mathrm{\Theta }}_{t+1}& ={\mathcal{F}}_{\rho }({\mathrm{\Psi }}_t-{\mathrm{\Lambda }}_t/\rho -\hat{\mathrm{\Sigma }}/\rho )\\ {\mathrm{\Psi }}_{t+1}& =\text{SoftThreshold}({\mathrm{\Theta }}_{t+1}+{\mathrm{\Lambda }}_t/\rho ,\lambda /\rho )\\ {\mathrm{\Lambda }}_{t+1}& ={\mathrm{\Lambda }}_t+\rho ({\mathrm{\Theta }}_{t+1}-{\mathrm{\Psi }}_{t+1})\end{array}$$

where for the spectral decomposition $X=UD{U}^{\mathrm{\top }}$, ${\mathcal{F}}_{\rho }(X)=U\mathrm{d}\mathrm{i}\mathrm{a}\mathrm{g}\{{\lambda }_i+\sqrt{{\lambda }_i+4/\rho }\}{U}^{\mathrm{\top }}$.

Example: Consensus Optimization

We design a divide-and-conquer ADMM for Lasso for massive data when $n$ is ultra-large. The Lasso:

$$\underset{\beta }{min}\sum _{i=1}^n(Y_i-X_i\beta {)}^2+\lambda \Vert \beta {\Vert }_1$$

can be formulated as the general block separable form:

$$\underset{x}{min}\sum _{i=1}^n{f}_i(x)⟺\underset{x_i,z}{min}\sum _{i=1}^n{f}_i(x)\text{ s.t. }{x}_i=z\text{ for all }i=1,\dots ,n$$

Applying ADMM, we get:

ADMM for Consensus Optimization

$$\begin{array}{rl}\text{Divide: }{x}_i^{t+1}& =\arg \underset{x_i}{min}{f}_i(x_i)+\frac{\rho }{2}\Vert x_i-z^t+{\lambda }_i^t/\rho {\Vert }_2^2,\mathrm{\forall }i=1,\dots ,n\\ \text{Gather: }{z}^{t+1}& =\frac{1}{n}\sum _{i=1}^n(x_i^{t+1}+{\lambda }_i^t/\rho )\\ \text{Broadcast: }{\lambda }_i^{t+1}& ={\lambda }_i^t+\rho (x_i^{t+1}-z^{t+1}),\mathrm{\forall }i=1,\dots ,n\end{array}$$

In the first step, we update each $x_i^{t+1}$ in parallel, then gather all local iterates, and finally broadcast the updated dual variables back to each core.