Automatic Differentiation and Computation Graphs

Automatic differentiation is the mechanism that lets PyTorch train neural networks without you manually deriving every partial derivative. You write the forward computation with tensors, PyTorch records the operations that produced gradient-tracked results, and a backward pass applies the chain rule from the final scalar objective back to the parameters. The outcome is practical: after this lesson you should be able to predict when a tensor will have a gradient, inspect a computation graph, fix common autograd errors, and use gradients correctly inside a training loop.

This chapter belongs early in the PyTorch foundations sequence because almost every later model, optimizer, and loss function depends on this machinery. A convolutional network, transformer block, or custom layer is still a composition of tensor operations. Autograd is the bridge between that composition and parameter updates.

What PyTorch Records

PyTorch uses dynamic computation graphs. Dynamic means the graph is built as Python executes the forward pass, not declared once in a separate graph-building phase. If an if statement, loop, or function call changes which tensor operations run, the graph for that pass changes too. This is why PyTorch feels like ordinary Python while still supporting gradient-based learning.

A tensor participates in autograd when it has requires_grad=True or is computed from tensors that do. Such a tensor may carry a grad_fn, which points to the backward function that knows how to propagate gradients through the operation that created it. Leaf tensors are tensors created directly by the user, such as model parameters. Non-leaf tensors are intermediate results. By default, PyTorch stores accumulated gradients in the .grad field of leaf tensors only.

During the forward pass, PyTorch records operation nodes and the edges between produced and consumed tensors. During backward(), it starts from a scalar output, gives that output an initial gradient of 1, and walks backward through the graph. Each operation contributes a local derivative, and autograd multiplies and sums these local derivatives according to the chain rule. If a value is used by two downstream branches, gradients from both branches are accumulated before reaching the shared ancestor.

API Anatomy

Autograd normally expects the backward root to be a scalar, such as a mean loss. If the root has multiple elements, you must provide the gradient of some outer objective with respect to that root. This is useful for vector-Jacobian products, but most training code avoids it by reducing losses with mean() or sum().

Example 1: A Scalar Graph

The smallest useful graph has one leaf input, two operations, and one scalar output. Here x is a leaf tensor. The multiplication and addition create intermediate tensors with backward functions. The derivative of y = x * x + 3 * x is 2x + 3, so at x = 2 the gradient is 7.

import torch

x = torch.tensor(2.0, requires_grad=True)
y = x * x + 3 * x

print(y.item())
y.backward()
print(x.grad.item())

The deterministic output is 10.0 followed by 7.0. Notice that y has a grad_fn, but x receives the stored gradient because it is the leaf tensor that autograd is tracking for optimization.

Example 2: Linear Regression by Hand

A training loss is just a larger graph. In this example, w and b are learnable leaf tensors. The prediction, error, squared error, and mean loss are intermediate graph nodes. After loss.backward(), the gradients tell us how to change w and b to reduce the mean squared error.

import torch

x = torch.tensor([1.0, 2.0, 3.0])
y_true = torch.tensor([2.0, 4.0, 6.0])
w = torch.tensor(0.0, requires_grad=True)
b = torch.tensor(0.0, requires_grad=True)

prediction = w * x + b
loss = ((prediction - y_true) ** 2).mean()
loss.backward()

print(round(loss.item(), 4))
print(round(w.grad.item(), 4), round(b.grad.item(), 4))

The output is 18.6667 and then -18.6667 -8.0. The negative gradients mean increasing both w and b would reduce loss for this batch. An optimizer such as stochastic gradient descent performs that update repeatedly, but the gradient calculation comes from the same graph traversal shown here.

Example 3: Training Step with Graph Reset

Gradients accumulate in PyTorch. Accumulation is useful for advanced techniques such as gradient accumulation across microbatches, but it surprises beginners. A normal training step clears old gradients, runs a fresh forward pass, calls backward, and then updates parameters without recording the update itself as part of the learning graph.

import torch

torch.manual_seed(0)
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = 3.0 * x + 1.0
w = torch.randn(1, 1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
learning_rate = 0.05

for step in range(40):
    prediction = x @ w + b
    loss = ((prediction - y) ** 2).mean()
    loss.backward()

    with torch.no_grad():
        w -= learning_rate * w.grad
        b -= learning_rate * b.grad
        w.grad.zero_()
        b.grad.zero_()

print(round(loss.item(), 4))
print(round(w.item(), 2), round(b.item(), 2))

The final loss should be small, and the learned values should be close to 3.0 and 1.0. The exact path depends on initialization and learning rate, but the expected behavior is stable: repeated backward passes produce gradients that move the parameters toward the line that generated the data.

Design Choices and Trade-offs

PyTorch uses reverse-mode automatic differentiation because deep learning usually has many parameters and one scalar loss. Reverse mode is efficient for this shape: one backward pass computes gradients for all parameters with respect to the loss. Forward-mode differentiation is useful for different shapes, such as directional derivatives with many outputs and fewer inputs, but reverse mode is the everyday training workhorse.

Dynamic graphs make debugging and Python control flow natural. You can print tensors, branch on values, and use ordinary functions. The trade-off is that graph construction happens each forward pass, so Python overhead can matter for very small operations or tight loops. Tools such as vectorization, compiled execution paths, and batching can reduce that overhead, but they do not change the underlying need to keep differentiable operations inside autograd-visible tensor code.

Another trade-off is memory. Backward functions often need saved forward values. For example, the backward pass for multiplication needs the multiplied operands. Large models can store many activations until backward completes. Techniques such as smaller batches, checkpointing, mixed precision, and careful use of no_grad() manage this cost.

Failure Modes and Troubleshooting

Symptom: RuntimeError: element 0 of tensors does not require grad. Cause: the loss was computed from tensors outside the graph, often through detach(), torch.no_grad(), .item(), or conversion to NumPy. Diagnose: inspect loss.requires_grad and intermediate grad_fn values. Correction: keep the forward computation in PyTorch tensors with gradient tracking enabled until after backward().

Symptom: gradients double or training diverges even though the formula looks right. Cause: gradients were not cleared between optimization steps. Diagnose: print parameter.grad before and after backward for two consecutive steps. Correction: call optimizer.zero_grad() or explicitly zero each leaf gradient before the next backward pass.

Symptom: RuntimeError about a variable modified by an in-place operation. Cause: an operation ending in an underscore, or indexed assignment, changed a tensor value that a backward function saved from the forward pass. Diagnose: enable anomaly detection with torch.autograd.set_detect_anomaly(True) around the failing step. Correction: replace the in-place operation with an out-of-place expression or move mutation into a no_grad() parameter update.

Symptom: calling backward() a second time on the same loss fails. Cause: PyTorch frees saved graph buffers after backward to reduce memory use. Diagnose: check whether the same graph object is reused instead of recomputing the forward pass. Correction: recompute the loss for the next step, or use retain_graph=True only when you truly need multiple backward traversals of one graph.

Reliability and Performance Implications

Autograd correctness is a reliability concern because silent gradient mistakes produce models that appear to train while learning the wrong relationship. Keep metric logging outside the graph when possible, but do not detach values needed for loss computation. During evaluation, use torch.no_grad() so PyTorch avoids saving activations and reduces memory pressure. During training, avoid unnecessary Python loops over individual examples; batched tensor operations create smaller, more efficient graphs than thousands of tiny scalar operations.

Gradient values can also expose information about training data in sensitive settings, so treat saved gradients, checkpoints, and debug dumps as model artifacts rather than harmless logs. Store only what you need, and clear intermediate tensors when debugging large or private batches.

Hands-on Lab: Inspect and Fix a Training Step

Prerequisites: Python with PyTorch installed, a terminal or notebook, and basic familiarity with tensors. The lab uses synthetic data, so no dataset download is required.

1. Create tensors x and y for the relationship y = 2x - 1. Create scalar leaf tensors w and b with requires_grad=True.
2. Compute prediction = w * x + b and loss = ((prediction - y) ** 2).mean(). Verify that loss.requires_grad is True and that loss.grad_fn is not None.
3. Run loss.backward(). Print w.grad and b.grad. Confirm that both are present and finite.
4. Update w and b inside with torch.no_grad():, then zero their gradients.
5. Repeat the forward, backward, update, and zeroing sequence for at least 50 steps. Verify that the loss decreases and that the parameters move toward 2 and -1.
6. Break the lab intentionally by computing loss_value = loss.item() and trying to call backward() on that Python float. Record the error, then restore the tensor loss path.

Verification: the loss should decrease across steps, w.grad and b.grad should be finite after backward, and the final parameters should be close to the generating values. Cleanup: delete temporary notebook cells or scripts, and restart the kernel if you enabled anomaly detection or kept large tensors alive during experimentation.

Assessment Exercises

1. A tensor has requires_grad=True, but after a computation the final loss has requires_grad=False. List two operations that could have broken the graph and explain how you would locate the break.
2. Why does PyTorch accumulate gradients instead of replacing them automatically? Describe one case where accumulation is useful and one case where it is a bug.
3. In the linear regression example, explain why the gradients for w and b are negative at initialization. What parameter movement does that imply?
4. Rewrite a parameter update so that it does not become part of the next computation graph. Explain the role of torch.no_grad().
5. Choose a model layer you already know, such as a linear layer or activation function, and describe what values its backward function likely needs from the forward pass.

Summary

PyTorch autograd records tensor operations during the forward pass, represents them as a dynamic computation graph, and applies reverse-mode differentiation during the backward pass. Leaf tensors with requires_grad=True receive accumulated gradients, while intermediate tensors carry the backward functions that connect the graph. Correct training code preserves the graph until backward, clears gradients deliberately, avoids unsafe in-place mutation, and uses no_grad() for evaluation and parameter updates.