Week 7: Proximal Methods and ADMM

Composite objectives#

Many machine learning problems have a composite structure: a smooth loss plus a nonsmooth regulariser:

$\min_x \; f(x) + g(x)$

where $f$ is smooth (e.g., squared loss, logistic loss) and $g$ is convex but possibly nonsmooth (e.g., $\ell_1$ norm, indicator of a convex set, nuclear norm). The prototypical example is the lasso:

$\min_\beta \; \tfrac{1}{2}\|X\beta - y\|_2^2 + \lambda \|\beta\|_1$

Gradient descent on the full objective fails because $g$ is not differentiable — the subgradient of $\|\beta\|_1$ at a coordinate where $\beta_j = 0$ is the entire interval $[-1, 1]$, giving no unique direction.

Proximal operator#

The proximal operator of a convex function $g$ is:

$\operatorname{prox}_{t g}(v) = \arg\min_x \left[ g(x) + \frac{1}{2t}\|x - v\|_2^2 \right]$

The prox evaluates $g$ at $x$ while penalising distance from $v$. Intuitively, it takes a step toward minimising $g$ but stays close to the reference point $v$, with $t$ controlling the trade-off. The prox is well-defined for any closed convex $g$, even nondifferentiable ones.

Key examples:

• $g(x) = \|x\|_1$ (lasso penalty): $\operatorname{prox}_{t\|\cdot\|_1}(v) = \operatorname{soft}(v, t)$, where the soft-thresholding operator is $\operatorname{soft}(v, t) = \operatorname{sign}(v) \cdot \max(|v| - t, 0)$ applied coordinate-wise.

• $g(x) = I_C(x)$ (indicator of closed convex set $C$): $\operatorname{prox}_{I_C}(v) = \operatorname{proj}_C(v)$, the Euclidean projection onto $C$.

• $g(x) = 0$: $\operatorname{prox}_{0}(v) = v$ — prox of zero is the identity.

Worked soft-threshold example. Let $v = 3.7$, $t = 2.0$. Then $\operatorname{soft}(3.7, 2.0) = \operatorname{sign}(3.7) \cdot (3.7 - 2.0) = 1.7$. For $v = 0.5$ with $t = 2.0$: $\operatorname{soft}(0.5, 2.0) = 0$ (shrunk to exactly zero). The threshold $t$ acts as a "kill zone": any $|v| \leq t$ is mapped to zero, implementing sparsity.

ISTA and FISTA#

The Iterative Shrinkage-Thresholding Algorithm (ISTA) is the proximal equivalent of gradient descent for composite problems:

$x_{k+1} = \operatorname{prox}_{\eta g}\big(x_k - \eta \nabla f(x_k)\big)$

Each iteration: take a gradient step on the smooth $f$, then apply the prox of the nonsmooth $g$. For lasso with $f(\beta) = \tfrac{1}{2}\|X\beta - y\|_2^2$ and $g(\beta) = \lambda\|\beta\|_1$:

$\beta_{k+1} = \operatorname{soft}\big(\beta_k - \eta X^T(X\beta_k - y),\; \eta\lambda\big)$

ISTA converges at rate $O(1/k)$ — the same as gradient descent on smooth problems. This is remarkable: the nondifferentiable $\ell_1$ term does not degrade the worst-case rate.

FISTA (Fast ISTA) adds Nesterov momentum to achieve $O(1/k^2)$ — the optimal rate for first-order methods on this class:

$y_{k+1} = x_k + \frac{t_k - 1}{t_{k+1}}(x_k - x_{k-1}), \quad x_{k+1} = \operatorname{prox}_{\eta g}(y_{k+1} - \eta \nabla f(y_{k+1}))$

where $t_{k+1} = (1 + \sqrt{1 + 4 t_k^2})/2$ is the Nesterov sequence. The acceleration comes from extrapolation — evaluating the gradient at a point beyond the current iterate, extrapolated in the direction of recent progress.

Why slow convergence is acceptable for lasso. Lasso is typically solved once per $\lambda$ for a fixed dataset of moderate dimension. $O(1/k^2)$ with a well-tuned step size converges to high accuracy in hundreds of iterations — seconds in practice. The simplicity of ISTA/FISTA (no linear system solve) makes them the go-to first-order choice for structured convex problems.

Operator splitting and ADMM#

ADMMAlternating Direction Method of Multipliers (Alternating Direction Method of Multipliers) solves problems of the form:

$\min_{x, z} \; f(x) + g(z) \quad \text{s.t.} \quad Ax + Bz = c$

where $f$ and $g$ are convex (possibly nonsmooth). The variables $x$ and $z$ are decoupled in the objective but coupled through a linear constraint. ADMM iterates:

where $\rho > 0$ is a penalty parameter and $u$ is the scaled dual variable. Each $x$- and $z$-subproblem involves only one of the two objective terms — this is the splitting that gives ADMM its power. When $f$ and $g$ have simple proximal operators (e.g., $f$ is a quadratic, $g$ is $\ell_1$), each subproblem is cheap.

Consensus ADMM. A common pattern is $f(x) = \sum_i f_i(x_i)$ and $g$ enforces consensus $x_i = z$ for all $i$. Each $f_i$ can be handled independently (in parallel), then the $z$-update averages the results. This maps directly to distributed optimisation.

When ADMM helps:

• Problems that decompose naturally into two blocks with cheap prox operators (e.g., lasso in consensus form across data shards).
• Distributed settings where the $x$-update can be parallelised.
• Problems where an IPM is too expensive but ISTA on the monolithic problem converges slowly due to coupling.

When ADMM is slower: For tightly coupled objectives where the $x$- and $z$-updates are as hard as the original problem, ADMM adds overhead (the penalty parameter $\rho$ and dual variable $u$) without reducing per-iteration cost.

Choosing Prox vs IPM vs SGD#

A quick decision guide for structured problems:

The patterns repeat across applications: find the structure, match it to the method, and exploit cheap proximal operators when they exist.

Knowledge Check#

Browser lab#

Compare GD on ridge regression (smooth) vs ISTA on lasso (composite) on a synthetic dataset. Plot objective value per iteration to contrast convergence behaviour.

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)
n, d = 100, 50
X = np.random.randn(n, d)
beta_true = np.zeros(d)
beta_true[:5] = [3, -2, 1.5, -1, 0.5]  # sparse
y = X @ beta_true + 0.5 * np.random.randn(n)

lam = 0.5

def soft_thresh(v, t):
    return np.sign(v) * np.maximum(np.abs(v) - t, 0)

# Ridge: smooth objective, GD
def ridge_grad(beta):
    return X.T @ (X @ beta - y) / n + lam * beta

# Lasso ISTA
def lasso_grad_smooth(beta):
    return X.T @ (X @ beta - y) / n

# GD on ridge
eta_ridge = 0.1
n_iter = 200
beta_ridge = np.zeros(d)
ridge_obj = []
for _ in range(n_iter):
    beta_ridge = beta_ridge - eta_ridge * ridge_grad(beta_ridge)
    obj = 0.5 * np.sum((X @ beta_ridge - y)**2) / n + 0.5 * lam * np.sum(beta_ridge**2)
    ridge_obj.append(obj)

# ISTA on lasso
eta_lasso = 0.1
beta_lasso = np.zeros(d)
lasso_obj = []
for _ in range(n_iter):
    beta_lasso = soft_thresh(beta_lasso - eta_lasso * lasso_grad_smooth(beta_lasso), eta_lasso * lam)
    obj = 0.5 * np.sum((X @ beta_lasso - y)**2) / n + lam * np.sum(np.abs(beta_lasso))
    lasso_obj.append(obj)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

ax1.semilogy(ridge_obj, "b-", linewidth=1.5, label="GD on Ridge")
ax1.semilogy(lasso_obj, "r-", linewidth=1.5, label="ISTA on Lasso")
ax1.set_xlabel("Iteration")
ax1.set_ylabel("Objective (log scale)")
ax1.set_title("Convergence: Ridge (GD) vs Lasso (ISTA)")
ax1.legend()
ax1.grid(True, alpha=0.3)

ax2.stem(range(d), beta_lasso, linefmt="r-", markerfmt="ro", basefmt="k-", label="Lasso (ISTA)")
ax2.stem(range(d), beta_true, linefmt="b--", markerfmt="b^", basefmt="k-", label="True")
ax2.set_xlabel("Coefficient index")
ax2.set_ylabel("Value")
ax2.set_title("Recovered vs True Coefficients (Lasso)")
ax2.legend()

plt.tight_layout()
plt.show()

print(f"Lasso nonzero coefficients: {np.sum(np.abs(beta_lasso) > 1e-4)}")
print(f"True nonzero coefficients: {np.sum(np.abs(beta_true) > 1e-4)}")
print(f"Estimation error ||beta - beta_true||: {np.linalg.norm(beta_lasso - beta_true):.4f}")

print("Notice: lasso selects exactly 5 nonzero coefficients, matching the 5 true ones, with estimation error 1.2149.")

Key takeaways#

• Composite problems $f + g$ with smooth $f$ and nonsmooth $g$ are handled by the proximal operator, which minimises $g$ while staying near a reference point.
• For $g = \|x\|_1$ the prox is soft thresholding, with a kill zone: any $|v| \leq t$ maps to exactly zero. That is the source of sparsity.
• The prox of an indicator function is Euclidean projection, which is why projected gradient descent is a proximal method in disguise.
• ISTA is a gradient step on $f$ followed by the prox of $g$, and it keeps the $O(1/k)$ rate despite the nonsmooth term. FISTA adds Nesterov momentum for $O(1/k^2)$.
• ADMM splits $\min f(x) + g(z)$ subject to $Ax + Bz = c$ into two cheap prox updates plus a dual ascent, and maps naturally to distributed and consensus settings.
• Match the structure: IPM for small high-accuracy cones, ISTA/FISTA for composite prox-friendly objectives, ADMM for multi-block or distributed problems, SGD/Adam for huge nonconvex ones.

Further reading#

• Neal Parikh & Stephen Boyd, Proximal Algorithms (Foundations and Trends in Optimization, 2014) — free online; the reference for prox operators and splitting.
• Stephen Boyd, Neal Parikh, Eric Chu, Borja Peleato & Jonathan Eckstein, "Distributed Optimization and Statistical Learning via ADMM" (Foundations and Trends in Machine Learning, 2011) — the ADMM paper.
• Amir Beck & Marc Teboulle, "A Fast Iterative Shrinkage-Thresholding Algorithm for Linear Inverse Problems" (SIAM Journal on Imaging Sciences, 2009) — ISTA and FISTA.