Calculus for Machine Learning: From Gradients to Production Optimization

Machine learning at its core is optimization: finding the set of parameters that minimizes a loss function. Calculus provides the mathematical machinery to compute how changes in parameters affect the loss, guiding us toward the minimum. Without calculus, training neural networks would be a blind search.

In 2026, with models growing to billions of parameters, efficient gradient computation is more critical than ever. Matrix calculus—the extension of single-variable calculus to vectors and matrices—is the language of backpropagation, the algorithm that makes deep learning feasible. Understanding it separates practitioners who tune hyperparameters by intuition from those who reason about optimization dynamics.

This article assumes you're comfortable with basic calculus (derivatives, chain rule) and linear algebra (vectors, matrices, matrix multiplication). We'll cover the core concepts of matrix calculus, their application in ML, and common pitfalls in production. You'll learn not just the math, but how to debug gradient issues and avoid costly mistakes.

We'll draw from the canonical reference "Mathematics for Machine Learning" by Deisenroth, Faisal, and Ong, and the Wikipedia article on matrix calculus, but go deeper into production framing. By the end, you'll be able to derive gradients for custom loss functions, interpret Jacobians in sensitivity analysis, and diagnose vanishing/exploding gradients in your models.

Why Calculus Matters in Machine Learning: Optimization as the Core Problem

Machine learning is fundamentally an optimization problem. Every model, from a linear regression to a 100-layer transformer, defines a loss function that measures prediction error. The goal is to find the model parameters that minimize this loss. Calculus provides the machinery to do this efficiently. Without derivatives, you'd be guessing parameter updates blindly, like tuning a radio by randomly twisting knobs. With gradients, you know exactly which direction to turn each knob to reduce static.

Consider training a simple linear model y = wx + b on 10,000 points. The mean squared error loss L = (1/N) Σ(y_i - (wx_i + b))² is a convex function of w and b. The derivative dL/dw tells you the instantaneous slope of the loss with respect to w. If the derivative is positive, increasing w increases loss, so you decrease w. This is gradient descent: w ← w - η * dL/dw. Without calculus, you'd need to evaluate the loss at many w values to approximate this slope, costing O(N) per evaluation. With the analytic derivative, you compute it in O(N) directly, a 100x speedup for typical batch sizes.

The core insight is that all ML training is just repeated application of the chain rule from calculus. The chain rule lets you decompose the derivative of a complex function (the loss through many layers) into a product of simpler derivatives. This is backpropagation. Every epoch of training computes millions of these derivatives automatically. Frameworks like PyTorch and TensorFlow are essentially automatic differentiation engines built on calculus. Understanding the underlying math lets you debug training issues, choose better optimizers, and design new architectures.

Production reality: when your model doesn't converge, it's often because gradients are vanishing, exploding, or incorrect. You can't diagnose gradient problems without understanding what a gradient is. Calculus isn't abstract theory here—it's the difference between a model that trains in 2 hours and one that never converges.

import numpy as np

# Generate synthetic data: y = 2x + 1 + noise
np.random.seed(42)
X = np.random.randn(1000, 1)
y_true = 2 * X + 1 + 0.1 * np.random.randn(1000, 1)

# Initialize parameters
w = np.random.randn()
b = np.random.randn()
lr = 0.1
n_epochs = 100

for epoch in range(n_epochs):
    # Forward pass: predict
    y_pred = X * w + b
    
    # Compute loss (MSE)
    loss = np.mean((y_pred - y_true) ** 2)
    
    # Compute gradients using calculus
    dw = np.mean(2 * X * (y_pred - y_true))
    db = np.mean(2 * (y_pred - y_true))
    
    # Update parameters
    w -= lr * dw
    b -= lr * db
    
    if epoch % 20 == 0:
        print(f"Epoch {epoch}: loss={loss:.6f}, w={w:.4f}, b={b:.4f}")

print(f"Final: w={w:.4f}, b={b:.4f} (true: w=2, b=1)")

From Scalar to Vector Calculus: Gradients, Jacobians, and Hessians

In practice, ML models have millions of parameters, not one. You can't compute a separate derivative for each parameter by hand. Vector calculus generalizes the scalar derivative to handle vectors and matrices as single objects. The gradient ∇f is a vector of partial derivatives: for f: ℝⁿ → ℝ, ∇f = [∂f/∂x₁, ∂f/∂x₂, ..., ∂f/∂xₙ]ᵀ. This is the fundamental object in first-order optimization. When you call loss.backward() in PyTorch, you're computing this gradient vector.

When your function maps vectors to vectors (f: ℝᵐ → ℝⁿ), the derivative is a Jacobian matrix J of shape n×m, where J_{ij} = ∂fᵢ/∂xⱼ. This is critical for understanding how transformations propagate. For example, a linear layer y = Wx + b has Jacobian ∂y/∂x = Wᵀ. The Jacobian of a softmax output with respect to logits is a full matrix, not a vector, because each output depends on all inputs. This matrix structure is what makes backpropagation through softmax nontrivial.

The Hessian H is the matrix of second derivatives for scalar functions: H_{ij} = ∂²f/∂xᵢ∂xⱼ. It describes the curvature of the loss landscape. A positive-definite Hessian means you're near a local minimum; negative-definite means a maximum; indefinite means a saddle point. Second-order methods like Newton's method use the Hessian to take more informed steps, but computing the full Hessian for a 100M-parameter model costs O(n²) memory—impossible. Practical approximations like Adam use diagonal estimates of the Hessian's eigenvalues (via gradient variance).

Understanding these objects prevents bugs. When implementing a custom layer, you need the correct Jacobian shape for backprop. A common mistake is treating the gradient of a vector function as a vector when it should be a matrix. The shape rules are: scalar→vector gives gradient (vector), vector→vector gives Jacobian (matrix), scalar→matrix gives gradient matrix. Get these wrong and your gradients silently corrupt training.

import torch
import torch.nn.functional as F

# Define a simple function: f(x) = [x1^2 + x2, x1 * x2]
def f(x):
    return torch.stack([x[0]**2 + x[1], x[0] * x[1]])

x = torch.tensor([2.0, 3.0], requires_grad=True)
y = f(x)

# Compute Jacobian manually using torch.autograd.functional
from torch.autograd.functional import jacobian
J = jacobian(f, x)
print("Jacobian:\n", J)
# Expected: [[2*x1, 1], [x2, x1]] = [[4, 1], [3, 2]]

# Hessian of a scalar function: g(x) = x1^2 + x2^2 + x1*x2
def g(x):
    return x[0]**2 + x[1]**2 + x[0]*x[1]

from torch.autograd.functional import hessian
H = hessian(g, x)
print("Hessian:\n", H)
# Expected: [[2, 1], [1, 2]]

The Chain Rule in Matrix Form: Backpropagation Foundation

Backpropagation is just the chain rule applied to computational graphs. For a scalar loss L that depends on a series of transformations, the chain rule says dL/dx = (dL/dy) (dy/dx). In matrix form, if y = f(x) and z = g(y), then the Jacobian of the composition is J_{z∘f} = J_g(y) J_f(x). This matrix multiplication is what makes backprop efficient: you compute gradients backward, multiplying Jacobians as you go, reusing intermediate results.

Consider a two-layer network: L = loss(softmax(W₂ ReLU(W₁x + b₁) + b₂)). The gradient dL/dW₁ involves the chain rule through the softmax, the linear layer, the ReLU, and the first linear layer. In matrix form, dL/dW₁ = (dL/dz₂) (dz₂/dh) (dh/dz₁) (dz₁/dW₁), where each term is a Jacobian. The key insight is that these Jacobians are sparse or have special structure: ReLU's Jacobian is a diagonal matrix of 0s and 1s; linear layers have Jacobians equal to the weight matrix.

Automatic differentiation frameworks like PyTorch implement this as a computational graph. Each operation stores a "gradient function" that computes the local Jacobian-vector product (VJP). Instead of building full Jacobian matrices, they compute Jᵀv for a given vector v, which is O(n) instead of O(n²). This is the key to efficiency: backprop never materializes the full Jacobian. For a layer with 4096 inputs and 4096 outputs, the Jacobian would be 16 million elements, but the VJP costs only 4096 operations.

Understanding the matrix chain rule lets you implement custom autograd functions correctly. If you define a custom operation, you must provide the backward function that computes the VJP. Getting the shapes wrong is the #1 bug. The rule: if forward takes input of shape (N, D_in) and produces output (N, D_out), then backward receives grad_output of shape (N, D_out) and must return grad_input of shape (N, D_in). The matrix multiplication is grad_input = grad_output @ J_f(x).T.

import torch

class CustomLinearFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, bias):
        ctx.save_for_backward(x, weight, bias)
        return x @ weight.T + bias
    
    @staticmethod
    def backward(ctx, grad_output):
        x, weight, bias = ctx.saved_tensors
        # Chain rule: dL/dx = dL/dy * dy/dx = grad_output @ weight
        grad_x = grad_output @ weight
        # dL/dweight = dL/dy * dy/dweight = grad_output.T @ x
        grad_weight = grad_output.T @ x
        # dL/dbias = sum of grad_output along batch dimension
        grad_bias = grad_output.sum(dim=0)
        return grad_x, grad_weight, grad_bias

# Test with random data
x = torch.randn(32, 64)
w = torch.randn(128, 64, requires_grad=True)
b = torch.randn(128, requires_grad=True)

# Use custom function
out = CustomLinearFunction.apply(x, w, b)
loss = out.sum()
loss.backward()

print(f"Gradient shapes: dx={x.grad.shape}, dw={w.grad.shape}, db={b.grad.shape}")
print(f"Gradient norms: ||dw||={w.grad.norm():.4f}")

Common Derivatives in ML: Linear Layers, Activation Functions, Loss Functions

Every ML practitioner should know the derivatives of common operations by heart. For a linear layer y = Wx + b, the gradients are: dL/dW = (dL/dy)ᵀ x, dL/dx = Wᵀ (dL/dy), dL/db = Σ dL/dy over batch. These are matrix multiplications with specific transpositions. Getting the transpose wrong is the most common gradient bug. For batch size N, input dim D_in, output dim D_out: dL/dW has shape (D_out, D_in), dL/dx has shape (N, D_in), dL/db has shape (D_out,).

Activation functions have simple element-wise derivatives. ReLU: f'(x) = 1 if x > 0 else 0. This means gradients only flow through positive activations—dead neurons have zero gradient forever. Sigmoid: f'(x) = f(x)(1 - f(x)), max at 0.25. Tanh: f'(x) = 1 - f(x)². These saturating activations cause vanishing gradients when inputs are large in magnitude. GELU, used in transformers, has derivative involving the normal CDF and PDF: f'(x) ≈ 0.5 (1 + erf(x/√2)) + x φ(x).

Loss functions have clean derivatives. MSE loss L = (1/N)Σ(y_pred - y_true)²: dL/dy_pred = (2/N)(y_pred - y_true). Cross-entropy loss with softmax: dL/dz = softmax(z) - y_true (where y_true is one-hot). This is the most elegant gradient in ML: the gradient is simply the difference between prediction and target. For binary classification with sigmoid: dL/dz = σ(z) - y. These simple forms make gradient computation trivial.

In production, knowing these derivatives helps you debug loss spikes. If loss suddenly jumps, check if your softmax logits have extreme values (causing numerical overflow in exp). If gradients vanish, check if you're using sigmoid/tanh in deep networks. If ReLU causes dead neurons, try LeakyReLU or ELU. These aren't theoretical concerns—they're daily debugging tools.

import torch
import torch.nn.functional as F

# Example: compute gradients for a small network manually
batch_size, d_in, d_out = 4, 3, 2
x = torch.randn(batch_size, d_in)
w = torch.randn(d_out, d_in, requires_grad=True)
b = torch.randn(d_out, requires_grad=True)
y_true = torch.randint(0, d_out, (batch_size,))

# Forward pass with manual gradient computation
z = x @ w.T + b  # linear layer
# Use sigmoid activation for binary-like case (simplified)
a = torch.sigmoid(z)
# Cross-entropy loss (manually for demonstration)
loss = F.cross_entropy(z, y_true)  # uses log-softmax internally

# Backward
loss.backward()

# Verify gradient shapes
print(f"dL/dw shape: {w.grad.shape}")  # (d_out, d_in)
print(f"dL/db shape: {b.grad.shape}")  # (d_out,)

# Manual gradient for cross-entropy with softmax:
# dL/dz = softmax(z) - one_hot(y_true)
softmax_out = F.softmax(z, dim=1)
manual_grad_z = softmax_out - F.one_hot(y_true, num_classes=d_out).float()
print(f"Manual dL/dz norm: {manual_grad_z.norm():.4f}")
print(f"AutoDiff dL/dz norm: {z.grad.norm():.4f}")

Layout Conventions: Numerator vs Denominator and Why Consistency Matters

Matrix calculus has two competing layout conventions: numerator layout and denominator layout. In numerator layout, the derivative of a scalar y with respect to a vector x (size n×1) yields a row vector of shape 1×n. In denominator layout, the same derivative yields a column vector of shape n×1. This seemingly minor choice cascades through every chain rule, gradient update, and backpropagation pass. Mixing conventions silently produces transposed or permuted results—your loss might decrease but your weights update in the wrong direction, or your Jacobian product yields a shape error that only surfaces at inference.

Consider a simple linear model y = Wx + b. Under denominator layout, ∂y/∂W has shape (output_dim, input_dim) and the gradient update is W ← W - η (∂L/∂y) x^T. Under numerator layout, the same derivative has shape (input_dim, output_dim) and the update becomes W ← W - η x (∂L/∂y)^T. Both are correct if applied consistently, but swapping them produces a shape mismatch or a rank-1 update that corrupts the weight matrix. PyTorch and TensorFlow use denominator layout (gradients are column vectors), while many textbooks and the original backpropagation papers use numerator layout. Always verify the convention of any paper or library before porting code.

In production systems, the convention affects not just training but also model serialization and inference. If you export gradients from one framework and apply them in another, a transposition bug can silently degrade accuracy by 5-10% over a few epochs. The fix is to write explicit shape assertions at every gradient application point: assert grad_w.shape == w.shape. This catches convention mismatches immediately rather than letting them accumulate. Standardize on one convention across your entire codebase—denominator layout is more common in modern ML frameworks—and document it in your style guide.

When implementing custom autograd functions, you must know which convention your framework uses. PyTorch's backward() receives gradients in the same shape as the input tensors (denominator layout). If you manually compute gradients using a numerator-layout formula, you must transpose before returning. A common mistake is to copy formulas from a paper without transposing, leading to silent training failures. Always test with a small random input and compare against finite differences to catch layout errors early.

import torch
import torch.nn.functional as F

def check_layout_convention():
    """Demonstrate denominator layout in PyTorch."""
    torch.manual_seed(42)
    x = torch.randn(3, requires_grad=True)
    W = torch.randn(2, 3, requires_grad=True)
    b = torch.randn(2, requires_grad=True)
    y = F.linear(x, W, b)
    loss = y.sum()
    loss.backward()
    
    # Under denominator layout, gradient of loss w.r.t. W has shape (2,3)
    assert W.grad.shape == (2, 3), f"Expected (2,3), got {W.grad.shape}"
    # Gradient w.r.t. x has shape (3,)
    assert x.grad.shape == (3,), f"Expected (3,), got {x.grad.shape}"
    print(f"W.grad shape: {W.grad.shape}")
    print(f"x.grad shape: {x.grad.shape}")
    print("Denominator layout confirmed: gradients match input shapes.")

if __name__ == "__main__":
    check_layout_convention()

Gradient Checking: Numerical Verification for Custom Operations

Gradient checking is the gold standard for verifying that your analytical gradients are correct. The idea is simple: approximate the gradient using finite differences and compare it to your analytical gradient. For a scalar function f(x), the central difference approximation is ∂f/∂x_i ≈ (f(x + ε e_i) - f(x - ε e_i)) / (2ε), where ε is a small perturbation (typically 1e-5 to 1e-7). The relative error between analytical and numerical gradients should be on the order of ε^2 for central differences, i.e., 1e-10 to 1e-14. If the error is larger than 1e-5, your gradient is likely wrong.

In practice, you only need to check a subset of parameters—checking every single weight is overkill and slow. Randomly sample 10-20 parameters from each layer and compare their gradients. Use the relative error formula: |analytical - numerical| / (|analytical| + |numerical| + 1e-8). This avoids division by zero when gradients are near zero. A relative error below 1e-7 is excellent, between 1e-7 and 1e-5 is acceptable, and above 1e-5 indicates a bug. For ReLU or other non-differentiable points, skip those inputs or use a subgradient-aware check.

When implementing custom autograd functions in PyTorch, always write a gradient check as part of your unit tests. Use torch.autograd.gradcheck() which handles the details: it perturbs inputs, computes numerical gradients, and compares against the analytical backward. However, be aware that gradcheck() can be slow for large tensors—limit the input size to 10-100 elements. Also, it uses double precision internally, so ensure your custom function works in float64. If your function uses float32, the numerical precision may degrade and you'll need to relax the tolerance.

A common pitfall is gradient checking with stochastic operations like dropout or batch normalization in training mode. These introduce randomness that makes numerical gradients noisy. Always set your model to evaluation mode and disable randomness (e.g., set dropout to 0, freeze batch norm statistics) before gradient checking. Another pitfall: checking gradients after a backward pass that already modified parameters. Always compute gradients on a fresh forward pass with the original parameters. Finally, remember that gradient checking only verifies correctness at a single point—it doesn't guarantee your gradients are correct everywhere, but it catches most implementation bugs.

import torch
import torch.nn as nn

def gradient_check(model, x, y, epsilon=1e-5, threshold=1e-5):
    """Check gradients for a random subset of parameters."""
    model.eval()
    # Forward pass
    loss_fn = nn.MSELoss()
    loss = loss_fn(model(x), y)
    loss.backward()
    
    errors = []
    for name, param in model.named_parameters():
        if param.grad is None:
            continue
        # Randomly sample 5 indices
        flat_grad = param.grad.view(-1)
        flat_param = param.data.view(-1)
        indices = torch.randint(0, flat_param.numel(), (5,))
        for idx in indices:
            original = flat_param[idx].item()
            # f(x + epsilon)
            flat_param[idx] = original + epsilon
            loss_plus = loss_fn(model(x), y)
            # f(x - epsilon)
            flat_param[idx] = original - epsilon
            loss_minus = loss_fn(model(x), y)
            # Numerical gradient
            numerical = (loss_plus - loss_minus) / (2 * epsilon)
            analytical = flat_grad[idx].item()
            # Relative error
            denominator = abs(analytical) + abs(numerical) + 1e-8
            rel_error = abs(analytical - numerical) / denominator
            errors.append(rel_error)
            # Restore parameter
            flat_param[idx] = original
    
    max_error = max(errors) if errors else 0.0
    print(f"Max relative error: {max_error:.2e}")
    if max_error > threshold:
        print("WARNING: Gradient check failed!")
    else:
        print("Gradient check passed.")
    return max_error

if __name__ == "__main__":
    model = nn.Linear(10, 5)
    x = torch.randn(4, 10)
    y = torch.randn(4, 5)
    gradient_check(model, x, y)

Production Pitfalls: Vanishing/Exploding Gradients, Shape Mismatches, and Debugging

Vanishing and exploding gradients are the most common production killers in deep learning. Vanishing gradients occur when gradients become exponentially small as they backpropagate through many layers, causing early layers to stop learning. This is typical with sigmoid/tanh activations in deep networks—the derivative of sigmoid maxes out at 0.25, so after 10 layers the gradient is scaled by 0.25^10 ≈ 9.5e-7. Exploding gradients are the opposite: gradients grow exponentially, causing NaN losses and training divergence. This happens with recurrent networks (especially LSTMs without gradient clipping) or poorly initialized deep networks. The fix is threefold: use ReLU or its variants (Leaky ReLU, ELU) to avoid saturation, apply batch normalization to keep activations in a reasonable range, and use gradient clipping (torch.nn.utils.clip_grad_norm_) to cap the global gradient norm at 1.0 or 5.0.

Shape mismatches are the second most frequent production issue. A common scenario: you reshape a tensor incorrectly, the forward pass works (PyTorch broadcasts silently), but the backward pass produces a shape error because gradients can't be broadcast back. For example, if you squeeze a dimension that shouldn't be squeezed, the gradient for that dimension becomes ambiguous. Always use .view() or .reshape() with explicit dimensions rather than relying on -1 inference. Add shape assertions at every reshape operation: assert x.shape[-1] == hidden_dim. Another shape pitfall is mixing batch-first and sequence-first conventions across layers—standardize on batch-first (batch_size, seq_len, features) throughout your model.

Debugging gradient issues requires systematic tooling. Start by logging gradient statistics (mean, std, min, max) for each layer every N steps. If the mean gradient for a layer is consistently below 1e-7, you have vanishing gradients. If it's above 1e3, you have exploding gradients. Use torch.autograd.set_detect_anomaly(True) to catch NaN gradients early—it slows training by 2x but is invaluable for debugging. When you hit a NaN, inspect the backward hook: register_hook on each tensor to print gradients as they flow. Common causes: division by zero (add epsilon), log of zero (use log_sigmoid or clamp), or in-place operations that break the autograd graph.

In production pipelines, gradient issues often manifest as sudden loss spikes after hours of training. The root cause is usually a data distribution shift that pushes activations into saturation regions. Monitor input statistics (mean, std per feature) and compare against training distribution. If inputs drift, gradients can vanish or explode even with proper initialization. Implement gradient health checks: if the global gradient norm exceeds a threshold (e.g., 100), log a warning and skip the update. This prevents a single bad batch from corrupting your model. Also, use gradient accumulation carefully—accumulating gradients over many micro-batches can amplify numerical errors if not normalized correctly.

import torch
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_

def train_with_gradient_monitoring(model, dataloader, epochs=5):
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    loss_fn = nn.CrossEntropyLoss()
    
    for epoch in range(epochs):
        for batch_idx, (x, y) in enumerate(dataloader):
            optimizer.zero_grad()
            output = model(x)
            loss = loss_fn(output, y)
            loss.backward()
            
            # Monitor gradient norms
            total_norm = 0.0
            for p in model.parameters():
                if p.grad is not None:
                    param_norm = p.grad.data.norm(2)
                    total_norm += param_norm.item() ** 2
            total_norm = total_norm ** 0.5
            
            if total_norm > 100.0:
                print(f"WARNING: Gradient norm {total_norm:.2f} at batch {batch_idx}")
                # Skip update or clip
                clip_grad_norm_(model.parameters(), max_norm=1.0)
            
            optimizer.step()
            
            if batch_idx % 100 == 0:
                print(f"Epoch {epoch}, Batch {batch_idx}, Loss {loss.item():.4f}, Grad norm {total_norm:.4f}")

if __name__ == "__main__":
    model = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 128),
        nn.ReLU(),
        nn.Linear(128, 10)
    )
    dummy_data = [(torch.randn(64, 784), torch.randint(0, 10, (64,)))]
    train_with_gradient_monitoring(model, dummy_data)

Advanced Topics: Second-Order Methods, Automatic Differentiation Internals

Second-order optimization methods use the Hessian matrix (matrix of second partial derivatives) to accelerate convergence. Newton's method updates parameters as θ ← θ - H⁻¹∇L, where H is the Hessian and ∇L is the gradient. For a quadratic loss, Newton's method converges in one step. In practice, computing the full Hessian is O(n²) memory and O(n³) time for n parameters—prohibitive for models with millions of parameters. Quasi-Newton methods like L-BFGS approximate the Hessian using gradient differences over the last m iterations, requiring only O(mn) memory. L-BFGS is popular for small-batch optimization in scientific computing but rarely used in deep learning due to its sensitivity to stochastic gradients.

Automatic differentiation (autodiff) is the engine behind all modern ML frameworks. There are two modes: forward mode and reverse mode. Forward mode computes derivatives alongside the forward pass, one input variable at a time. It's efficient for functions with few inputs and many outputs (e.g., f: R → Rⁿ). Reverse mode (backpropagation) computes gradients by traversing the computation graph backward, accumulating gradients from outputs to inputs. It's efficient for functions with many inputs and few outputs (e.g., f: Rⁿ → R), which is exactly the loss function scenario in ML. Reverse mode requires storing intermediate activations for the backward pass, leading to O(n) memory per layer—this is why training consumes more memory than inference.

PyTorch's autograd system builds a dynamic computation graph at runtime. Each tensor operation creates a Function node that records the operation and its inputs. During backward, it traverses this graph in topological order, applying the chain rule via the grad_fn attribute. The key insight is that autograd uses the vector-Jacobian product (VJP) rather than computing the full Jacobian. For a function y = f(x), the VJP computes v^T J where v is the incoming gradient and J is the Jacobian. This is O(n) instead of O(n²) for the full Jacobian. Custom autograd functions must implement forward() and backward() where backward() receives the gradient of the loss w.r.t. the output and returns the gradient w.r.t. each input.

Second-order information can be approximated without computing the full Hessian. The Gauss-Newton matrix (J^T J) approximates the Hessian for least-squares problems and is positive semi-definite by construction. Kronecker-factored approximate curvature (K-FAC) uses a block-diagonal approximation of the Fisher information matrix, with each block factored as a Kronecker product of two smaller matrices. K-FAC achieves faster convergence than SGD on some architectures but requires careful tuning and is sensitive to batch size. In production, second-order methods are rarely worth the complexity—Adam with learning rate scheduling matches or beats them on most benchmarks with far less engineering overhead.

import torch
import torch.nn as nn
import torch.optim as optim

def hessian_vector_product(model, loss, v):
    """Compute Hv where H is the Hessian of loss w.r.t. parameters."""
    grads = torch.autograd.grad(loss, model.parameters(), create_graph=True)
    flat_grads = torch.cat([g.view(-1) for g in grads])
    grad_v = torch.dot(flat_grads, v)
    hvp = torch.autograd.grad(grad_v, model.parameters(), retain_graph=True)
    return torch.cat([h.view(-1) for h in hvp])

def conjugate_gradient(Av_func, b, n_iter=10):
    """Solve Ax = b using conjugate gradient method."""
    x = torch.zeros_like(b)
    r = b - Av_func(x)
    p = r
    rsold = torch.dot(r, r)
    for _ in range(n_iter):
        Ap = Av_func(p)
        alpha = rsold / torch.dot(p, Ap)
        x = x + alpha * p
        r = r - alpha * Ap
        rsnew = torch.dot(r, r)
        if torch.sqrt(rsnew) < 1e-8:
            break
        p = r + (rsnew / rsold) * p
        rsold = rsnew
    return x

if __name__ == "__main__":
    model = nn.Linear(10, 1)
    x = torch.randn(5, 10)
    y = torch.randn(5, 1)
    loss_fn = nn.MSELoss()
    loss = loss_fn(model(x), y)
    
    # Compute Hv for a random vector v
    v = torch.randn(sum(p.numel() for p in model.parameters()))
    hvp = hessian_vector_product(model, loss, v)
    print(f"Hessian-vector product shape: {hvp.shape}")
    print(f"First 5 elements: {hvp[:5]}")