KBKnowledge Base
Machine Learning · 2.2.5

Coordinate & Proximal Methods

Coordinate descent, subgradients, and the soft-thresholding operator behind Lasso.

On this page
In plain English — beginner to advanced

Beginner: coordinate descent optimizes one variable at a time, holding every other variable fixed, cycling through all of them repeatedly — like solving a jigsaw puzzle by adjusting one piece at a time to its best position while leaving the rest untouched, then moving to the next piece.

Intermediate: why would you want this instead of gradient descent? Because sometimes optimizing just ONE coordinate, with everything else fixed, has an easy closed-form solution even when the full joint problem has no such formula. That's exactly the situation for L1-regularized ("Lasso-style") objectives, which is why this lesson exists.

Advanced: an L1 penalty term |θⱼ| is not differentiable at θⱼ = 0 — it has a sharp corner there, not a smooth curve — so plain gradient descent's "follow the gradient downhill" recipe breaks down exactly at the most important point, since the whole reason to use an L1 penalty is to push many coefficients to exactly zero. The subgradient generalizes "gradient" to make sense at such a corner (a whole range of valid slopes instead of one specific slope), and the proximal operator is a clean, general way to handle a non-differentiable penalty by inserting a periodic shrinkage step into an otherwise ordinary optimization loop.

Formula
θj={sign(θj)θj0[1,1]θj=0\partial|\theta_j| = \begin{cases}\text{sign}(\theta_j) & \theta_j \ne 0 \\ [-1,1] & \theta_j = 0\end{cases}
soft(θ,λ)=sign(θ)max(θλ,0)\text{soft}(\theta, \lambda) = \text{sign}(\theta)\max(|\theta|-\lambda, 0)

The subgradient set of the L1 penalty, and the soft-thresholding operator — the proximal operator of the L1 penalty — that this lesson derives below.

Derivation: the Lasso coordinate update, case by case

Fix every coefficient except θⱼ and minimize ½‖y − Xθ‖² + λ‖θ‖₁ with respect to θⱼ alone. Expanding the squared term and collecting everything that depends on θⱼ, the 1D sub-problem reduces to minimizing ½aθⱼ² − bθⱼ + λ|θⱼ|, where a = Σᵢ Xᵢⱼ² (column j's squared norm) and b = Σᵢ Xᵢⱼ(yᵢ − Σₖ≠ⱼ Xᵢₖθₖ) (column j's correlation with the current partial residual — everything that depends on the OTHER, currently-fixed coefficients).

Solve this 1D problem by cases:

  • Case θⱼ > 0: the penalty term is smooth here (+λθⱼ), so differentiate the whole thing, aθⱼ − b + λ, set it to zero, giving θⱼ = (b − λ)/a — valid only when this comes out positive, i.e. when b > λ.
  • Case θⱼ < 0: symmetric, the penalty is −λθⱼ here, giving θⱼ = (b + λ)/a — valid only when b < −λ.
  • Case θⱼ = 0: check whether zero is already a valid subgradient solution — it is exactly when some value in the subgradient set [−λ,λ] can make the total (sub)gradient −b + [−λ,λ] contain zero, i.e. whenever |b| ≤ λ.

These three cases combine into exactly one closed form: θⱼ* = soft(b, λ)/a — the soft-thresholding operator applied to b, then rescaled by 1/a. Every case above is a direct instance of this single formula: it returns 0 precisely when |b| ≤ λ, and the two smooth-region solutions otherwise.

Where this is used: this exact coordinate-wise soft-thresholding update, applied repeatedly across all coefficients until convergence, is literally the algorithm sklearn.linear_model.Lasso and the classical glmnet package use in production — coordinate descent isn't a theoretical curiosity here, it's the actual solver.

Soft-thresholding: the shrinkage operator behind Lasso

Drag any colored marker's input position and watch where it lands after soft-thresholding. Widen λ with the slider and watch the flat 'dead zone' around zero swallow more and more markers, snapping them to exactly 0.

Practical example — Lasso via coordinate descent
python
import numpy as np

def soft_threshold(theta, lam):
    return np.sign(theta) * np.maximum(np.abs(theta) - lam, 0)

def lasso_coordinate_descent(X, y, lam, iters=200):
    n, p = X.shape
    theta = np.zeros(p)
    col_sq = (X ** 2).sum(axis=0)
    for _ in range(iters):
        for j in range(p):
            residual = y - X @ theta + X[:, j] * theta[j]
            b = X[:, j] @ residual
            theta[j] = soft_threshold(b, lam) / col_sq[j]
    return theta

rng = np.random.default_rng(0)
n, p = 200, 12
X = rng.normal(size=(n, p))
true_theta = np.array([3.0, -2.0, 0, 0, 1.5, 0, 0, 0, 0, -1.0, 0, 0])
y = X @ true_theta + 0.3 * rng.normal(size=n)

fitted = lasso_coordinate_descent(X, y, lam=8.0)
print(np.round(fitted, 3))
# Several entries should land at exactly 0.0, matching the true zero coefficients.
Real-world examples
  • Genomics and bioinformatics — Lasso regression is a standard tool for automatic feature selection when there are thousands of candidate genes/markers and only a handful actually matter.
  • Finance and econometrics — sparse regression is often preferred over a dense model specifically because knowing which few factors matter is as important as the prediction itself.
  • Large-scale recommender system factorization — some matrix factorization solvers cycle through user and item factor updates in a pattern very similar to coordinate descent, for the same "each sub-step is cheap and exact" reason.
  • Group Lasso and Elastic Net extend this exact soft-thresholding idea to structured sparsity (whole groups of coefficients pushed to zero together) and to a blend of L1 and L2 penalties — covered fully in a later Regression module, but built on precisely the machinery derived here.
  • Coordinate descent is often faster in wall-clock time than gradient-based methods specifically for L1 problems, even though each individual step only touches one variable — because each step is cheap and exact rather than an approximate gradient step that still needs a learning rate.
Common mistakes
  • Applying plain gradient descent directly to an L1-penalized objective with no special handling — it won't reliably produce exact zeros, which defeats the entire point of choosing an L1 penalty in the first place.
  • Assuming coordinate descent's nice convergence guarantees carry over to non-convex problems — they're specific to convex objectives like Lasso's; the previous module's convexity lesson is exactly what makes this safe here.
  • Assuming the specific soft-thresholding formula generalizes automatically to every non-smooth penalty — it's tailored to the L1 penalty's particular corner shape; other penalties (group lasso, nuclear norm) have their own, different proximal operators.
Going deeper

The general proximal gradient framework, of which this coordinate-wise soft-thresholding update is a simple special case, extends to a much broader class of objectives: any smooth loss plus any non-smooth-but-"simple" penalty. Many other common penalties have their own closed-form proximal operators — group lasso, the nuclear norm for low-rank matrix problems, and indicator functions of constraint sets for projected gradient methods — all sharing the identical pattern: take an ordinary smooth gradient step, then apply a shrinkage or projection operator specific to the penalty in play.

Check yourself
Why does soft-thresholding produce coefficients that are exactly zero, rather than just very small, the way an L2 (ridge) penalty would?

Because the L1 penalty's subgradient at zero is an entire interval, [-λ, λ], not a single slope. Whenever the data term's pull, b, has magnitude at or below λ, zero itself is already a valid subgradient solution — the optimization has no incentive to move away from it. An L2 penalty's gradient at zero is exactly 0 with no such interval, so it always pulls coefficients toward zero proportionally but essentially never lands exactly on it for generic data.

Key takeaway

Coordinate descent plus the soft-thresholding proximal operator turns an objective that plain gradient descent can't handle cleanly — a non-differentiable corner exactly where the interesting behavior (sparsity) happens — into a sequence of trivial, closed-form 1D updates. The next lesson extends this same "handle the hard part separately" instinct to constrained problems via duality.

Newsletter

Stay in the loop

Subscribe to get new docs, diagrams, and engineering write-ups by Dharaneesh Boobalan delivered to your inbox.

  • Deep-dive write-ups on ML, inference, and systems.
  • New Draw.io diagrams & interactive canvases.
  • Agentic patterns and rocket-science notes.
  • No spam. One tasteful email when there's something new.

Crafted by Dharaneesh Boobalan

Newsletter

Get new docs, diagrams, and write-ups in your inbox.

We never share your details. Unsubscribe anytime.