Gradient-based optimization

If information about the gradient is available, it can accelerate convergence substantially by pointing in the direction of steepest change.

import numpy as np
from scipy.linalg import solve
import plotly.graph_objects as go

try:
    import executable_engineering as exe
except ImportError:
    %pip install -q executable_engineering
    import executable_engineering as exe

from executable_engineering import plot_optimization_path_3d

First order gradient information: Gradient descent methods¶

If you have information about the gradient, moving down the negative gradient is an intuitive approach to reach the minimum. While straightforward, steepest descent can suffer from severe zig-zagging in narrow valleys, which motivates constructing conjugate directions.

Each step is incremented as:

\mathbf{x}^{i+1} = \mathbf{x}^i + - a \nabla f(\mathbf{x}^i)

(1)

Steepest descent maximizes the decrease in $f$ locally, but a modified direction may allow longer steps. It is fast to execute and very parallelizable on GPUs.

Stochastic gradient descent¶

In machine learning training, optimization involves finding model parameters in very high-dimensional spaces where computing the true gradient over millions of data points is prohibitively expensive.

In Stochastic Gradient Descent (SGD), the gradient is approximated at each step by randomly sampling a mini-batch of data. The step length $a$ is termed the learning rate, and momentum terms are added to dampen oscillations across ravine-like loss surfaces.

Second order gradient information: Newton’s optimization method¶

Newton’s optimization routine aims to find the root of the gradient, $\nabla f(\mathbf{x}) = \mathbf{0}$ , which identifies the stationary points (minima, maxima, or saddle points). Since we are focused on a scalar objective $f(\mathbf{x})$ , the gradient is a vector $\nabla f$ , and we require the second-order derivative matrix—the Hessian $\mathbf{H}$ :

\mathbf{H} = \frac{\partial^2 f}{\partial \mathbf{x}^2}

(2)

The increment $\Delta \mathbf{x}$ is solved at each iteration from the linear system:

\mathbf{H} \Delta \mathbf{x} = -\nabla f(\mathbf{x})

(3)

where $\mathbf{H}$ must be symmetric. Newton’s optimization method has quadratic convergence near a minimum, but requires assembling and solving a linear system with the Hessian matrix at every step.

def newton_method(f, grad_f, hessian_f, x0, tol=1e-6, max_iter=100):
    x = np.asarray(x0, dtype=float)
    guesses = [x.copy()]
    for _ in range(max_iter):
        grad = grad_f(x)
        hess = hessian_f(x)
        if np.linalg.norm(grad) < tol:
            break
        delta_x = solve(hess, -grad, assume_a='sym')
        x = x + delta_x
        guesses.append(x.copy())
    return x, guesses

Example: Minimize $x^4 + 2y^4$ ¶

Let’s minimize $f(x, y) = x^4 + 2y^4$ starting from an initial guess of $\mathbf{x}_0 = (1.5, 2.0)$ :

def f(x):
    return x[0]**4 + 2 * x[1]**4

def grad_f(x):
    return np.array([4 * x[0]**3, 8 * x[1]**3])

def hessian_f(x):
    return np.array([[12 * x[0]**2, 0], [0, 24 * x[1]**2]])

x0 = np.array([1.5, 2.0])
x_sol, guesses = newton_method(f, grad_f, hessian_f, x0)
print("Converged to:", np.round(x_sol, 5), "in", len(guesses) - 1, "steps")

plot_optimization_path_3d(f, guesses, x_range=(-2, 2), y_range=(-2, 2), title="Newton's Method on f(x, y) = x^4 + 2y^4")

Example: Minimize $x^2 + y^2$ ¶

Now let’s apply Newton’s method to $f(x, y) = x^2 + y^2$ starting from $\mathbf{x}_0 = (1.5, 1.5)$ :

Why does this converge in a single step? Because $f(x, y)$ is an exact quadratic; the second-order Taylor expansion upon which Newton’s method is built is exact everywhere!

def f_quad(x):
    return x[0]**2 + x[1]**2

def grad_quad(x):
    return np.array([2 * x[0], 2 * x[1]])

def hessian_quad(x):
    return np.array([[2, 0], [0, 2]])

x0 = np.array([1.5, 1.5])
x_sol, guesses = newton_method(f_quad, grad_quad, hessian_quad, x0)
print("Converged to:", x_sol, "in", len(guesses) - 1, "step!")

plot_optimization_path_3d(f_quad, guesses, x_range=(-2, 2), y_range=(-2, 2), title="Newton's Method on Quadratic f(x, y) = x^2 + y^2")

Example: Saddle point in $x^2 - 6xy + y^2$ ¶

Let’s test Newton’s method on $f(x, y) = x^2 - 6xy + y^2$ with $\mathbf{x}_0 = (1.5, 2.0)$ :

Notice what happens: Newton’s step jumps to $(0, 0)$ because $\nabla f(0, 0) = \mathbf{0}$ . But is $(0, 0)$ a minimum? Along $y = 0$ , $f(x, 0) = x^2$ (curves up), but along $y = x$ , $f(x, x) = -4x^2$ (curves down!). Point $(0, 0)$ is a saddle point, and the Hessian is indefinite.

def f_saddle(x):
    return x[0]**2 + x[1]**2 - 6 * x[0] * x[1]

def grad_saddle(x):
    return np.array([2 * x[0] - 6 * x[1], 2 * x[1] - 6 * x[0]])

def hessian_saddle(x):
    return np.array([[2, -6], [-6, 2]])

x0 = np.array([1.5, 2.0])
x_sol, guesses = newton_method(f_saddle, grad_saddle, hessian_saddle, x0)
print("Stationary point found at:", x_sol)

plot_optimization_path_3d(f_saddle, guesses, x_range=(-3, 3), y_range=(-3, 3), title="Newton's Method on Indefinite System f(x, y) = x^2 - 6xy + y^2")

First order gradient information: Gradient descent methods¶

Stochastic gradient descent¶

Second order gradient information: Newton’s optimization method¶

Example: Minimize x4+2y4x^4 + 2y^4x4+2y4¶

Example: Minimize x2+y2x^2 + y^2x2+y2¶

Example: Saddle point in x2−6xy+y2x^2 - 6xy + y^2x2−6xy+y2¶

Example: Minimize $x^4 + 2y^4$ ¶

Example: Minimize $x^2 + y^2$ ¶

Example: Saddle point in $x^2 - 6xy + y^2$ ¶