LEARN · MATHEMATICS FOR MACHINE LEARNING
The small idea hiding inside enormous models
A neural network with a million trainable parameters sounds as though it should require a million separate calculus problems.
It does not.
The mathematical engine underneath backpropagation is still the chain rule you learned in introductory calculus. What changes at machine-learning scale is not the rule itself, but how systematically we organize and reuse intermediate derivatives.
That distinction is the key to understanding automatic differentiation.
A modern autodiff system does not stare at a giant symbolic formula and differentiate it from scratch. It records a computation as a graph of small operations—multiply, add, sine, matrix multiply, activation, normalization—and associates each operation with a local derivative rule. Then it propagates derivative information through that graph.
For the overwhelmingly common machine-learning situation where many parameters produce one scalar loss, reverse-mode automatic differentiation is especially efficient. A single reverse sweep can compute the derivative of that scalar loss with respect to every parameter participating in the graph. PyTorch’s autograd system is reverse-mode automatic differentiation, and its current Tensor.backward() API explicitly differentiates a tensor with respect to graph leaves using the chain rule.
That is how calc-101 becomes backpropagation.
Set up a current PyTorch environment
The examples below use the current torch.func and autograd APIs rather than the older standalone functorch interface. PyTorch’s documentation describes torch.func as the integrated home for composable transforms such as grad, vjp, jvp, jacrev, and jacfwd.
As of August 2026, PyPI lists PyTorch 2.13.0 as the latest release.
Create a virtual environment:
python -m venv .venv
Activate it on Linux or macOS:
source .venv/bin/activate
On Windows PowerShell:
.venv\Scripts\Activate.ps1
Install PyTorch:
python -m pip install --upgrade pip python -m pip install torch==2.13.0
Check the installation:
python -c "import torch; print(torch.__version__)"
The examples run on CPU, so you do not need a GPU to follow the derivations.
Start with the ordinary chain rule
Suppose one variable depends on another:
u = g(x) f = h(u)
Then the derivative of f with respect to x is:
df/dx = df/du · du/dx
The important word is local.
To differentiate the whole composition, we only need:
-
how
hchanges with respect to its immediate inputu; -
how
gchanges with respect to its immediate inputx; -
a multiplication connecting those local derivatives.
Now make the graph longer:
x → a → b → c → loss
The derivative becomes:
dloss/dx = dloss/dc · dc/db · db/da · da/dx
That multiplication of local sensitivities is the conceptual core of backpropagation.
Neural networks are larger computational graphs, but they do not require a new calculus rule.
A tiny function we can differentiate completely by hand
Consider this function of two variables:
u = x · y v = sin(u) w = v + x² f = w²
Equivalently:
f(x, y) = (sin(x · y) + x²)²
We want two gradients:
∂f/∂x ∂f/∂y
Instead of attacking the final expression directly, keep the intermediate variables. That is exactly the decomposition an autodiff engine benefits from.
Step 1: differentiate the final square
Because:
f = w²
we have:
∂f/∂w = 2w
Step 2: differentiate the addition
Because:
w = v + x²
the local derivatives are:
∂w/∂v = 1 ∂w/∂x = 2x
Be careful: w depends on x through two routes.
There is a direct route:
x → x² → w
and an indirect route:
x → u → v → w
Whenever a variable influences the output through multiple paths, the total derivative is the sum of the contributions from those paths.
This addition of gradient contributions becomes extremely important in residual networks, parameter sharing, recurrent structures, attention blocks, and any computational graph that branches and later rejoins.
Step 3: differentiate sine
Because:
v = sin(u)
we have:
∂v/∂u = cos(u)
Step 4: differentiate multiplication
Because:
u = x · y
we have:
∂u/∂x = y ∂u/∂y = x
Now we can assemble everything.
Hand-derive ∂f/∂x
There are two ways x affects w:
w = sin(x · y) + x²
So:
∂w/∂x = cos(x · y) · y + 2x
Then propagate through the final square:
∂f/∂x = ∂f/∂w · ∂w/∂x
Therefore:
∂f/∂x = 2w · (cos(x · y) · y + 2x)
Substitute w:
∂f/∂x = 2(sin(x · y) + x²) · (cos(x · y) · y + 2x)
Hand-derive ∂f/∂y
The variable y affects the result only through:
y → u → v → w → f
Following that path:
∂f/∂y = ∂f/∂w · ∂w/∂v · ∂v/∂u · ∂u/∂y
Substitute the local derivatives:
∂f/∂y = 2w · 1 · cos(u) · x
Therefore:
∂f/∂y = 2(sin(x · y) + x²) · cos(x · y) · x
That is the full chain rule calculation.
No machine learning yet. Just calculus.
Evaluate the derivatives numerically
Take:
x = 2 y = 3
The intermediate values are:
u = 6 v = sin(6) ≈ -0.2794154982 w ≈ 3.7205845018 f ≈ 13.8427490350
The gradients are approximately:
∂f/∂x ≈ 51.1990441400 ∂f/∂y ≈ 14.2895787504
We can calculate those values directly in ordinary Python:
import math x = 2.0 y = 3.0 u = x * y v = math.sin(u) w = v + x**2 f = w**2 df_dx = 2 * w * (math.cos(u) * y + 2 * x) df_dy = 2 * w * math.cos(u) * x print(f"f = {f:.12f}") print(f"df/dx = {df_dx:.12f}") print(f"df/dy = {df_dy:.12f}")
Expected output:
f = 13.842749035042 df/dx = 51.199044140016 df/dy = 14.289578750405
Now we have values against which an autodiff engine can be tested.
Let PyTorch build the same computation graph
Create x and y as tensors whose gradients should be tracked:
import torch x = torch.tensor( 2.0, dtype=torch.float64, requires_grad=True, ) y = torch.tensor( 3.0, dtype=torch.float64, requires_grad=True, ) u = x * y v = torch.sin(u) w = v + x**2 f = w**2 f.backward() print(f"f = {f.item():.12f}") print(f"df/dx = {x.grad.item():.12f}") print(f"df/dy = {y.grad.item():.12f}")
You should get the same numbers:
f = 13.842749035042 df/dx = 51.199044140016 df/dy = 14.289578750405
PyTorch records operations involving tensors that require gradients and constructs the information needed for the reverse pass. For leaf tensors with requires_grad=True, backward() accumulates the resulting gradients in their .grad fields.
Notice what we did not tell PyTorch.
We never supplied:
d(x · y)/dx = y
or:
d sin(u)/du = cos(u)
or:
d(w²)/dw = 2w
Those derivative rules belong to the operations themselves.
Autograd’s job is to compose them.
Think of backpropagation as passing messages backward
During the forward calculation, values flow left to right:
x ─┐ ├─ multiply → u → sin → v ─┐ y ─┘ ├─ add → w → square → f x ───────────────→ square ─────┘
During reverse mode, sensitivities flow right to left.
Start from:
∂f/∂f = 1
Then the square node receives that incoming sensitivity and sends:
1 · 2w
back toward w.
The addition node sends the incoming gradient down both branches.
The sine node multiplies its incoming gradient by:
cos(u)
The multiplication node sends different gradient messages to each input:
toward x: incoming · y toward y: incoming · x
At x, gradient contributions from multiple routes are added.
That is backpropagation in a compact form:
-
Receive an upstream gradient.
-
Multiply it by the operation’s local derivative.
-
Send the resulting gradients to the operation’s inputs.
-
Add contributions when multiple paths converge.
Why the backward pass starts with 1
We quietly started the reverse process with:
∂f/∂f = 1
That is why calling:
f.backward()
works naturally when f is scalar.
For a scalar output, PyTorch can implicitly use a scalar seed of 1. Its documentation notes that a non-scalar tensor instead requires a gradient argument describing the gradient of the quantity being differentiated with respect to that tensor.
For example:
import torch x = torch.tensor( [1.0, 2.0, 3.0], requires_grad=True, ) y = x**2 seed = torch.tensor( [1.0, 0.5, -1.0], ) y.backward(seed) print(x.grad)
The seed says, in effect, “differentiate this weighted combination of the outputs.”
The result is:
tensor([ 2., 2., -6.])
because the local derivative of x² is 2x, multiplied elementwise by the supplied upstream vector.
This is already pointing toward the deeper abstraction behind reverse-mode autodiff: the vector-Jacobian product.
Jacobians are the bridge from scalar calculus to tensors
For a function with vector input and vector output:
y = f(x)
the derivative is not one number.
It is the Jacobian matrix:
J[i, j] = ∂y[i]/∂x[j]
A naive strategy would build this entire matrix.
That is often unnecessary.
Suppose the input contains one million parameters and the output is one scalar loss:
f: R¹⁰⁰⁰⁰⁰⁰ → R
The Jacobian has shape:
1 × 1,000,000
That single row is the gradient.
Reverse mode is beautifully matched to this shape.
Rather than asking separately:
What is ∂loss/∂w₁? What is ∂loss/∂w₂? What is ∂loss/∂w₃? ...
it propagates one scalar sensitivity backward through the graph and collects all parameter sensitivities along the way.
The million-parameter experiment
Let’s make this literal.
Create one million trainable weights and define a scalar objective:
import torch torch.manual_seed(42) n_parameters = 1_000_000 weights = torch.randn( n_parameters, dtype=torch.float32, requires_grad=True, ) inputs = torch.randn( n_parameters, dtype=torch.float32, ) targets = torch.randn( n_parameters, dtype=torch.float32, ) predictions = weights * inputs loss = ((predictions - targets) ** 2).mean() loss.backward() print(f"loss: {loss.item():.6f}") print(f"gradient shape: {weights.grad.shape}") print(f"number of gradients: {weights.grad.numel():,}") print(f"first five gradients: {weights.grad[:5]}")
The important output is:
gradient shape: torch.Size([1000000]) number of gradients: 1,000,000
One call:
loss.backward()
produces one million gradient components.
This does not mean computing one million derivatives is free.
The reverse pass must still traverse the relevant computation and perform substantial work. A million parameters also means storing or otherwise processing a million gradient values.
The important scaling property is different: reverse mode does not require one independent full graph traversal for every input parameter.
For a scalar objective with a huge input dimension, that distinction is decisive.
PyTorch’s current Jacobian documentation describes the same geometry: reverse mode works row-by-row through a Jacobian, while forward mode works column-by-column. As a rule of thumb, reverse mode is attractive when outputs are fewer than inputs; forward mode becomes attractive when outputs greatly outnumber inputs.
Training neural networks usually has exactly the reverse-mode-friendly shape:
millions or billions of parameters → one scalar loss
Make reverse mode explicit with a VJP
backward() is convenient, but we can expose the underlying vector-Jacobian-product idea more directly with torch.func.vjp.
PyTorch’s current torch.func.vjp returns both the function result and a callable that computes reverse-mode vector-Jacobian products for supplied cotangents.
Try:
import torch from torch.func import vjp torch.manual_seed(42) weights = torch.randn( 1_000_000, dtype=torch.float32, ) inputs = torch.randn_like(weights) targets = torch.randn_like(weights) def loss_fn(w): predictions = w * inputs return ((predictions - targets) ** 2).mean() loss, vjp_fn = vjp(loss_fn, weights) gradient_seed = torch.ones_like(loss) (gradient,) = vjp_fn(gradient_seed) print(f"loss: {loss.item():.6f}") print(f"gradient shape: {gradient.shape}") print(f"number of gradients: {gradient.numel():,}")
Because loss is scalar, its Jacobian with respect to weights has one row.
Feeding the cotangent:
1
through the VJP returns that entire row.
That row is the one-million-component gradient.
Forward mode answers a different question
Forward-mode autodiff propagates a tangent in the same direction as the forward computation.
If:
y = f(x)
and v is a direction in input space, forward mode efficiently computes:
J · v
This is the Jacobian-vector product, or JVP.
PyTorch exposes this directly through the current torch.func.jvp API.
Here is a small example:
import torch from torch.func import jvp def f(x): return torch.stack( ( x[0] * x[1], torch.sin(x[0]) + x[1] ** 2, ) ) x = torch.tensor( [2.0, 3.0], dtype=torch.float64, ) direction = torch.tensor( [1.0, 0.0], dtype=torch.float64, ) value, directional_derivative = jvp( f, (x,), (direction,), ) print("f(x):") print(value) print("J @ direction:") print(directional_derivative)
Because the direction is:
[1, 0]
the JVP gives the Jacobian’s response in the first-coordinate direction.
Conceptually, it gives one Jacobian column’s worth of information.
Change the direction:
direction = torch.tensor( [0.0, 1.0], dtype=torch.float64, )
and you probe another direction.
Why forward mode would struggle with the million-weight gradient
Return to:
f: R¹⁰⁰⁰⁰⁰⁰ → R
Forward mode can efficiently tell you the directional derivative for one chosen direction:
J · v
But if you specifically want every partial derivative independently, you can imagine choosing the coordinate basis vectors:
e₁ = [1, 0, 0, ...] e₂ = [0, 1, 0, ...] e₃ = [0, 0, 1, ...] ...
There are one million such directions.
That makes a naive full-gradient calculation via forward mode require one million directional probes.
Reverse mode turns the geometry around.
The output is scalar, so there is only one output basis direction to seed:
1
That single reverse sweep propagates information back to every input.
This is one of the most important facts in deep learning:
Reverse-mode autodiff is not magical because it makes derivatives constant-time. It is powerful because its cost structure matches functions with many inputs and few outputs.
PyTorch’s current guidance captures the same rule: jacrev uses reverse mode and jacfwd uses forward mode; forward mode generally wins for very tall Jacobians with more outputs than inputs, while reverse mode is favored in the opposite regime.
Backpropagation is not gradient descent
These terms are often blurred together, but they describe different jobs.
Backpropagation computes derivatives.
It answers:
Given this loss, what is ∂loss/∂parameter for every parameter?
Gradient descent uses those derivatives to modify parameters.
It performs something like:
parameter_new = parameter_old - learning_rate · gradient
You can run backpropagation without taking an optimization step.
You can also use the calculated gradients in algorithms other than ordinary gradient descent.
Keeping these concepts separate makes framework code much easier to reason about.
A complete miniature optimization step
Suppose our million parameters represent coefficients in a deliberately simple synthetic fitting problem.
We can calculate gradients and update the weights manually:
import torch torch.manual_seed(42) n = 1_000_000 learning_rate = 0.1 weights = torch.randn( n, requires_grad=True, ) inputs = torch.randn(n) targets = 2.5 * inputs predictions = weights * inputs loss = ((predictions - targets) ** 2).mean() loss.backward() with torch.no_grad(): weights -= learning_rate * weights.grad weights.grad = None print(f"loss before update: {loss.item():.6f}")
The torch.no_grad() block matters because the parameter update itself should not become another differentiable part of the training graph.
And:
weights.grad = None
matters because PyTorch accumulates gradients into leaf tensors rather than automatically replacing them. The current documentation explicitly warns that repeated backward calls accumulate into .grad, and recommends clearing or setting gradients to None when appropriate.
Gradient accumulation: one of the first autograd surprises
Consider:
import torch x = torch.tensor( 3.0, requires_grad=True, ) loss = x**2 loss.backward() print(x.grad) loss = x**2 loss.backward() print(x.grad)
The first result is:
tensor(6.)
The second is:
tensor(12.)
Why?
Because the second gradient was added to the first.
Reset it before another independent calculation:
x.grad = None
Then:
loss = x**2 loss.backward() print(x.grad)
returns:
tensor(6.)
Gradient accumulation is intentional. It enables techniques such as accumulating gradients over multiple microbatches.
It also means that forgetting to clear gradients can silently alter training.
Leaf tensors, intermediate tensors, and .grad
Another common surprise is that not every intermediate value retains a .grad field after backpropagation.
For example:
import torch x = torch.tensor( 2.0, requires_grad=True, ) y = x**2 z = y**3 z.backward() print(x.grad) print(y.grad)
x is a leaf tensor requiring gradients, so its .grad is populated.
y is an intermediate, non-leaf tensor. PyTorch needs gradient information for it during backward, but does not normally retain that gradient in y.grad afterward. Current PyTorch documentation explicitly distinguishes gradient computation through intermediate nodes from .grad storage on leaf tensors.
If you genuinely need the intermediate gradient for debugging or analysis, request retention before backward:
import torch x = torch.tensor( 2.0, requires_grad=True, ) y = x**2 y.retain_grad() z = y**3 z.backward() print(f"x.grad = {x.grad}") print(f"y.grad = {y.grad}")
Do not retain every intermediate gradient casually in a large network. Retaining more data means retaining more memory.
Numerical gradient checking is still useful
Automatic differentiation is exact up to floating-point effects for the operations it differentiates, but your model implementation can still be wrong.
A finite-difference approximation provides an independent sanity check.
For one variable:
df/dx ≈ [f(x + ε) - f(x - ε)] / (2ε)
That is not how autograd computes gradients. It is a separate numerical approximation, which makes it useful as a check.
For our original function:
import math def f(x, y): return ( math.sin(x * y) + x**2 ) ** 2 x = 2.0 y = 3.0 epsilon = 1e-6 numeric_dx = ( f(x + epsilon, y) - f(x - epsilon, y) ) / (2 * epsilon) numeric_dy = ( f(x, y + epsilon) - f(x, y - epsilon) ) / (2 * epsilon) print(f"numeric df/dx: {numeric_dx:.12f}") print(f"numeric df/dy: {numeric_dy:.12f}")
You should see values close to:
numeric df/dx: 51.199044... numeric df/dy: 14.289579...
For custom differentiable operations, PyTorch also provides dedicated gradient-checking machinery, but the simple finite-difference idea is worth understanding because it gives you a framework-independent diagnostic.
The graph can be dynamic
A useful property of PyTorch autograd is that the graph reflects the tensor operations that actually happened during the forward computation.
That means ordinary Python control flow can determine which differentiable operations execute.
For example:
import torch def conditional_loss(x): if x.item() > 0: return x**2 return x**3 x = torch.tensor( 2.0, requires_grad=True, ) loss = conditional_loss(x) loss.backward() print(x.grad)
For x = 2, the executed branch is:
loss = x²
so the gradient is:
4
The backward computation follows the graph created by the actual forward operations.
This ability to work with dynamic computation is part of why autograd feels closer to ordinary Python programming than manual symbolic differentiation.
The deeper abstraction: local linear maps
There is another way to think about derivatives that scales better than memorizing formulas.
Near a point, a differentiable function behaves approximately like a linear transformation.
For:
y = f(x)
the Jacobian describes that local linear transformation.
Forward mode asks:
If I perturb the input in direction v, how does that perturbation move forward?
That is:
J · v
Reverse mode asks:
Given sensitivity at the output, how does that sensitivity pull backward onto the inputs?
That is a vector-Jacobian product.
This viewpoint explains why multiplication by local derivatives is the right operation at every graph node.
Each node receives a small linear sensitivity map from its neighbor and composes it with its own.
The chain rule is composition of these local linear maps.
Why you normally never construct the full Jacobian
Imagine a neural network with:
1,000,000 parameters
and:
1 scalar loss
Its Jacobian with respect to the parameters contains one million entries. That is manageable.
But now imagine an intermediate layer mapping:
10,000 inputs → 10,000 outputs
The explicit Jacobian would contain:
100,000,000 entries
For most training steps, building that whole matrix would be pointless.
Autodiff instead works with products involving the Jacobian.
Reverse mode computes quantities of the form:
vᵀ · J
without materializing J.
Forward mode computes:
J · v
without materializing J.
That ability to calculate the needed product rather than the entire derivative matrix is central to practical autodiff.
One pass does not mean one operation
The phrase “one backward pass computes all million gradients” can be misunderstood.
It does not mean:
one CPU instruction
or:
constant work regardless of model size
A backward pass performs many operations.
For a neural network it may include:
-
matrix multiplications;
-
elementwise multiplications;
-
reductions;
-
activation derivatives;
-
normalization derivatives;
-
communication between GPUs;
-
gradient accumulation;
-
memory reads and writes.
“One pass” means one coordinated reverse traversal of the computation graph, rather than one independent traversal per parameter.
That is the scaling breakthrough.
Backpropagation also has a memory cost
Reverse mode needs information from the forward computation to evaluate backward rules.
For example, if a local derivative depends on an activation computed earlier, the backward calculation may need that activation or enough information to reconstruct it.
This is why training generally consumes substantially more memory than pure inference.
Autograd frameworks therefore make tradeoffs between:
save intermediate values
and:
recompute intermediate values later
Techniques such as activation checkpointing exploit that tradeoff by recomputing parts of the forward pass during backward to save memory.
The chain rule stays the same. What changes is the engineering strategy for supplying the values needed to evaluate it.
Higher-order derivatives are chain rule again
Autodiff is not limited to first derivatives.
A gradient calculation can itself be represented as differentiable computation, enabling second derivatives, Hessian-vector products, and related quantities.
PyTorch exposes APIs for higher-order autodiff, including Hessian-vector-product functionality in its current autograd interface.
For a simple second derivative:
import torch x = torch.tensor( 2.0, requires_grad=True, ) y = x**3 (first_derivative,) = torch.autograd.grad( y, x, create_graph=True, ) (second_derivative,) = torch.autograd.grad( first_derivative, x, ) print(first_derivative) print(second_derivative)
For:
y = x³
we know:
dy/dx = 3x²
At x = 2:
dy/dx = 12
and:
d²y/dx² = 6x = 12
Expected output:
tensor(12., grad_fn=<MulBackward0>) tensor(12.)
The important option is:
create_graph=True
because we want the first derivative calculation itself to remain differentiable.
Cherry on the cake: even “weights-only” model loading had an RCE CVE
Autograd is a mathematical topic, but real ML systems also have a software supply-chain surface.
A particularly useful reminder arrived in January 2026.
PyTorch disclosed CVE-2026-24747, a high-severity vulnerability affecting versions through 2.9.1. A malicious checkpoint loaded through torch.load(..., weights_only=True) could exploit insufficient validation in the restricted unpickler, causing memory corruption and potentially arbitrary code execution. PyTorch lists version 2.10.0 and later as patched.
That story is surprising because weights_only=True exists specifically to restrict what checkpoint deserialization can reconstruct.
There had already been another critical torch.load(..., weights_only=True) remote-code-execution vulnerability, CVE-2025-32434, affecting PyTorch through 2.5.1 and patched in 2.6.0.
The practical lesson is broader than either CVE:
-
Treat model checkpoints as software artifacts, not harmless bags of numbers.
-
Do not load arbitrary checkpoints from unknown sources merely because their filename ends in
.ptor.pth. -
Keep ML frameworks patched.
-
Prefer current hardened serialization paths.
-
Maintain provenance, hashes, and access controls for model artifacts just as you would for executable dependencies.
The mathematics of backprop may be elegant, but production ML still lives inside ordinary security boundaries.
A compact mental model to keep
When you look at a large model, avoid imagining a giant derivative formula.
Imagine a graph.
Every node knows:
1. how to compute its forward value; 2. how an incoming sensitivity should be transformed into sensitivities for its inputs.
The forward pass computes values.
The reverse pass starts with:
∂loss/∂loss = 1
and repeatedly applies:
upstream gradient × local derivative
When paths merge, contributions add.
When millions of parameters feed one scalar objective, reverse mode lets that one output sensitivity flow backward to all of them in a coordinated sweep.
That is backpropagation.
What you should be able to do now
You should now be able to explain, rather than merely use, the essential mechanics:
-
The chain rule composes local derivatives.
-
A computational graph turns a huge function into manageable primitive operations.
-
Backpropagation is reverse-mode automatic differentiation applied to that graph.
-
.backward()propagates sensitivities from a scalar objective to all relevant leaf parameters. -
Reverse mode naturally suits functions with many inputs and few outputs.
-
Forward mode propagates input directions through JVPs and naturally suits different Jacobian shapes.
-
VJPs and JVPs let autodiff avoid constructing huge Jacobian matrices.
-
Gradient accumulation, graph retention, and intermediate storage are engineering consequences of the same underlying mathematics.
The next step is to make the graph slightly more realistic.
Take the million-parameter example and replace its elementwise prediction with a two-layer torch.nn.Module. Print every parameter’s shape, call loss.backward(), and inspect the matching .grad tensors. Then hand-derive the gradient of one weight in one neuron and verify that autograd gives the same value.
Once you can move comfortably between one derivative on paper, one node in a computation graph, and a million gradients produced by reverse mode, backpropagation stops being a black box.
Build that example next—and make the chain rule earn its keep at scale.