Module 02: Autograd

Introduction

Automatic differentiation enables neural network training. It computes gradients through any computation - the foundation of backpropagation.

Autograd (automatic differentiation) computes gradients for you. When you call loss.backward() in PyTorch, autograd figures out how to adjust every parameter to reduce the loss.

Why is this essential for LLMs?

  • Millions of parameters: An LLM has millions (or billions) of numbers to adjust. Only automatic computation can handle this scale.
  • Complex computations: Attention, embeddings, layer norms - the gradient must flow through all of them.
  • Training loop: Every training step computes gradients to update weights.

Autograd makes deep learning trainable.

What You’ll Learn

After this module, you can:

  • Explain how computational graphs enable gradient computation
  • Build a working scalar autograd engine from scratch
  • Extend these principles to tensor operations
  • Use PyTorch’s autograd for gradient computation
  • Identify and avoid common autograd pitfalls

Intuition: The Computational Graph

Every computation builds a directed acyclic graph (DAG) dynamically as operations execute. Frameworks call this approach “define-by-run” because the code path determines graph structure—enabling conditionals and loops to create different graphs each pass.

Think of a computation as a graph of operations:

Forward pass: Compute values flowing down (x -> a -> b -> c)

Backward pass: Compute gradients flowing up (c -> b -> a -> x)

  • dc/dc = 1 (output gradient is always 1)
  • dc/db = 2b = 14 (derivative of b^2)
  • dc/da = dc/db x db/da = 14 x 1 = 14
  • dc/dx = dc/da x da/dx = 14 x 3 = 42

The chain rule connects everything: multiply local gradients as you go back.

Computational Graph: Forward and Backward

An example shows both passes:

The Math

Chain Rule

For composed functions f(g(x)):

d/dx[f(g(x))] = f'(g(x)) x g'(x)

This extends to any depth - just multiply local derivatives along the path.

Common Gradients

Operation Forward Local Gradient
c = a + b sum dc/da = 1, dc/db = 1
c = a * b product dc/da = b, dc/db = a
c = a ** n power dc/da = n x a^(n-1)
c = exp(a) exp dc/da = exp(a)
c = log(a) log dc/da = 1/a
c = tanh(a) tanh dc/da = 1 - tanh^2(a)
c = relu(a) ReLU dc/da = 1 if a > 0 else 0

Matrix Multiplication Gradients

For C = A @ B:

  • dL/dA = dL/dC @ B.T
  • dL/dB = A.T @ dL/dC

This is why matrix shapes matter so much!

Gradient Accumulation

When a value is used multiple times, gradients ADD:

Code Walkthrough

Explore autograd interactively:

import torch

print(f"PyTorch version: {torch.__version__}")
PyTorch version: 2.10.0+cu128

Building a Simple Autograd Engine

This simple autograd engine (inspired by Andrej Karpathy’s micrograd) handles scalar values only - PyTorch’s autograd extends these same principles to tensors of any shape, which is what makes it powerful for real neural networks.

class Value:
    """A scalar value that tracks its gradient."""

    def __init__(self, data, children=(), op='', label=''):
        self.data = data
        self.grad = 0.0
        self._backward = lambda: None
        self._prev = set(children)
        self._op = op
        self.label = label

    def __repr__(self):
        return f"Value(data={self.data}, grad={self.grad})"

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')

        def _backward():
            self.grad += out.grad   # d(a+b)/da = 1
            other.grad += out.grad  # d(a+b)/db = 1
        out._backward = _backward

        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), '*')

        def _backward():
            self.grad += other.data * out.grad  # d(a*b)/da = b
            other.grad += self.data * out.grad  # d(a*b)/db = a
        out._backward = _backward

        return out

    def __pow__(self, other):
        assert isinstance(other, (int, float)), "only supporting int/float powers"
        out = Value(self.data ** other, (self,), f'**{other}')

        def _backward():
            self.grad += (other * self.data ** (other - 1)) * out.grad
        out._backward = _backward

        return out

    def __neg__(self):
        return self * -1

    def __sub__(self, other):
        return self + (-other)

    def __radd__(self, other):
        return self + other

    def __rmul__(self, other):
        return self * other

    def tanh(self):
        import math
        t = math.tanh(self.data)
        out = Value(t, (self,), 'tanh')

        def _backward():
            self.grad += (1 - t ** 2) * out.grad
        out._backward = _backward

        return out

    def relu(self):
        out = Value(max(0, self.data), (self,), 'relu')

        def _backward():
            self.grad += (self.data > 0) * out.grad
        out._backward = _backward

        return out

    def exp(self):
        import math
        out = Value(math.exp(self.data), (self,), 'exp')

        def _backward():
            self.grad += out.data * out.grad
        out._backward = _backward

        return out

    def log(self):
        import math
        out = Value(math.log(self.data), (self,), 'log')

        def _backward():
            self.grad += (1.0 / self.data) * out.grad
        out._backward = _backward

        return out

    def __truediv__(self, other):
        return self * (other ** -1)

    def backward(self):
        # Topological sort to process in correct order
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)
        build_topo(self)

        # Go backwards, applying chain rule
        self.grad = 1.0
        for v in reversed(topo):
            v._backward()

Testing Our Value Class

# Create some Values
a = Value(2.0, label='a')
b = Value(3.0, label='b')

print(f"a = {a}")
print(f"b = {b}")
print(f"\nInitially, gradients are 0.0")
a = Value(data=2.0, grad=0.0)
b = Value(data=3.0, grad=0.0)

Initially, gradients are 0.0
# Perform a computation
c = a * b
c.label = 'c'

print(f"c = a * b = {c.data}")
print(f"\nNow let's compute gradients with backward():")

c.backward()

print(f"\ndc/da = {a.grad}  (should be b = 3.0)")
print(f"dc/db = {b.grad}  (should be a = 2.0)")
c = a * b = 6.0

Now let's compute gradients with backward():

dc/da = 3.0  (should be b = 3.0)
dc/db = 2.0  (should be a = 2.0)

Verifying Gradients Numerically

# Let's verify: increase 'a' by a tiny amount
epsilon = 0.0001

a_original = 2.0
b_val = 3.0

c_original = a_original * b_val
c_perturbed = (a_original + epsilon) * b_val

numerical_gradient = (c_perturbed - c_original) / epsilon

print(f"Numerical gradient dc/da = {numerical_gradient:.4f}")
print(f"Our computed gradient: {a.grad}")
print(f"\nThey match! (The tiny difference is numerical precision)")
Numerical gradient dc/da = 3.0000
Our computed gradient: 3.0

They match! (The tiny difference is numerical precision)

The Chain Rule in Action

# Let's trace through: c = (a + b) * b
a = Value(2.0, label='a')
b = Value(3.0, label='b')

sum_ab = a + b
sum_ab.label = 'sum'
c = sum_ab * b
c.label = 'c'

c.backward()

print(f"Expression: c = (a + b) * b")
print(f"a = {a.data}, b = {b.data}")
print(f"sum = a + b = {sum_ab.data}")
print(f"c = sum * b = {c.data}")
print(f"\nGradients:")
print(f"dc/da = {a.grad}  (expected: b = 3)")
print(f"dc/db = {b.grad}  (expected: sum + a = 5 + 3 = 8)")
Expression: c = (a + b) * b
a = 2.0, b = 3.0
sum = a + b = 5.0
c = sum * b = 15.0

Gradients:
dc/da = 3.0  (expected: b = 3)
dc/db = 8.0  (expected: sum + a = 5 + 3 = 8)

Training a Neuron

Train a neuron to output a target value:

# Training loop
x_val = 2.0
target_val = 0.8

# Learnable parameters (start with small random values)
w = 0.3
b = 0.1

learning_rate = 0.5
losses = []

print("Training a neuron to output 0.8 when input is 2.0")
print("=" * 50)

for step in range(20):
    # Create fresh Values (gradients reset)
    x = Value(x_val, label='x')
    w_v = Value(w, label='w')
    b_v = Value(b, label='b')
    target = Value(target_val, label='target')

    # Forward pass
    y = (w_v * x + b_v).tanh()
    loss = (y - target) ** 2

    # Backward pass
    loss.backward()

    # Gradient descent update
    w = w - learning_rate * w_v.grad
    b = b - learning_rate * b_v.grad

    losses.append(loss.data)

    if step % 4 == 0:
        print(f"Step {step:2d}: y={y.data:.4f}, loss={loss.data:.6f}, w={w:.4f}, b={b:.4f}")

print(f"\nFinal output: {y.data:.4f}")
print(f"Target: {target_val}")
print("Pretty close!")
Training a neuron to output 0.8 when input is 2.0
==================================================
Step  0: y=0.6044, loss=0.038272, w=0.5484, b=0.2242
Step  4: y=0.8121, loss=0.000146, w=0.4650, b=0.1825
Step  8: y=0.8002, loss=0.000000, w=0.4595, b=0.1798
Step 12: y=0.8000, loss=0.000000, w=0.4594, b=0.1797
Step 16: y=0.8000, loss=0.000000, w=0.4594, b=0.1797

Final output: 0.8000
Target: 0.8
Pretty close!

The loss decreases because gradients tell us which way to adjust w and b!

The Evolutionary Leap: Scalar to Tensor

We’ve now mastered scalar autograd with our Value class. The chain rule, computational graphs, and backward passes all work the same way at any scale. But real neural networks operate on tensors with thousands or millions of elements. Let’s build a second autograd engine — this time for tensors — to see how the same principles scale up.

Why Scalars Don’t Scale

Our Value class works beautifully for understanding autograd. But try to imagine training a real neural network with it:

  • A small MLP with 10,000 parameters creates 10,000 Value objects
  • Each forward pass creates thousands more intermediate Value nodes
  • Python loops for every dot product (catastrophically slow)
  • Memory explodes with millions of tiny objects

The fix isn’t “optimize Python.” The fix is stop pretending a neural network is a pile of scalars.

# The problem: dot product with scalar Values
def slow_dot_product(values_a, values_b):
    """This is what we're doing now - Python loop over scalars."""
    result = values_a[0] * values_b[0]
    for a, b in zip(values_a[1:], values_b[1:]):
        result = result + a * b
    return result

# With 1000-dimensional vectors, that's 1000 Python operations
# A GPU can do this in ONE operation

The solution: upgrade from scalar Value to tensor Tensor.

Building a Tensor Autograd

The evolutionary leap: a Tensor class backed by NumPy arrays instead of Python floats.

import numpy as np
from typing import Optional, Tuple, Callable, Set

class Tensor:
    """
    A NumPy-backed tensor with reverse-mode autodiff.

    This is what PyTorch's tensors do internally (but in C++/CUDA).
    """
    def __init__(self, data, requires_grad: bool = False, _prev: Set["Tensor"] = None, _op: str = ""):
        # Convert to numpy array if needed
        if isinstance(data, np.ndarray):
            self.data = data.astype(np.float32)
        else:
            self.data = np.array(data, dtype=np.float32)

        self.requires_grad = requires_grad
        self.grad: Optional[np.ndarray] = None
        self._prev = _prev or set()
        self._op = _op
        self._backward: Callable[[], None] = lambda: None

    def __repr__(self) -> str:
        return f"Tensor(shape={self.data.shape}, requires_grad={self.requires_grad})"

    @property
    def shape(self) -> Tuple[int, ...]:
        return self.data.shape

    def zero_grad(self) -> None:
        self.grad = None

Already we see the key difference: self.data is a NumPy array, not a float. This means operations work on entire arrays at once.

The Tricky Part: Broadcasting Gradients

When you add a (batch, dim) tensor to a (dim,) bias, NumPy broadcasts the bias across all batches. But during backprop, we need to undo that broadcasting - the gradient for the bias should be summed across the batch dimension.

def _unbroadcast(grad: np.ndarray, target_shape: Tuple[int, ...]) -> np.ndarray:
    """
    Undo NumPy broadcasting for gradients.

    Example:
      Forward: y = x + b where x is (B, D) and b is (D,)
      Backward: grad wrt b should be sum over batch axis -> (D,)

    Rules:
      1. While grad has extra leading dims, sum them away
      2. For dims where target had size 1, sum over that axis
    """
    g = grad

    # Remove leading dims added by broadcasting
    while len(g.shape) > len(target_shape):
        g = g.sum(axis=0)

    # Sum over axes where target had size 1
    for axis, (gdim, tdim) in enumerate(zip(g.shape, target_shape)):
        if tdim == 1 and gdim != 1:
            g = g.sum(axis=axis, keepdims=True)

    return g.reshape(target_shape)

# Test it
grad = np.ones((3, 4))  # Gradient has batch dimension
bias_shape = (4,)       # Bias was (4,) before broadcasting
result = _unbroadcast(grad, bias_shape)
print(f"Input grad shape: {grad.shape}")
print(f"Target shape: {bias_shape}")
print(f"Result shape: {result.shape}")  # (4,) - summed over batch
print(f"Result values: {result}")       # [3, 3, 3, 3] - sum of 3 ones per position
Input grad shape: (3, 4)
Target shape: (4,)
Result shape: (4,)
Result values: [3. 3. 3. 3.]

This is the key insight PyTorch handles automatically. When you see RuntimeError: grad shape doesn't match, it’s usually a broadcasting gradient issue.

Tensor Arithmetic with Gradients

Now we add operations. Each operation stores a _backward function that knows how to push gradients to its inputs.

# Add these methods to our Tensor class (shown separately for clarity)

def tensor_add(self, other) -> "Tensor":
    """
    Addition: c = a + b
    Gradient: dc/da = 1, dc/db = 1 (with unbroadcasting)
    """
    other = other if isinstance(other, Tensor) else Tensor(other)
    out = Tensor(
        self.data + other.data,
        requires_grad=self.requires_grad or other.requires_grad,
        _prev={self, other},
        _op="+"
    )

    def _backward():
        if out.grad is None:
            return
        if self.requires_grad:
            g = _unbroadcast(out.grad, self.data.shape)
            self.grad = g if self.grad is None else (self.grad + g)
        if other.requires_grad:
            g = _unbroadcast(out.grad, other.data.shape)
            other.grad = g if other.grad is None else (other.grad + g)

    out._backward = _backward
    return out

def tensor_mul(self, other) -> "Tensor":
    """
    Multiplication: c = a * b
    Gradient: dc/da = b, dc/db = a (with unbroadcasting)
    """
    other = other if isinstance(other, Tensor) else Tensor(other)
    out = Tensor(
        self.data * other.data,
        requires_grad=self.requires_grad or other.requires_grad,
        _prev={self, other},
        _op="*"
    )

    def _backward():
        if out.grad is None:
            return
        if self.requires_grad:
            g = _unbroadcast(out.grad * other.data, self.data.shape)
            self.grad = g if self.grad is None else (self.grad + g)
        if other.requires_grad:
            g = _unbroadcast(out.grad * self.data, other.data.shape)
            other.grad = g if other.grad is None else (other.grad + g)

    out._backward = _backward
    return out

# Attach to Tensor class
Tensor.__add__ = tensor_add
Tensor.__radd__ = lambda self, other: tensor_add(self, other)
Tensor.__mul__ = tensor_mul
Tensor.__rmul__ = lambda self, other: tensor_mul(self, other)
Tensor.__neg__ = lambda self: tensor_mul(self, -1.0)
Tensor.__sub__ = lambda self, other: tensor_add(self, -other)

This pattern matches our scalar Value - with arrays and _unbroadcast.

Matrix Multiplication: The Core of Neural Networks

Matrix multiplication is where tensors really shine. This single operation replaces thousands of scalar multiplications and additions.

def tensor_matmul(self, other) -> "Tensor":
    """
    Matrix multiplication: C = A @ B

    Forward: C[i,j] = sum_k A[i,k] * B[k,j]

    Backward:
      dA = dC @ B.T   (gradient flows back through B transposed)
      dB = A.T @ dC   (gradient flows back through A transposed)
    """
    other = other if isinstance(other, Tensor) else Tensor(other)
    out = Tensor(
        np.matmul(self.data, other.data),
        requires_grad=self.requires_grad or other.requires_grad,
        _prev={self, other},
        _op="matmul"
    )

    def _backward():
        if out.grad is None:
            return
        if self.requires_grad:
            # dA = dC @ B.T
            dA = np.matmul(out.grad, np.swapaxes(other.data, -1, -2))
            dA = _unbroadcast(dA, self.data.shape)
            self.grad = dA if self.grad is None else (self.grad + dA)
        if other.requires_grad:
            # dB = A.T @ dC
            dB = np.matmul(np.swapaxes(self.data, -1, -2), out.grad)
            dB = _unbroadcast(dB, other.data.shape)
            other.grad = dB if other.grad is None else (other.grad + dB)

    out._backward = _backward
    return out

Tensor.matmul = tensor_matmul

# Test: simple 2x2 matmul
A = Tensor([[1, 2], [3, 4]], requires_grad=True)
B = Tensor([[5, 6], [7, 8]], requires_grad=True)
C = A.matmul(B)
print(f"A @ B =\n{C.data}")

# Backward pass
C.grad = np.ones_like(C.data)  # Seed gradient
C._backward()
print(f"\ndA =\n{A.grad}")
print(f"\ndB =\n{B.grad}")
A @ B =
[[19. 22.]
 [43. 50.]]

dA =
[[11. 15.]
 [11. 15.]]

dB =
[[4. 4.]
 [6. 6.]]

Matrix multiplication is the fundamental operation of neural networks. Every linear layer computes x @ W + b.

Backward Pass: Topological Sort

Just like our scalar Value, we need to traverse the graph in reverse order.

def tensor_backward(self) -> None:
    """
    Reverse-mode autodiff: topologically sort and backprop.
    """
    # Build topological order
    topo = []
    visited = set()

    def build(v: Tensor):
        if v not in visited:
            visited.add(v)
            for p in v._prev:
                build(p)
            topo.append(v)

    build(self)

    # Seed gradient and propagate
    self.grad = np.ones_like(self.data, dtype=np.float32)
    for v in reversed(topo):
        v._backward()

Tensor.backward = tensor_backward

Activation Functions

Neural networks need nonlinearities. Here are the common ones:

def tensor_relu(self) -> "Tensor":
    """ReLU: max(0, x). Gradient: 1 if x > 0, else 0."""
    out = Tensor(
        np.maximum(self.data, 0.0),
        requires_grad=self.requires_grad,
        _prev={self},
        _op="relu"
    )

    def _backward():
        if out.grad is None or not self.requires_grad:
            return
        g = out.grad * (self.data > 0.0)
        self.grad = g if self.grad is None else (self.grad + g)

    out._backward = _backward
    return out

def tensor_tanh(self) -> "Tensor":
    """Tanh: gradient is (1 - tanh^2)."""
    t = np.tanh(self.data)
    out = Tensor(t, requires_grad=self.requires_grad, _prev={self}, _op="tanh")

    def _backward():
        if out.grad is None or not self.requires_grad:
            return
        g = out.grad * (1.0 - t * t)
        self.grad = g if self.grad is None else (self.grad + g)

    out._backward = _backward
    return out

Tensor.relu = tensor_relu
Tensor.tanh = tensor_tanh

Reduction Operations

We need sum and mean for computing losses.

def tensor_sum(self, axis=None, keepdims=False) -> "Tensor":
    """Sum over axis. Gradient broadcasts back to input shape."""
    out = Tensor(
        self.data.sum(axis=axis, keepdims=keepdims),
        requires_grad=self.requires_grad,
        _prev={self},
        _op="sum"
    )

    def _backward():
        if out.grad is None or not self.requires_grad:
            return
        g = out.grad
        # Expand reduced axes back to input shape
        if axis is None:
            g = np.ones_like(self.data) * g
        else:
            axes = axis if isinstance(axis, tuple) else (axis,)
            if not keepdims:
                for ax in sorted(axes):
                    g = np.expand_dims(g, axis=ax)
            g = np.ones_like(self.data) * g
        self.grad = g if self.grad is None else (self.grad + g)

    out._backward = _backward
    return out

def tensor_mean(self, axis=None, keepdims=False) -> "Tensor":
    """Mean = sum / count."""
    if axis is None:
        denom = self.data.size
    else:
        axes = axis if isinstance(axis, tuple) else (axis,)
        denom = np.prod([self.data.shape[ax] for ax in axes])
    return self.sum(axis=axis, keepdims=keepdims) * (1.0 / float(denom))

Tensor.sum = tensor_sum
Tensor.mean = tensor_mean

# Test
x = Tensor([[1, 2, 3], [4, 5, 6]], requires_grad=True)
loss = x.mean()
print(f"Mean: {loss.data}")
loss.backward()
print(f"Gradient (each element contributes 1/6): \n{x.grad}")
Mean: 3.5
Gradient (each element contributes 1/6): 
[[0.16666667 0.16666667 0.16666667]
 [0.16666667 0.16666667 0.16666667]]

Putting It Together: A Tiny Neural Network

Use our tensor autograd to train a small network:

# Complete example: train a 2-layer network on XOR
np.random.seed(42)

# XOR dataset
X = Tensor([[0, 0], [0, 1], [1, 0], [1, 1]], requires_grad=False)
y = Tensor([[0], [1], [1], [0]], requires_grad=False)

# Weights (small random init)
W1 = Tensor(np.random.randn(2, 8) * 0.5, requires_grad=True)
b1 = Tensor(np.zeros((1, 8)), requires_grad=True)
W2 = Tensor(np.random.randn(8, 1) * 0.5, requires_grad=True)
b2 = Tensor(np.zeros((1, 1)), requires_grad=True)

params = [W1, b1, W2, b2]
lr = 0.5

for step in range(200):
    # Forward
    h = X.matmul(W1) + b1      # (4, 8)
    h = h.tanh()               # activation
    out = h.matmul(W2) + b2    # (4, 1)

    # MSE loss
    diff = out + (y * -1.0)    # out - y
    loss = (diff * diff).mean()

    # Backward
    for p in params:
        p.grad = None
    loss.backward()

    # SGD update
    for p in params:
        p.data -= lr * p.grad

    if step % 50 == 0:
        print(f"Step {step}: loss = {loss.data:.4f}")

print(f"\nFinal predictions:")
print(f"  [0,0] -> {out.data[0,0]:.3f} (target: 0)")
print(f"  [0,1] -> {out.data[1,0]:.3f} (target: 1)")
print(f"  [1,0] -> {out.data[2,0]:.3f} (target: 1)")
print(f"  [1,1] -> {out.data[3,0]:.3f} (target: 0)")
Step 0: loss = 1.0835
Step 50: loss = 0.1430
Step 100: loss = 655805535748096.0000
Step 150: loss = nan

Final predictions:
  [0,0] -> nan (target: 0)
  [0,1] -> nan (target: 1)
  [1,0] -> nan (target: 1)
  [1,1] -> nan (target: 0)
/opt/hostedtoolcache/Python/3.11.14/x64/lib/python3.11/site-packages/numpy/_core/_methods.py:49: RuntimeWarning: overflow encountered in reduce
  return umr_sum(a, axis, dtype, out, keepdims, initial, where)
/tmp/ipykernel_3006/919648715.py:36: RuntimeWarning: overflow encountered in multiply
  self.data * other.data,
/tmp/ipykernel_3006/744146672.py:24: RuntimeWarning: overflow encountered in matmul
  dA = np.matmul(out.grad, np.swapaxes(other.data, -1, -2))
/tmp/ipykernel_3006/616841562.py:27: RuntimeWarning: invalid value encountered in multiply
  g = out.grad * (1.0 - t * t)

This trained neural network uses only NumPy and 100 lines of autograd code.

PyTorch’s Tensor Autograd

The same XOR problem in PyTorch:

import torch

# Same XOR problem
X_pt = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=torch.float32)
y_pt = torch.tensor([[0], [1], [1], [0]], dtype=torch.float32)

torch.manual_seed(42)
# Note: create tensor, multiply, THEN set requires_grad to keep as leaf tensors
W1_pt = (torch.randn(2, 8) * 0.5).requires_grad_(True)
b1_pt = torch.zeros(1, 8, requires_grad=True)
W2_pt = (torch.randn(8, 1) * 0.5).requires_grad_(True)
b2_pt = torch.zeros(1, 1, requires_grad=True)

params_pt = [W1_pt, b1_pt, W2_pt, b2_pt]

for step in range(200):
    # Forward - identical logic!
    h = (X_pt @ W1_pt + b1_pt).tanh()
    out = h @ W2_pt + b2_pt
    loss = ((out - y_pt) ** 2).mean()

    # Backward - one line
    loss.backward()

    # Update - use .data to modify in-place, keeping tensors as leaves
    with torch.no_grad():
        for p in params_pt:
            p.data -= 0.5 * p.grad
            p.grad = None

    if step % 50 == 0:
        print(f"Step {step}: loss = {loss.item():.4f}")
Step 0: loss = 0.3617
Step 50: loss = nan
Step 100: loss = nan
Step 150: loss = nan

Key Insight

The logic is identical. PyTorch:

  • Handles _unbroadcast automatically
  • Uses optimized C++/CUDA kernels
  • Provides convenient loss.backward() without manual graph traversal
  • Has torch.no_grad() context for updates

You now understand what happens inside requires_grad=True and backward().

PyTorch Autograd (Scalar Examples)

See how PyTorch does the same thing with scalars:

# Create tensors that track gradients
x = torch.tensor(2.0, requires_grad=True)
w = torch.tensor(3.0, requires_grad=True)
b = torch.tensor(1.0, requires_grad=True)

# Forward pass
y = w * x + b    # Linear function: y = 3*2 + 1 = 7
loss = y ** 2    # Square loss: loss = 49

# Backward pass - PyTorch computes all gradients automatically
loss.backward()

# Derivation using chain rule:
# dloss/dy = 2*y = 2*7 = 14
# dloss/dx = dloss/dy * dy/dx = 14 * w = 14 * 3 = 42
# dloss/dw = dloss/dy * dy/dw = 14 * x = 14 * 2 = 28
# dloss/db = dloss/dy * dy/db = 14 * 1 = 14
print(f"y = {y.item():.1f}, loss = {loss.item():.1f}")
print(f"dloss/dx = {x.grad.item():.1f} (= 2*y*w = 2*7*3)")
print(f"dloss/dw = {w.grad.item():.1f} (= 2*y*x = 2*7*2)")
print(f"dloss/db = {b.grad.item():.1f} (= 2*y*1 = 2*7)")
y = 7.0, loss = 49.0
dloss/dx = 42.0 (= 2*y*w = 2*7*3)
dloss/dw = 28.0 (= 2*y*x = 2*7*2)
dloss/db = 14.0 (= 2*y*1 = 2*7)
# With vectors
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
w = torch.tensor([0.1, 0.2, 0.3], requires_grad=True)

# Forward pass
y = (x * w).sum()  # Dot product
print(f"y = x . w = {y.item():.4f}")

# Backward pass
y.backward()

print(f"\ndy/dx = {x.grad}")  # Should be w
print(f"dy/dw = {w.grad}")  # Should be x
y = x . w = 1.4000

dy/dx = tensor([0.1000, 0.2000, 0.3000])
dy/dw = tensor([1., 2., 3.])
# With matrices (like in attention)
Q = torch.randn(2, 3, requires_grad=True)  # Query
K = torch.randn(2, 3, requires_grad=True)  # Key

# Attention scores: Q @ K^T
scores = Q @ K.T
loss = scores.sum()

loss.backward()

print(f"Q shape: {Q.shape}")
print(f"K shape: {K.shape}")
print(f"scores shape: {scores.shape}")
print(f"\ndloss/dQ shape: {Q.grad.shape}")
print(f"dloss/dK shape: {K.grad.shape}")
print("\nGradients have the same shape as the original tensors!")
Q shape: torch.Size([2, 3])
K shape: torch.Size([2, 3])
scores shape: torch.Size([2, 2])

dloss/dQ shape: torch.Size([2, 3])
dloss/dK shape: torch.Size([2, 3])

Gradients have the same shape as the original tensors!

Interactive Exploration

Now that you’ve seen gradient computation in code, explore it interactively. The widget below lets you modify input values and see how gradients change in real-time — demonstrating the chain rule in action.

```{ojs}
//| echo: false

// Theme colors - uses CSS variables from _diagram-variables.scss
// Falls back to diagramTheme values when specific variables aren't available
theme = {
  // Primary colors from CSS variables
  bgPrimary: getCSSVar('--diagram-bg', diagramTheme.bg),
  textPrimary: getCSSVar('--diagram-node-text', diagramTheme.nodeText),
  textMuted: getCSSVar('--diagram-edge-stroke', diagramTheme.edgeStroke),

  // Input nodes - use accent color system
  inputBg: getCSSVar('--diagram-node-fill', diagramTheme.nodeFill),
  inputBorder: getCSSVar('--diagram-accent', diagramTheme.accent),
  inputText: getCSSVar('--diagram-accent', diagramTheme.accent),

  // Operation nodes - use highlight color system
  opBg: getCSSVar('--diagram-node-fill-alt', diagramTheme.nodeFillHover),
  opBorder: getCSSVar('--diagram-highlight', diagramTheme.highlight),
  opText: getCSSVar('--diagram-highlight', diagramTheme.highlight),

  // Output nodes - blend of node colors
  outputBg: getCSSVar('--diagram-node-fill', diagramTheme.nodeFill),
  outputBorder: diagramTheme.success,
  outputText: diagramTheme.success,

  // Gradient text - use error color for visibility
  gradientText: diagramTheme.error,

  // Arrow color
  arrowColor: getCSSVar('--diagram-edge-stroke', diagramTheme.edgeStroke)
}
```
ImportantOJS Syntax Error (line 1683, column 14)Unexpected token
TipTry This

  1. Gradient depends on values: Change a from 2 to 4. Watch ∂c/∂b change (it equals a + 2b).

  2. b has two paths: Notice ∂c/∂b is the sum of two contributions - one from the multiplication (= sum) and one via the addition (= b).

  3. Zero gradients: Set b = 0. Now ∂c/∂a = 0 because the multiplication by b kills the gradient.

  4. Negative gradients: Try negative values. Gradients can be negative, indicating the output decreases when the input increases.

  5. Verify the formula: For any a, b values: ∂c/∂a should equal b, and ∂c/∂b should equal a + 2b.

Exercises

Exercise 1: Verify the Gradient of x^3

# Verify the gradient of x^3 at x=2
# Expected: d/dx[x^3] = 3x^2 = 12

x = Value(2.0, label='x')
y = x ** 3
y.backward()

print(f"x = {x.data}")
print(f"y = x^3 = {y.data}")
print(f"dy/dx = {x.grad} (expected: 12.0)")
x = 2.0
y = x^3 = 8.0
dy/dx = 12.0 (expected: 12.0)

Exercise 2: Softmax Gradient

# Compute gradient of softmax numerator
# If y = exp(x) / (exp(x) + exp(z)), what is dy/dx?

x = Value(1.0, label='x')
z = Value(2.0, label='z')

exp_x = x.exp()
exp_z = z.exp()
y = exp_x / (exp_x + exp_z)

y.backward()

print(f"x = {x.data}, z = {z.data}")
print(f"y = softmax(x)[0] = {y.data:.4f}")
print(f"dy/dx = {x.grad:.4f}")
print(f"\nThis is the gradient that flows back through softmax!")
x = 1.0, z = 2.0
y = softmax(x)[0] = 0.2689
dy/dx = 0.1966

This is the gradient that flows back through softmax!

Exercise 3: ReLU vs Tanh Gradients

# Compare ReLU vs tanh gradients

for x_val in [-2.0, -0.5, 0.5, 2.0]:
    x_relu = Value(x_val)
    x_tanh = Value(x_val)

    y_relu = x_relu.relu()
    y_tanh = x_tanh.tanh()

    y_relu.backward()
    y_tanh.backward()

    print(f"x={x_val:5.1f}: ReLU grad={x_relu.grad:.3f}, tanh grad={x_tanh.grad:.3f}")

print("\nNotice: ReLU has 0 or 1, tanh is smooth but saturates for large |x|")
x= -2.0: ReLU grad=0.000, tanh grad=0.071
x= -0.5: ReLU grad=0.000, tanh grad=0.786
x=  0.5: ReLU grad=1.000, tanh grad=0.786
x=  2.0: ReLU grad=1.000, tanh grad=0.071

Notice: ReLU has 0 or 1, tanh is smooth but saturates for large |x|

Backpropagation in Neural Networks

Gradients flow backward through a neural network layer:

Key insights:

  1. Gradients flow backward: Starting from the loss, we trace back through every operation
  2. Chain rule connects layers: Multiply local gradients along the path
  3. Accumulation: If a value is used multiple times, gradients add up
  4. Shape matters: dL/dW must have the same shape as W for the update

Detaching Tensors and Stopping Gradients

Sometimes you want to stop gradient flow. PyTorch provides several mechanisms:

Using detach()

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x * 2
z = y.detach()  # z has no gradient history

print(f"y.requires_grad: {y.requires_grad}")
print(f"z.requires_grad: {z.requires_grad}")

# z is now a "constant" - no gradients flow through it
loss = z.sum()
# loss.backward() would NOT compute gradients for x
y.requires_grad: True
z.requires_grad: False

Using no_grad()

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x * 2

# Compute without building graph (saves memory)
with torch.no_grad():
    z = y * 3
    print(f"Inside no_grad, z.requires_grad: {z.requires_grad}")

# Common use: inference
model_output = y  # pretend this is model output
with torch.no_grad():
    prediction = model_output.argmax()  # no gradients needed
Inside no_grad, z.requires_grad: False

When to Stop Gradients

  • Inference: No training, no gradients needed
  • Frozen layers: Transfer learning with some layers fixed
  • Metrics: Computing accuracy, loss for logging (not backprop)
  • Target values: In losses like MSE, the target is a constant

Memory Implications

Why Autograd Uses Memory

During the forward pass, autograd stores intermediate values needed for gradient computation:

# Each operation stores data for backward pass
x = torch.randn(1000, 1000, requires_grad=True)
y = x @ x.T          # Stores x for gradient computation
z = y.relu()         # Stores y (to know which elements > 0)
out = z.mean()       # Stores z

# All of x, y, z stay in memory until backward() completes
out.backward()

# Now intermediate tensors can be freed

Memory Grows with Depth

For a model with L layers: - Forward pass: O(L) memory for activations - Each activation tensor can be large (batch_size x hidden_dim) - LLMs with 100+ layers and large hidden dims = huge memory

Gradient Checkpointing

Trade compute for memory by recomputing activations during backward:

from torch.utils.checkpoint import checkpoint

def expensive_layer(x):
    """A layer we want to checkpoint."""
    return x.relu().pow(2)

x = torch.randn(100, 100, requires_grad=True)

# Without checkpointing: stores intermediate activations
y_normal = expensive_layer(x)

# With checkpointing: discards intermediates, recomputes during backward
y_checkpoint = checkpoint(expensive_layer, x, use_reentrant=False)

print(f"Both produce same result: {torch.allclose(y_normal, y_checkpoint)}")
Both produce same result: True

In practice, checkpoint every few layers to reduce memory by ~sqrt(L).

Common Pitfalls

1. Forgetting to Zero Gradients

x = torch.tensor(2.0, requires_grad=True)

# First backward
y = x ** 2
y.backward()
print(f"After first backward: x.grad = {x.grad}")

# Second backward without zeroing - gradients ACCUMULATE!
y = x ** 2
y.backward()
print(f"After second backward: x.grad = {x.grad}")  # 8.0, not 4.0!

# Always zero gradients before backward in training loops
x.grad.zero_()
y = x ** 2
y.backward()
print(f"After zeroing: x.grad = {x.grad}")  # Back to 4.0
After first backward: x.grad = 4.0
After second backward: x.grad = 8.0
After zeroing: x.grad = 4.0

2. In-place Operations

x = torch.tensor([1.0, 2.0], requires_grad=True)
y = x * 2

# This would break the graph (commented to avoid error):
# y.add_(1)  # In-place operation on a tensor needed for gradient

# Instead, use out-of-place operations:
y = y + 1  # Creates new tensor, graph preserved
y.sum().backward()
print(f"x.grad: {x.grad}")
x.grad: tensor([2., 2.])

3. Losing requires_grad

x = torch.tensor([1.0, 2.0], requires_grad=True)

# Can't call .numpy() directly on a tensor requiring grad:
# y = x.numpy()  # RuntimeError: Can't call numpy() on tensor that requires grad

# Must detach first - this explicitly breaks the graph
y = x.detach().numpy()
print(f"y is now numpy array: {type(y)}")

# Converting back loses gradient tracking:
z = torch.from_numpy(y)
print(f"z.requires_grad: {z.requires_grad}")  # False

# Be careful when mixing numpy and autograd
y is now numpy array: <class 'numpy.ndarray'>
z.requires_grad: False

4. Backward Through Non-Scalar

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x ** 2  # Vector output

# This fails (commented to avoid error):
# y.backward()  # RuntimeError: need gradient argument

# For non-scalar outputs, provide a gradient tensor:
y.backward(torch.ones_like(y))  # Equivalent to y.sum().backward()
print(f"x.grad: {x.grad}")
x.grad: tensor([2., 4., 6.])

Summary

Key takeaways:

  1. Gradients flow backward through the computational graph
  2. Chain rule connects everything: multiply local gradients
  3. Gradients accumulate when a value is used multiple times
  4. Training = gradient descent: use gradients to adjust parameters
  5. PyTorch autograd does this automatically for any computation
  6. Memory matters: intermediate activations consume memory; use detach(), no_grad(), or gradient checkpointing
  7. Zero gradients: always zero gradients before each backward pass in training loops

What’s Next

Module 03: Tokenization covers converting text into numbers that our model can process. Autograd works behind the scenes during training, but we need input data first!