What backpropagation does before gradient descent takes a step

A typical PyTorch training loop calls loss.backward() and then optimizer.step(). You can stop between those calls and inspect the difference: the gradients have been calculated, while the parameters still have their old values.

Backpropagation computes the gradient of the loss with respect to each trainable parameter. Plain gradient descent uses those gradients to update the parameters. Other optimizers use different update rules, but occupy the same place in the PyTorch training loop, after the backward pass.

We will follow one batch through an affine layer using matrices throughout. Two examples, two input features, and two outputs keep the arithmetic small while showing how a batch produces a gradient for an entire weight matrix. The worked layer has no nonlinear activation; later, we connect its backward calculation to a multilayer network.

Start with a prediction and a loss

Put examples in the rows of an input matrix $X$. For a batch of $m$ examples with $d$ input features and $k$ outputs, the prediction is

$$\hat Y = XW + \mathbf{1}_m b^\top.$$

Every entry of $W$ is a separate parameter. For example, $W_{12}$ connects input feature 1 to output 2. All rows in the batch use the same $W$ and $b$.

Here are our inputs and targets:

$$\begin{gathered} X = \begin{bmatrix}1 & 2\\2 & 1\end{bmatrix}, \\[6pt] Y = \begin{bmatrix}0 & 1\\1 & 0\end{bmatrix}. \end{gathered}$$

Start with an identity weight matrix and zero biases:

$$\begin{gathered} W = \begin{bmatrix}1 & 0\\0 & 1\end{bmatrix}, \\[6pt] b = \begin{bmatrix}0\\0\end{bmatrix}. \end{gathered}$$

Multiplication by the identity leaves $X$ unchanged, so the predictions and residuals are

$$\begin{gathered} \hat Y = \begin{bmatrix}1 & 2\\2 & 1\end{bmatrix}, \\[6pt] R = \hat Y-Y = \begin{bmatrix}1 & 1\\1 & 1\end{bmatrix}. \end{gathered}$$

Use half the average squared error per example, summing over that example's outputs:

$$L = \frac{1}{2m}\|R\|_F^2 = \frac{1}{2m}\sum_{i=1}^{m}\sum_{j=1}^{k}R_{ij}^2.$$

The squared Frobenius norm, $\|R\|_F^2$, means square every entry and add. Here $m=2$ and all four residuals equal one, giving $L=4/4=1$. The loss is a scalar even though its inputs and parameters are matrices. This normalization averages over examples, not over all matrix entries; the code below uses exactly this definition.

Work backward to find the gradients

Start at the loss. Its gradient with respect to the prediction matrix has one entry for every prediction:

$$G = \nabla_{\hat Y}L = \frac{R}{m} = \begin{bmatrix}0.5 & 0.5\\0.5 & 0.5\end{bmatrix}.$$

Here $G_{ij}=\partial L/\partial\hat Y_{ij}$ is a partial derivative. It records how sensitive the loss is to prediction $j$ for example $i$, holding the other prediction entries fixed.

For a scalar-valued loss, a gradient with respect to a matrix collects those partial derivatives in an array with the matrix's shape. Flattening that array gives the gradient vector for the flattened inputs. A Jacobian uses a different layout: one row per output and one column per input. The loss has just one output, so its Jacobian with respect to the flattened predictions would have one row; $G$ keeps the derivatives in the prediction matrix's original layout.

The backward pass through the affine layer gives

$$\begin{gathered} \nabla_W L = X^\top G, \\[6pt] \nabla_b L = G^\top\mathbf{1}_m. \end{gathered}$$

The matrix multiplication collects each example's contribution to every weight. In component form, $(\nabla_W L)_{aj}=\sum_i X_{ia}G_{ij}$. The bias derivative sums the incoming gradients over examples because the same bias was added to every row. Each gradient has the same shape as its parameter: $d\times k$ for $W$, and $k\times1$ for $b$.

For our batch,

$$\begin{gathered} \nabla_W L = \begin{bmatrix}1 & 2\\2 & 1\end{bmatrix} \begin{bmatrix}0.5 & 0.5\\0.5 & 0.5\end{bmatrix} \\[6pt]= \begin{bmatrix}1.5 & 1.5\\1.5 & 1.5\end{bmatrix}. \end{gathered}$$ $$\begin{gathered} \nabla_b L = \begin{bmatrix}0.5 & 0.5\\0.5 & 0.5\end{bmatrix} \begin{bmatrix}1\\1\end{bmatrix} \\[6pt]= \begin{bmatrix}1\\1\end{bmatrix}. \end{gathered}$$

For example, the derivative for $W_{12}$ is $1\times0.5+2\times0.5=1.5$. Both examples contribute to that one parameter. The equal entries in these gradients come from the symmetric numbers chosen for this example; gradients in a trained network generally differ from entry to entry.

We have calculated six derivatives: four for the weights and two for the biases. The parameter matrices have not changed. Backpropagation has evaluated the chain rule through the graph. It has not chosen a learning rate or taken an optimization step. Our chain rule lesson explains the scalar case; the multivariable chain rule lesson shows how derivative matrices compose through intermediate values.

Use the gradients to update the parameters

The gradient descent update subtracts a scaled gradient from each parameter array. With learning rate $\eta$,

$$\begin{gathered} W_{\text{new}} = W-\eta\nabla_W L, \\[6pt] b_{\text{new}} = b-\eta\nabla_b L. \end{gathered}$$

Using $\eta=0.1$ gives

$$\begin{gathered} W_{\text{new}}= \begin{bmatrix}0.85 & -0.15\\-0.15 & 0.85\end{bmatrix}, \\[6pt] b_{\text{new}}= \begin{bmatrix}-0.1\\-0.1\end{bmatrix}. \end{gathered}$$

Every derivative came from the same starting state. We calculate both gradients before changing either parameter array.

A new forward pass gives

$$\begin{gathered} \hat Y_{\text{new}} = XW_{\text{new}}+\mathbf{1}_m b_{\text{new}}^\top \\[6pt]= \begin{bmatrix}0.45 & 1.45\\1.45 & 0.45\end{bmatrix}. \end{gathered}$$

Each prediction is now $0.45$ above its target. The new loss is $4\times0.45^2/4=0.2025$, down from $1$.

The second row is the distinction you can inspect between loss.backward() and optimizer.step(). The backward pass adds derivative information; the optimization step changes the model.

For a classification loss, the starting derivative changes. Our cross-entropy derivative example derives the softmax logit gradient $p-y$ and checks it numerically. The same chain rule then carries that gradient through the network.

How this extends to a multilayer network

An affine layer inside a larger network still computes $Z=XW+\mathbf{1}_m b^\top$. Its incoming gradient $G=\nabla_Z L$ now comes from the layers after it. Once that incoming gradient is known, the local backward rules remain

$$\begin{gathered} \nabla_W L=X^\top G, \\[6pt] \nabla_b L=G^\top\mathbf{1}_m, \\[6pt] \nabla_X L=GW^\top. \end{gathered}$$

The last expression passes a gradient to the preceding layer when $X$ is itself a learned intermediate activation. Our vector–Jacobian product lesson explains how an incoming output sensitivity travels backward through a vector operation. These matrix formulas perform that same propagation while preserving the batch and feature axes.

If the affine layer is followed by an elementwise activation $A=\phi(Z)$, the backward pass through that activation first computes $G=(\nabla_A L)\odot\phi'(Z)$, wherever the activation is differentiable. The symbol $\odot$ means entry-by-entry multiplication. Branches add contributions when multiple paths lead back to the same value. This is how the autograd computation graph organizes the chain rule across a network.

You may also see all weights and biases collected into one parameter vector $\theta$. In that notation, the same plain gradient descent update is $\theta_{t+1}=\theta_t-\eta\nabla_\theta L(\theta_t)$. Keeping separate matrices simply preserves the shapes of the layers.

Changing the learning rate does not change this starting gradient

Return to the original $W$ and $b$, keep the same batch and loss, and choose $\eta=0.5$. The starting gradients are unchanged. The larger update gives

$$\begin{gathered} W_{\text{new}}= \begin{bmatrix}0.25 & -0.75\\-0.75 & 0.25\end{bmatrix}, \\[6pt] b_{\text{new}}= \begin{bmatrix}-0.5\\-0.5\end{bmatrix}. \end{gathered}$$

Now every residual is $-1.75$, and the loss rises to $3.0625$. Even a direction that reduces the loss locally can make it worse if the step is too large.

For these starting matrices, each residual after one update is $1-5.5\eta$. The resulting loss is

$$L_{\text{after}}(\eta)=(1-5.5\eta)^2.$$

Figure link ↗SVG ↓PNG ↓Data ↓

This plot varies one scalar, the learning rate, while the underlying update still changes all six parameters. The rate $\eta=2/11\approx0.1818$ happens to fit this particular batch exactly. That is a consequence of our chosen numbers, not a recommended learning rate for a neural network. Our gradient descent oscillation article examines how step size affects repeated updates.

Reproduce the whole step in Python

The NumPy code below uses @ for matrix multiplication. b is stored as a column; adding b.T broadcasts the bias row across the batch. Both gradients are calculated before either parameter is updated.

import numpy as np

X = np.array([[1., 2.], [2., 1.]])
Y = np.array([[0., 1.], [1., 0.]])
W = np.eye(2)
b = np.zeros((2, 1))
m = X.shape[0]
eta = 0.1

# Forward pass: half the mean per-example squared error.
prediction = X @ W + b.T
residual = prediction - Y
loss = np.sum(residual ** 2) / (2 * m)

# Backward pass: each gradient matches its parameter's shape.
G = residual / m
grad_W = X.T @ G
grad_b = G.T @ np.ones((m, 1))

# Plain gradient descent uses gradients from the same state.
W_new = W - eta * grad_W
b_new = b - eta * grad_b
new_prediction = X @ W_new + b_new.T
new_loss = np.sum((new_prediction - Y) ** 2) / (2 * m)

print("weight gradient:\n", grad_W)
print("bias gradient:\n", grad_b)
print("new weights:\n", W_new)
print("new biases:\n", b_new)
print(f"loss: {loss:.4f} -> {new_loss:.4f}")

It prints the gradient and parameter matrices derived above, and loss: 1.0000 -> 0.2025. This code implements the local backward rules explicitly. In PyTorch, loss.backward() performs the differentiation and accumulates the results in parameter .grad attributes. A usual loop clears old gradients before the next backward pass unless it deliberately accumulates them across batches. The autograd tutorial explains that behavior.

What changes when you use a different optimizer?

The gradient calculation can stay the same when the update rule changes. SGD with momentum or Adam uses gradients together with optimizer state to determine an update. Its parameter changes therefore differ from the plain subtraction used in this article.

Gradient clipping adds another possible operation between the backward pass and the optimizer step. It limits the calculated gradients; the optimizer still determines how those gradients change the parameters.

When debugging, inspect the loss and computation graph if derivatives are wrong. If the derivatives are right but the loss grows after an update, examine the optimizer and step size. In our matrix example, changing only the learning rate was enough to turn an improving step into a harmful one.