Augmented Lagrangian and ADMM

Why This Matters

Many important ML problems have structure that a single monolithic optimizer cannot exploit. The Lasso has a smooth least-squares term plus a non-smooth $\ell_1$ penalty. Distributed training needs to split a problem across machines. Matrix completion has a rank constraint that must be handled carefully.

ADMM (Alternating Direction Method of Multipliers) handles all of these by splitting the problem into subproblems, each of which is easy to solve. It combines the decomposability of dual ascent with the convergence properties of the augmented Lagrangian method. Since Boyd et al.'s 2011 monograph, ADMM has become a standard tool for structured optimization in machine learning, signal processing, and statistics.

Mental Model

You want to minimize $f(\mathbf{x}) + g(\mathbf{z})$ subject to $\mathbf{A}\mathbf{x} + \mathbf{B}\mathbf{z} = \mathbf{c}$. You cannot minimize over $\mathbf{x}$ and $\mathbf{z}$ jointly because $f$ and $g$ have different structure (e.g., one is smooth, the other is non-smooth). ADMM alternates: first optimize over $\mathbf{x}$ (holding $\mathbf{z}$ fixed), then optimize over $\mathbf{z}$ (holding $\mathbf{x}$ fixed), then update a dual variable that enforces the constraint. Each step is easy; the alternation converges to the joint solution.

Formal Setup

The Penalty Method and Its Problem

The simplest approach to constrained optimization: penalize constraint violations. For the problem $\min_{\mathbf{x}} f(\mathbf{x})$ subject to $\mathbf{A}\mathbf{x} = \mathbf{b}$, the penalty method minimizes:

$f(\mathbf{x}) + \frac{\rho}{2}\|\mathbf{A}\mathbf{x} - \mathbf{b}\|^2$

As $\rho \to \infty$, the solution approaches the constrained optimum. But large $\rho$ makes the problem ill-conditioned: the Hessian has eigenvalues proportional to $\rho$, and optimization becomes slow.

The Augmented Lagrangian

Definition

Augmented Lagrangian

The augmented Lagrangian for $\min_{\mathbf{x}} f(\mathbf{x})$ subject to $\mathbf{A}\mathbf{x} = \mathbf{b}$ is:

$\mathcal{L}_\rho(\mathbf{x}, \boldsymbol{\lambda}) = f(\mathbf{x}) + \boldsymbol{\lambda}^T(\mathbf{A}\mathbf{x} - \mathbf{b}) + \frac{\rho}{2}\|\mathbf{A}\mathbf{x} - \mathbf{b}\|^2$

It combines the ordinary Lagrangian (with dual variable $\boldsymbol{\lambda}$) with a quadratic penalty (weight $\rho > 0$).

Proposition

Augmented Lagrangian Exactness

Statement

Unlike the pure penalty method, the augmented Lagrangian method achieves the exact constrained optimum at a finite penalty parameter $\rho$, provided the dual variable $\boldsymbol{\lambda}$ is updated correctly. The method of multipliers alternates:

$\mathbf{x}^{k+1} = \arg\min_{\mathbf{x}} \; \mathcal{L}_\rho(\mathbf{x}, \boldsymbol{\lambda}^k)$

$\boldsymbol{\lambda}^{k+1} = \boldsymbol{\lambda}^k + \rho(\mathbf{A}\mathbf{x}^{k+1} - \mathbf{b})$

Under convexity and the Slater condition, $\mathbf{x}^k$ converges to the primal optimum and $\boldsymbol{\lambda}^k$ converges to the dual optimum.

Intuition

The dual variable $\boldsymbol{\lambda}$ adapts to enforce the constraint. If the constraint is violated at iteration $k$, $\boldsymbol{\lambda}$ increases, penalizing future violations more. This adaptive mechanism means you do not need $\rho \to \infty$; a moderate $\rho$ suffices because $\boldsymbol{\lambda}$ does the fine-tuning.

Why It Matters

The augmented Lagrangian method avoids the ill-conditioning of the pure penalty method while maintaining convergence. ADMM is the augmented Lagrangian method applied to a split (decomposed) problem.

ADMM: The Algorithm

Consider the structured problem:

$\min_{\mathbf{x}, \mathbf{z}} \; f(\mathbf{x}) + g(\mathbf{z}) \quad \text{subject to} \quad \mathbf{A}\mathbf{x} + \mathbf{B}\mathbf{z} = \mathbf{c}$

Definition

ADMM Updates

The ADMM algorithm applies the augmented Lagrangian method with alternating minimization:

x-update: $\mathbf{x}^{k+1} = \arg\min_{\mathbf{x}} \left[f(\mathbf{x}) + \frac{\rho}{2}\|\mathbf{A}\mathbf{x} + \mathbf{B}\mathbf{z}^k - \mathbf{c} + \mathbf{u}^k\|^2\right]$

z-update: $\mathbf{z}^{k+1} = \arg\min_{\mathbf{z}} \left[g(\mathbf{z}) + \frac{\rho}{2}\|\mathbf{A}\mathbf{x}^{k+1} + \mathbf{B}\mathbf{z} - \mathbf{c} + \mathbf{u}^k\|^2\right]$

Dual update: $\mathbf{u}^{k+1} = \mathbf{u}^k + \mathbf{A}\mathbf{x}^{k+1} + \mathbf{B}\mathbf{z}^{k+1} - \mathbf{c}$

where $\mathbf{u} = \boldsymbol{\lambda}/\rho$ is the scaled dual variable.

The power of ADMM is that the $\mathbf{x}$-update only involves $f$ (plus a quadratic) and the $\mathbf{z}$-update only involves $g$ (plus a quadratic). If $f$ and $g$ have different structure, each subproblem can be solved with the appropriate specialized method.

Convergence

Theorem

ADMM Convergence

Statement

Under convexity and the existence of a saddle point, ADMM converges:

1. Residual convergence: $\mathbf{A}\mathbf{x}^k + \mathbf{B}\mathbf{z}^k - \mathbf{c} \to \mathbf{0}$ (primal feasibility)
2. Objective convergence: $f(\mathbf{x}^k) + g(\mathbf{z}^k) \to p^*$ (optimal value)
3. Dual convergence: $\mathbf{u}^k \to \mathbf{u}^*$ (dual optimality)

The convergence rate is $O(1/k)$ for the primal and dual residuals.

Intuition

ADMM makes progress in two ways each iteration: the primal updates reduce the objective (or at least the augmented Lagrangian), and the dual update enforces feasibility. The $O(1/k)$ rate means it reaches moderate accuracy quickly but converges slowly to high accuracy. This is acceptable for many ML applications where approximate solutions suffice.

Proof Sketch

Define the primal residual $\mathbf{r}^k = \mathbf{A}\mathbf{x}^k + \mathbf{B}\mathbf{z}^k - \mathbf{c}$ and the dual residual $\mathbf{s}^k = \rho \mathbf{A}^T \mathbf{B}(\mathbf{z}^k - \mathbf{z}^{k-1})$. Show that the augmented Lagrangian $\mathcal{L}_\rho$ decreases sufficiently each iteration: the decrease is lower bounded by $(\rho/2)\|\mathbf{r}^k\|^2 + (1/(2\rho))\|\mathbf{s}^k\|^2$. Since $\mathcal{L}_\rho$ is bounded below, the residuals must converge to zero. Telescoping the bound gives the $O(1/k)$ rate.

Why It Matters

The convergence guarantee holds under very mild conditions (convexity, saddle point existence) and does not require smoothness of $f$ or $g$. This is why ADMM handles non-smooth objectives like the $\ell_1$ norm naturally.

Failure Mode

The $O(1/k)$ rate is slow for problems requiring high precision. ADMM is not competitive with specialized methods (interior point, proximal gradient with acceleration) when high-accuracy solutions are needed. The penalty parameter $\rho$ significantly affects practical convergence speed and may need adaptive tuning.

Applications

Lasso via ADMM

The Lasso problem $\min_{\mathbf{x}} \frac{1}{2}\|\mathbf{A}\mathbf{x} - \mathbf{b}\|^2 + \lambda\|\mathbf{x}\|_1$ splits as $f(\mathbf{x}) = \frac{1}{2}\|\mathbf{A}\mathbf{x} - \mathbf{b}\|^2$ and $g(\mathbf{z}) = \lambda\|\mathbf{z}\|_1$ with the constraint $\mathbf{x} = \mathbf{z}$. The x-update is a ridge regression (closed-form). The z-update is soft-thresholding: $z_i^{k+1} = S_{\lambda/\rho}(x_i^{k+1} + u_i^k)$ where $S_\kappa(a) = \text{sign}(a)\max(|a|-\kappa, 0)$.

Consensus Optimization

For distributed ML, split data across $N$ machines. Each machine $i$ has a local objective $f_i(\mathbf{x}_i)$. The consensus problem is:

$\min \sum_{i=1}^{N} f_i(\mathbf{x}_i) \quad \text{subject to} \quad \mathbf{x}_1 = \mathbf{x}_2 = \cdots = \mathbf{x}_N$

ADMM introduces a global variable $\mathbf{z}$ and the constraint $\mathbf{x}_i = \mathbf{z}$. Each local update involves only $f_i$ and can be computed independently on machine $i$. The z-update is just averaging: $\mathbf{z}^{k+1} = \frac{1}{N}\sum_i (\mathbf{x}_i^{k+1} + \mathbf{u}_i^k)$. Communication is limited to gathering local variables and broadcasting the average.

Matrix Completion

For matrix completion (Netflix problem), the objective involves a low-rank constraint or nuclear norm regularization. ADMM splits the problem so that one subproblem is a least-squares fit to observed entries and the other is a singular value thresholding step that enforces low rank.

Common Confusions

Watch Out

ADMM is not just a penalty method

The dual variable update is what makes ADMM work. Without it, you have a penalty method that requires $\rho \to \infty$ for exactness and suffers from ill-conditioning. The dual update corrects for constraint violations at finite $\rho$. This is the key insight of the augmented Lagrangian approach.

Watch Out

Rho is not a regularization parameter

The penalty weight $\rho$ affects convergence speed but not the solution. All values of $\rho > 0$ converge to the same answer. In practice, too small $\rho$ gives slow primal convergence, too large $\rho$ gives slow dual convergence. Adaptive schemes adjust $\rho$ during optimization.

Summary

• Augmented Lagrangian = Lagrangian + quadratic penalty. Achieves exact solutions at finite $\rho$ (unlike pure penalty methods)
• ADMM splits a problem into an x-update, a z-update, and a dual update. Each subproblem has simpler structure than the original
• Convergence is guaranteed under convexity at $O(1/k)$ rate
• Lasso via ADMM: ridge regression + soft thresholding
• Consensus ADMM distributes optimization across machines with minimal communication (local compute + global averaging)
• ADMM excels at moderate accuracy; use specialized methods when high accuracy is required

Exercises

ExerciseCore

Problem

Write out the ADMM updates for the Lasso problem $\min_{\mathbf{x}} \frac{1}{2}\|\mathbf{A}\mathbf{x} - \mathbf{b}\|^2 + \lambda\|\mathbf{x}\|_1$. What is the closed-form solution for the x-update? What is the z-update?

Hint

Reveal Solution

ExerciseAdvanced

Problem

In consensus ADMM with $N$ machines, each local x-update is $\mathbf{x}_i^{k+1} = \arg\min_{\mathbf{x}_i} [f_i(\mathbf{x}_i) + (\rho/2)\|\mathbf{x}_i - \mathbf{z}^k + \mathbf{u}_i^k\|^2]$. Show that the z-update is the average of all $\mathbf{x}_i^{k+1} + \mathbf{u}_i^k$.

Hint

Reveal Solution

ExerciseResearch

Problem

ADMM converges at $O(1/k)$ for convex problems. Can ADMM be applied to non-convex problems (e.g., neural network training)? What guarantees exist, and what can go wrong?

Hint

Reveal Solution

References

Canonical:

• Boyd et al., "Distributed Optimization and Statistical Learning via ADMM" (2011)
• Bertsekas & Tsitsiklis, Parallel and Distributed Computation (1989)

Current:

• Parikh & Boyd, "Proximal Algorithms" (Foundations and Trends, 2014)
• Hong et al., "Convergence Analysis of ADMM for Nonconvex Optimization" (2016)

Next Topics

The natural next steps from ADMM:

• Trust region methods: another approach to constrained optimization where the constraint is a region of trust around the current iterate