Calculus for ML
Gradients, chain rule, Hessian, convexity
You are training logistic regression with binary cross-entropy loss: L = -[y log(ŷ) + (1-y) log(1-ŷ)] where ŷ = σ(Wx + b). You want to update W to reduce L. How much does L change when you nudge W₁₁ — the weight from feature 1 to the output — by a tiny ε? This is the partial derivative ∂L/∂W₁₁. Computing all such partial derivatives for all weights gives the gradient ∇L — the direction of steepest ascent in loss space. You step in the opposite direction. This is gradient descent.
The chain rule makes gradient computation tractable for deep networks. For the composition L(ŷ(z(W))), the chain rule gives ∂L/∂W = (∂L/∂ŷ) × (∂ŷ/∂z) × (∂z/∂W) where z = Wx + b. Each arrow in the computation graph corresponds to a derivative. The product of derivatives along a path is the derivative of the composed function. For the sigmoid: ∂σ(z)/∂z = σ(z)(1 - σ(z)) — a closed form that depends only on the output value.
Key derivatives in ML: for MSE loss L = (y - Wx)²/2, the gradient is -(y - Wx)x = -residual × feature. For cross-entropy with sigmoid output, the gradient is (ŷ - y)x — the residual times the feature. These clean gradient formulas are not accidents. They are chosen precisely because the derivative of the cross-entropy loss through the sigmoid produces a numerically stable, interpretable update.
Taylor expansion: f(x + ε) ≈ f(x) + f'(x)ε + f''(x)ε²/2. Gradient descent uses the first-order approximation — it assumes the loss surface is locally linear at each step. Second-order methods (Newton's method) use the quadratic term via the Hessian. They take more accurate steps but require O(n²) memory to store the Hessian — infeasible for large models.
NOT this. Calculus in ML is not just gradient descent. Calculus is why certain loss and activation combinations work together and others do not. Using MSE loss with a sigmoid output produces gradient saturation — in the saturated region where σ(z) ≈ 0 or 1, the derivative σ(z)(1-σ(z)) ≈ 0, and the gradient of the MSE loss through this near-zero value nearly vanishes. The network cannot learn from these examples. Cross-entropy with sigmoid was chosen specifically because the saturating term cancels algebraically, leaving a non-saturating gradient. Every loss-activation combination is a calculus decision.
Key points
- When designing a new loss function or output layer, compute the gradient analytically first. If the gradient saturates — approaches 0 — in regions of the output space, the model cannot learn from examples that fall there. This is why ReLU replaced sigmoid in hidden layers: ReLU's gradient is 1 for all positive inputs, never saturates. Sigmoid's gradient σ(z)(1-σ(z)) peaks at 0.25 and goes to 0 for large |z|.
- Trap: using finite-difference gradient checking as a routine verification step on large models. Computing (f(x+ε) - f(x-ε)) / 2ε requires 2 forward passes per parameter. For a network with 10 million parameters, that is 20 million forward passes. Use gradient checking only for debugging specific custom layers, not the full network — use automatic differentiation for everything else.
- Diagnostic: if a layer's gradient magnitude is near 0 during backpropagation, the gradient is vanishing. Check the activation function in that layer — if it is sigmoid or tanh, the layer inputs may be in the saturated region. Plot the distribution of pre-activation values for that layer. If they are large in magnitude (|z| > 4), replace the activation with ReLU or use batch normalization to keep activations in the non-saturated range.
The choice of loss function and activation is a calculus decision, not an aesthetic one. The clean gradient (ŷ - y)x that makes logistic regression easy to train exists because cross-entropy and sigmoid were chosen together precisely to cancel the saturation term.
Recap
- Gradient ∇L = direction of steepest ascent; gradient descent steps the opposite way.
- Chain rule makes deep nets tractable: ∂L/∂W = (∂L/∂ŷ)(∂ŷ/∂z)(∂z/∂W) — one derivative per graph edge.
- Sigmoid derivative σ'(z)=σ(z)(1−σ(z)), peaks at 0.25, → 0 for large |z| (saturation).
- Clean gradients are chosen, not accidental: cross-entropy + sigmoid gives (ŷ−y)x by cancelling the saturation term.
- MSE + sigmoid saturates — near σ≈0 or 1 the gradient vanishes and the net stops learning. ReLU's gradient is 1, never saturates.
- Taylor: GD uses the 1st-order (linear) approximation; Newton uses the Hessian but costs O(n²) memory.
- Vanishing gradient near 0 in a layer → check the activation; if sigmoid/tanh, inputs are saturated (|z|>4).
Check your understanding
Q1. What is the gradient of f(x) = ‖Ax − b‖₂² with respect to x? Derive it step by step.
- A) ∇_x f = 2Ax − 2b. Derive: f(x) = (Ax−b)ᵀ(Ax−b). The chain rule gives ∂f/∂x = 2(Ax−b)·A, but since we differentiate with respect to x (not Ax), the A term does not transpose: ∇_x f = 2A(Ax−b) expanded term-by-term gives 2Ax − 2b. Setting to zero: Ax = b, which is only valid when A is square — the gradient directly gives the solution without the normal equations.
- B) ∇_x f = 2A(Ax−b). Derive: f(x) = (Ax−b)ᵀ(Ax−b). Let r = Ax−b, so f = rᵀr. ∂f/∂x = 2r·(∂r/∂x) = 2(Ax−b)·A = 2A(Ax−b), treating the Jacobian ∂r/∂x as A itself rather than requiring the transpose Aᵀ. Setting to zero gives AAx = Ab — a distinct linear system from the true normal equations, valid only when A happens to be symmetric and invertible.
- C) f(x) = (Ax−b)ᵀ(Ax−b) = xᵀAᵀAx − 2bᵀAx + bᵀb. Taking the gradient term by term: ∇_x(xᵀAᵀAx) = 2AᵀAx (using ∇_x(xᵀMx) = (M+Mᵀ)x with M=AᵀA symmetric), ∇_x(−2bᵀAx) = −2Aᵀb, ∇_x(bᵀb) = 0. Total: ∇_x f = 2AᵀAx − 2Aᵀb = 2Aᵀ(Ax−b). Setting to zero gives the normal equations AᵀAx = Aᵀb — the update direction is Aᵀ times the residual.
- D) ∇_x f = (Ax−b). Derive: f(x) = ‖r‖² where r = Ax−b. The gradient of ‖r‖² with respect to r is 2r, and by the chain rule ∂r/∂x = I, treating Ax−b as depending directly on x with a unit Jacobian rather than through the linear map A. So ∇_x f = 2(Ax−b) = 2Ax − 2b, skipping the transpose correction that the true chain rule through A actually requires.
Q2. The chain rule for ∂L/∂W at layer l requires the upstream gradient ∂L/∂z_{l+1}. Which two of the following correctly explain why backpropagation computes gradients with a backward pass rather than a forward pass?
- A) Gradients propagate backward because the loss is computed at the output layer (forward end), not the input layer. The gradient information must physically travel from where the loss is measured back to where the parameters are. Forward propagation moves data forward; backward propagation moves gradient information the other way, and both directions take exactly the same total compute time regardless of network depth.
- B) Backward propagation is required by the chain rule: ∂L/∂Wₗ = (∂L/∂zₗ)·(∂zₗ/∂Wₗ), and ∂L/∂zₗ can only be computed after ∂L/∂zₙ is known. Forward accumulation would compute ∂zₙ/∂Wₗ, but this requires a separate forward pass per parameter and does not benefit from the scalar loss structure. Backprop processes all layers simultaneously in one backward pass due to dynamic programming — it reuses intermediate activations stored during the forward pass.
- C) The chain rule for f(g(x)): df/dx=(df/dg)(dg/dx). For L=L(z_n(...z_1(x))), ∂L/∂W_l is a product of Jacobians, computable right-to-left or left-to-right. Forward accumulation needs one pass per input dimension; backward accumulation (from ∂L/∂z_n) needs one pass per output. Since L is scalar, backward mode gets every parameter's gradient in one pass — forward mode needs one pass per parameter, often millions.
- D) Gradients flow backward because neurons are connected by directed edges from input to output, and gradient flow must respect edge directionality. In the computational graph, edges point forward (input→output), so gradient information naturally flows in the reverse direction along the same edges — if the network had bidirectional connections, gradients could instead flow both ways simultaneously.
Q3. You want to find the minimum of f(x₁,x₂) = (x₁−3)² + 2(x₂+1)². What are ∂f/∂x₁ and ∂f/∂x₂, and what is the minimum?
- A) ∂f/∂x₁ = 2(x₁−3). ∂f/∂x₂ = 4(x₂+1). Setting both to zero: x₁=3, x₂=−1. Minimum value f(3,−1) = 0. The Hessian H = [[2,0],[0,4]] is positive definite, confirming a global minimum. Gradient descent must use learning rate < 1/4 (half the inverse of the largest eigenvalue). The x₂ dimension converges faster because it has higher curvature (4 vs. 2), reaching the optimum in fewer steps.
- B) ∂f/∂x₁ = 2(x₁−3). ∂f/∂x₂ = 4(x₂+1). Setting both to zero: x₁=3, x₂=−1. The Hessian H = [[2,0],[0,4]] is positive definite (eigenvalues 2, 4), confirming a strict local minimum that is also global since f is convex. Minimum value f(3,−1) = 0. Gradient descent needs learning rate < 1/4; x₂ (curvature 4) takes smaller effective steps than x₁ (curvature 2) — condition number 2.
- C) ∂f/∂x₁ = (x₁−3). ∂f/∂x₂ = 2(x₂+1). Setting both to zero: x₁=3, x₂=−1. Minimum value f(3,−1) = 0. The Hessian H = [[1,0],[0,2]] — positive definite, having dropped the constant multipliers 2 and 4 from the true second derivatives. The condition number is 2/1 = 2, so gradient descent converges with at most 2× more steps in the x₁ direction than the x₂ direction under this curvature estimate.
- D) ∂f/∂x₁ = 2x₁−6. ∂f/∂x₂ = 4x₂+4. Setting to zero: x₁=3, x₂=−1. Minimum f(3,−1) = 9+2 = 11, since expanding (x₁−3)² = x₁²−6x₁+9 and 2(x₂+1)² = 2x₂²+4x₂+2 gives f(3,−1) = 0+0+9+2 = 11 by summing only the expanded constant terms, ignoring that the linear and quadratic terms in x₁ and x₂ also vanish exactly at this same stationary point.
Try it interactively
ML Systems Lab is a free interview-prep platform for ML engineers — work through the full interactive module, quizzes, and drills.
Open ML Systems Lab →