Automatic Differentiation
The Problem with Manual Gradients
In the previous modules, you computed gradients by hand. For simple functions like , that’s easy. But what about this? That’s a 2-layer neural network. Now imagine 100 layers. Nobody computes these gradients by hand. Automatic differentiation (autodiff) does it for you - and that’s what powers PyTorch and TensorFlow.Difficulty: Intermediate
Prerequisites: Chain Rule module
What You’ll Build: Your own mini autodiff system!
Three Ways to Compute Gradients
Think of computing derivatives as three different approaches to getting driving directions:- Symbolic differentiation is like working out the route with pen and paper, applying road rules algebraically. You get a perfect formula, but it gets impossibly complex for a cross-country trip with thousands of turns.
- Numerical differentiation is like driving the route twice with slightly different starting positions and comparing where you end up. Simple but slow (you literally run the function twice per parameter), and small measurement errors accumulate.
- Automatic differentiation is like having a co-pilot who records every turn you make and can replay the route backward, noting exactly which turns contributed to going north vs. south. One forward trip plus one backward replay, and you know the derivative with respect to every single starting condition.
The Core Idea: Computational Graphs
Every computation can be broken into a graph of simple operations.Example:
Backward Pass: Chain Rule Through the Graph
Build Your Own Autodiff System
Let’s implement a simple autodiff system from scratch!Forward Mode vs Reverse Mode
There are two ways to do autodiff:PyTorch: Autodiff in Action
Common Pitfalls
1. Gradient Accumulation
2. Detaching from Graph
3. In-place Operations
4. Numerical Precision in Gradient Checking
Visualizing Computation Graphs
Practice Exercises
Exercise 1: Add More Operations
Exercise 1: Add More Operations
Value class with:tanh()activation functionexp()function- Division operator
Exercise 2: Softmax Gradient
Exercise 2: Softmax Gradient
Exercise 3: Implement from Scratch
Exercise 3: Implement from Scratch
Value class:- Define network with random weights
- Forward pass on XOR data
- Compute MSE loss
- Backward pass
- Update weights with gradient descent
Summary
Interview Deep-Dive
Explain the difference between symbolic differentiation, numerical differentiation, and automatic differentiation. Why did deep learning frameworks choose autodiff over the other two?
Explain the difference between symbolic differentiation, numerical differentiation, and automatic differentiation. Why did deep learning frameworks choose autodiff over the other two?
- Symbolic differentiation applies algebraic rules to produce an exact derivative formula. For f(x) = x^2 * sin(x), it gives f’(x) = 2xsin(x) + x^2cos(x). The problem is expression swell: for complex compositions, the symbolic derivative can be exponentially larger than the original expression. A 10-layer neural network’s loss function, symbolically differentiated, would produce an unmanageably large expression that is slow to evaluate even if you could derive it.
- Numerical differentiation uses finite differences: (f(x+h) - f(x-h))/(2h). It is simple and works for any function you can evaluate. But it has two fatal flaws for deep learning: it requires O(n) function evaluations for n parameters (one per parameter), and it suffers from the truncation-cancellation trade-off where no choice of h gives both accuracy and stability simultaneously.
- Automatic differentiation computes exact derivatives (to machine precision) at a cost proportional to the original function evaluation. It works by decomposing the computation into elementary operations and applying the chain rule through them. Reverse-mode autodiff (backpropagation) computes the gradient of a scalar output with respect to ALL inputs in a single backward pass, regardless of the number of inputs.
- Deep learning chose autodiff because it is the only method that scales. With millions of parameters and a scalar loss, reverse-mode autodiff gives exact gradients in O(1) backward passes. Symbolic differentiation would produce an expression too large to store. Numerical differentiation would require millions of forward passes per gradient step. The cost ratio is not 2x or 10x — it is millions-to-one. This is literally what makes deep learning computationally feasible.
PyTorch uses dynamic computational graphs while TensorFlow 1.x used static graphs. What are the trade-offs, and why did the industry converge toward dynamic graphs?
PyTorch uses dynamic computational graphs while TensorFlow 1.x used static graphs. What are the trade-offs, and why did the industry converge toward dynamic graphs?
- In a static graph system (TensorFlow 1.x), you first define the entire computation as a graph object, then execute it in a separate “session.” The graph is compiled once and reused. This enables aggressive ahead-of-time optimizations: operation fusion, memory planning, dead code elimination, and cross-device placement. The downside is that Python control flow (if/else, loops) cannot be used naturally — you need special graph operations like tf.cond and tf.while_loop.
- In a dynamic graph system (PyTorch, JAX in eager mode), the graph is built on-the-fly as Python code executes. Each line of Python immediately computes a value and appends a node to the graph. This means standard Python debugging tools (print statements, pdb, breakpoints) work naturally. Control flow is just regular Python.
- The industry converged on dynamic graphs for a simple reason: researcher productivity. In ML research, you spend most of your time writing and debugging new model architectures. The ability to set a breakpoint, inspect intermediate tensors, and step through the computation in a standard debugger dramatically accelerates the research cycle. The performance overhead of dynamic graphs (typically 10-20% slower than optimized static graphs) is acceptable because researcher time is more expensive than GPU time.
- The convergence is not absolute. TensorFlow 2.0 adopted eager execution by default but added tf.function for compiling hot paths into static graphs. PyTorch added torch.compile (PyTorch 2.0) which traces the dynamic graph and compiles it for performance. JAX takes a hybrid approach: eager by default, with jax.jit to compile functions. The modern consensus is: develop dynamically, deploy with compilation.
You need to implement a custom backward pass for a non-standard operation in PyTorch. Walk me through how torch.autograd.Function works and what pitfalls to watch for.
You need to implement a custom backward pass for a non-standard operation in PyTorch. Walk me through how torch.autograd.Function works and what pitfalls to watch for.
- You subclass torch.autograd.Function and implement two static methods: forward() and backward(). forward() receives input tensors and a context object, computes the output, and saves any tensors needed for the backward pass using ctx.save_for_backward(). backward() receives the upstream gradient (grad_output) and the context, and must return one gradient per forward input.
- The contract is: backward must return tensors with the same shape as the corresponding forward inputs. If an input does not need a gradient (like an integer parameter), return None for that position. Getting the number and order of returned gradients wrong is one of the most common bugs.
- Pitfall one: saving too much or too little in the context. If you forget to save an intermediate tensor and try to use a forward-pass local variable in backward, it will be garbage-collected and you get a crash or wrong result. If you save too many large tensors, you waste memory.
- Pitfall two: in-place modification of saved tensors. If you save tensor A in the forward pass and then modify A in-place before backward runs, the saved reference points to the modified data, not the original. Always save copies if there is any chance of in-place modification.
- Pitfall three: not handling batched inputs. Your backward must work for arbitrary batch sizes. A common bug is implementing the gradient for a single sample and then getting shape errors with batches.
- Always validate with torch.autograd.gradcheck, which compares your analytical backward against numerical finite differences. Run this with float64 tensors for maximum precision.
What is the 'gradient accumulation' bug in PyTorch, and why does the framework not zero gradients automatically?
What is the 'gradient accumulation' bug in PyTorch, and why does the framework not zero gradients automatically?
- In PyTorch, calling loss.backward() ADDS the computed gradients to the .grad attribute of each parameter rather than replacing them. If you forget to call optimizer.zero_grad() before computing the next batch’s gradients, the gradients from the previous batch are still there, and the new batch’s gradients are added on top. The result is that your effective gradient is the sum of all batches since the last zero_grad(), which is mathematically wrong for standard SGD.
- This is the single most common PyTorch bug for beginners. The symptom is that the model appears to train but converges to a worse solution or oscillates erratically. It is insidious because the code does not crash — it just silently computes wrong updates.
- Why PyTorch does not zero automatically: gradient accumulation is a deliberate feature, not a bug. It enables training with effective batch sizes larger than what fits in GPU memory. If your GPU can hold batch size 8 but you want the gradient quality of batch size 32, you run 4 forward-backward passes (each with batch 8) without zeroing gradients, then call optimizer.step(). The accumulated gradient is mathematically equivalent to computing the gradient over all 32 samples at once.
- The accumulation semantics also enable multi-task and multi-loss training. If you have two losses (classification loss + reconstruction loss), you can call loss1.backward() and loss2.backward() separately, and the gradients accumulate correctly.
- Best practice: always call optimizer.zero_grad() at the start of each training iteration (before loss.backward()), not after optimizer.step().