Roadmap

1. Automatic differentiation, the idea under every framework
2. PyTorch and JAX, two designs for the same mathematics
3. Tensors: shapes, dtypes, indices, and a first story
4. Models that looked fine: two more stories
5. The training loop, and what Adam's knobs do
6. Accelerators, and when they lose

Automatic differentiation

• Training needs for every parameter, every step
• Millions of parameters, one scalar loss
• Three ways to get a derivative:
  • Symbolic algebra: exact, but expressions blow up
  • Finite differences: 2 evaluations per parameter, and approximate
  • Automatic differentiation: exact, about 2 forward passes total

Automatic differentiation, a network small enough to do by hand

One hidden layer. Data: input , target . Parameters: . Loss: squared error.
Chain rule, one step at a time, each reusing the last:

is the derivative of , and it needs from the forward pass

Automatic differentiation, the graph

Automatic differentiation, reverse and forward mode

Training is reverse mode. A 3-variable simulator with 1,000 outputs is forward mode.

Automatic differentiation, against finite differences

• Best finite difference over 40 steps: 1.4 × 10⁻¹², and the best step is unknowable in advance
• It also costs 2 evaluations per parameter: use it only to spot-check autodiff

PyTorch and JAX, two designs

Same numbers either way: the two agree to 4.4 × 10⁻¹⁶, float64 epsilon.

Docs: autograd mechanics, JAX key concepts

PyTorch and JAX, the PyTorch tape

Read off a real loss.grad_fn. Rebuilt from scratch on every call: define-by-run.

PyTorch and JAX, .grad accumulates

loss = loss_fn(model(x), y)
loss.backward()                  # AccumulateGrad: W.grad += dL/dW

• The leaves of the tape add into .grad; they do not overwrite it
• Deliberate: split a big batch into micro-batches, backward() each, step once
• The cost: every loop must clear it, optimizer.zero_grad()

PyTorch and JAX, a question

PyTorch and JAX, why switch the tape off

The tape keeps every intermediate until backward() uses it. No backward(), no release.

total = 0
for x, y in val_loader:
    loss = loss_fn(model(x), y)
    total += loss          # keeps this batch's whole graph alive

PyTorch and JAX, when to switch the tape off

torch.inference_mode() is a stricter, slightly faster no_grad() for serving.

PyTorch and JAX, switching the tape off

model.eval()                 # dropout off, batch norm uses running statistics
with torch.no_grad():        # record nothing
    total = sum(loss_fn(model(x), y).item() for x, y in val_loader)
model.train()                # back to training behavior

• eval() and no_grad() are different switches: layers versus the tape. Evaluation needs both
• Forget eval(): a validation score another script cannot reproduce
• JAX: nothing is recorded outside jax.grad, and arrays are immutable, so only freezing needs code: jax.lax.stop_gradient(x)

PyTorch and JAX, a JAX gradient is a function

def loss(params, x, y):
    a = jnp.tanh(params["W"] @ x + params["b"])
    return (params["v"] @ a + params["c"] - y) ** 2

grads = jax.grad(loss)(params, x, y)   # same structure as params

• No tape, no .grad, nothing to zero
• jax.grad(jax.grad(f)) differentiates twice; jax.jit compiles it through XLA

PyTorch and JAX, the jaxpr

g = dot_general a e      # W @ x
h = add g b              # + b
i = tanh h
j = dot_general d i      # v · a
k = add j c
l = sub k f              # - y
m = integer_pow[y=2] l

• jax.make_jaxpr(loss) prints this: 7 equations, the same 7 nodes as PyTorch's tape
• jax.grad rewrites it into a 16-equation program: forward and backward together

Docs: Understanding jaxprs

PyTorch and JAX, limits of tracing

• JAX calls your function once with placeholders and writes down one path
• A Python if on an array's value cannot be traced: use jax.lax.cond
• A data-dependent loop: jax.lax.while_loop
• The function must be pure: a print or list append runs once, at trace time
• Arrays are immutable: x[0] = 1.0 is a TypeError, use x = x.at[0].set(1.0)

PyTorch pays none of this. Define-by-run means any Python works.

PyTorch and JAX, vmap and jit

vmap: write it for one example, map the batch axis. Agrees with a Python loop over 256 examples to 2.2 × 10⁻¹⁵.

Compilation pays where there are many small operations to fuse.

PyTorch and JAX, side by side

The gap is now mostly defaults: PyTorch starts eager and opts in; JAX starts functional.

Tensors, what a tensor is

• Batch first: (N, features), (N, channels, time)
• Broadcasting follows NumPy: trailing dimensions line up, size 1 stretches
• Broadcasting raises no error when shapes are compatible, even when the result is not what you meant

Tensors, a question

Tensors, reading a shape error

model = nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 1))
model[2](x)    # x: 32 mixes x 8 features, but the first layer was skipped

mat1 and mat2 shapes cannot be multiplied (32x8 and 64x1)

Fix the model, not the input: a skipped layer, or a stale in_features.

Tensors, default dtypes disagree

• Mixing float32 with float64 gives float64, so every dtype downstream reads float64
• JAX: jax.config.update("jax_enable_x64", True) before any array exists

Tensors, JAX will not raise

• Out-of-bounds read clamps: jnp.arange(3.0)[10] returns 2.0
• Out-of-bounds write .at[10].set(v) is silently dropped
• Raising from compiled accelerator code is expensive, so JAX does not

Read before writing any JAX: the Sharp Bits

Story 1, the gradient that disagreed

A story: code that ran, a number that looked fine, then the cause. Three today.

How it looked

• Autodiff checked against a hand-derived gradient, in float64
• JAX agrees to 4.4 × 10⁻¹⁶
• PyTorch disagrees at 7.5 × 10⁻⁸
• Every tensor.dtype printed: float64
• Conclusion drafted: "PyTorch's autograd is less precise"

Story 1, a question

Story 1, debugged

• Two scalars were Python floats: 1.234, and a bare rng.normal()
• torch.tensor stored both as float32
• The fingerprint: relative error 1.25 × 10⁻⁷, float32 epsilon to two figures
• Fixed, dtype=torch.float64 on both: 1.7 × 10⁻¹⁸

Check dtypes where data enters, not after. An error near 10⁻⁷ means float32.

Models that looked fine, the setup

Predict concrete strength from the mix and its age: 1,030 rows, 8 inputs, a crush test that takes weeks

• 428 mixes; 76% of rows belong to a mix tested at more than one age
• Each story, like story 1: the code as written, the number, why it looked fine, the cause, the fix

Story 2, the tree that won

folds = KFold(5, shuffle=True, random_state=0).split(X)
for seed in range(5):
    for tr, va in folds:
        tree = HistGradientBoostingRegressor(random_state=seed).fit(X[tr], y[tr])
        net  = train(X[tr], y[tr], X[va], y[va], seed=seed)

• Gradient boosting 4.44 MPa, MLP 4.87: the tree wins by 0.43 ± 0.09, nearly five standard errors
• Five seeds, five folds, matches the received wisdom. Where is the bug?

Story 2, debugged

• Under KFold, 75% of validation rows share a mix with training: the same concrete at another age
• GroupKFold by mix: tree 6.25, MLP 6.48, a tie; the leak was worth 1.81 MPa to the tree, 1.61 to the net

Story 2, what about the leaky scaler?

The textbook leak: StandardScaler fitted on all rows before splitting (Lecture 7)

• Eight means and eight standard deviations barely move. Still a bug
• The leak that changed the conclusion was the split
• Grinsztajn et al., 45 datasets: trees do lead on tabular data. Check the split before crediting the model family

Story 3, the noisy run

for xb, yb in loader:
    loss = loss_fn(model(xb), yb)
    loss.backward()
    opt.step()

• No NaN, no exception; validation RMSE wanders between 10 and 17 MPa, under the 17.9 baseline
• Reads as "noisy, lower the learning rate". Where is the bug?

Story 3, debugged

No opt.zero_grad(): .grad sums every past step, so each update follows the sum of all past gradients. Ends at 23.2 MPa, worse than the mean. Fixed: 5.5.

Models that looked fine, what caught each

None raised an error. All four models, rerun and taken apart: l11-four-models.ipynb

The training loop, PyTorch

for epoch in range(n_epochs):
    model.train()
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        loss = loss_fn(model(xb), yb)   # forward
        optimizer.zero_grad()           # clear the accumulator
        loss.backward()                 # backward
        optimizer.step()                # update

• One optimizer step per mini-batch; an epoch is one pass over the rows
• Four objects: DataLoader, nn.Module, loss function, optimizer. Story 3 was one missing line here

The training loop, JAX and optax

@jax.jit
def step(params, opt_state, xb, yb):
    loss, grads = jax.value_and_grad(loss_fn)(params, xb, yb)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, loss

• No zero_grad, no .to(device): state goes in and comes out
• Same network, same fold, 5 seeds: PyTorch 5.70 ± 0.29, JAX 5.66 ± 0.24 MPa

The training loop, Adam

gradient, learning rate, and running averages (this is Adam's, not the network's). Kingma and Ba, 2015

The training loop, what Adam's knobs do

Narrow bowl, 40× steeper in . Adam's first step is in both coordinates; SGD's are 0.18 and 2.7.

The training loop, and

Gradient drops 100× at step 200: with the step collapses about 60× for hundreds of steps. Large turns Adam back into SGD.

The training loop, defaults that differ

• AdamW decays the weights outside the adaptive scaling (Loshchilov and Hutter)
• Port with defaults and the regularization changes 100×, silently

The training loop, the learning rate

SGD on the same fold, 120 epochs:

• At lr = 1.0, six seeds: nan in two, 91 to 1,235 MPa in the other four
• Stable is a property of this initialization. One surviving run proves little

Accelerators, the rules

• PyTorch: model.to(device), x.to(device); mismatch is a loud error, the good case
• JAX: arrays go to the default device; jax.devices() shows it
• .cpu() or print inside the loop forces a sync and serializes everything
• GPU calls return when queued: time with torch.cuda.synchronize() or x.block_until_ready()
• A GPU has a fixed cost per kernel launch; it pays only with enough arithmetic behind each launch

Accelerators, a question

Accelerators, the crossover

Today's model: 8.6 ms per epoch on the CPU, 21.1 ms on the GPU. The GPU wins past a few hundred hidden units.

Trade-offs

And on 1,030 rows of concrete, gradient boosting ties the MLP in under a second, with no learning rate.

Recap

• Reverse-mode AD: every gradient, exactly, for about two forward passes
• PyTorch records a tape and accumulates into .grad; JAX transforms pure functions
• Default float32 in both; check dtypes where data enters
• Three stories, no errors raised; each was caught by a comparison: an exact reference, a grouped split, a correct loop
• Adam gives each parameter its own step size; it tolerates unscaled inputs, and so hides them
• The GPU is 2.5× slower on today's model; measure before you migrate

Standings

Nicknames only. Everyone who skipped one still counted in every bar you saw.

Before you go

Practice module for this session, for participation credit
Reading PyTorch: Learn the Basics, JAX Sharp Bits, Grinsztajn et al. 2022

Full notes, with all sources: lectures/l11/notes.md