Optimization & Gradient Descent

An optimization problem asks: given a function $f: \mathbb{R}^n \to \mathbb{R}$ (the objective or loss) and a set $\mathcal{C} \subseteq \mathbb{R}^n$ (the feasible set), find $$x^* = \arg\min_{x \in \mathcal{C}} f(x).$$ When $\mathcal{C} = \mathbb{R}^n$ the problem is unconstrained; otherwise it is constrained.

Critical points are where $\nabla f(x) = 0$. A critical point can be:

• A local minimum: $f(x^) \le f(x)$ for all $x$ near $x^$.
• A local maximum: $f(x^) \ge f(x)$ for all $x$ near $x^$.
• A saddle point: neither — some directions go up, others go down.

The second-order test uses the Hessian $H = \nabla^2 f(x^*)$:

• $H \succ 0$ (positive definite): strict local minimum.
• $H \prec 0$ (negative definite): strict local maximum.
• $H$ has mixed eigenvalues: saddle point.

Analogy: imagine hiking in a mountainous landscape. A valley floor is a local minimum; a mountain peak is a local maximum; a mountain pass (high in two directions, low in the other two) is a saddle. Gradient descent will roll you into some valley, but not necessarily the deepest one — the global minimum. The central challenge of non-convex optimization is that the landscape may have exponentially many valleys.

In high-dimensional machine learning, the landscape has billions of dimensions. Empirically, most critical points are saddle points (mixed Hessian eigenvalues), not bad local minima — a surprising insight that makes gradient descent work better in practice than theory predicts.

import numpy as np

# f(x,y) = x^2 - y^2 (saddle at origin)
def f(x, y): return x**2 - y**2
def grad_f(x, y): return np.array([2*x, -2*y])
def hessian_f(): return np.array([[2, 0], [0, -2]])

H = hessian_f()
eigenvalues = np.linalg.eigvalsh(H)
print("Eigenvalues:", eigenvalues)  # [-2.  2.] => saddle point

Definition. A set $\mathcal{C}$ is convex if for all $x, y \in \mathcal{C}$ and $\lambda \in [0,1]$, $$\lambda x + (1-\lambda)y \in \mathcal{C}.$$ Every point on the line segment between any two members is also in the set. Think of a bowl versus a donut: a bowl is convex; a donut is not (the segment from one side to the other passes through the hole).

Definition. A function $f: \mathcal{C} \to \mathbb{R}$ is convex if for all $x, y \in \mathcal{C}$ and $\lambda \in [0,1]$, $$f(\lambda x + (1-\lambda)y) \le \lambda f(x) + (1-\lambda)f(y).$$ The chord between any two points on the graph lies above the graph. Intuitively: the function curves upward everywhere.

First-order condition. If $f$ is differentiable, it is convex if and only if $$f(y) \ge f(x) + \nabla f(x)^\top (y-x) \quad \forall x,y.$$ The tangent hyperplane is a global underestimate. This is extraordinary: knowing $f$ and $\nabla f$ at any single point gives a lower bound on $f$ everywhere.

Second-order condition. If $f$ is twice differentiable, it is convex if and only if $\nabla^2 f(x) \succeq 0$ everywhere (the Hessian is positive semidefinite).

The key theorem: For a convex function on a convex set, every local minimum is a global minimum. This collapses the whole search problem: you cannot get stuck in a bad valley because there is only one valley.

Strict vs. strong convexity. $f$ is strictly convex if the inequality above is strict for $x \ne y$ — then the minimizer is unique. $f$ is $\mu$-strongly convex ($\mu > 0$) if $$f(y) \ge f(x) + \nabla f(x)^\top(y-x) + \frac{\mu}{2}|y-x|^2.$$ Strong convexity means $f$ grows at least as fast as a quadratic — it rules out flat regions and gives fast convergence rates.

Examples:

• $f(x) = x^2$: convex (Hessian = 2 > 0).
• $f(x) = |x|$: convex but not differentiable at 0.
• $f(x) = \log(x)$: concave (Hessian = $-1/x^2 < 0$).
• $f(x,y) = x^2 - y^2$: not convex (Hessian has eigenvalue $-2$).
• Linear regression MSE: convex in the weights.
• Neural network loss: not convex — yet gradient descent often works!

import numpy as np

def is_convex_numerically(f, domain, n_samples=1000):
    """Monte Carlo check of convexity on a 1D domain."""
    xs = np.random.uniform(*domain, n_samples)
    ys = np.random.uniform(*domain, n_samples)
    lam = np.random.uniform(0, 1, n_samples)
    lhs = f(lam * xs + (1 - lam) * ys)
    rhs = lam * f(xs) + (1 - lam) * f(ys)
    return np.all(lhs <= rhs + 1e-9)

print(is_convex_numerically(lambda x: x**2, (-5, 5)))   # True
print(is_convex_numerically(lambda x: -x**2, (-5, 5)))  # False

The gradient $\nabla f(x) \in \mathbb{R}^n$ is the vector of partial derivatives: $$\nabla f(x) = \left(\frac{\partial f}{\partial x_1}, \ldots, \frac{\partial f}{\partial x_n}\right).$$ The directional derivative of $f$ at $x$ in direction $v$ (unit vector) is $$D_v f(x) = \nabla f(x)^\top v.$$ By the Cauchy-Schwarz inequality, $D_v f(x) \le |\nabla f(x)|$ with equality when $v = \nabla f(x) / |\nabla f(x)|$. The gradient points in the direction of steepest ascent. The direction of steepest descent is $-\nabla f(x)$.

Analogy: $\nabla f(x)$ is the slope vector on a topographic map. It tells you which way is uphill, and how steep. Gradient descent says: always walk in the direction that goes downhill fastest.

The gradient is perpendicular to level sets. A level set is the set ${x : f(x) = c}$ — a contour line on the map. The gradient at a point is exactly perpendicular to the contour through that point. This is why gradient descent zig-zags: it always goes perpendicular to the current contour, which is not necessarily straight toward the minimum.

Computing gradients in practice is done by automatic differentiation (as in PyTorch/JAX): chain-rule applications tracked symbolically through a computation graph. Understanding the gradient analytically is essential for designing and debugging those systems.

import numpy as np

def f(x):    # f(x1, x2) = x1^2 + 4*x2^2
    return x[0]**2 + 4*x[1]**2

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

# Verify numerically via finite differences
def numerical_grad(f, x, eps=1e-5):
    g = np.zeros_like(x, dtype=float)
    for i in range(len(x)):
        e = np.zeros_like(x, dtype=float)
        e[i] = eps
        g[i] = (f(x + e) - f(x - e)) / (2 * eps)
    return g

x = np.array([3.0, 1.0])
print("Analytic:", grad_f(x))    # [6. 8.]
print("Numeric: ", numerical_grad(f, x))  # [6. 8.]

The algorithm is beautifully simple. Starting from an initial point $x_0$, repeat: $$x_{t+1} = x_t - \alpha \nabla f(x_t)$$ where $\alpha > 0$ is the learning rate (or step size). We walk a small step in the direction of steepest descent.

Why does this work? By Taylor expansion, $$f(x_t - \alpha \nabla f(x_t)) \approx f(x_t) - \alpha |\nabla f(x_t)|^2.$$ For small enough $\alpha$, the function value decreases at each step (as long as $\nabla f \ne 0$). The approximation is exact when $\alpha \to 0$; the art is choosing $\alpha$ large enough to make real progress but small enough for the approximation to hold.

The step size is critical:

• Too large: the update overshoots the minimum, and the loss can increase or oscillate wildly (divergence).
• Too small: convergence is very slow (many small steps to reach the valley).
• Just right: fast, steady descent.

Convergence rates. Let $f$ be $L$-smooth (i.e., $|\nabla f(x) - \nabla f(y)| \le L|x-y|$) and convex. With step size $\alpha = 1/L$:

• Convex: $f(x_T) - f^* \le \frac{L|x_0 - x^*|^2}{2T}$ — this is $O(1/T)$ convergence.
• $\mu$-strongly convex: $|x_T - x^|^2 \le (1 - \mu/L)^T |x_0 - x^|^2$ — this is linear (geometric) convergence, exponentially fast.

The ratio $\kappa = L/\mu$ is called the condition number. Ill-conditioned problems ($\kappa \gg 1$) converge very slowly and cause the zig-zagging behavior seen on elongated ellipsoidal loss surfaces.

Stopping criteria: (1) $|\nabla f(x_t)| < \epsilon$, (2) $|f(x_t) - f(x_{t-1})| < \epsilon$, or (3) a fixed number of iterations.

Line search adaptively chooses $\alpha$ at each step by finding the minimizer along the descent direction. The Armijo (sufficient decrease) and Wolfe conditions provide practical rules that guarantee convergence without a fixed schedule.

import numpy as np

def gradient_descent(f, grad_f, x0, alpha=0.1, tol=1e-6, max_iter=1000):
    x = np.array(x0, dtype=float)
    history = [f(x)]
    for t in range(max_iter):
        g = grad_f(x)
        if np.linalg.norm(g) < tol:
            break
        x = x - alpha * g
        history.append(f(x))
    return x, history

# Minimize f(x1,x2) = x1^2 + 4*x2^2  (minimum at origin)
f     = lambda x: x[0]**2 + 4*x[1]**2
grad  = lambda x: np.array([2*x[0], 8*x[1]])

x_star, hist = gradient_descent(f, grad, x0=[3.0, 2.0], alpha=0.1)
print(f"Minimizer: {x_star}")   # ~ [0, 0]
print(f"Steps: {len(hist)}")    # converges quickly

In machine learning, the objective is almost always a sum over training examples: $$f(w) = \frac{1}{n} \sum_{i=1}^{n} \ell(w; x_i, y_i).$$ Computing the full gradient costs $O(n)$ per step, which is prohibitive when $n$ is millions. Stochastic Gradient Descent (SGD) replaces the true gradient with a noisy estimate using a single randomly chosen example: $$w_{t+1} = w_t - \alpha_t \nabla_w \ell(w_t; x_{i_t}, y_{i_t}).$$ This is an unbiased estimator: $\mathbb{E}[\nabla \ell(w; x_i, y_i)] = \nabla f(w)$. Each step is $O(1)$ instead of $O(n)$, making it $n$ times cheaper per update.

Mini-batch SGD averages over a small batch $\mathcal{B}$ of $b$ examples: $$w_{t+1} = w_t - \frac{\alpha_t}{|\mathcal{B}|} \sum_{i \in \mathcal{B}} \nabla \ell(w_t; x_i, y_i).$$ Typically $b \in {32, 64, 128, 256}$. Mini-batches reduce gradient variance, leverage GPU parallelism, and give smoother convergence than pure SGD.

The noise is a feature, not just a bug. SGD's noise lets it escape sharp local minima and saddle points that full-batch gradient descent would get stuck in. Flat minima (low curvature at $x^*$) tend to generalize better, and SGD's noise biases it toward finding them — a key empirical reason why SGD often outperforms full-batch methods in deep learning.

Learning rate schedules. For SGD to converge, the learning rate must decrease over time. The classical Robbins-Monro conditions require: $$\sum_t \alpha_t = \infty \quad \text{and} \quad \sum_t \alpha_t^2 < \infty.$$ Common schedules: step decay (halve every $k$ epochs), exponential decay, cosine annealing.

Convergence rate. For convex $f$, mini-batch SGD achieves $O(1/\sqrt{T})$ convergence — slower than full-batch's $O(1/T)$, but each of the $T$ steps costs $O(b)$ instead of $O(n)$, so the total work to reach $\epsilon$ accuracy is typically much less.

import numpy as np

def sgd(X, y, w0, alpha=0.01, epochs=50, batch_size=32):
    """Mini-batch SGD for linear regression: min sum (w^T x_i - y_i)^2."""
    w = w0.copy().astype(float)
    n = len(y)
    for epoch in range(epochs):
        idx = np.random.permutation(n)
        X_s, y_s = X[idx], y[idx]
        for start in range(0, n, batch_size):
            Xb = X_s[start:start+batch_size]
            yb = y_s[start:start+batch_size]
            grad = 2 * Xb.T @ (Xb @ w - yb) / len(yb)
            w -= alpha * grad
    return w

np.random.seed(0)
n, d = 1000, 5
X = np.random.randn(n, d)
true_w = np.ones(d)
y = X @ true_w + 0.1 * np.random.randn(n)

w_hat = sgd(X, y, np.zeros(d))
print("Recovered w:", np.round(w_hat, 2))  # close to [1,1,1,1,1]

Plain gradient descent zig-zags on ill-conditioned surfaces because it only uses the current gradient. Momentum methods accumulate a velocity across steps, smoothing the trajectory.

Momentum (Polyak, 1964): $$v_{t+1} = \beta v_t + \nabla f(x_t), \qquad x_{t+1} = x_t - \alpha v_{t+1}.$$ With $\beta \approx 0.9$, the velocity is an exponential moving average of past gradients. In directions where gradients consistently point the same way, the velocity builds up (a ball rolling downhill accelerates); in oscillating directions, the gradients cancel and slow the update. Momentum reduces zig-zagging and can cut convergence from $O(\kappa)$ to $O(\sqrt{\kappa})$ iterations.

Nesterov Accelerated Gradient (NAG): Evaluate the gradient at the lookahead point $x_t - \beta v_t$ rather than $x_t$. This gives optimal $O(1/T^2)$ convergence for smooth convex functions — the best possible first-order rate.

Adaptive methods — RMSProp and Adam. Different coordinates of the gradient can have vastly different magnitudes. Adaptive methods maintain per-coordinate learning rates:

Adam (Kingma & Ba, 2015) keeps running estimates of the first moment $m_t$ (mean) and second moment $v_t$ (uncentered variance) of the gradient: $$m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t, \quad v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$$ $$\hat m_t = m_t/(1-\beta_1^t), \quad \hat v_t = v_t/(1-\beta_2^t)$$ $$x_{t+1} = x_t - \frac{\alpha}{\sqrt{\hat v_t} + \epsilon} \hat m_t$$ With defaults $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\epsilon = 10^{-8}$. The denominator $\sqrt{\hat v_t}$ normalizes each coordinate by its recent gradient magnitude — directions with large, consistent gradients get small effective steps; sparse or small-gradient directions get large steps. This makes Adam extremely robust to learning rate choice and dominant in modern deep learning.

Caveat: Adam does not always converge to the global optimum on convex problems (AMSGrad was proposed to fix this). For final fine-tuning of large models, SGD with momentum sometimes outperforms Adam by converging to flatter (better-generalizing) minima.

import numpy as np

def adam(grad_f, x0, alpha=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, steps=500):
    x = np.array(x0, dtype=float)
    m = np.zeros_like(x)
    v = np.zeros_like(x)
    for t in range(1, steps + 1):
        g = grad_f(x)
        m = beta1 * m + (1 - beta1) * g
        v = beta2 * v + (1 - beta2) * g**2
        m_hat = m / (1 - beta1**t)
        v_hat = v / (1 - beta2**t)
        x -= alpha * m_hat / (np.sqrt(v_hat) + eps)
    return x

# Rosenbrock: notoriously hard for plain GD
f    = lambda x: (1 - x[0])**2 + 100*(x[1] - x[0]**2)**2
grad = lambda x: np.array([
    -2*(1-x[0]) - 400*x[0]*(x[1]-x[0]**2),
     200*(x[1]-x[0]**2)
])

result = adam(grad, x0=[-1.0, 1.0])
print(f"Minimizer: {result}")  # close to [1, 1]

Now suppose the feasible set is defined by equality constraints: minimize $f(x)$ subject to $g_i(x) = 0$ for $i = 1, \ldots, m$. We cannot just ignore the constraints and minimize $f$ freely.

The geometric insight: At a constrained minimum $x^$, the gradient $\nabla f(x^)$ must be perpendicular to the constraint surface — otherwise we could move along the surface and decrease $f$. The constraint surface has tangent plane spanned by directions $v$ with $\nabla g_i(x^)^\top v = 0$. The condition is: $$\nabla f(x^) = \sum_i \lambda_i \nabla g_i(x^*).$$ $\nabla f$ must lie in the span of the constraint gradients.

The Lagrangian. Define $$\mathcal{L}(x, \lambda) = f(x) - \sum_i \lambda_i g_i(x).$$ The scalars $\lambda_i$ are Lagrange multipliers. The necessary conditions for an optimum are: $$\nabla_x \mathcal{L} = 0 \implies \nabla f(x) = \sum_i \lambda_i \nabla g_i(x)$$ $$\nabla_{\lambda_i} \mathcal{L} = 0 \implies g_i(x) = 0 \quad \text{(constraints satisfied)}.$$ So a constrained problem with $n$ variables and $m$ constraints is converted to an unconstrained system of $n + m$ equations in $n + m$ unknowns $(x, \lambda)$.

Example: Maximize $f(x,y) = xy$ subject to $x + y = 10$ (split a fixed sum to maximize product). $$\mathcal{L} = xy - \lambda(x + y - 10).$$ $\partial \mathcal{L}/\partial x = y - \lambda = 0$, $\partial \mathcal{L}/\partial y = x - \lambda = 0$, $\partial \mathcal{L}/\partial \lambda = -(x+y-10) = 0$. Solving: $x = y = 5$, $\lambda = 5$. Maximum product is 25. This is the AM-GM inequality in action.

Interpretation of $\lambda$. The multiplier $\lambda_i$ measures the sensitivity of the optimal value to the constraint: $df^*/db_i$ where $g_i(x) = b_i$. In economics, $\lambda$ is the "shadow price" of the constraint.

Linear constraints yield the classic method of Lagrange. More general equality-constrained optimization uses this as the core, solved numerically by Newton's method on the KKT system.

from scipy.optimize import minimize
import numpy as np

# Maximize xy subject to x + y = 10
# Equivalently: minimize -xy subject to x + y - 10 = 0
f      = lambda x: -x[0] * x[1]
constr = {'type': 'eq', 'fun': lambda x: x[0] + x[1] - 10}

result = minimize(f, x0=[3.0, 7.0], constraints=constr)
print(f"x={result.x[0]:.4f}, y={result.x[1]:.4f}")  # 5.0000, 5.0000
print(f"max xy = {-result.fun:.4f}")                  # 25.0000

Inequality constraints — minimize $f(x)$ subject to $h_j(x) \le 0$ for $j = 1, \ldots, p$ — require a generalization: the Karush-Kuhn-Tucker (KKT) conditions.

The Lagrangian for inequality constraints: $$\mathcal{L}(x, \mu) = f(x) + \sum_j \mu_j h_j(x), \quad \mu_j \ge 0.$$ Note the sign: for inequality constraints the multipliers must be non-negative (intuitively, they can only push the solution against the constraint, not away from it).

KKT conditions (necessary for optimality under constraint qualification):

1. Stationarity: $\nabla_x \mathcal{L}(x^, \mu^) = 0$
2. Primal feasibility: $h_j(x^*) \le 0$ for all $j$
3. Dual feasibility: $\mu_j^* \ge 0$ for all $j$
4. Complementary slackness: $\mu_j^* h_j(x^*) = 0$ for all $j$

Condition 4 is the key new ingredient. It says: either the constraint is active ($h_j = 0$) or the multiplier is zero — you cannot have both $\mu_j > 0$ and the constraint inactive. This captures the intuition that constraints only matter when they are "tight".

For convex $f$ and convex constraints: the KKT conditions are sufficient as well as necessary. This is the crown jewel of convex optimization: a set of equations that fully characterizes the global optimum.

Example: Minimize $x^2 + y^2$ (distance to origin) subject to $x + y \ge 1$ (i.e., $1 - x - y \le 0$). KKT gives $\mu(1-x-y) = 0$ and $\nabla(x^2+y^2) = \mu \nabla(-(x+y-1))$. The solution is $x^* = y^* = 1/2$, $\mu^* = 1 > 0$ (constraint is active). Without the constraint the minimum is at the origin, but the constraint pushes us to the nearest point on the line $x+y=1$.

Support vector machines are entirely built on KKT conditions — the support vectors are exactly the training points where the constraints are active (complementary slackness).

# KKT conditions solved by hand for: min x^2+y^2 s.t. x+y >= 1
# Lagrangian: L = x^2 + y^2 + mu*(1 - x - y),  mu >= 0
# Stationarity: 2x - mu = 0,  2y - mu = 0  =>  x = y = mu/2
# Complementary slackness: mu*(1 - x - y) = 0
# Try mu > 0: then 1 - x - y = 0 => 1 - mu/2 - mu/2 = 0 => mu = 1
# Solution: x = y = 0.5, mu = 1 (constraint active)
print("x* = 0.5, y* = 0.5, mu* = 1")
print("Distance to origin:", (0.5**2 + 0.5**2)**0.5)  # 0.707

Every constrained optimization problem (the primal) has a companion problem called the dual. Duality theory is one of the most beautiful and useful ideas in all of optimization.

The dual function. Given the Lagrangian $\mathcal{L}(x, \lambda, \mu) = f(x) + \lambda^\top g(x) + \mu^\top h(x)$, define: $$d(\lambda, \mu) = \inf_x \mathcal{L}(x, \lambda, \mu).$$ This is the dual function. It is always concave (as an infimum of linear functions of $\lambda, \mu$).

Weak duality: For any dual-feasible $(\lambda, \mu)$ (i.e., $\mu \ge 0$), $d(\lambda, \mu) \le f^$. Every dual point gives a lower bound on the primal optimum. Maximizing the dual gives the tightest such lower bound: $$\max_{\lambda, \mu \ge 0} d(\lambda, \mu) \le f^.$$ The gap $f^* - d^*$ is the duality gap.

Strong duality: Under mild conditions (Slater's condition: there exists a strictly feasible $x$), the gap is zero: $d^* = f^*$. Primal and dual have the same optimal value. This means we can solve the primal by solving the dual (which is always convex!).

Practical importance:

• SVMs: The primal SVM is a QP in $n$ variables (number of features). Its dual is a QP in $m$ variables (number of training points), allowing the kernel trick.
• Linear programming: The dual of an LP is another LP; LP duality is the foundation of the simplex method and interior-point methods.
• Duality gap as a certificate: A small duality gap certifies nearoptimality even when you cannot compute the true optimum.

Analogy: the primal is building a fence to enclose land; the dual is finding the most efficient way to price the land so the fence value equals the land value. At optimality, they agree.

# Dual of: min c^T x s.t. Ax = b, x >= 0  (LP primal)
# Dual:    max b^T y s.t. A^T y <= c  (LP dual)
# Strong duality: same optimal value

from scipy.optimize import linprog
import numpy as np

c = np.array([1.0, 2.0])         # minimize x1 + 2*x2
A = np.array([[1.0, 1.0]])       # x1 + x2 = 3
b = np.array([3.0])

# Primal
res_p = linprog(c, A_eq=A, b_eq=b, bounds=[(0,None),(0,None)])
print(f"Primal optimum: {res_p.fun:.4f}")  # 3.0  (x1=3, x2=0)

# Dual: max 3*y  s.t.  y<=1, y<=2  =>  y*=1, obj=3.0
print("Dual optimum: 3.0 (y*=1)")  # strong duality confirmed

Gradient descent uses a linear (first-order) Taylor approximation to choose the next step. Newton's method uses a quadratic (second-order) approximation, incorporating curvature information from the Hessian.

Derivation. Approximate $f$ near $x_t$ as: $$f(x) \approx f(x_t) + \nabla f(x_t)^\top (x - x_t) + \frac{1}{2}(x-x_t)^\top H_t (x-x_t),$$ where $H_t = \nabla^2 f(x_t)$. Minimize this quadratic over $x$: $$x_{t+1} = x_t - H_t^{-1} \nabla f(x_t).$$ This is the Newton step. Note there is no learning rate — the curvature automatically scales the step.

Convergence. Newton's method has quadratic convergence near the optimum: $$|x_{t+1} - x^| \le C |x_t - x^|^2.$$ If you are within distance $10^{-3}$, the next step gets you to $10^{-6}$, then $10^{-12}$. Extremely fast in theory.

Cost: Each step requires computing and inverting the $n \times n$ Hessian at cost $O(n^3)$ — completely infeasible for neural networks with millions of parameters.

Quasi-Newton methods approximate $H^{-1}$ cheaply. BFGS builds a low-rank approximation updated at each step using gradient differences. L-BFGS (limited-memory BFGS) stores only the last $k$ gradient-difference vectors (usually $k \in [5, 20]$), giving $O(kn)$ cost per step. L-BFGS is the workhorse for smooth optimization in moderate dimensions and is the default in scipy.optimize.minimize.

Natural gradient uses the Fisher information matrix as the curvature metric — the second-order method for models over distributions. It is invariant to reparameterization and underpins natural policy gradient methods in reinforcement learning.

import numpy as np

def newton_method(f, grad_f, hess_f, x0, tol=1e-8, max_iter=50):
    x = np.array(x0, dtype=float)
    for i in range(max_iter):
        g = grad_f(x)
        if np.linalg.norm(g) < tol:
            break
        H = hess_f(x)
        # Newton step: solve H * d = -g
        d = np.linalg.solve(H, -g)
        x = x + d   # full step (no line search for illustration)
    return x, i

# f(x) = x^4 - 4x^2 + 3  (two local minima: x = +-sqrt(2))
f      = lambda x: x**4 - 4*x**2 + 3
grad_f = lambda x: np.array([4*x[0]**3 - 8*x[0]])
hess_f = lambda x: np.array([[12*x[0]**2 - 8]])

x_star, iters = newton_method(f, grad_f, hess_f, x0=[1.5])
print(f"Converged to x={x_star[0]:.6f} in {iters} iterations")  # 1.414213...