Gradient Descent for Model Training

PyTorch's autograd system computes gradients in dynamic computational graphs automatically, and it records the computation that produces the loss so gradients can be found later.

Programming

Why Training Needs Gradients

Training a model means adjusting its parameters so that the loss gets smaller. To adjust a parameter sensibly, you need to know how much a small change to that parameter would change the loss. That quantity is a gradient. For a model with many parameters, working out these gradients by hand for every calculation becomes impractical. This is the problem PyTorch's autograd system is built to handle.

The key idea behind the rest of this article is that PyTorch does not only compute your loss. It also keeps track of how the loss was produced, so that the gradients can be worked out afterward. Before looking at that bookkeeping, ask yourself a simple question: what would you need to know about a calculation in order to find its gradients without redoing it from scratch?

What do you think happens?

If you compute a loss from a weight in PyTorch, does PyTorch record every number in your program along the way?

Yes, every number is always recorded
No, recording depends on whether the values involved are flagged as needing gradients

Reveal answer

Answer:

No, recording depends on whether the values involved are flagged as needing gradients.

PyTorch builds a computation graph internally by default when the relevant tensors have requires_grad set to True. The recording is not automatic for every number in the program.

Tracing a Forward Pass

PyTorch builds a computation graph in the background. A computation graph is a record of the operations that turned inputs into outputs, arranged so that each result points to the values it was calculated from. Its purpose here is specific: it lets PyTorch calculate gradients of a loss function with respect to the model parameters, which is what training needs.

Which operations create which intermediate values, and how do the inputs, weight, bias, and loss connect as nodes and edges?

Read the diagram from the bottom up as the forward pass. The weight and the bias feed into the intermediate result z, and the loss is then computed from z. The graph is the record of this structure. Each node knows which values it was calculated from, so the path from the loss back to the weight and the bias is preserved and can be walked later.

When PyTorch Starts Recording

The trigger for building the graph is a flag on the tensors involved. When you carry out computations in PyTorch, it builds a computational graph internally by default as long as one of the terminal nodes has its requires_grad attribute set to True. The terminal nodes are the values at the end of the computation. The requires_grad attribute marks the quantities for which you want gradients. So the recording depends on whether the values you care about are flagged, not on every number in your program.

python

import torch

w = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
x = torch.tensor(3.0)

z = w * x + b
loss = z ** 2

loss.backward()
print(w.grad)
print(b.grad)

import torch

w = torch.tensor(2.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)
x = torch.tensor(3.0)

z = w * x + b
loss = z ** 2

loss.backward()
print(w.grad)
print(b.grad)

Output (expected)

tensor(42.)
tensor(14.)

Notes on the code

w and b have requires_grad=True, so PyTorch tracks the operations that use them
x has no requires_grad flag, so it is not marked as a quantity needing gradients
loss.backward() walks the recorded graph from the loss back to the flagged tensors
w.grad and b.grad hold the gradients of the loss with respect to each parameter

Following the Chain Rule Backward

Once the graph exists, gradients are found by applying the chain rule from calculus. Backpropagation can be considered an implementation of the chain rule for neural networks. The most common way of computing loss gradients in a computation graph applies the chain rule from right to left. This is also called reverse-mode automatic differentiation or backpropagation.

The direction of travel is the important idea. Calculation moves forward, from inputs toward the loss. Gradient calculation moves in the opposite direction. It starts at the loss and works backward through the graph toward the parameters. This lets the procedure compute the gradient of the loss with respect to each parameter, meaning each weight and bias. Those gradients then inform how the parameters are updated during training.

Starting at the loss, in what order are local derivatives multiplied as the gradient travels back through each operation to reach each weight?

Each arrow in the backward pass carries one local partial derivative. To reach the weight, the gradient travels from the loss to z, and then from z to w. The contributions along that path are chained together. In TeX form, the chain for this single path is shown below.

Rule

The same procedure is repeated for each further link. So every parameter receives its own gradient by the same backward procedure, even when the graph is much larger than this one.

Partial Derivatives Versus the Gradient

Partial derivatives are the building blocks of this process. Each one measures how a single quantity changes with respect to one of its inputs while the other inputs are held in place. In a computation graph, each link between two nodes contributes one partial derivative, and the chain rule combines these pieces along a path.

How does one partial derivative for a single input differ from the full gradient that collects a partial derivative for every parameter?

A gradient is what you obtain for the parameters when these pieces are assembled. The gradient of the loss with respect to the model parameters tells the training procedure how each parameter affects the loss, and that is the information used to update the parameters. The distinction is one of scope. A partial derivative describes one link in the graph. The gradient describes the effect on the loss for each parameter, produced by chaining those links.

Worked Example: Tracing One Weight

Worked example

Gradient of the loss with respect to a weight

Let z = w x + b and loss L = z^{^2}, with w = 2, x = 3, and b = 1. Find the gradient of L with respect to w by tracing the graph backward.

1
Forward pass
Compute the intermediate value: $z = wx + b = 2 \cdot 3 + 1 = 7$. Then compute the loss: $L = z^{^2} = 49$.
2
Local derivative at the loss
The partial derivative of the loss with respect to z is $\frac{\partial L}{\partial z} = 2z = 14$.
3
Local derivative at the weight
The partial derivative of z with respect to w is $\frac{\partial z}{\partial w} = x = 3$.
4
Chain the links
Multiply along the path from the loss back to w: $\frac{\partial L}{\partial w} = 14 \cdot 3 = 42$.

Result

The gradient of the loss with respect to w is 42, which matches the w.grad value printed by the code example.

Mistakes Beginners Make with Graphs

Common mistake
Assuming PyTorch records every operation in the program
What it looks like
Writing a calculation with plain tensors and expecting a gradient later, when none of the involved tensors has requires_grad set to True.
Why it fails
The graph is built by default only when the relevant tensors are flagged. Without the flag, there is nothing marked as needing gradients.
Do this instead
Set requires_grad=True on the parameters whose gradients you need before running the forward computation.
Common mistake
Treating a partial derivative as the full gradient
What it looks like
Reporting dz/dw as the answer for how the loss changes with the weight.
Why it fails
A partial derivative describes one link. The loss depends on the weight through the whole path, so the chain of links must be combined.
Do this instead
Chain the partial derivatives along the path from the loss back to the parameter, then collect one gradient per parameter.
Common mistake
Reading the chain rule as running forward
What it looks like
Starting the gradient calculation at the inputs and moving toward the loss.
Why it fails
The gradient calculation runs from right to left, starting at the loss and moving backward toward the parameters.

Practice Tracing a Graph

MEDIUM

Let z = w x + b and L = z^{^2}, with w = 1, x = 4, and b = 2. Trace the graph backward and find the gradient of L with respect to w. Write each partial derivative you use, then chain them together.

Hints

Compute z first and use it to find the partial derivative of L with respect to z.
The partial derivative of z with respect to w is the input x.
Multiply the two partial derivatives along the path from L back to w.

Your attempt (not saved)

Key Takeaways

PyTorch's autograd system records the operations that produce a loss, so gradients can be computed afterward instead of being worked out by hand.
A computation graph is built by default when the tensors involved are flagged with requires_grad set to True; unflagged values are not marked as needing gradients.
Gradients are found by applying the chain rule from right to left, starting at the loss and moving backward, which is reverse-mode automatic differentiation or backpropagation.
A partial derivative describes one link in the graph, while the gradient collects the chained effects on the loss for every parameter.

Related concepts

More concepts starting with G