3. How networks learn: gradients, autograd, optimizers

Why an inference book teaches training

You’ll mostly run models in this book. But every weight you load was produced by the process in this chapter, and Parts II and V make you train and fine-tune models yourself. More immediately, the backward-pass machinery you build here is exactly what PyTorch runs when you call loss.backward(). Engine builders who know it can read memory reports, debug fine-tuning, and understand why training needs so much more memory than inference.

Learning is adjusting knobs to reduce a number

A neural network is a function with adjustable numbers inside it, called parameters or weights. Feed it an input, and it produces an output. A loss function turns “how wrong was that output?” into a single number. Learning means changing the parameters to make that number smaller.

The smallest possible example has one parameter w. It predicts w * x, and we want the prediction for x = 3 to be 8. The squared-error loss is

$$ L(w) = (w \cdot x - \text{target})^2 = (3w - 8)^2 . $$

At w = 2, the prediction is 6 and the loss is $(6-8)^2 = 4$. Should we increase or decrease w, and by how much?

The derivative says which way is downhill

The derivative $dL/dw$ is the slope of the loss as a function of w: how much the loss changes per small change in w. For our loss,

$$ \frac{dL}{dw} = 2x(wx - \text{target}) = 2 \cdot 3 \cdot (6 - 8) = -12 . $$

The slope is negative, so increasing w decreases the loss. Gradient descent takes a small step against the slope:

$$ w \leftarrow w - \eta \frac{dL}{dw} $$

where $\eta$ (eta) is the learning rate. With $\eta = 0.01$, the first steps are:

The loss falls, and the steps get smaller as the slope flattens near the minimum at w = 8/3. Try a learning rate twelve times larger ($\eta = 0.12$) and the loss grows: w goes 2.0 → 3.44 → 1.77 → 3.71 and the loss goes 4.0 → 5.38 → 7.24 → 9.75, because every step overshoots. Choosing the learning rate is the first hyperparameter you’ll tune.

Real networks have millions or billions of parameters. The gradient is the vector of all their partial derivatives, $\partial L / \partial w_i$, one per parameter, and gradient descent updates all of them at once. The idea is the same; only the bookkeeping grows.

The chain rule, and computing gradients mechanically

Nobody derives gradients for a billion-parameter model by hand. A network is built from simple operations (add, multiply, exp, …), and each one knows its own local derivative. The chain rule combines them: if $L$ depends on $d$, and $d$ depends on $a$, then

$$ \frac{\partial L}{\partial a} = \frac{\partial L}{\partial d}\cdot\frac{\partial d}{\partial a}. $$

Here’s a small expression worked through by hand. Let $a = 2$, $b = -3$, $c = 10$, $f = -2$, and compute

$$ e = a \cdot b = -6,\qquad d = e + c = 4,\qquad L = d \cdot f = -8 . $$

Work backwards from $L$, multiplying local derivatives as you go:

That procedure is backpropagation: a forward pass computes values and remembers how each was made, then a backward pass visits every operation once, in reverse order, multiplying the incoming gradient by the local derivative and passing it on. The cost of the backward pass is about twice the forward pass, regardless of the number of parameters. That’s why training billion-parameter models is possible at all.

Two details matter when you implement it:

1. Order. A node’s gradient is complete only after every node that uses it has passed its contribution back. Sorting the graph topologically and walking it in reverse guarantees this.
2. Accumulation. If a value is used twice (like x in x * x + x), contributions from each use add up. Gradients are accumulated with +=, never assigned with =.

Build an autograd engine

The Value class below wraps one number, remembers the values it was computed from, and stores a small closure that applies its local chain-rule step. This design follows Andrej Karpathy’s micrograd. Each operation creates the output, then defines how gradient flows from the output to its inputs:

class Value:
    def __init__(self, data, children=(), op=""):
        self.data, self.grad = float(data), 0.0
        self._children, self._op = tuple(children), op
        self._backward = lambda: None

    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(ab)/da = b
            other.grad += self.data * out.grad     # d(ab)/db = a
        out._backward = backward
        return out

    def backward(self):
        order, seen = [], set()
        def visit(node):                           # topological sort
            if id(node) not in seen:
                seen.add(id(node))
                for child in node._children:
                    visit(child)
                order.append(node)
        visit(self)
        self.grad = 1.0
        for node in reversed(order):
            node._backward()

// A scalar autograd node. shared_ptr because one value can feed many later operations.
struct Node {
    double data, grad = 0;
    std::vector<std::shared_ptr<Node>> children;
    std::function<void(Node&)> backward = [](Node&) {};
    explicit Node(double d) : data(d) {}
};
using Value = std::shared_ptr<Node>;

inline Value make(double d) { return std::make_shared<Node>(d); }
inline Value add(Value a, Value b) {
    auto out = make(a->data + b->data);
    out->children = {a, b};
    out->backward = [a, b](Node& self) { a->grad += self.grad; b->grad += self.grad; };
    return out;
}
inline Value mul(Value a, Value b) {
    auto out = make(a->data * b->data);
    out->children = {a, b};
    out->backward = [a, b](Node& self) { a->grad += b->data * self.grad; b->grad += a->data * self.grad; };
    return out;
}
inline Value power(Value a, double n) {
    auto out = make(std::pow(a->data, n));
    out->children = {a};
    out->backward = [a, n](Node& self) { a->grad += n * std::pow(a->data, n - 1) * self.grad; };
    return out;
}
inline void backward(const Value& root) {
    std::vector<Value> order;
    std::set<Node*> seen;
    std::function<void(const Value&)> visit = [&](const Value& v) {
        if (!seen.insert(v.get()).second) return;
        for (auto& c : v->children) visit(c);
        order.push_back(v);
    };
    visit(root);
    root->grad = 1;
    for (auto it = order.rbegin(); it != order.rend(); ++it) (*it)->backward(**it);
}

#![allow(unused)]
fn main() {
#[derive(Clone)]
pub struct Value(Rc<RefCell<Node>>);

struct Node {
    data: f64,
    grad: f64,
    children: Vec<Value>,
    // Given this node's gradient, add the chain-rule contribution to each child.
    backward: Option<Box<dyn Fn(f64, &[Value])>>,
}

impl Value {
    pub fn new(data: f64) -> Self {
        Value(Rc::new(RefCell::new(Node { data, grad: 0.0, children: vec![], backward: None })))
    }
    fn op(data: f64, children: Vec<Value>, backward: impl Fn(f64, &[Value]) + 'static) -> Self {
        Value(Rc::new(RefCell::new(Node { data, grad: 0.0, children, backward: Some(Box::new(backward)) })))
    }
    pub fn data(&self) -> f64 { self.0.borrow().data }
    pub fn grad(&self) -> f64 { self.0.borrow().grad }
    fn add_grad(&self, g: f64) { self.0.borrow_mut().grad += g; }

    pub fn add(&self, other: &Value) -> Value {
        Value::op(self.data() + other.data(), vec![self.clone(), other.clone()], |g, c| {
            c[0].add_grad(g);
            c[1].add_grad(g);
        })
    }
    pub fn mul(&self, other: &Value) -> Value {
        Value::op(self.data() * other.data(), vec![self.clone(), other.clone()], |g, c| {
            let (a, b) = (c[0].data(), c[1].data());
            c[0].add_grad(b * g);
            c[1].add_grad(a * g);
        })
    }
    pub fn powf(&self, n: f64) -> Value {
        Value::op(self.data().powf(n), vec![self.clone()], move |g, c| {
            c[0].add_grad(n * c[0].data().powf(n - 1.0) * g);
        })
    }
    pub fn tanh(&self) -> Value {
        let t = self.data().tanh();
        Value::op(t, vec![self.clone()], move |g, c| c[0].add_grad((1.0 - t * t) * g))
    }

    /// Visit every node after all nodes that use it, then apply each local rule once.
    pub fn backward(&self) {
        let mut order = vec![];
        let mut seen = HashSet::new();
        fn visit(v: &Value, seen: &mut HashSet<usize>, order: &mut Vec<Value>) {
            if seen.insert(Rc::as_ptr(&v.0) as usize) {
                for child in &v.0.borrow().children {
                    visit(child, seen, order);
                }
                order.push(v.clone());
            }
        }
        visit(self, &mut seen, &mut order);
        self.0.borrow_mut().grad = 1.0;
        for v in order.iter().rev() {
            let node = v.0.borrow();
            if let Some(f) = &node.backward {
                f(node.grad, &node.children);
            }
        }
    }
}
}

With it, the worked example above becomes:

from engine.autograd import Value
a, b, c, f = Value(2.0), Value(-3.0), Value(10.0), Value(-2.0)
L = (a * b + c) * f
L.backward()
print(L.data, a.grad, b.grad, c.grad, f.grad)    # -8.0 6.0 -4.0 -2.0 4.0

Neurons, layers and why nonlinearity matters

A neuron computes a weighted sum of its inputs plus a bias, then applies a nonlinear function: $\tanh(w \cdot x + b)$. A layer is many neurons reading the same inputs, which is exactly the linear layer from Chapter 2 followed by an elementwise nonlinearity. A multilayer perceptron (MLP) stacks layers.

Why the nonlinearity? Without it, two linear layers collapse into one: $W_2(W_1 x) = (W_2 W_1)x$, a single matrix. You can stack a hundred linear layers and still only represent linear functions, which can’t even compute XOR. The nonlinearity between layers lets the network bend space. A classic demonstration is XOR, where the output is high when exactly one input is on:

No straight line separates the +1s from the −1s, but a 2-input MLP with two hidden tanh layers of 8 neurons learns it in a few dozen steps. This is the output of python run.py autograd --steps 200 with the reference engine:

{"step": 1, "loss": 5.63926}
{"step": 26, "loss": 2.49858}
{"step": 51, "loss": 0.00013}
{"step": 200, "loss": 0.0}
{"predictions": [-1.0, 1.0, 1.0, -1.0], "targets": [-1.0, 1.0, 1.0, -1.0]}

Transformer MLPs (Chapter 6) are exactly this structure, scaled up to thousands of neurons per layer, with smoother nonlinearities (GELU, SiLU) than tanh.

From scalars to tensors: PyTorch autograd

Your Value engine works on single numbers. PyTorch’s autograd does the same thing on whole tensors, so one node represents a million multiplications, and the local backward rules are themselves tensor operations running on the GPU. The interface is the same idea:

import torch
w = torch.tensor(2.0, requires_grad=True)   # track operations on w
loss = (w * 3.0 - 8.0) ** 2
loss.backward()                             # fills w.grad
print(w.grad)                               # tensor(-12.)

Three switches control it, and they’re easy to confuse:

• requires_grad=True on a tensor (parameters have it by default) means “record operations involving me”.
• torch.no_grad() and torch.inference_mode() stop recording. Every inference path in your engine runs under one of them. Recording costs memory, because each operation keeps its inputs alive for the backward pass.
• model.train() / model.eval() switch layer behavior (dropout on or off). They do not turn gradient recording on or off. That’s a common misconception.

Warning

Gradients accumulate across backward() calls, by design (it lets you sum gradients over several small batches). Forgetting optimizer.zero_grad() makes every step use the sum of all previous gradients. The loss usually explodes after a few steps.

The training loop

Every training script you’ll ever read has the same five lines at its core:

for batch in data:
    loss = loss_fn(model(batch.x), batch.y)   # 1. forward
    optimizer.zero_grad()                     # 2. clear old gradients
    loss.backward()                           # 3. backward: fill .grad for every parameter
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)   # 4. (optional) cap the gradient size
    optimizer.step()                          # 5. update the parameters

The optimizer decides how to turn gradients into parameter updates.

Optimizers: SGD, momentum, Adam, AdamW

SGD (stochastic gradient descent) is the update rule from earlier, applied to gradients estimated on a random batch of data. It works, but one learning rate must suit every parameter, and noisy gradients make it zigzag.

Momentum keeps a running average of recent gradients and steps along it, which smooths the noise:

$$ m \leftarrow \beta m + (1-\beta) g, \qquad w \leftarrow w - \eta, m . $$

Adam also tracks a running average of squared gradients, $v$, and divides each parameter’s step by $\sqrt{v}$. Parameters with consistently large gradients take smaller steps, and rarely updated parameters take relatively larger ones:

$$ m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t,\quad v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2,\quad w \leftarrow w - \eta,\frac{\hat m_t}{\sqrt{\hat v_t} + \epsilon} $$

where $\hat m_t = m_t/(1-\beta_1^t)$ and $\hat v_t = v_t/(1-\beta_2^t)$ correct for starting both averages at zero. Here’s a worked first step with $g_1 = 2$, $\beta_1 = 0.9$, $\beta_2 = 0.999$: $m_1 = 0.2$ and $v_1 = 0.004$, and after correction $\hat m = 2$ and $\hat v = 4$. The step is $\eta \cdot 2/\sqrt 4 = \eta$. Adam’s first step has size equal to the learning rate, whatever the gradient’s scale. (PyTorch agrees: with $\eta = 0.1$, one Adam step moves the weight from 0 to −0.1.)

AdamW adds weight decay, a gentle pull of every weight toward zero, applied separately from the adaptive gradient step: $w \leftarrow (1 - \eta\lambda)w - \eta,\hat m/(\sqrt{\hat v} + \epsilon)$. It’s the default optimizer for transformers.

Two facts matter for engine builders:

1. Optimizer state is big. AdamW stores $m$ and $v$ for every parameter, usually in FP32. Training a 0.6B model needs the weights (1.2 GB in BF16), plus gradients (another 1.2 GB), plus $m$ and $v$ (4.8 GB), plus FP32 master weights in mixed precision (2.4 GB), plus activations. Inference needs only the 1.2 GB. Chapter 22 shows how LoRA avoids most of that.
2. Schedules. Learning rates usually ramp up from near zero for the first few hundred steps (warmup), then decay (often along a cosine curve). Your training loop in Chapter 7 does both.

Build it

pytest tests/test_ch03_autograd.py
python run.py autograd --impl engine      # trains XOR with YOUR engine

Tip

If a gradient is wrong, compare it with the finite-difference estimate $\big(L(w+\epsilon) - L(w-\epsilon)\big)/2\epsilon$ (numerical_gradient in the module). It’s slow but independent of your backward code, which makes it the most useful debugging tool for gradients.

Stretch exercises

1. ★ Add sigmoid to Value. Its derivative is $\sigma(x)(1-\sigma(x))$. Verify it with numerical_gradient. Where: add Value.sigmoid in engine/autograd.py.
2. ★★ Train the XOR MLP with learning rates 0.005, 0.05 and 0.5. Plot the loss curves. Which converges, which is slow, and which oscillates? Where: experiments/ch03.py (create it), adapting run.py’s cmd_autograd.
3. ★★ Implement momentum and Adam for lists of Values and compare the steps-to-convergence on XOR with plain SGD. Where: add optimizer helpers beside sgd_step in engine/autograd.py; compare them in experiments/ch03.py (create it).
4. ★★★ Re-implement the XOR network in PyTorch with nn.Linear and torch.optim.AdamW. Then count the parameters, gradients and optimizer state in bytes for a 1B-parameter model trained this way. Where: experiments/ch03.py (create it).

Check your understanding

1. Why does gradient descent move against the gradient?
2. Why must gradients be accumulated with += in the backward pass?
3. What is lost if you remove every nonlinearity from an MLP?
4. Why is model.eval() not a substitute for torch.no_grad()?
5. Roughly how many bytes does AdamW training need per parameter, compared with BF16 inference?

Going deeper

• BALLM Appendix A §§A.3-A.7 (pp. 258-277): computation graphs, autograd and training loops in PyTorch. Appendix D (pp. 313-321): warmup, cosine decay and gradient clipping.
• Andrej Karpathy, The spelled-out intro to neural networks and backpropagation: building micrograd (video and code), the inspiration for this chapter’s engine.
• Kingma and Ba, Adam: A Method for Stochastic Optimization (2015); Loshchilov and Hutter, Decoupled Weight Decay Regularization (2019), the paper behind AdamW.
• GPU Mode L6 (Jane Xu, Optimizing PyTorch Optimizers): how optimizer steps are fused into a few kernels.