Calculus & Optimization for AI

Linear algebra tells you how data flows through a model. Calculus tells you how to change the model so it gets better. Every training run, at its core, is an optimization loop built entirely on the ideas in this page.

Intuition: Slope, Curvature, and the Chain Rule

Three ideas, layered: the gradient tells you which way is uphill from where you're standing (first derivative — slope). The Hessian tells you whether the ground curves like a bowl, a dome, or a saddle beneath you (second derivative — curvature) — the entire reason a critical point (flat gradient) isn't automatically a minimum. And the chain rule is the mechanical trick that lets you compute a gradient through an arbitrarily deep composition of functions — which is all a neural network is — one local derivative at a time. Everything else on this page (Newton's method, Adam, KKT conditions) is a specific, practical consequence of one of these three ideas.

Derivatives and Gradients

A derivative measures how much a function's output changes when you nudge its input: $\dfrac{df}{dx}$. In ML, our functions take many inputs (millions of weights), so we use the gradient — a vector of partial derivatives, one per parameter: $\nabla f = \left[\dfrac{\partial f}{\partial w_1}, \dfrac{\partial f}{\partial w_2}, \dots\right]$.

The key idea: the gradient points in the direction of steepest increase of the loss function. So to reduce the loss, you step in the opposite direction of the gradient. Drag the point below and watch both directions recompute live, on a real bowl:

That single sentence — step opposite the gradient — is the entire idea behind training every neural network that exists.

Jacobians & Hessians: Derivatives of Vector-Valued and Multivariate Functions

A plain gradient is the right tool when a function has many inputs but one scalar output (a loss function). Two extensions handle the cases that come up just as often in practice.

The Jacobian — when a function has many inputs and many outputs, e.g. $\mathbf{y} = f(\mathbf{x})$ with $\mathbf{x} \in \mathbb{R}^n$, $\mathbf{y} \in \mathbb{R}^m$. The Jacobian $J$ is the $m \times n$ matrix of every output's partial derivative w.r.t. every input: $J_{ij} = \dfrac{\partial y_i}{\partial x_j}$. For a linear layer $\mathbf{y} = W\mathbf{x}$, the Jacobian is simply $W$ itself — which is exactly why backprop through a layer with multiple outputs (basically every layer) multiplies the incoming gradient by that layer's Jacobian, a direct generalization of the single-number chain rule below to the vector case.

The Hessian — the matrix of second partial derivatives of a scalar function: $H_{ij} = \dfrac{\partial^2 f}{\partial x_i \partial x_j}$. Where the gradient tells you the slope, the Hessian tells you the curvature. Toggle between the three real critical-point types below and watch the real eigenvalue classification match the shading:

If the Hessian is positive definite at a critical point (all eigenvalues positive — see Linear Algebra), that point is a local minimum. If it's negative definite, a local maximum. If eigenvalues have mixed signs, it's a saddle point — flat in the gradient ($\nabla f = 0$) but not actually a minimum, which is why high-dimensional non-convex optimization (deep learning) is dominated by escaping saddle points, not just avoiding local minima.

Newton's method uses the Hessian directly: instead of a fixed-size gradient step, it steps to $\theta - H^{-1}\nabla f$ — jumping straight to the minimum of the local quadratic approximation. Real trajectories, same surface, same starting point:

This converges dramatically faster than gradient descent near a minimum, but computing and inverting an $n \times n$ Hessian for a model with billions of parameters is completely infeasible — which is precisely why deep learning uses first-order methods (SGD, Adam) and only approximates second-order information cheaply, if at all.

Taylor expansion ties gradient and Hessian together: near a point $\theta_0$, $f(\theta) \approx f(\theta_0) + \nabla f(\theta_0)^T(\theta - \theta_0) + \frac{1}{2}(\theta-\theta_0)^T H (\theta-\theta_0)$ — a linear term (the gradient) plus a quadratic correction (the Hessian). This is the justification for why loss surfaces look locally like the elongated bowl in the optimizer-race chart below — near any point, a smooth loss function genuinely is approximately a quadratic form, which is why understanding quadratic optimization (above) transfers so directly to understanding real training dynamics.

The Chain Rule → Backpropagation

If $y = f(g(x))$, then $\dfrac{dy}{dx} = \dfrac{dy}{dg} \cdot \dfrac{dg}{dx}$. A neural network is a chain of functions (layer 1 → layer 2 → ... → loss). Step through a real forward pass, then a real backward pass, on a genuine 2-layer network:

Backpropagation is nothing more than applying the chain rule repeatedly, from the loss backward to each weight, to compute how much each individual weight contributed to the error. This is why it's called backward propagation — you compute the output forward, then walk the chain rule backward, exactly as shown above.

Convexity

A function is convex if a line segment between any two points on its graph never dips below the graph itself — it has a single global minimum, no false valleys. Toggle between a convex and a genuinely non-convex function below, and watch where real gradient descent actually lands depending on where it starts:

Linear regression's loss (MSE) is convex, which is why it can be solved exactly. Neural network loss surfaces are non-convex — full of local minima and saddle points, exactly like the function above. This is why deep learning relies on iterative, gradient-based search rather than a closed-form solution, and why techniques like good initialization, momentum, and learning rate schedules matter so much: they're all ways of navigating a bumpy, high-dimensional surface without getting stuck.

Gradient Descent and Its Variants

Batch gradient descent: compute the gradient using the entire dataset, then take one step. Accurate but painfully slow for large datasets. Stochastic Gradient Descent (SGD): compute the gradient using a single example, then step — fast and noisy, and the noise actually helps escape shallow local minima. Mini-batch SGD: the practical middle ground — compute the gradient over a small batch (e.g. 32-512 examples), what "batch size" refers to in every training run you'll configure.

Update rule: $\theta \leftarrow \theta - \eta \nabla L(\theta)$, where $\eta$ is the learning rate — the single most important hyperparameter to get right. Too large and training diverges; too small and training crawls.

Adaptive Optimizers

Plain SGD treats every parameter the same. Modern optimizers adapt the step size per parameter — the classic demonstration is a "ravine" loss surface, steep in one direction and shallow in the other, exactly the shape that breaks plain SGD:

• Momentum: accumulates a running average of past gradients, so the optimizer keeps moving in a consistent direction and dampens oscillation — like a ball rolling downhill picking up speed.
• AdaGrad: scales down the learning rate for parameters that get frequent large updates — good for sparse features, but its accumulated sum only grows, eventually shrinking the learning rate to near zero.
• RMSProp: fixes AdaGrad's decay problem by using a moving average of squared gradients instead of a running sum.
• Adam: combines momentum (first moment) and RMSProp-style adaptive scaling (second moment). The default choice for training almost every modern neural network, including LLMs.

Learning Rate Schedules & Warmup

Real formulas, evaluated at every step, not a diagram of the concept:

• Warmup: start with a tiny learning rate and ramp up over the first several hundred/thousand steps. Prevents early, noisy gradients (before the model has learned anything sensible) from causing instability.
• Decay schedules (cosine, linear, step): reduce the learning rate over training so the model can settle into a sharper minimum near the end, rather than bouncing around it.
• Gradient clipping: cap the gradient's magnitude before applying it, to prevent a single bad batch from blowing up the weights — essential when training large Transformers.

Constrained Optimization: Lagrange Multipliers & KKT Conditions

Equality constraints: used when optimizing a function subject to a constraint, e.g. "minimize $f(x)$ such that $g(x) = 0$." You form the Lagrangian $\mathcal{L}(x, \lambda) = f(x) - \lambda g(x)$ and optimize jointly over $x$ and $\lambda$ — at the solution, $\nabla f$ and $\nabla g$ point in the same (or opposite) direction, which is exactly what setting $\nabla_x \mathcal{L} = 0$ enforces.

Inequality constraints — the KKT conditions: real constrained problems are usually inequalities, not equalities — "minimize $f(x)$ such that $g(x) \leq 0$." Toggle between a constraint that binds and one that doesn't, and watch the real solution (and real $\lambda$) respond:

This is the actual formulation Support Vector Machines need (maximize the margin subject to every point being correctly classified, $y_i(w^Tx_i + b) \geq 1$ — an inequality, not an equality), so plain Lagrange multipliers aren't quite enough on their own. The Karush-Kuhn-Tucker (KKT) conditions generalize the Lagrangian approach to inequalities, requiring at the optimum:

1. Stationarity: $\nabla_x \mathcal{L} = 0$, same as the equality case.
2. Primal feasibility: the original constraint $g(x) \leq 0$ actually holds.
3. Dual feasibility: the multiplier $\lambda \geq 0$ (unlike equality constraints, the sign matters here — it enforces that you're pushing against the constraint boundary from the right side).
4. Complementary slackness: $\lambda \cdot g(x) = 0$ — either the constraint is exactly tight ($g(x) = 0$, an "active" constraint) or its multiplier is zero (the constraint wasn't binding at all), exactly the two states the diagram above toggles between.

That last condition is why SVM has support vectors in the first place: complementary slackness means only the points sitting exactly on the margin boundary get a nonzero multiplier — every other point's constraint is slack (not binding), so it contributes nothing to the final decision boundary. The KKT framework is the same machinery behind constrained RL formulations (e.g. constraining a policy update's KL divergence from the previous policy, as in PPO).

Code: Newton's Method, For Real

The exact update rule the Newton-vs-gradient-descent diagram above races against plain gradient descent:

import numpy as np

def newtons_method(grad_fn, hessian_fn, x0: np.ndarray, steps: int = 10) -> np.ndarray:
    x = x0.copy()
    for _ in range(steps):
        g = grad_fn(x)
        H = hessian_fn(x)
        x = x - np.linalg.solve(H, g)   # H^-1 @ g, solved directly rather than inverting H
    return x

Where this shows up in the rest of the curriculum

Next: Probability & Statistics — the math of uncertainty, which underlies loss functions, evaluation, and how LLMs generate text token by token.