Week 10: Nonsmooth and Composite Objectives

When gradients don't exist#

Not every useful loss function is differentiable everywhere. Common nondifferentiable functions in ML:

• $f(x) = |x|$: nondifferentiable at $x = 0$ (the kink)
• $\text{ReLU}(x) = \max(0, x)$: nondifferentiable at $x = 0$
• $f(x) = \|x\|_1$, the $\ell_1$ norm, nondifferentiable on the coordinate axes
• $f(x) = \max(x_1, x_2)$: nondifferentiable when $x_1 = x_2$

These functions are convex (Week 2's Jensen inequality still holds), so we can extend the notion of "gradient" to a set-valued generalisation.

Subgradients#

For a convex function $f$, a vector $g$ is a subgradient at $x$ if:

$f(y) \geq f(x) + g^T (y - x) \quad \text{for all } y$

This is the same first-order condition from Week 2, but now $g$ is not unique. The set of all subgradients at $x$ is the subdifferential $\partial f(x)$.

Worked examples:

• $f(x) = |x|$: $\partial f(0) = [-1, 1]$ (any value between $-1$ and $1$ works). For $x > 0$, $\partial f(x) = \{1\}$; for $x < 0$, $\partial f(x) = \{-1\}$.

• $\text{ReLU}(x) = \max(0, x)$: $\partial f(0) = [0, 1]$. For $x > 0$, $\partial f(x) = \{1\}$; for $x < 0$, $\partial f(x) = \{0\}$.

• $f(x) = \max(x_1, x_2)$: At $x_1 = x_2$, $\partial f(x) = \{(\theta, 1-\theta) \mid \theta \in [0, 1]\}$: the convex combinations of the two active gradients.

The subgradient method is the direct extension of GD:

$x_{k+1} = x_k - \eta_k g_k \quad \text{where } g_k \in \partial f(x_k)$

But it has a major drawback: the subgradient method is not a descent method. Unlike GD, taking a subgradient step does not guarantee $f(x_{k+1}) < f(x_k)$. The function value can increase. For this reason, the subgradient method is less popular in ML than proximal methods.

Composite objectives#

Many important problems have the form:

$\min_x \; f(x) = \underbrace{g(x)}_{\text{smooth}} + \underbrace{h(x)}_{\text{nonsmooth}}$

The classic example is the lasso:

$\min_w \; \frac{1}{2}\|Xw - y\|_2^2 + \lambda \|w\|_1$

where $g(w) = \frac{1}{2}\|Xw - y\|_2^2$ is $L$-smooth (least squares) and $h(w) = \lambda \|w\|_1$ is nonsmooth (sparsity-inducing regulariser).

The $\ell_1$ penalty encourages sparse solutions: many $w_i$ become exactly zero, which performs automatic feature selection.

Proximal gradient (lite)#

The proximal gradient method (also called ISTAIterative Shrinkage-Thresholding Algorithm) elegantly handles composite objectives:

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

where the proximal operator of $h$ is:

$\text{prox}_{\eta h}(v) = \arg\min_x \left\{ h(x) + \frac{1}{2\eta}\|x - v\|_2^2 \right\}$

This takes a gradient step on the smooth part $g$, then "projects" the result through the proximal operator of $h$.

The key example: soft thresholding for $\ell_1$. For $h(x) = \lambda \|x\|_1$ (vector case):

$\text{prox}_{\eta \lambda \|\cdot\|_1}(v)_i = \text{sign}(v_i) \cdot \max(|v_i| - \eta\lambda, 0)$

This shrinks each coordinate toward zero by $\eta\lambda$, and sets it to exactly zero if it was within that threshold. This is why ISTA produces sparse solutions: coordinates smaller than the threshold vanish.

Worked example: 1D lasso step.

Take $g(w) = \frac{1}{2}(w - 3)^2$ (gradient $\nabla g(w) = w - 3$), $h(w) = 2|w|$, $\eta = 0.5$.

• Gradient step: $v = w_k - \eta \nabla g(w_k)$. If $w_k = 1$, $v = 1 - 0.5 \cdot (-2) = 2$.
• Proximal step: soft-threshold $v = 2$ with $\eta\lambda = 0.5 \cdot 2 = 1$. The result is $\text{sign}(2) \cdot \max(2-1, 0) = 1$.

The iterate moved from $w_k = 1$ to $w_{k+1} = 1$, stopped at the threshold. If the gradient were stronger, the iterate would cross the threshold and move.

Early stopping as regularization#

An intriguing phenomenon: running GD on the unregularised least squares problem

$\min_w \frac{1}{2}\|Xw - y\|_2^2$

and stopping early (after a small number of iterations) produces a solution that behaves similarly to ridge regression. The early-stopped GD solution has smaller norm and better generalisation than the fully-converged solution.

Intuition: GD started from $w_0 = 0$ explores the parameter space from the origin outward. The first few iterations capture large-scale, low-curvature directions (principal components); later iterations fit fine-grained, high-curvature directions that often correspond to noise. Stopping early implicitly truncates the high-frequency fitting.

This is a special case of implicit regularisation: the optimisation algorithm itself, not an explicit penalty term, biases the solution. This theme returns in Week 13 (nonconvex deep learning).

Knowledge check#

Browser lab: subgradient vs proximal gradient#

Compare the subgradient method vs proximal gradient (ISTA) on L1-regularised least squares (lasso). Plot objective vs iteration.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
n, d = 100, 50
w_true = rng.normal(size=d) * (rng.random(d) > 0.7)
X = rng.normal(size=(n, d))
y = X @ w_true + rng.normal(scale=0.3, size=n)

lam = 0.5
L = np.linalg.eigvalsh(X.T @ X / n).max()
eta = 1.0 / L
K = 200

def soft_thresh(v, thresh):
    return np.sign(v) * np.maximum(np.abs(v) - thresh, 0.0)

def objective(w):
    return 0.5 * np.mean((X @ w - y)**2) + lam * np.sum(np.abs(w))

w_sg = np.zeros(d)
loss_sg = [objective(w_sg)]
for k in range(1, K+1):
    g = X.T @ (X @ w_sg - y) / n
    subg = g + lam * np.sign(w_sg)
    subg[np.abs(w_sg) < 1e-10] = g[np.abs(w_sg) < 1e-10] + lam * np.sign(g[np.abs(w_sg) < 1e-10])
    w_sg = w_sg - (eta / np.sqrt(k)) * subg
    loss_sg.append(objective(w_sg))

w_ista = np.zeros(d)
loss_ista = [objective(w_ista)]
for _ in range(K):
    grad_g = X.T @ (X @ w_ista - y) / n
    w_ista = soft_thresh(w_ista - eta * grad_g, eta * lam)
    loss_ista.append(objective(w_ista))

plt.figure(figsize=(8, 5))
plt.plot(loss_sg, label="Subgradient method (η_k = η₀/√k)", alpha=0.8)
plt.plot(loss_ista, label="ISTA (proximal gradient)", linewidth=2)
plt.yscale("log"); plt.xlabel("Iteration"); plt.ylabel("Objective f(w)")
plt.title("Subgradient Method vs ISTA on Lasso")
plt.legend(); plt.grid(True, alpha=0.3); plt.show()

print(f"Final objective  subgradient={loss_sg[-1]:.4f}  ISTA={loss_ista[-1]:.4f}")
print("Notice: ISTA (proximal gradient) descends monotonically, while the subgradient method is not a descent method and is noisier.")

Key takeaways#

• Convex functions without a gradient still have subgradients, defined by the same linear lower bound; the subdifferential is the set of all of them.
• The subgradient method is not a descent method: the objective can increase step to step, which is why it is rarely the tool of choice.
• Composite objectives $g(x) + h(x)$ with smooth $g$ and nonsmooth $h$ are handled by proximal gradient: a gradient step on $g$, then $\mathrm{prox}_{\eta h}$.
• For $h = \lambda\|w\|_1$ the proximal operator is soft thresholding, which sets small coordinates to exactly zero, the source of sparsity in the lasso.
• Subgradients are what make ReLU, hinge, $\ell_1$, and max-based losses trainable, kinks and all.
• Early stopping of gradient descent on unregularized least squares approximates ridge regression: the algorithm itself acts as a regularizer.

Further reading#

• Neal Parikh & Stephen Boyd, Proximal Algorithms (Foundations and Trends in Optimization, 2014). The reference for proximal operators and soft thresholding. Freely available at stanford.edu/~boyd/papers/prox_algs.html.
• Amir Beck & Marc Teboulle, "A Fast Iterative Shrinkage-Thresholding Algorithm for Linear Inverse Problems" (SIAM Journal on Imaging Sciences, 2009). ISTA and its accelerated variant.
• Robert Tibshirani, "Regression Shrinkage and Selection via the Lasso" (Journal of the Royal Statistical Society B, 1996). The $\ell_1$ problem and why it produces sparsity.