Neural network forward and backward pass on two records
Trace a two-layer ReLU network, derive one gradient with the chain rule and compare it with PyTorch autograd.
A neural network forward pass computes predictions from inputs and parameters. The backward pass applies the chain rule to compute how the loss changes with each parameter. A two-record example is enough to inspect both without hiding the calculation inside a training framework.
We will use a two-feature, two-hidden-unit network with ReLU and one output. The inputs, weights and targets are original teaching fixtures. We choose values away from the ReLU kink so an analytic gradient check is unambiguous.
Trace the tensor shapes
Write Z = X @ W1 + b1, H = relu(Z) and prediction = H @ W2 + b2. With two records and two features, X, Z and H are (2, 2). W2 is (2, 1), so the prediction and target are (2, 1).
The loss is mean((prediction-target)²) / 2. Its prediction derivative is the residual divided by the number of records. Multiplying by W2.T sends that signal back to the hidden units; the ReLU derivative masks units whose pre-activation is negative.
For W1, the manual gradient is X.T @ ((residual/n @ W2.T) * (Z > 0)). Parentheses matter: every factor represents a particular axis and chain-rule operation. In larger networks, autodiff performs this bookkeeping, but you still need the correct forward computation and loss definition.
Compare the manual and automatic gradients
import torch
torch.set_num_threads(1)
dtype = torch.float64
X = torch.tensor([[1., 2.], [2., -1.]], dtype=dtype)
y = torch.tensor([[1.], [0.]], dtype=dtype)
W1 = torch.tensor([[.5, -.4], [.3, .2]], dtype=dtype, requires_grad=True)
b1 = torch.tensor([.1, -.2], dtype=dtype, requires_grad=True)
W2 = torch.tensor([[.7], [-.5]], dtype=dtype, requires_grad=True)
b2 = torch.tensor([.05], dtype=dtype, requires_grad=True)
Z = X @ W1 + b1
H = torch.relu(Z)
prediction = H @ W2 + b2
loss = ((prediction-y)**2).mean()/2
loss.backward()
with torch.no_grad():
residual_gradient = (prediction-y) / len(X)
manual = X.T @ ((residual_gradient @ W2.T) * (Z > 0))
print('pre-activation:', Z.detach().round(decimals=6).tolist())
print('predictions:', prediction.detach().flatten().round(decimals=6).tolist())
print(f'loss={loss.item():.6f}')
print('W1 gradient:', W1.grad.round(decimals=6).tolist())
print('manual matches:', torch.allclose(manual, W1.grad))
assert torch.allclose(manual, W1.grad)
assert prediction.shape == y.shape
Checked output:
pre-activation: [[1.2, -0.2], [0.8, -1.2]]
predictions: [0.89, 0.61]
loss=0.096050
W1 gradient: [[0.3885, 0.0], [-0.2905, 0.0]]
manual matches: True
Read the inactive-unit gradient
The second hidden unit is negative for both records in this fixture, so ReLU returns zero there and its contribution to the weight gradient is zero. This demonstrates the local derivative convention, not a claim that the unit can never become active under different data or parameters.
The first hidden unit receives a weighted error signal from both records. A weight gradient is a sum over the batch: one record can push a parameter upward while another pushes it downward. Looking at a single activation does not determine the final update direction.
At exactly zero, ReLU is not differentiable in the ordinary sense; PyTorch uses a defined backward convention. Avoid placing a finite-difference check on that boundary and then treating a mismatch as proof that the chain rule is wrong.
Keep backward separate from updating
loss.backward() accumulates gradients into parameter .grad fields. It does not update the weights. An optimizer step is a separate operation. Clear gradients before the next training iteration when you do not intend accumulation; otherwise successive backward calls add to the stored values.
Use torch.no_grad() for a manual parameter update or an inspection calculation that should not build another graph. Do not detach the prediction before calculating the training loss, because that would disconnect the gradient path you need.
Extend the network
Change the second hidden bias so that one record activates the second unit. Predict which zero entries in the gradient can become nonzero, then rerun the calculation. Add a third observation and confirm that the loss and gradient still average over records.
The Deep Learning syllabus covers neural networks and training; lessons are being published. Review matrix multiplication shapes for the forward map, derivative checks for numerical verification and train, eval and no_grad for the separate execution controls.
References
PyTorch autograd mechanics — Graph recording, backward behavior and gradients at nondifferentiable points.